fix: Remove active connections when RTC connection is lost

coder · kylecarbs · Jul 8, 2021 · Jul 7, 2021 · Jul 8, 2021 · Jul 8, 2021
commit 8d5db5d9adc6c63db5679f45e69e655666019219
diff --git a/wsnet/dial.go b/wsnet/dial.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"sync"
 	"time"
 
 	"github.com/pion/datachannel"
@@ -81,9 +82,10 @@ func Dial(conn net.Conn, iceServers []webrtc.ICEServer) (*Dialer, error) {
 	flushCandidates()
 
 	dialer := &Dialer{
-		conn: conn,
-		ctrl: ctrl,
-		rtc:  rtc,
+		conn:        conn,
+		ctrl:        ctrl,
+		rtc:         rtc,
+		connClosers: make([]io.Closer, 0),
 	}
 
 	return dialer, dialer.negotiate()
@@ -97,6 +99,9 @@ type Dialer struct {
 	ctrl   *webrtc.DataChannel
 	ctrlrw datachannel.ReadWriteCloser
 	rtc    *webrtc.PeerConnection
+
+	connClosers    []io.Closer
+	connClosersMut sync.Mutex
 }
 
 func (d *Dialer) negotiate() (err error) {
@@ -117,6 +122,19 @@ func (d *Dialer) negotiate() (err error) {
 			errCh <- err
 			return
 		}
+		d.rtc.OnConnectionStateChange(func(pcs webrtc.PeerConnectionState) {
+			if pcs == webrtc.PeerConnectionStateConnected {
+				return
+			}
+
+			// Close connections opened while the RTC was alive.
+			d.connClosersMut.Lock()
+			defer d.connClosersMut.Unlock()
+			for _, connCloser := range d.connClosers {
+				_ = connCloser.Close()
+			}
+			d.connClosers = make([]io.Closer, 0)
+		})
 		go func() {
 			// Closing this connection took 30ms+.
 			_ = d.conn.Close()
@@ -252,6 +270,10 @@ func (d *Dialer) DialContext(ctx context.Context, network, address string) (net.
 		return nil, ctx.Err()
 	}
 
+	d.connClosersMut.Lock()
+	defer d.connClosersMut.Unlock()
+	d.connClosers = append(d.connClosers, rw)
+
 	c := &conn{
 		addr: &net.UnixAddr{
 			Name: address,

diff --git a/wsnet/dial_test.go b/wsnet/dial_test.go
@@ -5,11 +5,13 @@ import (
 	"context"
 	"crypto/rand"
 	"errors"
+	"fmt"
 	"io"
 	"net"
 	"strconv"
 	"testing"
 
+	"github.com/pion/ice/v2"
 	"github.com/pion/webrtc/v3"
 )
 
@@ -184,6 +186,46 @@ func TestDial(t *testing.T) {
 			t.Error(err)
 		}
 	})
+
+	t.Run("Disconnect DialContext", func(t *testing.T) {
+		tcpListener, err := net.Listen("tcp", "0.0.0.0:0")
+		if err != nil {
+			t.Error(err)
+			return
+		}
+		go tcpListener.Accept()
+
+		connectAddr, listenAddr := createDumbBroker(t)
+		_, err = Listen(context.Background(), listenAddr)
+		if err != nil {
+			t.Error(err)
+			return
+		}
+		turnAddr, closeTurn := createTURNServer(t, ice.SchemeTypeTURN, "test")
+		dialer, err := DialWebsocket(context.Background(), connectAddr, []webrtc.ICEServer{{
+			URLs:           []string{fmt.Sprintf("turn:%s", turnAddr)},
+			Username:       "example",
+			Credential:     "test",
+			CredentialType: webrtc.ICECredentialTypePassword,
+		}})
+		if err != nil {
+			t.Error(err)
+			return
+		}
+		conn, err := dialer.DialContext(context.Background(), "tcp", tcpListener.Addr().String())
+		if err != nil {
+			t.Error(err)
+			return
+		}
+		// Close the TURN server before reading...
+		// WebRTC connections take a few seconds to timeout.
+		closeTurn()
+		_, err = conn.Read(make([]byte, 16))
+		if err != io.EOF {
+			t.Error(err)
+			return
+		}
+	})
 }
 
 func BenchmarkThroughput(b *testing.B) {

diff --git a/wsnet/rtc.go b/wsnet/rtc.go
@@ -164,6 +164,8 @@ func newPeerConnection(servers []webrtc.ICEServer) (*webrtc.PeerConnection, erro
 	lf.DefaultLogLevel = logging.LogLevelDisabled
 	se.LoggerFactory = lf
 
+	transportPolicy := webrtc.ICETransportPolicyAll
+
 	// If one server is provided and we know it's TURN, we can set the
 	// relay acceptable so the connection starts immediately.
 	if len(servers) == 1 {
@@ -174,12 +176,18 @@ func newPeerConnection(servers []webrtc.ICEServer) (*webrtc.PeerConnection, erro
 				se.SetNetworkTypes([]webrtc.NetworkType{webrtc.NetworkTypeTCP4, webrtc.NetworkTypeTCP6})
 				se.SetRelayAcceptanceMinWait(0)
 			}
+			if err == nil && (url.Scheme == ice.SchemeTypeTURN || url.Scheme == ice.SchemeTypeTURNS) {
+				// Local peers will connect if they discover they live on the same host.
+				// For testing purposes, it's simpler if they cannot peer on the same host.
+				transportPolicy = webrtc.ICETransportPolicyRelay
+			}
 		}
 	}
 	api := webrtc.NewAPI(webrtc.WithSettingEngine(se))
 
 	return api.NewPeerConnection(webrtc.Configuration{
-		ICEServers: servers,
+		ICEServers:         servers,
+		ICETransportPolicy: transportPolicy,
 	})
 }
 

diff --git a/wsnet/rtc_test.go b/wsnet/rtc_test.go
@@ -16,7 +16,7 @@ func TestDialICE(t *testing.T) {
 	t.Run("TURN with TLS", func(t *testing.T) {
 		t.Parallel()
 
-		addr := createTURNServer(t, ice.SchemeTypeTURNS, "test")
+		addr, _ := createTURNServer(t, ice.SchemeTypeTURNS, "test")
 		err := DialICE(webrtc.ICEServer{
 			URLs:           []string{fmt.Sprintf("turns:%s", addr)},
 			Username:       "example",
@@ -34,7 +34,7 @@ func TestDialICE(t *testing.T) {
 	t.Run("Protocol mismatch", func(t *testing.T) {
 		t.Parallel()
 
-		addr := createTURNServer(t, ice.SchemeTypeTURNS, "test")
+		addr, _ := createTURNServer(t, ice.SchemeTypeTURNS, "test")
 		err := DialICE(webrtc.ICEServer{
 			URLs:           []string{fmt.Sprintf("turn:%s", addr)},
 			Username:       "example",
@@ -52,7 +52,7 @@ func TestDialICE(t *testing.T) {
 	t.Run("Invalid auth", func(t *testing.T) {
 		t.Parallel()
 
-		addr := createTURNServer(t, ice.SchemeTypeTURNS, "test")
+		addr, _ := createTURNServer(t, ice.SchemeTypeTURNS, "test")
 		err := DialICE(webrtc.ICEServer{
 			URLs:           []string{fmt.Sprintf("turns:%s", addr)},
 			Username:       "example",

diff --git a/wsnet/wsnet_test.go b/wsnet/wsnet_test.go
@@ -86,7 +86,7 @@ func createDumbBroker(t testing.TB) (connectAddr string, listenAddr string) {
 }
 
 // createTURNServer allocates a TURN server and returns the address.
-func createTURNServer(t *testing.T, server ice.SchemeType, pass string) string {
+func createTURNServer(t *testing.T, server ice.SchemeType, pass string) (string, func()) {
 	var (
 		listeners   []turn.ListenerConfig
 		pcListeners []turn.PacketConnConfig
@@ -143,17 +143,18 @@ func createTURNServer(t *testing.T, server ice.SchemeType, pass string) string {
 	if err != nil {
 		t.Error(err)
 	}
-	t.Cleanup(func() {
+	closeFunc := func() {
 		for _, l := range listeners {
 			l.Listener.Close()
 		}
 		for _, l := range pcListeners {
 			l.PacketConn.Close()
 		}
 		srv.Close()
-	})
+	}
+	t.Cleanup(closeFunc)
 
-	return listenAddr.String()
+	return listenAddr.String(), closeFunc
 }
 
 func generateTLSConfig(t testing.TB) *tls.Config {