From af76c3868ab0dab531648d31608633eff94dd20f Mon Sep 17 00:00:00 2001
From: Simon Ser <contact@emersion.fr>
Date: Mon, 16 Mar 2020 15:05:24 +0100
Subject: [PATCH] Add CAP support for downstream connections

---
 downstream.go | 111 +++++++++++++++++++++++++++++++++++++++++++++++++-
 irc.go        |  27 ++++++------
 2 files changed, 124 insertions(+), 14 deletions(-)

diff --git a/downstream.go b/downstream.go
index eb51e6a..ef61d34 100644
--- a/downstream.go
+++ b/downstream.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"strconv"
 	"strings"
 	"sync"
 	"time"
@@ -71,6 +72,10 @@ type downstreamConn struct {
 	password    string   // empty after authentication
 	network     *network // can be nil
 
+	negociatingCaps bool
+	capVersion      int
+	caps            map[string]bool
+
 	lock        sync.Mutex
 	ourMessages map[*irc.Message]struct{}
 }
@@ -84,6 +89,7 @@ func newDownstreamConn(srv *Server, netConn net.Conn) *downstreamConn {
 		outgoing:     make(chan *irc.Message, 64),
 		ringMessages: make(chan ringMessage),
 		closed:       make(chan struct{}),
+		caps:         make(map[string]bool),
 		ourMessages:  make(map[*irc.Message]struct{}),
 	}
 
@@ -328,16 +334,119 @@ func (dc *downstreamConn) handleMessageUnregistered(msg *irc.Message) error {
 		if err := parseMessageParams(msg, &dc.password); err != nil {
 			return err
 		}
+	case "CAP":
+		var subCmd string
+		if err := parseMessageParams(msg, &subCmd); err != nil {
+			return err
+		}
+		subCmd = strings.ToUpper(subCmd)
+		if err := dc.handleCapCommand(subCmd, msg.Params[1:]); err != nil {
+			return err
+		}
 	default:
 		dc.logger.Printf("unhandled message: %v", msg)
 		return newUnknownCommandError(msg.Command)
 	}
-	if dc.rawUsername != "" && dc.nick != "" {
+	if dc.rawUsername != "" && dc.nick != "" && !dc.negociatingCaps {
 		return dc.register()
 	}
 	return nil
 }
 
+func (dc *downstreamConn) handleCapCommand(cmd string, args []string) error {
+	replyTo := dc.nick
+	if !dc.registered {
+		replyTo = "*"
+	}
+
+	switch cmd {
+	case "LS":
+		if len(args) > 0 {
+			var err error
+			if dc.capVersion, err = strconv.Atoi(args[0]); err != nil {
+				return err
+			}
+		}
+
+		var caps []string
+		/*if dc.capVersion >= 302 {
+			caps = append(caps, "sasl=PLAIN")
+		} else {
+			caps = append(caps, "sasl")
+		}*/
+
+		// TODO: multi-line replies
+		dc.SendMessage(&irc.Message{
+			Prefix:  dc.srv.prefix(),
+			Command: "CAP",
+			Params:  []string{replyTo, "LS", strings.Join(caps, " ")},
+		})
+
+		if !dc.registered {
+			dc.negociatingCaps = true
+		}
+	case "LIST":
+		var caps []string
+		for name := range dc.caps {
+			caps = append(caps, name)
+		}
+
+		// TODO: multi-line replies
+		dc.SendMessage(&irc.Message{
+			Prefix:  dc.srv.prefix(),
+			Command: "CAP",
+			Params:  []string{replyTo, "LIST", strings.Join(caps, " ")},
+		})
+	case "REQ":
+		if len(args) == 0 {
+			return ircError{&irc.Message{
+				Command: err_invalidcapcmd,
+				Params:  []string{replyTo, cmd, "Missing argument in CAP REQ command"},
+			}}
+		}
+
+		caps := strings.Fields(args[0])
+		ack := true
+		for _, name := range caps {
+			name = strings.ToLower(name)
+			enable := !strings.HasPrefix(name, "-")
+			if !enable {
+				name = strings.TrimPrefix(name, "-")
+			}
+
+			enabled := dc.caps[name]
+			if enable == enabled {
+				continue
+			}
+
+			switch name {
+			/*case "sasl":
+				dc.caps[name] = enable*/
+			default:
+				ack = false
+			}
+		}
+
+		reply := "NAK"
+		if ack {
+			reply = "ACK"
+		}
+		dc.SendMessage(&irc.Message{
+			Prefix:  dc.srv.prefix(),
+			Command: "CAP",
+			Params:  []string{replyTo, reply, args[0]},
+		})
+	case "END":
+		dc.negociatingCaps = false
+	default:
+		return ircError{&irc.Message{
+			Command: err_invalidcapcmd,
+			Params:  []string{replyTo, cmd, "Unknown CAP command"},
+		}}
+	}
+	return nil
+}
+
 func sanityCheckServer(addr string) error {
 	dialer := net.Dialer{Timeout: 30 * time.Second}
 	conn, err := tls.DialWithDialer(&dialer, "tcp", addr, nil)
diff --git a/irc.go b/irc.go
index 001c582..e523856 100644
--- a/irc.go
+++ b/irc.go
@@ -8,19 +8,20 @@ import (
 )
 
 const (
-	rpl_statsping    = "246"
-	rpl_localusers   = "265"
-	rpl_globalusers  = "266"
-	rpl_topicwhotime = "333"
-	rpl_loggedin     = "900"
-	rpl_loggedout    = "901"
-	err_nicklocked   = "902"
-	rpl_saslsuccess  = "903"
-	err_saslfail     = "904"
-	err_sasltoolong  = "905"
-	err_saslaborted  = "906"
-	err_saslalready  = "907"
-	rpl_saslmechs    = "908"
+	rpl_statsping     = "246"
+	rpl_localusers    = "265"
+	rpl_globalusers   = "266"
+	rpl_topicwhotime  = "333"
+	err_invalidcapcmd = "410"
+	rpl_loggedin      = "900"
+	rpl_loggedout     = "901"
+	err_nicklocked    = "902"
+	rpl_saslsuccess   = "903"
+	err_saslfail      = "904"
+	err_sasltoolong   = "905"
+	err_saslaborted   = "906"
+	err_saslalready   = "907"
+	rpl_saslmechs     = "908"
 )
 
 type modeSet string