From 194ba43a50f4f3f33d9d7e1255219f9fe1b9da5c Mon Sep 17 00:00:00 2001
From: Eric Daniels <eric@erdaniels.com>
Date: Mon, 12 Aug 2024 16:48:57 -0400
Subject: [PATCH] Unblock handshake ch on close

---
 association.go      | 23 +++++++++++++++++++----
 association_test.go |  4 ++++
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/association.go b/association.go
index 2e8708fc..de658fd4 100644
--- a/association.go
+++ b/association.go
@@ -1376,7 +1376,11 @@ func (a *Association) handleCookieEcho(c *chunkCookieEcho) []*packet {
 
 		a.setState(established)
 		// Note: This is a future place where the user could be notified (COMMUNICATION UP)
-		a.handshakeCompletedCh <- nil
+		select {
+		case a.handshakeCompletedCh <- nil:
+		case <-a.closeWriteLoopCh: // check the write side since we are in a read path
+			return nil
+		}
 	}
 
 	p := &packet{
@@ -1405,7 +1409,10 @@ func (a *Association) handleCookieAck() {
 
 	a.setState(established)
 	// Note: This is a future place where the user could be notified (COMMUNICATION UP)
-	a.handshakeCompletedCh <- nil
+	select {
+	case a.handshakeCompletedCh <- nil:
+	case <-a.closeWriteLoopCh: // check the write side since we are in a read path
+	}
 }
 
 // The caller should hold the lock.
@@ -2698,13 +2705,21 @@ func (a *Association) onRetransmissionFailure(id int) {
 
 	if id == timerT1Init {
 		a.log.Errorf("[%s] retransmission failure: T1-init", a.name)
-		a.handshakeCompletedCh <- ErrHandshakeInitAck
+		select {
+		case a.handshakeCompletedCh <- ErrHandshakeInitAck:
+		case <-a.closeWriteLoopCh: // check the read/write sides
+		case <-a.readLoopCloseCh:
+		}
 		return
 	}
 
 	if id == timerT1Cookie {
 		a.log.Errorf("[%s] retransmission failure: T1-cookie", a.name)
-		a.handshakeCompletedCh <- ErrHandshakeCookieEcho
+		select {
+		case a.handshakeCompletedCh <- ErrHandshakeCookieEcho:
+		case <-a.closeWriteLoopCh: // check the read/write sides
+		case <-a.readLoopCloseCh:
+		}
 		return
 	}
 
diff --git a/association_test.go b/association_test.go
index 483be4b5..1a512149 100644
--- a/association_test.go
+++ b/association_test.go
@@ -3222,6 +3222,10 @@ func TestAssociation_Abort(t *testing.T) {
 
 // TestAssociation_createClientWithContext tests that the client is closed when the context is canceled.
 func TestAssociation_createClientWithContext(t *testing.T) {
+	// Limit runtime in case of deadlocks
+	lim := test.TimeOut(time.Second * 5)
+	defer lim.Stop()
+
 	checkGoroutineLeaks(t)
 
 	udp1, udp2 := createUDPConnPair()