From de3c4062196056a55b244af367fe29ef94404946 Mon Sep 17 00:00:00 2001
From: Eric Daniels <eric@erdaniels.com>
Date: Wed, 14 Aug 2024 18:35:51 -0400
Subject: [PATCH] set close state before unregisterStream

---
 association.go      |  1 +
 association_test.go | 22 ++++++++++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/association.go b/association.go
index 5be8838..29e9978 100644
--- a/association.go
+++ b/association.go
@@ -573,6 +573,7 @@ func (a *Association) readLoop() {
 		a.closeWriteLoopOnce.Do(func() { close(a.closeWriteLoopCh) })
 
 		a.lock.Lock()
+		a.setState(closed)
 		for _, s := range a.streams {
 			a.unregisterStream(s, closeErr)
 		}
diff --git a/association_test.go b/association_test.go
index 1a51214..141daa1 100644
--- a/association_test.go
+++ b/association_test.go
@@ -3475,3 +3475,25 @@ func TestAssociation_OpenStreamAfterClose(t *testing.T) {
 	_, err = a2.OpenStream(1, PayloadTypeWebRTCString)
 	require.ErrorIs(t, err, ErrAssociationClosed)
 }
+
+// https://github.com/pion/sctp/pull/350
+// may need to run with a high test count to reproduce if there
+// is ever a regression.
+func TestAssociation_OpenStreamAfterInternalClose(t *testing.T) {
+	checkGoroutineLeaks(t)
+
+	a1, a2, err := createAssocs()
+	require.NoError(t, err)
+
+	a1.netConn.Close()
+	a2.netConn.Close()
+
+	a1.OpenStream(1, PayloadTypeWebRTCString)
+	a2.OpenStream(1, PayloadTypeWebRTCString)
+
+	require.NoError(t, a1.Close())
+	require.NoError(t, a2.Close())
+
+	require.Equal(t, 0, len(a1.streams))
+	require.Equal(t, 0, len(a2.streams))
+}