From 3b0639bacc848c2b703f64aceffc598023ed6184 Mon Sep 17 00:00:00 2001
From: Simon Ser <contact@emersion.fr>
Date: Thu, 6 Feb 2020 22:25:32 +0100
Subject: [PATCH] Keep track of upstream connections in a list

---
 server.go | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/server.go b/server.go
index f231908..dc0a193 100644
--- a/server.go
+++ b/server.go
@@ -45,6 +45,7 @@ type Server struct {
 
 	lock            sync.Mutex
 	downstreamConns []*downstreamConn
+	upstreamConns   []*upstreamConn
 }
 
 func (s *Server) prefix() *irc.Prefix {
@@ -61,9 +62,23 @@ func (s *Server) Run() {
 				s.Logger.Printf("failed to connect to upstream server %q: %v", upstream.Addr, err)
 				return
 			}
+
+			s.lock.Lock()
+			s.upstreamConns = append(s.upstreamConns, conn)
+			s.lock.Unlock()
+
 			if err := conn.readMessages(); err != nil {
 				conn.logger.Printf("failed to handle messages: %v", err)
 			}
+
+			s.lock.Lock()
+			for i, c := range s.upstreamConns {
+				if c == conn {
+					s.upstreamConns = append(s.upstreamConns[:i], s.upstreamConns[i+1:]...)
+					break
+				}
+			}
+			s.lock.Unlock()
 		}()
 	}
 }
@@ -80,9 +95,11 @@ func (s *Server) Serve(ln net.Listener) error {
 			s.lock.Lock()
 			s.downstreamConns = append(s.downstreamConns, conn)
 			s.lock.Unlock()
+
 			if err := conn.readMessages(); err != nil {
 				conn.logger.Printf("failed to handle messages: %v", err)
 			}
+
 			s.lock.Lock()
 			for i, c := range s.downstreamConns {
 				if c == conn {