From 582b3e6ab1d77f323e5dd7aca6847dd70b2c9524 Mon Sep 17 00:00:00 2001
From: Kyle Carberry <kyle@coder.com>
Date: Wed, 21 Jul 2021 07:35:19 +0000
Subject: [PATCH 1/2] fix: Increase ICE timeouts

---
 wsnet/dial.go | 27 ++++++++++++++++++++++-----
 wsnet/rtc.go  |  4 +++-
 2 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/wsnet/dial.go b/wsnet/dial.go
index af4b422c..74b2f5b6 100644
--- a/wsnet/dial.go
+++ b/wsnet/dial.go
@@ -137,6 +137,7 @@ type Dialer struct {
 	closedChan     chan struct{}
 	connClosers    []io.Closer
 	connClosersMut sync.Mutex
+	pingMut        sync.Mutex
 }
 
 func (d *Dialer) negotiate() (err error) {
@@ -160,7 +161,7 @@ func (d *Dialer) negotiate() (err error) {
 			return
 		}
 		d.rtc.OnConnectionStateChange(func(pcs webrtc.PeerConnectionState) {
-			if pcs == webrtc.PeerConnectionStateConnected {
+			if pcs != webrtc.PeerConnectionStateDisconnected {
 				return
 			}
 
@@ -178,6 +179,7 @@ func (d *Dialer) negotiate() (err error) {
 			default:
 			}
 			close(d.closedChan)
+			_ = d.rtc.Close()
 		})
 	}()
 
@@ -263,7 +265,7 @@ func (d *Dialer) Ping(ctx context.Context) error {
 	// Since we control the client and server we could open this
 	// data channel with `Negotiated` true to reduce traffic being
 	// sent when the RTC connection is opened.
-	err := waitForDataChannelOpen(context.Background(), d.ctrl)
+	err := waitForDataChannelOpen(ctx, d.ctrl)
 	if err != nil {
 		return err
 	}
@@ -273,13 +275,28 @@ func (d *Dialer) Ping(ctx context.Context) error {
 			return err
 		}
 	}
+	d.pingMut.Lock()
+	defer d.pingMut.Unlock()
 	_, err = d.ctrlrw.Write([]byte{'a'})
 	if err != nil {
 		return fmt.Errorf("write: %w", err)
 	}
-	b := make([]byte, 4)
-	_, err = d.ctrlrw.Read(b)
-	return err
+	errCh := make(chan error)
+	go func() {
+		// There's a race in which connections can get lost-mid ping
+		// in which case this would block forever.
+		defer close(errCh)
+		_, err = d.ctrlrw.Read(make([]byte, 1))
+		errCh <- err
+	}()
+	ctx, cancelFunc := context.WithTimeout(ctx, time.Second*5)
+	defer cancelFunc()
+	select {
+	case err := <-errCh:
+		return err
+	case <-ctx.Done():
+		return ctx.Err()
+	}
 }
 
 // DialContext dials the network and address on the remote listener.
diff --git a/wsnet/rtc.go b/wsnet/rtc.go
index 05c04f1b..0b93c57c 100644
--- a/wsnet/rtc.go
+++ b/wsnet/rtc.go
@@ -160,7 +160,9 @@ func newPeerConnection(servers []webrtc.ICEServer, dialer proxy.Dialer) (*webrtc
 	se.SetNetworkTypes([]webrtc.NetworkType{webrtc.NetworkTypeUDP4})
 	se.SetSrflxAcceptanceMinWait(0)
 	se.DetachDataChannels()
-	se.SetICETimeouts(time.Second*3, time.Second*3, time.Second*2)
+	// If the disconnect and keep-alive timeouts are too closely related, we'll
+	// experience "random" connection failures.
+	se.SetICETimeouts(time.Second*5, time.Second*25, time.Second*2)
 	lf := logging.NewDefaultLoggerFactory()
 	lf.DefaultLogLevel = logging.LogLevelDisabled
 	se.LoggerFactory = lf

From e2bcf1a28308dfa07951ce3119680e8bf223aebd Mon Sep 17 00:00:00 2001
From: Kyle Carberry <kyle@coder.com>
Date: Wed, 21 Jul 2021 15:59:32 +0000
Subject: [PATCH 2/2] Fix ping times

---
 wsnet/dial.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/wsnet/dial.go b/wsnet/dial.go
index 74b2f5b6..283bedf4 100644
--- a/wsnet/dial.go
+++ b/wsnet/dial.go
@@ -286,10 +286,10 @@ func (d *Dialer) Ping(ctx context.Context) error {
 		// There's a race in which connections can get lost-mid ping
 		// in which case this would block forever.
 		defer close(errCh)
-		_, err = d.ctrlrw.Read(make([]byte, 1))
+		_, err = d.ctrlrw.Read(make([]byte, 4))
 		errCh <- err
 	}()
-	ctx, cancelFunc := context.WithTimeout(ctx, time.Second*5)
+	ctx, cancelFunc := context.WithTimeout(ctx, time.Second*15)
 	defer cancelFunc()
 	select {
 	case err := <-errCh: