From 519cdd9e38936f0cc40d5d241e0cdf364c9d1bf1 Mon Sep 17 00:00:00 2001
From: Simon Ser <contact@emersion.fr>
Date: Fri, 7 Feb 2020 12:42:24 +0100
Subject: [PATCH] Close connection from writer goroutine

Connections were being closed from the reader goroutine, causing issues
when sending messages and immediately closing the connection.
---
 downstream.go | 10 +++++++---
 upstream.go   | 12 ++++++------
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/downstream.go b/downstream.go
index 4420b40..72445ff 100644
--- a/downstream.go
+++ b/downstream.go
@@ -70,6 +70,11 @@ func newDownstreamConn(srv *Server, netConn net.Conn) *downstreamConn {
 				conn.logger.Printf("failed to write message: %v", err)
 			}
 		}
+		if err := conn.net.Close(); err != nil {
+			conn.logger.Printf("failed to close connection: %v", err)
+		} else {
+			conn.logger.Printf("connection closed")
+		}
 	}()
 
 	return conn
@@ -85,7 +90,6 @@ func (c *downstreamConn) prefix() *irc.Prefix {
 
 func (c *downstreamConn) readMessages() error {
 	c.logger.Printf("new connection")
-	defer c.Close()
 
 	for {
 		msg, err := c.irc.ReadMessage()
@@ -108,7 +112,7 @@ func (c *downstreamConn) readMessages() error {
 		}
 	}
 
-	return c.Close()
+	return nil
 }
 
 func (c *downstreamConn) Close() error {
@@ -129,7 +133,7 @@ func (c *downstreamConn) Close() error {
 	close(c.messages)
 	c.closed = true
 
-	return c.net.Close()
+	return nil
 }
 
 func (c *downstreamConn) handleMessage(msg *irc.Message) error {
diff --git a/upstream.go b/upstream.go
index e3bd487..6c6a9a1 100644
--- a/upstream.go
+++ b/upstream.go
@@ -71,6 +71,11 @@ func connectToUpstream(u *user, upstream *Upstream) (*upstreamConn, error) {
 				conn.logger.Printf("failed to write message: %v", err)
 			}
 		}
+		if err := conn.net.Close(); err != nil {
+			conn.logger.Printf("failed to close connection: %v", err)
+		} else {
+			conn.logger.Printf("connection closed")
+		}
 	}()
 
 	return conn, nil
@@ -80,9 +85,6 @@ func (c *upstreamConn) Close() error {
 	if c.closed {
 		return fmt.Errorf("upstream connection already closed")
 	}
-	if err := c.net.Close(); err != nil {
-		return err
-	}
 	close(c.messages)
 	c.closed = true
 	return nil
@@ -330,8 +332,6 @@ func (c *upstreamConn) register() {
 }
 
 func (c *upstreamConn) readMessages() error {
-	defer c.Close()
-
 	for {
 		msg, err := c.irc.ReadMessage()
 		if err == io.EOF {
@@ -345,5 +345,5 @@ func (c *upstreamConn) readMessages() error {
 		}
 	}
 
-	return c.Close()
+	return nil
 }