Close websocket stream when client closes

pkg/util/wsstream/stream.go

@@ -20,6 +20,7 @@ import (
 	"encoding/base64"
 	"io"
 	"net/http"
+	"sync"
 	"time"
 
 	"golang.org/x/net/websocket"
@@ -107,11 +108,28 @@ func (r *Reader) Copy(w http.ResponseWriter, req *http.Request) error {
 
 // handle implements a WebSocket handler.
 func (r *Reader) handle(ws *websocket.Conn) {
+	// Close the connection when the client requests it, or when we finish streaming, whichever happens first
+	closeConnOnce := &sync.Once{}
+	closeConn := func() {
+		closeConnOnce.Do(func() {
+			ws.Close()
+		})
+	}
+
 	negotiated := ws.Config().Protocol
 	r.selectedProtocol = negotiated[0]
 	defer close(r.err)
-	defer ws.Close()
-	go IgnoreReceives(ws, r.timeout)
+	defer closeConn()
+
+	go func() {
+		defer runtime.HandleCrash()
+		// This blocks until the connection is closed.
+		// Client should not send anything.
+		IgnoreReceives(ws, r.timeout)
+		// Once the client closes, we should also close
+		closeConn()
+	}()
+
 	r.err <- messageCopy(ws, r.r, !r.protocols[r.selectedProtocol].Binary, r.ping, r.timeout)
 }