fix: Loop Dial for reconnects

coder · kylecarbs · May 9, 2021 · May 9, 2021 · May 9, 2021 · May 9, 2021
commit a9a8b4dc5870e874ad1500033e8f163b37e7038e
diff --git a/wsnet/dial_test.go b/wsnet/dial_test.go
@@ -48,10 +48,12 @@ func TestDial(t *testing.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
 		}
 		err = dialer.Ping(context.Background())
 		if err != nil {
@@ -64,6 +66,7 @@ func TestDial(t *testing.T) {
 		_, err := Listen(context.Background(), listenAddr)
 		if err != nil {
 			t.Error(err)
+			return
 		}
 		dialer, err := DialWebsocket(context.Background(), connectAddr, nil)
 		if err != nil {
@@ -100,10 +103,12 @@ func TestDial(t *testing.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
 		}
 		conn, err := dialer.DialContext(context.Background(), listener.Addr().Network(), listener.Addr().String())
 		if err != nil {
@@ -132,6 +137,7 @@ func TestDial(t *testing.T) {
 		srv, err := Listen(context.Background(), listenAddr)
 		if err != nil {
 			t.Error(err)
+			return
 		}
 		dialer, err := DialWebsocket(context.Background(), connectAddr, nil)
 		if err != nil {

diff --git a/wsnet/listen.go b/wsnet/listen.go
@@ -9,6 +9,7 @@ import (
 	"net"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/hashicorp/yamux"
 	"github.com/pion/webrtc/v3"
@@ -20,7 +21,57 @@ import (
 // Listen connects to the broker proxies connections to the local net.
 // Close will end all RTC connections.
 func Listen(ctx context.Context, broker string) (io.Closer, error) {
-	conn, resp, err := websocket.Dial(ctx, broker, nil)
+	l := &listener{
+		broker:      broker,
+		connClosers: make([]io.Closer, 0),
+	}
+	// We do a one-off dial outside of the loop to ensure the initial
+	// connection is successful. If not, there's likely an error the
+	// user needs to act on.
+	ch, err := l.dial(ctx)
+	if err != nil {
+		return nil, err
+	}
+	go func() {
+		for {
+			err := <-ch
+			if errors.Is(err, io.EOF) {
+				// If we hit an EOF, then the connection to the broker
+				// was interrupted. We'll take a short break then dial
+				// again.
+				time.Sleep(time.Second)
+				ch, err = l.dial(ctx)
+			}
+			if err != nil {
+				l.acceptError = err
+				_ = l.Close()
+				break
+			}
+		}
+
+		err := <-ch
+
+		if err != nil {
+			l.acceptError = err
+		}
+	}()
+	return l, nil
+}
+
+type listener struct {
+	broker string
+
+	acceptError    error
+	ws             *websocket.Conn
+	connClosers    []io.Closer
+	connClosersMut sync.Mutex
+}
+
+func (l *listener) dial(ctx context.Context) (<-chan error, error) {
+	if l.ws != nil {
+		_ = l.ws.Close(websocket.StatusNormalClosure, "new connection inbound")
+	}
+	conn, resp, err := websocket.Dial(ctx, l.broker, nil)
 	if err != nil {
 		if resp != nil {
 			return nil, &coder.HTTPError{
@@ -29,37 +80,25 @@ func Listen(ctx context.Context, broker string) (io.Closer, error) {
 		}
 		return nil, err
 	}
+	l.ws = conn
 	nconn := websocket.NetConn(ctx, conn, websocket.MessageBinary)
 	session, err := yamux.Server(nconn, nil)
 	if err != nil {
 		return nil, fmt.Errorf("create multiplex: %w", err)
 	}
-	l := &listener{
-		ws:          conn,
-		connClosers: make([]io.Closer, 0),
-	}
+	errCh := make(chan error)
 	go func() {
+		defer close(errCh)
 		for {
 			conn, err := session.Accept()
 			if err != nil {
-				if errors.Is(err, io.EOF) {
-					continue
-				}
-				l.acceptError = err
-				l.Close()
-				return
+				errCh <- err
+				break
 			}
 			go l.negotiate(conn)
 		}
 	}()
-	return l, nil
-}
-
-type listener struct {
-	acceptError    error
-	ws             *websocket.Conn
-	connClosers    []io.Closer
-	connClosersMut sync.Mutex
+	return errCh, nil
 }
 
 // Negotiates the handshake protocol over the connection provided.