From abfcc6aac2d081744d04b65ca0600dfde197f9c1 Mon Sep 17 00:00:00 2001
From: Stefano <stefano.scafiti96@gmail.com>
Date: Sun, 6 Mar 2022 23:20:06 +0100
Subject: [PATCH] Negotiate capabilities during registration

---
 client/connection.go | 51 +++++++++++++++++++++--------
 client/handlers.go   | 78 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 116 insertions(+), 13 deletions(-)

diff --git a/client/connection.go b/client/connection.go
index 43bcc24..ce1b355 100644
--- a/client/connection.go
+++ b/client/connection.go
@@ -45,6 +45,12 @@ type Conn struct {
 	out         chan string
 	connected   bool
 
+	// Capabilities supported by the server
+	supportedCaps *capSet
+
+	// Capabilites currently enabled
+	currCaps *capSet
+
 	// CancelFunc and WaitGroup for goroutines
 	die context.CancelFunc
 	wg  sync.WaitGroup
@@ -89,6 +95,12 @@ type Config struct {
 	// Passed through to https://golang.org/pkg/net/#Dialer
 	DualStack bool
 
+	// Enable IRCv3 capability negotiation.
+	EnableCapabilityNegotiation bool
+
+	// A list of capabilities to request to the server during registration.
+	Capabilites []string
+
 	// Replaceable function to customise the 433 handler's new nick.
 	// By default an underscore "_" is appended to the current nick.
 	NewNick func(string) string
@@ -125,12 +137,13 @@ type Config struct {
 // name, but these are optional.
 func NewConfig(nick string, args ...string) *Config {
 	cfg := &Config{
-		Me:       &state.Nick{Nick: nick},
-		PingFreq: 3 * time.Minute,
-		NewNick:  DefaultNewNick,
-		Recover:  (*Conn).LogPanic, // in dispatch.go
-		SplitLen: defaultSplit,
-		Timeout:  60 * time.Second,
+		Me:                          &state.Nick{Nick: nick},
+		PingFreq:                    3 * time.Minute,
+		NewNick:                     DefaultNewNick,
+		Recover:                     (*Conn).LogPanic, // in dispatch.go
+		SplitLen:                    defaultSplit,
+		Timeout:                     60 * time.Second,
+		EnableCapabilityNegotiation: true,
 	}
 	cfg.Me.Ident = "goirc"
 	if len(args) > 0 && args[0] != "" {
@@ -204,13 +217,15 @@ func Client(cfg *Config) *Conn {
 	}
 
 	conn := &Conn{
-		cfg:         cfg,
-		dialer:      dialer,
-		intHandlers: handlerSet(),
-		fgHandlers:  handlerSet(),
-		bgHandlers:  handlerSet(),
-		stRemovers:  make([]Remover, 0, len(stHandlers)),
-		lastsent:    time.Now(),
+		cfg:           cfg,
+		dialer:        dialer,
+		intHandlers:   handlerSet(),
+		fgHandlers:    handlerSet(),
+		bgHandlers:    handlerSet(),
+		stRemovers:    make([]Remover, 0, len(stHandlers)),
+		lastsent:      time.Now(),
+		supportedCaps: capabilitySet(),
+		currCaps:      capabilitySet(),
 	}
 	conn.addIntHandlers()
 	return conn
@@ -289,6 +304,16 @@ func (conn *Conn) DisableStateTracking() {
 	}
 }
 
+// SupportsCapability returns true if the server supports the given capability.
+func (conn *Conn) SupportsCapability(cap string) bool {
+	return conn.currCaps.Has(cap)
+}
+
+// HasCapability returns true if the given capability has been acked by the server during negotiation.
+func (conn *Conn) HasCapability(cap string) bool {
+	return conn.currCaps.Has(cap)
+}
+
 // Per-connection state initialisation.
 func (conn *Conn) initialise() {
 	conn.io = nil
diff --git a/client/handlers.go b/client/handlers.go
index 24165ad..b799788 100644
--- a/client/handlers.go
+++ b/client/handlers.go
@@ -5,6 +5,7 @@ package client
 
 import (
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/fluffle/goirc/logging"
@@ -18,6 +19,7 @@ var intHandlers = map[string]HandlerFunc{
 	CTCP:     (*Conn).h_CTCP,
 	NICK:     (*Conn).h_NICK,
 	PING:     (*Conn).h_PING,
+	CAP:      (*Conn).h_CAP,
 }
 
 func (conn *Conn) addIntHandlers() {
@@ -35,6 +37,10 @@ func (conn *Conn) h_PING(line *Line) {
 
 // Handler for initial registration with server once tcp connection is made.
 func (conn *Conn) h_REGISTER(line *Line) {
+	if conn.cfg.EnableCapabilityNegotiation {
+		conn.Cap(CAP_LS)
+	}
+
 	if conn.cfg.Pass != "" {
 		conn.Pass(conn.cfg.Pass)
 	}
@@ -42,6 +48,78 @@ func (conn *Conn) h_REGISTER(line *Line) {
 	conn.User(conn.cfg.Me.Ident, conn.cfg.Me.Name)
 }
 
+func (conn *Conn) negotiateCapabilities(supportedCaps []string) {
+	for _, cap := range supportedCaps {
+		conn.supportedCaps.Add(cap)
+	}
+
+	// intersect capabilities supported by server with those requested
+	reqCaps := make([]string, 0)
+	for _, cap := range conn.cfg.Capabilites {
+		if conn.supportedCaps.Has(cap) {
+			reqCaps = append(reqCaps, cap)
+		}
+	}
+
+	if len(reqCaps) > 0 {
+		conn.Cap(CAP_REQ, reqCaps...)
+	} else {
+		conn.Cap(CAP_END)
+	}
+}
+
+func (conn *Conn) handleCapAck(caps []string) {
+	for _, cap := range caps {
+		conn.currCaps.Add(cap)
+	}
+	conn.Cap(CAP_END)
+}
+
+const (
+	CAP_LS  = "LS"
+	CAP_REQ = "REQ"
+	CAP_ACK = "ACK"
+	CAP_NAK = "NAK"
+	CAP_END = "END"
+)
+
+type capSet struct {
+	caps map[string]struct{}
+	mu   sync.RWMutex
+}
+
+func capabilitySet() *capSet {
+	return &capSet{
+		caps: make(map[string]struct{}),
+	}
+}
+
+func (c *capSet) Add(cap string) {
+	c.mu.Lock()
+	c.caps[cap] = struct{}{}
+	c.mu.Unlock()
+}
+
+func (c *capSet) Has(cap string) bool {
+	c.mu.RLock()
+	_, ok := c.caps[cap]
+	c.mu.RUnlock()
+	return ok
+}
+
+// Handler for capability negotiation commands.
+func (conn *Conn) h_CAP(line *Line) {
+	subcommand := line.Args[1]
+
+	switch subcommand {
+	case CAP_LS:
+		conn.negotiateCapabilities(strings.Fields(line.Text()))
+	case CAP_ACK:
+		conn.handleCapAck(strings.Fields(line.Text()))
+	case CAP_NAK:
+	}
+}
+
 // Handler to trigger a CONNECTED event on receipt of numeric 001
 // :<server> 001 <nick> :Welcome message <nick>!<user>@<host>
 func (conn *Conn) h_001(line *Line) {