coder · deansheather · Jul 27, 2021 · Jul 26, 2021 · Jul 26, 2021 · Jul 26, 2021
diff --git a/wsnet/dial_test.go b/wsnet/dial_test.go
@@ -293,39 +293,55 @@ func TestDial(t *testing.T) {
 
 	t.Run("Close Listeners on Disconnect", func(t *testing.T) {
 		t.Parallel()
+		log := slogtest.Make(t, nil)
 
-		tcpListener, err := net.Listen("tcp", "0.0.0.0:0")
+		listener, err := net.Listen("tcp", "0.0.0.0:0")
 		require.NoError(t, err)
 		go func() {
-			_, _ = tcpListener.Accept()
+			for {
+				c, _ := listener.Accept()
+
+				go func() {
+					b := make([]byte, 5)
+					_, err := c.Read(b)
+					if err != nil {
+						return
+					}
+					_, err = c.Write(b)
+					require.NoError(t, err)
+				}()
+			}
 		}()
-
 		connectAddr, listenAddr := createDumbBroker(t)
-		l, err := Listen(context.Background(), slogtest.Make(t, nil), listenAddr, "")
+		_, err = Listen(context.Background(), slogtest.Make(t, nil), listenAddr, "")
 		require.NoError(t, err)
 
-		turnAddr, closeTurn := createTURNServer(t, ice.SchemeTypeTURN)
-		dialer, err := DialWebsocket(context.Background(), connectAddr, &DialOptions{
-			ICEServers: []webrtc.ICEServer{{
-				URLs:           []string{fmt.Sprintf("turn:%s", turnAddr)},
-				Username:       "example",
-				Credential:     testPass,
-				CredentialType: webrtc.ICECredentialTypePassword,
-			}},
+		d1, err := DialWebsocket(context.Background(), connectAddr, &DialOptions{
+			Log: &log,
 		}, nil)
 		require.NoError(t, err)
+		_, err = d1.DialContext(context.Background(), listener.Addr().Network(), listener.Addr().String())
+		require.NoError(t, err)
 
-		_, err = dialer.DialContext(context.Background(), "tcp", tcpListener.Addr().String())
+		d2, err := DialWebsocket(context.Background(), connectAddr, &DialOptions{
+			Log: &log,
+		}, nil)
+		require.NoError(t, err)
+		conn, err := d2.DialContext(context.Background(), listener.Addr().Network(), listener.Addr().String())
+		require.NoError(t, err)
+		err = d1.Close()
 		require.NoError(t, err)
 
-		closeTurn()
+		// TODO: This needs to be longer than the KeepAlive timeout for the RTC connection.
+		// Once the listener stores RTC connections instead of io.Closer we can directly
+		// reference the RTC connection to ensure it's properly closed.
+		time.Sleep(time.Second * 10)
 
-		list := l.(*listener)
-		assert.Eventually(t, func() bool {
-			list.connClosersMut.Lock()
-			defer list.connClosersMut.Unlock()
-			return len(list.connClosers) == 0
-		}, time.Second*15, time.Millisecond*100)
+		b := []byte("hello")
+		_, err = conn.Write(b)
+		require.NoError(t, err)
+		_, err = conn.Read(b)
+		require.NoError(t, err)
 	})
 }
 

diff --git a/wsnet/listen.go b/wsnet/listen.go
@@ -163,9 +163,11 @@ func (l *listener) negotiate(ctx context.Context, conn net.Conn) {
 	ctx = slog.With(ctx, slog.F("conn_id", id))
 
 	var (
-		err     error
-		decoder = json.NewDecoder(conn)
-		rtc     *webrtc.PeerConnection
+		err            error
+		decoder        = json.NewDecoder(conn)
+		rtc            *webrtc.PeerConnection
+		connClosers    = make([]io.Closer, 0)
+		connClosersMut sync.Mutex
 		// If candidates are sent before an offer, we place them here.
 		// We currently have no assurances to ensure this can't happen,
 		// so it's better to buffer and process than fail.
@@ -255,6 +257,9 @@ func (l *listener) negotiate(ctx context.Context, conn net.Conn) {
 				closeError(err)
 				return
 			}
+			l.connClosersMut.Lock()
+			l.connClosers = append(l.connClosers, rtc)
+			l.connClosersMut.Unlock()
 			rtc.OnConnectionStateChange(func(pcs webrtc.PeerConnectionState) {
 				l.log.Info(ctx, "connection state change", slog.F("state", pcs.String()))
 				switch pcs {
@@ -267,16 +272,16 @@ func (l *listener) negotiate(ctx context.Context, conn net.Conn) {
 				}
 
 				// Close connections opened when RTC was alive.
-				l.connClosersMut.Lock()
-				defer l.connClosersMut.Unlock()
-				for _, connCloser := range l.connClosers {
+				connClosersMut.Lock()
+				defer connClosersMut.Unlock()
+				for _, connCloser := range connClosers {
 					_ = connCloser.Close()
 				}
-				l.connClosers = make([]io.Closer, 0)
+				connClosers = make([]io.Closer, 0)
 			})
 
 			flushCandidates := proxyICECandidates(rtc, conn)
-			rtc.OnDataChannel(l.handle(ctx, msg))
+			rtc.OnDataChannel(l.handle(ctx, msg, &connClosers, &connClosersMut))
 
 			l.log.Debug(ctx, "set remote description", slog.F("offer", *msg.Offer))
 			err = rtc.SetRemoteDescription(*msg.Offer)
@@ -329,7 +334,7 @@ func (l *listener) negotiate(ctx context.Context, conn net.Conn) {
 }
 
 // nolint:gocognit
-func (l *listener) handle(ctx context.Context, msg BrokerMessage) func(dc *webrtc.DataChannel) {
+func (l *listener) handle(ctx context.Context, msg BrokerMessage, connClosers *[]io.Closer, connClosersMut *sync.Mutex) func(dc *webrtc.DataChannel) {
 	return func(dc *webrtc.DataChannel) {
 		if dc.Protocol() == controlChannel {
 			// The control channel handles pings.
@@ -430,9 +435,9 @@ func (l *listener) handle(ctx context.Context, msg BrokerMessage) func(dc *webrt
 				dc:   dc,
 				rw:   rw,
 			}
-			l.connClosersMut.Lock()
-			l.connClosers = append(l.connClosers, co)
-			l.connClosersMut.Unlock()
+			connClosersMut.Lock()
+			*connClosers = append(*connClosers, co)
+			connClosersMut.Unlock()
 			co.init()
 			defer nc.Close()
 			defer co.Close()