From 58c9607dfbf015b1db117a8a900546771399a550 Mon Sep 17 00:00:00 2001
From: Luca Bigliardi <shammash@google.com>
Date: Sat, 27 Mar 2021 11:35:19 +0100
Subject: [PATCH] Fix connection cleanup when context is canceled

Signed-off-by: Luca Bigliardi <shammash@google.com>
---
 client/connection.go      |  8 ++++++--
 client/connection_test.go | 22 ++++++++++++++++++++++
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/client/connection.go b/client/connection.go
index f98bc68..43bcc24 100644
--- a/client/connection.go
+++ b/client/connection.go
@@ -503,13 +503,17 @@ func (conn *Conn) ping(ctx context.Context) {
 // It pulls Lines from the input channel and dispatches them to any
 // handlers that have been registered for that IRC verb.
 func (conn *Conn) runLoop(ctx context.Context) {
-	defer conn.wg.Done()
 	for {
 		select {
 		case line := <-conn.in:
 			conn.dispatch(line)
 		case <-ctx.Done():
-			// control channel closed, bail out
+			// control channel closed, trigger Cancel() to clean
+			// things up properly and bail out
+
+			// We can't defer this, because Close() waits for it.
+			conn.wg.Done()
+			conn.Close()
 			return
 		}
 	}
diff --git a/client/connection_test.go b/client/connection_test.go
index c6a3b95..50d32bc 100644
--- a/client/connection_test.go
+++ b/client/connection_test.go
@@ -103,6 +103,28 @@ func TestEOF(t *testing.T) {
 	}
 }
 
+func TestCleanupOnContextDone(t *testing.T) {
+	c, s := setUp(t)
+	// Since we're not using tearDown() here, manually call Finish()
+	defer s.ctrl.Finish()
+
+	// Close() triggers DISCONNECT handler after cleaning up the state
+	// use this as a proxy to check that Close() was indeed called
+	dcon := callCheck(t)
+	c.Handle(DISCONNECTED, dcon)
+
+	// Simulate context cancelation using our cancel func
+	c.die()
+
+	// Verify that disconnected handler was called
+	dcon.assertWasCalled("Conn did not call disconnected handlers.")
+
+	// Verify that the connection no longer thinks it's connected
+	if c.Connected() {
+		t.Errorf("Conn still thinks it's connected to the server.")
+	}
+}
+
 func TestClientAndStateTracking(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	st := state.NewMockTracker(ctrl)