chore: add support for one-way websockets to backend (#16853)

Closes https://github.com/coder/coder/issues/16775

## Changes made
- Added `OneWayWebSocket` function that establishes WebSocket
connections that don't allow client-to-server communication
- Added tests for the new function
- Updated API endpoints to make new WS-based endpoints, and mark
previous SSE-based endpoints as deprecated
- Updated existing SSE handlers to use the same core logic as the new WS
handlers

## Notes
- Frontend changes handled via #16855

coderd/httpapi/httpapi.go

@@ -16,6 +16,9 @@ import (
 	"github.com/go-playground/validator/v10"
 	"golang.org/x/xerrors"
 
+	"github.com/coder/websocket"
+	"github.com/coder/websocket/wsjson"
+
 	"github.com/coder/coder/v2/coderd/httpapi/httpapiconstraints"
 	"github.com/coder/coder/v2/coderd/tracing"
 	"github.com/coder/coder/v2/codersdk"
@@ -282,7 +285,25 @@ func WebsocketCloseSprintf(format string, vars ...any) string {
 	return msg
 }
 
-func ServerSentEventSender(rw http.ResponseWriter, r *http.Request) (sendEvent func(ctx context.Context, sse codersdk.ServerSentEvent) error, closed chan struct{}, err error) {
+type EventSender func(rw http.ResponseWriter, r *http.Request) (
+	sendEvent func(sse codersdk.ServerSentEvent) error,
+	done <-chan struct{},
+	err error,
+)
+
+// ServerSentEventSender establishes a Server-Sent Event connection and allows
+// the consumer to send messages to the client.
+//
+// The function returned allows you to send a single message to the client,
+// while the channel lets you listen for when the connection closes.
+//
+// As much as possible, this function should be avoided in favor of using the
+// OneWayWebSocket function. See OneWayWebSocket for more context.
+func ServerSentEventSender(rw http.ResponseWriter, r *http.Request) (
+	func(sse codersdk.ServerSentEvent) error,
+	<-chan struct{},
+	error,
+) {
 	h := rw.Header()
 	h.Set("Content-Type", "text/event-stream")
 	h.Set("Cache-Control", "no-cache")
@@ -294,7 +315,8 @@ func ServerSentEventSender(rw http.ResponseWriter, r *http.Request) (sendEvent f
 		panic("http.ResponseWriter is not http.Flusher")
 	}
 
-	closed = make(chan struct{})
+	ctx := r.Context()
+	closed := make(chan struct{})
 	type sseEvent struct {
 		payload []byte
 		errC    chan error
@@ -304,16 +326,13 @@ func ServerSentEventSender(rw http.ResponseWriter, r *http.Request) (sendEvent f
 	// Synchronized handling of events (no guarantee of order).
 	go func() {
 		defer close(closed)
-
-		// Send a heartbeat every 15 seconds to avoid the connection being killed.
-		ticker := time.NewTicker(time.Second * 15)
+		ticker := time.NewTicker(HeartbeatInterval)
 		defer ticker.Stop()
 
 		for {
 			var event sseEvent
-
 			select {
-			case <-r.Context().Done():
+			case <-ctx.Done():
 				return
 			case event = <-eventC:
 			case <-ticker.C:
@@ -333,21 +352,21 @@ func ServerSentEventSender(rw http.ResponseWriter, r *http.Request) (sendEvent f
 		}
 	}()
 
-	sendEvent = func(ctx context.Context, sse codersdk.ServerSentEvent) error {
+	sendEvent := func(newEvent codersdk.ServerSentEvent) error {
 		buf := &bytes.Buffer{}
-		enc := json.NewEncoder(buf)
-
-		_, err := buf.WriteString(fmt.Sprintf("event: %s\n", sse.Type))
+		_, err := buf.WriteString(fmt.Sprintf("event: %s\n", newEvent.Type))
 		if err != nil {
 			return err
 		}
 
-		if sse.Data != nil {
+		if newEvent.Data != nil {
 			_, err = buf.WriteString("data: ")
 			if err != nil {
 				return err
 			}
-			err = enc.Encode(sse.Data)
+
+			enc := json.NewEncoder(buf)
+			err = enc.Encode(newEvent.Data)
 			if err != nil {
 				return err
 			}
@@ -364,8 +383,6 @@ func ServerSentEventSender(rw http.ResponseWriter, r *http.Request) (sendEvent f
 		}
 
 		select {
-		case <-r.Context().Done():
-			return r.Context().Err()
 		case <-ctx.Done():
 			return ctx.Err()
 		case <-closed:
@@ -375,8 +392,6 @@ func ServerSentEventSender(rw http.ResponseWriter, r *http.Request) (sendEvent f
 			// for early exit. We don't check closed here because it
 			// can't happen while processing the event.
 			select {
-			case <-r.Context().Done():
-				return r.Context().Err()
 			case <-ctx.Done():
 				return ctx.Err()
 			case err := <-event.errC:
@@ -387,3 +402,90 @@ func ServerSentEventSender(rw http.ResponseWriter, r *http.Request) (sendEvent f
 
 	return sendEvent, closed, nil
 }
+
+// OneWayWebSocketEventSender establishes a new WebSocket connection that
+// enforces one-way communication from the server to the client.
+//
+// The function returned allows you to send a single message to the client,
+// while the channel lets you listen for when the connection closes.
+//
+// We must use an approach like this instead of Server-Sent Events for the
+// browser, because on HTTP/1.1 connections, browsers are locked to no more than
+// six HTTP connections for a domain total, across all tabs. If a user were to
+// open a workspace in multiple tabs, the entire UI can start to lock up.
+// WebSockets have no such limitation, no matter what HTTP protocol was used to
+// establish the connection.
+func OneWayWebSocketEventSender(rw http.ResponseWriter, r *http.Request) (
+	func(event codersdk.ServerSentEvent) error,
+	<-chan struct{},
+	error,
+) {
+	ctx, cancel := context.WithCancel(r.Context())
+	r = r.WithContext(ctx)
+	socket, err := websocket.Accept(rw, r, nil)
+	if err != nil {
+		cancel()
+		return nil, nil, xerrors.Errorf("cannot establish connection: %w", err)
+	}
+	go Heartbeat(ctx, socket)
+
+	eventC := make(chan codersdk.ServerSentEvent)
+	socketErrC := make(chan websocket.CloseError, 1)
+	closed := make(chan struct{})
+	go func() {
+		defer cancel()
+		defer close(closed)
+
+		for {
+			select {
+			case event := <-eventC:
+				writeCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
+				err := wsjson.Write(writeCtx, socket, event)
+				cancel()
+				if err == nil {
+					continue
+				}
+				_ = socket.Close(websocket.StatusInternalError, "Unable to send newest message")
+			case err := <-socketErrC:
+				_ = socket.Close(err.Code, err.Reason)
+			case <-ctx.Done():
+				_ = socket.Close(websocket.StatusNormalClosure, "Connection closed")
+			}
+			return
+		}
+	}()
+
+	// We have some tools in the UI code to help enforce one-way WebSocket
+	// connections, but there's still the possibility that the client could send
+	// a message when it's not supposed to. If that happens, the client likely
+	// forgot to use those tools, and communication probably can't be trusted.
+	// Better to just close the socket and force the UI to fix its mess
+	go func() {
+		_, _, err := socket.Read(ctx)
+		if errors.Is(err, context.Canceled) {
+			return
+		}
+		if err != nil {
+			socketErrC <- websocket.CloseError{
+				Code:   websocket.StatusInternalError,
+				Reason: "Unable to process invalid message from client",
+			}
+			return
+		}
+		socketErrC <- websocket.CloseError{
+			Code:   websocket.StatusProtocolError,
+			Reason: "Clients cannot send messages for one-way WebSockets",
+		}
+	}()
+
+	sendEvent := func(event codersdk.ServerSentEvent) error {
+		select {
+		case eventC <- event:
+		case <-ctx.Done():
+			return ctx.Err()
+		}
+		return nil
+	}
+
+	return sendEvent, closed, nil
+}

coderd/httpapi/httpapi_test.go

@@ -1,14 +1,18 @@
 package httpapi_test
 
 import (
+	"bufio"
 	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
+	"net"
 	"net/http"
 	"net/http/httptest"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -16,6 +20,7 @@ import (
 
 	"github.com/coder/coder/v2/coderd/httpapi"
 	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/testutil"
 )
 
 func TestInternalServerError(t *testing.T) {
@@ -155,3 +160,436 @@ func TestWebsocketCloseMsg(t *testing.T) {
 		assert.Equal(t, len(trunc), 123)
 	})
 }
+
+// Our WebSocket library accepts any arbitrary ResponseWriter at the type level,
+// but the writer must also implement http.Hijacker for long-lived connections.
+type mockOneWaySocketWriter struct {
+	serverRecorder   *httptest.ResponseRecorder
+	serverConn       net.Conn
+	clientConn       net.Conn
+	serverReadWriter *bufio.ReadWriter
+	testContext      *testing.T
+}
+
+func (m mockOneWaySocketWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	return m.serverConn, m.serverReadWriter, nil
+}
+
+func (m mockOneWaySocketWriter) Flush() {
+	err := m.serverReadWriter.Flush()
+	require.NoError(m.testContext, err)
+}
+
+func (m mockOneWaySocketWriter) Header() http.Header {
+	return m.serverRecorder.Header()
+}
+
+func (m mockOneWaySocketWriter) Write(b []byte) (int, error) {
+	return m.serverReadWriter.Write(b)
+}
+
+func (m mockOneWaySocketWriter) WriteHeader(code int) {
+	m.serverRecorder.WriteHeader(code)
+}
+
+type mockEventSenderWrite func(b []byte) (int, error)
+
+func (w mockEventSenderWrite) Write(b []byte) (int, error) {
+	return w(b)
+}
+
+func TestOneWayWebSocketEventSender(t *testing.T) {
+	t.Parallel()
+
+	newBaseRequest := func(ctx context.Context) *http.Request {
+		url := "ws://www.fake-website.com/logs"
+		req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+		require.NoError(t, err)
+
+		h := req.Header
+		h.Add("Connection", "Upgrade")
+		h.Add("Upgrade", "websocket")
+		h.Add("Sec-WebSocket-Version", "13")
+		h.Add("Sec-WebSocket-Key", "dGhlIHNhbXBsZSBub25jZQ==") // Just need any string
+
+		return req
+	}
+
+	newOneWayWriter := func(t *testing.T) mockOneWaySocketWriter {
+		mockServer, mockClient := net.Pipe()
+		recorder := httptest.NewRecorder()
+
+		var write mockEventSenderWrite = func(b []byte) (int, error) {
+			serverCount, err := mockServer.Write(b)
+			if err != nil {
+				return 0, err
+			}
+			recorderCount, err := recorder.Write(b)
+			if err != nil {
+				return 0, err
+			}
+			return min(serverCount, recorderCount), nil
+		}
+
+		return mockOneWaySocketWriter{
+			testContext:    t,
+			serverConn:     mockServer,
+			clientConn:     mockClient,
+			serverRecorder: recorder,
+			serverReadWriter: bufio.NewReadWriter(
+				bufio.NewReader(mockServer),
+				bufio.NewWriter(write),
+			),
+		}
+	}
+
+	t.Run("Produces error if the socket connection could not be established", func(t *testing.T) {
+		t.Parallel()
+
+		incorrectProtocols := []struct {
+			major int
+			minor int
+			proto string
+		}{
+			{0, 9, "HTTP/0.9"},
+			{1, 0, "HTTP/1.0"},
+		}
+		for _, p := range incorrectProtocols {
+			ctx := testutil.Context(t, testutil.WaitShort)
+			req := newBaseRequest(ctx)
+			req.ProtoMajor = p.major
+			req.ProtoMinor = p.minor
+			req.Proto = p.proto
+
+			writer := newOneWayWriter(t)
+			_, _, err := httpapi.OneWayWebSocketEventSender(writer, req)
+			require.ErrorContains(t, err, p.proto)
+		}
+	})
+
+	t.Run("Returned callback can publish new event to WebSocket connection", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := testutil.Context(t, testutil.WaitShort)
+		req := newBaseRequest(ctx)
+		writer := newOneWayWriter(t)
+		send, _, err := httpapi.OneWayWebSocketEventSender(writer, req)
+		require.NoError(t, err)
+
+		serverPayload := codersdk.ServerSentEvent{
+			Type: codersdk.ServerSentEventTypeData,
+			Data: "Blah",
+		}
+		err = send(serverPayload)
+		require.NoError(t, err)
+
+		// The client connection will receive a little bit of additional data on
+		// top of the main payload. Have to make sure check has tolerance for
+		// extra data being present
+		serverBytes, err := json.Marshal(serverPayload)
+		require.NoError(t, err)
+		clientBytes, err := io.ReadAll(writer.clientConn)
+		require.NoError(t, err)
+		require.True(t, bytes.Contains(clientBytes, serverBytes))
+	})
+
+	t.Run("Signals to outside consumer when socket has been closed", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := context.WithCancel(testutil.Context(t, testutil.WaitShort))
+		req := newBaseRequest(ctx)
+		writer := newOneWayWriter(t)
+		_, done, err := httpapi.OneWayWebSocketEventSender(writer, req)
+		require.NoError(t, err)
+
+		successC := make(chan bool)
+		ticker := time.NewTicker(testutil.WaitShort)
+		go func() {
+			select {
+			case <-done:
+				successC <- true
+			case <-ticker.C:
+				successC <- false
+			}
+		}()
+
+		cancel()
+		require.True(t, <-successC)
+	})
+
+	t.Run("Socket will immediately close if client sends any message", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := testutil.Context(t, testutil.WaitShort)
+		req := newBaseRequest(ctx)
+		writer := newOneWayWriter(t)
+		_, done, err := httpapi.OneWayWebSocketEventSender(writer, req)
+		require.NoError(t, err)
+
+		successC := make(chan bool)
+		ticker := time.NewTicker(testutil.WaitShort)
+		go func() {
+			select {
+			case <-done:
+				successC <- true
+			case <-ticker.C:
+				successC <- false
+			}
+		}()
+
+		type JunkClientEvent struct {
+			Value string
+		}
+		b, err := json.Marshal(JunkClientEvent{"Hi :)"})
+		require.NoError(t, err)
+		_, err = writer.clientConn.Write(b)
+		require.NoError(t, err)
+		require.True(t, <-successC)
+	})
+
+	t.Run("Renders the socket inert if the request context cancels", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := context.WithCancel(testutil.Context(t, testutil.WaitShort))
+		req := newBaseRequest(ctx)
+		writer := newOneWayWriter(t)
+		send, done, err := httpapi.OneWayWebSocketEventSender(writer, req)
+		require.NoError(t, err)
+
+		successC := make(chan bool)
+		ticker := time.NewTicker(testutil.WaitShort)
+		go func() {
+			select {
+			case <-done:
+				successC <- true
+			case <-ticker.C:
+				successC <- false
+			}
+		}()
+
+		cancel()
+		require.True(t, <-successC)
+		err = send(codersdk.ServerSentEvent{
+			Type: codersdk.ServerSentEventTypeData,
+			Data: "Didn't realize you were closed - sorry! I'll try coming back tomorrow.",
+		})
+		require.Equal(t, err, ctx.Err())
+		_, open := <-done
+		require.False(t, open)
+		_, err = writer.serverConn.Write([]byte{})
+		require.Equal(t, err, io.ErrClosedPipe)
+		_, err = writer.clientConn.Read([]byte{})
+		require.Equal(t, err, io.EOF)
+	})
+
+	t.Run("Sends a heartbeat to the socket on a fixed internal of time to keep connections alive", func(t *testing.T) {
+		t.Parallel()
+
+		// Need add at least three heartbeats for something to be reliably
+		// counted as an interval, but also need some wiggle room
+		heartbeatCount := 3
+		hbDuration := time.Duration(heartbeatCount) * httpapi.HeartbeatInterval
+		timeout := hbDuration + (5 * time.Second)
+
+		ctx := testutil.Context(t, timeout)
+		req := newBaseRequest(ctx)
+		writer := newOneWayWriter(t)
+		_, _, err := httpapi.OneWayWebSocketEventSender(writer, req)
+		require.NoError(t, err)
+
+		type Result struct {
+			Err     error
+			Success bool
+		}
+		resultC := make(chan Result)
+		go func() {
+			err := writer.
+				clientConn.
+				SetReadDeadline(time.Now().Add(timeout))
+			if err != nil {
+				resultC <- Result{err, false}
+				return
+			}
+			for range heartbeatCount {
+				pingBuffer := make([]byte, 1)
+				pingSize, err := writer.clientConn.Read(pingBuffer)
+				if err != nil || pingSize != 1 {
+					resultC <- Result{err, false}
+					return
+				}
+			}
+			resultC <- Result{nil, true}
+		}()
+
+		result := <-resultC
+		require.NoError(t, result.Err)
+		require.True(t, result.Success)
+	})
+}
+
+// ServerSentEventSender accepts any arbitrary ResponseWriter at the type level,
+// but the writer must also implement http.Flusher for long-lived connections
+type mockServerSentWriter struct {
+	serverRecorder *httptest.ResponseRecorder
+	serverConn     net.Conn
+	clientConn     net.Conn
+	buffer         *bytes.Buffer
+	testContext    *testing.T
+}
+
+func (m mockServerSentWriter) Flush() {
+	b := m.buffer.Bytes()
+	_, err := m.serverConn.Write(b)
+	require.NoError(m.testContext, err)
+	m.buffer.Reset()
+
+	// Must close server connection to indicate EOF for any reads from the
+	// client connection; otherwise reads block forever. This is a testing
+	// limitation compared to the one-way websockets, since we have no way to
+	// frame the data and auto-indicate EOF for each message
+	err = m.serverConn.Close()
+	require.NoError(m.testContext, err)
+}
+
+func (m mockServerSentWriter) Header() http.Header {
+	return m.serverRecorder.Header()
+}
+
+func (m mockServerSentWriter) Write(b []byte) (int, error) {
+	return m.buffer.Write(b)
+}
+
+func (m mockServerSentWriter) WriteHeader(code int) {
+	m.serverRecorder.WriteHeader(code)
+}
+
+func TestServerSentEventSender(t *testing.T) {
+	t.Parallel()
+
+	newBaseRequest := func(ctx context.Context) *http.Request {
+		url := "ws://www.fake-website.com/logs"
+		req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+		require.NoError(t, err)
+		return req
+	}
+
+	newServerSentWriter := func(t *testing.T) mockServerSentWriter {
+		mockServer, mockClient := net.Pipe()
+		return mockServerSentWriter{
+			testContext:    t,
+			serverRecorder: httptest.NewRecorder(),
+			clientConn:     mockClient,
+			serverConn:     mockServer,
+			buffer:         &bytes.Buffer{},
+		}
+	}
+
+	t.Run("Mutates response headers to support SSE connections", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := testutil.Context(t, testutil.WaitShort)
+		req := newBaseRequest(ctx)
+		writer := newServerSentWriter(t)
+		_, _, err := httpapi.ServerSentEventSender(writer, req)
+		require.NoError(t, err)
+
+		h := writer.Header()
+		require.Equal(t, h.Get("Content-Type"), "text/event-stream")
+		require.Equal(t, h.Get("Cache-Control"), "no-cache")
+		require.Equal(t, h.Get("Connection"), "keep-alive")
+		require.Equal(t, h.Get("X-Accel-Buffering"), "no")
+	})
+
+	t.Run("Returned callback can publish new event to SSE connection", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := testutil.Context(t, testutil.WaitShort)
+		req := newBaseRequest(ctx)
+		writer := newServerSentWriter(t)
+		send, _, err := httpapi.ServerSentEventSender(writer, req)
+		require.NoError(t, err)
+
+		serverPayload := codersdk.ServerSentEvent{
+			Type: codersdk.ServerSentEventTypeData,
+			Data: "Blah",
+		}
+		err = send(serverPayload)
+		require.NoError(t, err)
+
+		clientBytes, err := io.ReadAll(writer.clientConn)
+		require.NoError(t, err)
+		require.Equal(
+			t,
+			string(clientBytes),
+			"event: data\ndata: \"Blah\"\n\n",
+		)
+	})
+
+	t.Run("Signals to outside consumer when connection has been closed", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := context.WithCancel(testutil.Context(t, testutil.WaitShort))
+		req := newBaseRequest(ctx)
+		writer := newServerSentWriter(t)
+		_, done, err := httpapi.ServerSentEventSender(writer, req)
+		require.NoError(t, err)
+
+		successC := make(chan bool)
+		ticker := time.NewTicker(testutil.WaitShort)
+		go func() {
+			select {
+			case <-done:
+				successC <- true
+			case <-ticker.C:
+				successC <- false
+			}
+		}()
+
+		cancel()
+		require.True(t, <-successC)
+	})
+
+	t.Run("Sends a heartbeat to the client on a fixed internal of time to keep connections alive", func(t *testing.T) {
+		t.Parallel()
+
+		// Need add at least three heartbeats for something to be reliably
+		// counted as an interval, but also need some wiggle room
+		heartbeatCount := 3
+		hbDuration := time.Duration(heartbeatCount) * httpapi.HeartbeatInterval
+		timeout := hbDuration + (5 * time.Second)
+
+		ctx := testutil.Context(t, timeout)
+		req := newBaseRequest(ctx)
+		writer := newServerSentWriter(t)
+		_, _, err := httpapi.ServerSentEventSender(writer, req)
+		require.NoError(t, err)
+
+		type Result struct {
+			Err     error
+			Success bool
+		}
+		resultC := make(chan Result)
+		go func() {
+			err := writer.
+				clientConn.
+				SetReadDeadline(time.Now().Add(timeout))
+			if err != nil {
+				resultC <- Result{err, false}
+				return
+			}
+			for range heartbeatCount {
+				pingBuffer := make([]byte, 1)
+				pingSize, err := writer.clientConn.Read(pingBuffer)
+				if err != nil || pingSize != 1 {
+					resultC <- Result{err, false}
+					return
+				}
+			}
+			resultC <- Result{nil, true}
+		}()
+
+		result := <-resultC
+		require.NoError(t, result.Err)
+		require.True(t, result.Success)
+	})
+}

coderd/httpapi/websocket.go

@@ -11,11 +11,13 @@ import (
 	"github.com/coder/websocket"
 )
 
+const HeartbeatInterval time.Duration = 15 * time.Second
+
 // Heartbeat loops to ping a WebSocket to keep it alive.
 // Default idle connection timeouts are typically 60 seconds.
 // See: https://docs.aws.amazon.com/elasticloadbalancing/latest/application/application-load-balancers.html#connection-idle-timeout
 func Heartbeat(ctx context.Context, conn *websocket.Conn) {
-	ticker := time.NewTicker(15 * time.Second)
+	ticker := time.NewTicker(HeartbeatInterval)
 	defer ticker.Stop()
 	for {
 		select {
@@ -33,8 +35,7 @@ func Heartbeat(ctx context.Context, conn *websocket.Conn) {
 // Heartbeat loops to ping a WebSocket to keep it alive. It calls `exit` on ping
 // failure.
 func HeartbeatClose(ctx context.Context, logger slog.Logger, exit func(), conn *websocket.Conn) {
-	interval := 15 * time.Second
-	ticker := time.NewTicker(interval)
+	ticker := time.NewTicker(HeartbeatInterval)
 	defer ticker.Stop()
 
 	for {
@@ -43,7 +44,7 @@ func HeartbeatClose(ctx context.Context, logger slog.Logger, exit func(), conn *
 			return
 		case <-ticker.C:
 		}
-		err := pingWithTimeout(ctx, conn, interval)
+		err := pingWithTimeout(ctx, conn, HeartbeatInterval)
 		if err != nil {
 			// context.DeadlineExceeded is expected when the client disconnects without sending a close frame
 			if !errors.Is(err, context.DeadlineExceeded) {