fix: WebSocketConnection: Timeout on Read shuts down writing connection

webSocketConnection Read and Write had separate
timeouts which, when expired, mutually closed the
websocket. In case only one party of the connection
was sending, the websocket was closed after the
timeout despite the fact that the connection was alive.
This comes from misunderstanding the impact of the Read/Write context parameter. See coder/websocket#242

Author: philippseith

websocketconnection.go

@@ -4,44 +4,41 @@ import (
 	"bytes"
 	"context"
 	"fmt"
-
 	"github.com/teivah/onecontext"
 	"nhooyr.io/websocket"
+	"time"
 )
 
 type webSocketConnection struct {
 	ConnectionBase
 	conn         *websocket.Conn
 	transferMode TransferMode
+	watchDogChan chan dogFood
 }
 
 func newWebSocketConnection(parentContext context.Context, requestContext context.Context, connectionID string, conn *websocket.Conn) *webSocketConnection {
 	ctx, _ := onecontext.Merge(parentContext, requestContext)
 	w := &webSocketConnection{
-		conn: conn,
+		conn:         conn,
+		watchDogChan: make(chan dogFood, 1),
 		ConnectionBase: ConnectionBase{
 			ctx:          ctx,
 			connectionID: connectionID,
 		},
 	}
+	go w.watchDog(ctx)
 	return w
 }
 
 func (w *webSocketConnection) Write(p []byte) (n int, err error) {
 	if err := w.Context().Err(); err != nil {
 		return 0, fmt.Errorf("webSocketConnection canceled: %w", w.ctx.Err())
 	}
-	ctx := w.ctx
-	if w.timeout > 0 {
-		var cancel context.CancelFunc
-		ctx, cancel = context.WithTimeout(w.ctx, w.Timeout())
-		defer cancel() // has no effect because timeoutCtx is either done or not used anymore after websocket returns. But it keeps lint quiet
-	}
 	messageType := websocket.MessageText
 	if w.transferMode == BinaryTransferMode {
 		messageType = websocket.MessageBinary
 	}
-	err = w.conn.Write(ctx, messageType, p)
+	err = w.conn.Write(w.resetWatchDog(), messageType, p)
 	if err != nil {
 		return 0, err
 	}
@@ -52,19 +49,70 @@ func (w *webSocketConnection) Read(p []byte) (n int, err error) {
 	if err := w.Context().Err(); err != nil {
 		return 0, fmt.Errorf("webSocketConnection canceled: %w", w.ctx.Err())
 	}
-	ctx := w.ctx
-	if w.timeout > 0 {
-		var cancel context.CancelFunc
-		ctx, cancel = context.WithTimeout(w.ctx, w.Timeout())
-		defer cancel() // has no effect because timeoutCtx is either done or not used anymore after websocket returns. But it keeps lint quiet
-	}
-	_, data, err := w.conn.Read(ctx)
+	_, data, err := w.conn.Read(w.resetWatchDog())
 	if err != nil {
 		return 0, err
 	}
 	return bytes.NewReader(data).Read(p)
 }
 
+// resetWatchDog resets the common watchDog for Read and Write.
+// the watchDog will stop waiting for the last set timeout and wait for the new timeout.
+func (w *webSocketConnection) resetWatchDog() context.Context {
+	ctx := w.ctx
+	food := dogFood{timeout: w.timeout}
+	if w.timeout > 0 {
+		ctx, food.bark = context.WithCancel(w.ctx)
+	}
+	w.watchDogChan <- food
+	return ctx
+}
+
+// dogFood is used to reset the watchDog
+type dogFood struct {
+	// After this, the dog will bark
+	timeout time.Duration
+	bark    context.CancelFunc
+}
+
+// watchDog is the common watchDog for Read and Write. It stops the connection (aka closes the Websocket)
+// when the last timeout has elapsed. If resetWatchDog is called before the last timeout has elapsed,
+// the watchDog will restart waiting for the new timeout. If timeout is set to 0, it will not wait at all.
+func (w *webSocketConnection) watchDog(ctx context.Context) {
+	var timer *time.Timer
+	var cancelTimeoutChan chan struct{}
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case food := <-w.watchDogChan:
+			if timer != nil {
+				if !timer.Stop() {
+					go func() {
+						<-timer.C
+					}()
+				}
+				go func() {
+					cancelTimeoutChan <- struct{}{}
+				}()
+			}
+			if food.timeout != 0 {
+				timer = time.NewTimer(food.timeout)
+				cancelTimeoutChan = make(chan struct{}, 1)
+				go func() {
+					select {
+					case <-cancelTimeoutChan:
+					case <-timer.C:
+						food.bark()
+					}
+				}()
+			} else {
+				timer = nil
+			}
+		}
+	}
+}
+
 func (w *webSocketConnection) TransferMode() TransferMode {
 	return w.transferMode
 }