From 18c401513a010e9989689ae4af54d154a44d24ca Mon Sep 17 00:00:00 2001
From: Eric Daniels <eric@erdaniels.com>
Date: Mon, 12 Aug 2024 15:00:40 -0400
Subject: [PATCH] Do not add streams after association is closed

---
 association.go      |  6 ++++++
 association_test.go | 16 ++++++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/association.go b/association.go
index d498e155..2e8708fc 100644
--- a/association.go
+++ b/association.go
@@ -33,6 +33,7 @@ var (
 	ErrChunk                         = errors.New("abort chunk, with following errors")
 	ErrShutdownNonEstablished        = errors.New("shutdown called in non-established state")
 	ErrAssociationClosedBeforeConn   = errors.New("association closed before connecting")
+	ErrAssociationClosed             = errors.New("association closed")
 	ErrSilentlyDiscard               = errors.New("silently discard")
 	ErrInitNotStoredToSend           = errors.New("the init not stored to send")
 	ErrCookieEchoNotStoredToSend     = errors.New("cookieEcho not stored to send")
@@ -1505,6 +1506,11 @@ func (a *Association) OpenStream(streamIdentifier uint16, defaultPayloadType Pay
 	a.lock.Lock()
 	defer a.lock.Unlock()
 
+	switch a.getState() {
+	case shutdownAckSent, shutdownPending, shutdownReceived, shutdownSent, closed:
+		return nil, ErrAssociationClosed
+	}
+
 	return a.getOrCreateStream(streamIdentifier, false, defaultPayloadType), nil
 }
 
diff --git a/association_test.go b/association_test.go
index dd833351..483be4b5 100644
--- a/association_test.go
+++ b/association_test.go
@@ -3455,3 +3455,19 @@ func TestAssociation_ReconfigRequestsLimited(t *testing.T) {
 	require.NoError(t, a1.Close())
 	require.NoError(t, a2.Close())
 }
+
+func TestAssociation_OpenStreamAfterClose(t *testing.T) {
+	checkGoroutineLeaks(t)
+
+	a1, a2, err := createAssocs()
+	require.NoError(t, err)
+
+	require.NoError(t, a1.Close())
+	require.NoError(t, a2.Close())
+
+	_, err = a1.OpenStream(1, PayloadTypeWebRTCString)
+	require.ErrorIs(t, err, ErrAssociationClosed)
+
+	_, err = a2.OpenStream(1, PayloadTypeWebRTCString)
+	require.ErrorIs(t, err, ErrAssociationClosed)
+}