From f8fcdd5109f180dae7b035d4426ab202e14b0e8f Mon Sep 17 00:00:00 2001
From: Kyle Carberry <kyle@carberry.com>
Date: Thu, 8 Jul 2021 23:32:44 -0500
Subject: [PATCH 1/3] feat: Add Closed func to Dialer

---
 wsnet/dial.go      | 14 ++++++++++++++
 wsnet/dial_test.go | 18 ++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/wsnet/dial.go b/wsnet/dial.go
index 637bc5fd..b7359113 100644
--- a/wsnet/dial.go
+++ b/wsnet/dial.go
@@ -85,6 +85,8 @@ func Dial(conn net.Conn, iceServers []webrtc.ICEServer) (*Dialer, error) {
 		conn:        conn,
 		ctrl:        ctrl,
 		rtc:         rtc,
+		closed:      false,
+		closedChan:  make(chan struct{}),
 		connClosers: make([]io.Closer, 0),
 	}
 
@@ -100,6 +102,8 @@ type Dialer struct {
 	ctrlrw datachannel.ReadWriteCloser
 	rtc    *webrtc.PeerConnection
 
+	closed         bool
+	closedChan     chan struct{}
 	connClosers    []io.Closer
 	connClosersMut sync.Mutex
 }
@@ -136,6 +140,10 @@ func (d *Dialer) negotiate() (err error) {
 				_ = connCloser.Close()
 			}
 			d.connClosers = make([]io.Closer, 0)
+			if !d.closed {
+				d.closed = true
+				close(d.closedChan)
+			}
 		})
 	}()
 
@@ -184,6 +192,12 @@ func (d *Dialer) negotiate() (err error) {
 	return <-errCh
 }
 
+// Closed returns a channel that closes when
+// the connection is closed.
+func (d *Dialer) Closed() <-chan struct{} {
+	return d.closedChan
+}
+
 // Close closes the RTC connection.
 // All data channels dialed will be closed.
 func (d *Dialer) Close() error {
diff --git a/wsnet/dial_test.go b/wsnet/dial_test.go
index 71fdc8c7..1a794012 100644
--- a/wsnet/dial_test.go
+++ b/wsnet/dial_test.go
@@ -229,6 +229,24 @@ func TestDial(t *testing.T) {
 			return
 		}
 	})
+
+	t.Run("Closed", func(t *testing.T) {
+		connectAddr, listenAddr := createDumbBroker(t)
+		_, err := Listen(context.Background(), listenAddr)
+		if err != nil {
+			t.Error(err)
+			return
+		}
+		dialer, err := DialWebsocket(context.Background(), connectAddr, nil)
+		if err != nil {
+			t.Error(err)
+			return
+		}
+		go func() {
+			_ = dialer.Close()
+		}()
+		<-dialer.Closed()
+	})
 }
 
 func BenchmarkThroughput(b *testing.B) {

From c230faa1a4cd2841267aba514a1aa6196b3cf93a Mon Sep 17 00:00:00 2001
From: Kyle Carberry <kyle@carberry.com>
Date: Thu, 8 Jul 2021 23:36:01 -0500
Subject: [PATCH 2/3] Fix gocyclo

---
 wsnet/dial_test.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/wsnet/dial_test.go b/wsnet/dial_test.go
index 1a794012..85e19d7a 100644
--- a/wsnet/dial_test.go
+++ b/wsnet/dial_test.go
@@ -46,7 +46,7 @@ func ExampleDial_basic() {
 	// You now have access to the proxied remote port in `conn`.
 }
 
-// nolint:gocognit
+// nolint:gocognit,gocyclo
 func TestDial(t *testing.T) {
 	t.Run("Ping", func(t *testing.T) {
 		connectAddr, listenAddr := createDumbBroker(t)

From 6c7c23bfb0656fed4f6b3a4c22f91e209480b7f6 Mon Sep 17 00:00:00 2001
From: Kyle Carberry <kyle@carberry.com>
Date: Thu, 8 Jul 2021 23:39:36 -0500
Subject: [PATCH 3/3] Cleanup closed bool

---
 wsnet/dial.go      | 11 ++++++-----
 wsnet/dial_test.go |  7 ++++++-
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/wsnet/dial.go b/wsnet/dial.go
index b7359113..c72dc513 100644
--- a/wsnet/dial.go
+++ b/wsnet/dial.go
@@ -85,7 +85,6 @@ func Dial(conn net.Conn, iceServers []webrtc.ICEServer) (*Dialer, error) {
 		conn:        conn,
 		ctrl:        ctrl,
 		rtc:         rtc,
-		closed:      false,
 		closedChan:  make(chan struct{}),
 		connClosers: make([]io.Closer, 0),
 	}
@@ -102,7 +101,6 @@ type Dialer struct {
 	ctrlrw datachannel.ReadWriteCloser
 	rtc    *webrtc.PeerConnection
 
-	closed         bool
 	closedChan     chan struct{}
 	connClosers    []io.Closer
 	connClosersMut sync.Mutex
@@ -140,10 +138,13 @@ func (d *Dialer) negotiate() (err error) {
 				_ = connCloser.Close()
 			}
 			d.connClosers = make([]io.Closer, 0)
-			if !d.closed {
-				d.closed = true
-				close(d.closedChan)
+
+			select {
+			case <-d.closedChan:
+				return
+			default:
 			}
+			close(d.closedChan)
 		})
 	}()
 
diff --git a/wsnet/dial_test.go b/wsnet/dial_test.go
index 85e19d7a..9b412a3e 100644
--- a/wsnet/dial_test.go
+++ b/wsnet/dial_test.go
@@ -10,6 +10,7 @@ import (
 	"net"
 	"strconv"
 	"testing"
+	"time"
 
 	"github.com/pion/ice/v2"
 	"github.com/pion/webrtc/v3"
@@ -245,7 +246,11 @@ func TestDial(t *testing.T) {
 		go func() {
 			_ = dialer.Close()
 		}()
-		<-dialer.Closed()
+		select {
+		case <-dialer.Closed():
+		case <-time.NewTimer(time.Second).C:
+			t.Error("didn't close in time")
+		}
 	})
 }