From 525734eeb346f4051eeac16be7ee34d181caab2d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?P=C5=99emysl=20Janouch?= <p@janouch.name>
Date: Sun, 22 Jul 2018 13:45:32 +0200
Subject: [PATCH] tls-autodetect: fix client-initiated shutdown

---
 prototypes/tls-autodetect.go | 57 ++++++++++++++++++------------------
 1 file changed, 29 insertions(+), 28 deletions(-)

diff --git a/prototypes/tls-autodetect.go b/prototypes/tls-autodetect.go
index 50017d5..aaff3d9 100644
--- a/prototypes/tls-autodetect.go
+++ b/prototypes/tls-autodetect.go
@@ -80,14 +80,15 @@ type connCloseWrite interface {
 }
 
 type client struct {
-	transport  net.Conn       // underlying connection
-	tls        *tls.Conn      // TLS, if detected
-	conn       connCloseWrite // high-level connection
-	inQ        []byte         // unprocessed input
-	outQ       []byte         // unprocessed output
-	writing    bool           // whether a writing goroutine is running
-	inShutdown bool           // whether we're closing connection
-	killTimer  *time.Timer    // timeout
+	transport net.Conn       // underlying connection
+	tls       *tls.Conn      // TLS, if detected
+	conn      connCloseWrite // high-level connection
+	inQ       []byte         // unprocessed input
+	outQ      []byte         // unprocessed output
+	reading   bool           // whether a reading goroutine is running
+	writing   bool           // whether a writing goroutine is running
+	closing   bool           // whether we're closing the connection
+	killTimer *time.Timer    // timeout
 }
 
 type preparedEvent struct {
@@ -164,7 +165,7 @@ func forceShutdown(reason string) {
 // --- Client ------------------------------------------------------------------
 
 func (c *client) send(line string) {
-	if !c.inShutdown {
+	if !c.closing {
 		c.outQ = append(c.outQ, (line + "\r\n")...)
 		c.flushOutQ()
 	}
@@ -172,7 +173,7 @@ func (c *client) send(line string) {
 
 // Tear down the client connection, trying to do so in a graceful manner.
 func (c *client) closeLink() {
-	if c.inShutdown {
+	if c.closing {
 		return
 	}
 	if c.conn == nil {
@@ -180,13 +181,13 @@ func (c *client) closeLink() {
 		return
 	}
 
-	// Since we send this goodbye, we don't need to call CloseWrite.
+	// Since we send this goodbye, we don't need to call CloseWrite here.
 	c.send("Goodbye")
 	c.killTimer = time.AfterFunc(3*time.Second, func() {
 		timeouts <- c
 	})
 
-	c.inShutdown = true
+	c.closing = true
 }
 
 // Close the connection and forget about the client.
@@ -218,6 +219,7 @@ func (c *client) onPrepared(host string, isTLS bool) {
 
 	// TODO: Save the host in the client structure.
 	go read(c)
+	c.reading = true
 }
 
 // Handle the results from trying to read from the client connection.
@@ -236,28 +238,34 @@ func (c *client) onRead(data []byte, readErr error) {
 	}
 
 	// TODO: Inform the client about the inQ overrun in the farewell message.
-	// TODO: We should stop receiving any more data from this client.
+	// TODO: We should stop receiving any more data from this client, or at
+	// least stop extending the inQ if we don't want to signal tho goroutine.
 	if len(c.inQ) > 8192 {
 		c.closeLink()
 		return
 	}
 
-	if readErr == io.EOF {
-		if c.inShutdown {
+	if readErr != nil {
+		c.reading = false
+
+		if readErr != io.EOF {
+			log.Println(readErr)
+			c.destroy()
+		} else if c.closing {
+			// Disregarding whether a clean shutdown has happened or not.
+			log.Println("client finished shutdown")
 			c.destroy()
 		} else {
+			log.Println("client EOF")
 			c.closeLink()
 		}
-	} else if readErr != nil {
-		log.Println(readErr)
-		c.destroy()
 	}
 }
 
 // Spawn a goroutine to flush the outQ if possible and necessary.  If the
 // connection is not ready yet, it needs to be retried as soon as it becomes.
 func (c *client) flushOutQ() {
-	if c.conn != nil && !c.writing {
+	if !c.writing && c.conn != nil {
 		go write(c, c.outQ)
 		c.writing = true
 	}
@@ -273,14 +281,8 @@ func (c *client) onWrite(written int, writeErr error) {
 		c.destroy()
 	} else if len(c.outQ) > 0 {
 		c.flushOutQ()
-	} else if c.inShutdown {
-		if c.conn != nil {
-			// FIXME: This is only correct for when /we/ initiate the shutdown,
-			// otherwise we should perhaps just Close. Though even if we
-			// Close, there's a/ no writer to fail on it, and b/ the reader
-			// has already exited, too, which is why the client stays alive
-			// up until the timeout. It seems that in that case we need to
-			// call c.destroy().
+	} else if c.closing {
+		if c.reading {
 			c.conn.CloseWrite()
 		} else {
 			c.destroy()
@@ -341,7 +343,6 @@ func prepare(client *client) {
 	}
 
 	prepared <- preparedEvent{client, host, isTLS}
-
 }
 
 func read(client *client) {