From c22d1aea058b6d7ef1ee1f4cfa9508df172e48d3 Mon Sep 17 00:00:00 2001
From: Kyle Carberry <kyle@coder.com>
Date: Thu, 22 Jul 2021 00:23:07 +0000
Subject: [PATCH 1/2] fix: Close conn if context deadline is exceeded

Previously the error occurred but never propogated due to blocking fromm
the read loop in negotiate. Now, connection timeouts can properly occur.
---
 wsnet/dial.go       |  4 +++-
 wsnet/dial_test.go  | 12 ++++++++++++
 wsnet/rtc.go        |  2 +-
 wsnet/wsnet_test.go |  2 +-
 4 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/wsnet/dial.go b/wsnet/dial.go
index d6f45461..3880b12c 100644
--- a/wsnet/dial.go
+++ b/wsnet/dial.go
@@ -147,7 +147,6 @@ func (d *Dialer) negotiate(ctx context.Context) (err error) {
 		// so it's better to buffer and process than fail.
 		pendingCandidates = []webrtc.ICECandidateInit{}
 	)
-
 	go func() {
 		defer close(errCh)
 		defer func() {
@@ -155,6 +154,9 @@ func (d *Dialer) negotiate(ctx context.Context) (err error) {
 		}()
 		err := waitForConnectionOpen(ctx, d.rtc)
 		if err != nil {
+			if errors.Is(err, context.DeadlineExceeded) {
+				_ = d.conn.Close()
+			}
 			errCh <- err
 			return
 		}
diff --git a/wsnet/dial_test.go b/wsnet/dial_test.go
index 5dd11b58..aad497d3 100644
--- a/wsnet/dial_test.go
+++ b/wsnet/dial_test.go
@@ -9,6 +9,7 @@ import (
 	"net"
 	"strconv"
 	"testing"
+	"time"
 
 	"cdr.dev/slog/sloggers/slogtest"
 	"github.com/pion/ice/v2"
@@ -51,6 +52,17 @@ func ExampleDial_basic() {
 }
 
 func TestDial(t *testing.T) {
+	t.Run("Timeout", func(t *testing.T) {
+		t.Parallel()
+
+		connectAddr, _ := createDumbBroker(t)
+
+		ctx, cancelFunc := context.WithTimeout(context.Background(), time.Millisecond*50)
+		defer cancelFunc()
+		_, err := DialWebsocket(ctx, connectAddr, nil, nil)
+		require.True(t, errors.Is(err, context.DeadlineExceeded))
+	})
+
 	t.Run("Ping", func(t *testing.T) {
 		t.Parallel()
 
diff --git a/wsnet/rtc.go b/wsnet/rtc.go
index 79702743..32a089a2 100644
--- a/wsnet/rtc.go
+++ b/wsnet/rtc.go
@@ -256,7 +256,7 @@ func waitForConnectionOpen(ctx context.Context, conn *webrtc.PeerConnection) err
 	})
 	<-ctx.Done()
 	if ctx.Err() == context.DeadlineExceeded {
-		return ctx.Err()
+		return context.DeadlineExceeded
 	}
 	return nil
 }
diff --git a/wsnet/wsnet_test.go b/wsnet/wsnet_test.go
index ad9ac381..5a6a9ecd 100644
--- a/wsnet/wsnet_test.go
+++ b/wsnet/wsnet_test.go
@@ -68,7 +68,7 @@ func createDumbBroker(t testing.TB) (connectAddr string, listenAddr string) {
 		mut.Lock()
 		defer mut.Unlock()
 		if sess == nil {
-			t.Error("listen not called")
+			_, _ = io.Copy(io.Discard, nc)
 			return
 		}
 		oc, err := sess.Open()

From 6d7d47626614e4282414eaabd46452e6c3663469 Mon Sep 17 00:00:00 2001
From: Kyle Carberry <kyle@coder.com>
Date: Thu, 22 Jul 2021 00:28:44 +0000
Subject: [PATCH 2/2] Fixes

---
 wsnet/dial_test.go  | 3 ++-
 wsnet/wsnet_test.go | 2 ++
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/wsnet/dial_test.go b/wsnet/dial_test.go
index aad497d3..5d75cfd4 100644
--- a/wsnet/dial_test.go
+++ b/wsnet/dial_test.go
@@ -59,8 +59,9 @@ func TestDial(t *testing.T) {
 
 		ctx, cancelFunc := context.WithTimeout(context.Background(), time.Millisecond*50)
 		defer cancelFunc()
-		_, err := DialWebsocket(ctx, connectAddr, nil, nil)
+		dialer, err := DialWebsocket(ctx, connectAddr, nil, nil)
 		require.True(t, errors.Is(err, context.DeadlineExceeded))
+		require.Error(t, dialer.conn.Close(), "already wrote close")
 	})
 
 	t.Run("Ping", func(t *testing.T) {
diff --git a/wsnet/wsnet_test.go b/wsnet/wsnet_test.go
index 5a6a9ecd..20aa7699 100644
--- a/wsnet/wsnet_test.go
+++ b/wsnet/wsnet_test.go
@@ -68,6 +68,8 @@ func createDumbBroker(t testing.TB) (connectAddr string, listenAddr string) {
 		mut.Lock()
 		defer mut.Unlock()
 		if sess == nil {
+			// We discard inbound to emulate a pubsub where we don't know if anyone
+			// is listening on the other side.
 			_, _ = io.Copy(io.Discard, nc)
 			return
 		}