From 8bbba42aefd811fe42e93d28b36fd98bc0f8d2ef Mon Sep 17 00:00:00 2001
From: Simon Ser <contact@emersion.fr>
Date: Thu, 6 Feb 2020 21:11:35 +0100
Subject: [PATCH] Maintain a list of downstream connections

---
 downstream.go | 75 +++++++++++++++++++++++++++++----------------------
 server.go     | 10 ++++---
 2 files changed, 50 insertions(+), 35 deletions(-)

diff --git a/downstream.go b/downstream.go
index 7c0268b..828595a 100644
--- a/downstream.go
+++ b/downstream.go
@@ -42,6 +42,8 @@ type downstreamConn struct {
 	net        net.Conn
 	irc        *irc.Conn
 	srv        *Server
+	logger     Logger
+
 	registered bool
 	closed     bool
 	nick       string
@@ -49,6 +51,45 @@ type downstreamConn struct {
 	realname   string
 }
 
+func newDownstreamConn(srv *Server, netConn net.Conn) *downstreamConn {
+	return &downstreamConn{
+		net: netConn,
+		irc: irc.NewConn(netConn),
+		srv: srv,
+		logger: &prefixLogger{srv.Logger, fmt.Sprintf("downstream %q: ", netConn.RemoteAddr())},
+	}
+}
+
+func (c *downstreamConn) readMessages() error {
+	c.logger.Printf("new connection")
+	defer c.Close()
+
+	for {
+		msg, err := c.irc.ReadMessage()
+		if err == io.EOF {
+			break
+		} else if err != nil {
+			return fmt.Errorf("failed to read IRC command: %v", err)
+		}
+
+		err = c.handleMessage(msg)
+		if ircErr, ok := err.(ircError); ok {
+			ircErr.Message.Prefix = c.srv.prefix()
+			if err := c.WriteMessage(ircErr.Message); err != nil {
+				return fmt.Errorf("failed to write IRC reply: %v", err)
+			}
+		} else if err != nil {
+			return fmt.Errorf("failed to handle IRC command %q: %v", msg.Command, err)
+		}
+
+		if c.closed {
+			return nil
+		}
+	}
+
+	return c.Close()
+}
+
 func (c *downstreamConn) Close() error {
 	if err := c.net.Close(); err != nil {
 		return err
@@ -95,6 +136,7 @@ func (c *downstreamConn) handleMessageUnregistered(msg *irc.Message) error {
 	case "QUIT":
 		return c.Close()
 	default:
+		c.logger.Printf("unhandled message: %v", msg)
 		return newUnknownCommandError(msg.Command)
 	}
 	if c.username != "" && c.nick != "" {
@@ -162,38 +204,7 @@ func (c *downstreamConn) handleMessageRegistered(msg *irc.Message) error {
 	case "QUIT":
 		return c.Close()
 	default:
+		c.logger.Printf("unhandled message: %v", msg)
 		return newUnknownCommandError(msg.Command)
 	}
 }
-
-func handleConn(s *Server, netConn net.Conn) error {
-	s.Logger.Printf("Handling connection from %v", netConn.RemoteAddr())
-
-	c := downstreamConn{net: netConn, irc: irc.NewConn(netConn), srv: s}
-	defer c.Close()
-	for {
-		msg, err := c.irc.ReadMessage()
-		if err == io.EOF {
-			break
-		} else if err != nil {
-			return fmt.Errorf("failed to read IRC command: %v", err)
-		}
-		s.Logger.Printf("Downstream message: %v", msg)
-
-		err = c.handleMessage(msg)
-		if ircErr, ok := err.(ircError); ok {
-			ircErr.Message.Prefix = s.prefix()
-			if err := c.WriteMessage(ircErr.Message); err != nil {
-				return fmt.Errorf("failed to write IRC reply: %v", err)
-			}
-		} else if err != nil {
-			return fmt.Errorf("failed to handle IRC command %q: %v", msg.Command, err)
-		}
-
-		if c.closed {
-			return nil
-		}
-	}
-
-	return c.Close()
-}
diff --git a/server.go b/server.go
index ec56e0f..986ac65 100644
--- a/server.go
+++ b/server.go
@@ -41,6 +41,8 @@ type Server struct {
 	Hostname  string
 	Logger    Logger
 	Upstreams []Upstream // TODO: per-user
+
+	downstreamConns []*downstreamConn
 }
 
 func (s *Server) prefix() *irc.Prefix {
@@ -61,14 +63,16 @@ func (s *Server) Run() {
 
 func (s *Server) Serve(ln net.Listener) error {
 	for {
-		c, err := ln.Accept()
+		netConn, err := ln.Accept()
 		if err != nil {
 			return fmt.Errorf("failed to accept connection: %v", err)
 		}
 
+		conn := newDownstreamConn(s, netConn)
+		s.downstreamConns = append(s.downstreamConns, conn)
 		go func() {
-			if err := handleConn(s, c); err != nil {
-				s.Logger.Printf("Error handling connection: %v", err)
+			if err := conn.readMessages(); err != nil {
+				conn.logger.Printf("Error handling messages: %v", err)
 			}
 		}()
 	}