From 398855b30985e0324fe1feb731b27fa576938ff1 Mon Sep 17 00:00:00 2001
From: sukun <sukunrt@gmail.com>
Date: Wed, 14 Aug 2024 19:40:01 +0530
Subject: [PATCH] check conn close on OpenStream

---
 go.mod                                 |  4 +---
 go.sum                                 |  4 ++--
 p2p/transport/webrtc/connection.go     | 10 ++++++++--
 p2p/transport/webrtc/transport.go      |  1 +
 p2p/transport/webrtc/transport_test.go |  8 ++++----
 test-plans/go.mod                      |  2 +-
 test-plans/go.sum                      |  4 ++--
 7 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/go.mod b/go.mod
index b1ff398586..bc17644c92 100644
--- a/go.mod
+++ b/go.mod
@@ -50,7 +50,7 @@ require (
 	github.com/pion/logging v0.2.2
 	github.com/pion/sctp v1.8.20
 	github.com/pion/stun v0.6.1
-	github.com/pion/webrtc/v3 v3.2.50
+	github.com/pion/webrtc/v3 v3.2.52-0.20240813151442-cef1db8adbf1
 	github.com/prometheus/client_golang v1.19.1
 	github.com/prometheus/client_model v0.6.1
 	github.com/quic-go/quic-go v0.45.2
@@ -130,5 +130,3 @@ require (
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 	lukechampine.com/blake3 v1.3.0 // indirect
 )
-
-replace github.com/pion/webrtc/v3 => github.com/pion/webrtc/v3 v3.2.52-0.20240812164622-d0c92a5c4578
diff --git a/go.sum b/go.sum
index 3be88d9a5c..a32962c3af 100644
--- a/go.sum
+++ b/go.sum
@@ -315,8 +315,8 @@ github.com/pion/transport/v3 v3.0.6/go.mod h1:HvJr2N/JwNJAfipsRleqwFoR3t/pWyHeZU
 github.com/pion/turn/v2 v2.1.3/go.mod h1:huEpByKKHix2/b9kmTAM3YoX6MKP+/D//0ClgUYR2fY=
 github.com/pion/turn/v2 v2.1.6 h1:Xr2niVsiPTB0FPtt+yAWKFUkU1eotQbGgpTIld4x1Gc=
 github.com/pion/turn/v2 v2.1.6/go.mod h1:huEpByKKHix2/b9kmTAM3YoX6MKP+/D//0ClgUYR2fY=
-github.com/pion/webrtc/v3 v3.2.52-0.20240812164622-d0c92a5c4578 h1:0qPS1mVONwjpIahWlkxlWmso8FKVxVh990C2hR43Sl8=
-github.com/pion/webrtc/v3 v3.2.52-0.20240812164622-d0c92a5c4578/go.mod h1:hVmrDJvwhEertRWObeb1xzulzHGeVUoPlWvxdGzcfU0=
+github.com/pion/webrtc/v3 v3.2.52-0.20240813151442-cef1db8adbf1 h1:T6Lwii20XHU7vxP+Q4Z5NeisV6bFJkxiR+7sK8+rqRM=
+github.com/pion/webrtc/v3 v3.2.52-0.20240813151442-cef1db8adbf1/go.mod h1:hVmrDJvwhEertRWObeb1xzulzHGeVUoPlWvxdGzcfU0=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
diff --git a/p2p/transport/webrtc/connection.go b/p2p/transport/webrtc/connection.go
index f40e4460ae..2fba37a970 100644
--- a/p2p/transport/webrtc/connection.go
+++ b/p2p/transport/webrtc/connection.go
@@ -16,6 +16,7 @@ import (
 
 	ma "github.com/multiformats/go-multiaddr"
 	"github.com/pion/datachannel"
+	"github.com/pion/sctp"
 	"github.com/pion/webrtc/v3"
 )
 
@@ -57,8 +58,7 @@ type connection struct {
 	streams      map[uint16]*stream
 	nextStreamID atomic.Int32
 
-	acceptQueue            chan dataChannel
-	peerConnectionClosedCh chan struct{}
+	acceptQueue chan dataChannel
 
 	ctx    context.Context
 	cancel context.CancelFunc
@@ -173,6 +173,12 @@ func (c *connection) OpenStream(ctx context.Context) (network.MuxedStream, error
 	}
 	rwc, err := c.detachChannel(ctx, dc)
 	if err != nil {
+		// There's a race between webrtc.SCTP.OnClose callback and the underlying
+		// association closing. It's nicer to close the connection here.
+		if errors.Is(err, sctp.ErrStreamClosed) {
+			c.closeWithError(errConnClosed)
+			return nil, c.closeErr
+		}
 		dc.Close()
 		return nil, fmt.Errorf("detach channel failed for stream(%d): %w", streamID, err)
 	}
diff --git a/p2p/transport/webrtc/transport.go b/p2p/transport/webrtc/transport.go
index 3f23be5b80..c4c16fd402 100644
--- a/p2p/transport/webrtc/transport.go
+++ b/p2p/transport/webrtc/transport.go
@@ -269,6 +269,7 @@ func (t *WebRTCTransport) dial(ctx context.Context, scope network.ConnManagement
 			}
 			if tConn != nil {
 				_ = tConn.Close()
+				tConn = nil
 			}
 		}
 	}()
diff --git a/p2p/transport/webrtc/transport_test.go b/p2p/transport/webrtc/transport_test.go
index abd3d7dcf6..d603d610da 100644
--- a/p2p/transport/webrtc/transport_test.go
+++ b/p2p/transport/webrtc/transport_test.go
@@ -1016,10 +1016,10 @@ func TestConnectionClosedWhenRemoteCloses(t *testing.T) {
 	require.NoError(t, err)
 
 	dialer, _ := getTransport(t)
-	var done sync.Mutex
-	done.Lock()
+	var wg sync.WaitGroup
+	wg.Add(1)
 	go func() {
-		defer done.Unlock()
+		defer wg.Done()
 		c, err := listener.Accept()
 		if !assert.NoError(t, err) {
 			return
@@ -1032,5 +1032,5 @@ func TestConnectionClosedWhenRemoteCloses(t *testing.T) {
 	c, err := dialer.Dial(context.Background(), listener.Multiaddr(), p)
 	require.NoError(t, err)
 	c.Close()
-	done.Lock()
+	wg.Wait()
 }
diff --git a/test-plans/go.mod b/test-plans/go.mod
index f654cdaded..dc9224ca29 100644
--- a/test-plans/go.mod
+++ b/test-plans/go.mod
@@ -80,7 +80,7 @@ require (
 	github.com/pion/stun v0.6.1 // indirect
 	github.com/pion/transport/v2 v2.2.10 // indirect
 	github.com/pion/turn/v2 v2.1.6 // indirect
-	github.com/pion/webrtc/v3 v3.2.50 // indirect
+	github.com/pion/webrtc/v3 v3.2.52-0.20240813151442-cef1db8adbf1 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/prometheus/client_golang v1.19.1 // indirect
diff --git a/test-plans/go.sum b/test-plans/go.sum
index 44fa241a25..88b753456a 100644
--- a/test-plans/go.sum
+++ b/test-plans/go.sum
@@ -261,8 +261,8 @@ github.com/pion/transport/v3 v3.0.6/go.mod h1:HvJr2N/JwNJAfipsRleqwFoR3t/pWyHeZU
 github.com/pion/turn/v2 v2.1.3/go.mod h1:huEpByKKHix2/b9kmTAM3YoX6MKP+/D//0ClgUYR2fY=
 github.com/pion/turn/v2 v2.1.6 h1:Xr2niVsiPTB0FPtt+yAWKFUkU1eotQbGgpTIld4x1Gc=
 github.com/pion/turn/v2 v2.1.6/go.mod h1:huEpByKKHix2/b9kmTAM3YoX6MKP+/D//0ClgUYR2fY=
-github.com/pion/webrtc/v3 v3.2.50 h1:C/rwL2mBfCxHv6tlLzDAO3krJpQXfVx8A8WHnGJ2j34=
-github.com/pion/webrtc/v3 v3.2.50/go.mod h1:dytYYoSBy7ZUWhJMbndx9UckgYvzNAfL7xgVnrIKxqo=
+github.com/pion/webrtc/v3 v3.2.52-0.20240813151442-cef1db8adbf1 h1:T6Lwii20XHU7vxP+Q4Z5NeisV6bFJkxiR+7sK8+rqRM=
+github.com/pion/webrtc/v3 v3.2.52-0.20240813151442-cef1db8adbf1/go.mod h1:hVmrDJvwhEertRWObeb1xzulzHGeVUoPlWvxdGzcfU0=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=