coder/coderd/httpapi/websocket_internal_test.go

package httpapi

import (
	"context"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"cdr.dev/slog/v3"
	"github.com/coder/coder/v2/testutil"
	"github.com/coder/quartz"
	"github.com/coder/websocket"
)

// websocketPair sets up an httptest server with a websocket endpoint and
// returns the server-side conn. The server handler stays alive until ctx
// is done.
func websocketPair(ctx context.Context, t *testing.T) *websocket.Conn {
	t.Helper()
	serverConnCh := make(chan *websocket.Conn, 1)

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := websocket.Accept(w, r, nil)
		if err != nil {
			return
		}
		serverConnCh <- conn
		// Keep the handler alive so the HTTP server doesn't close
		// the connection from under us.
		<-ctx.Done()
	}))
	t.Cleanup(srv.Close)

	//nolint:bodyclose
	clientConn, _, err := websocket.Dial(ctx, srv.URL, nil)
	require.NoError(t, err)
	_ = clientConn.CloseRead(ctx) // Needed to handle pings/pongs.
	t.Cleanup(func() {
		_ = clientConn.Close(websocket.StatusNormalClosure, "test cleanup")
	})

	select {
	case sc := <-serverConnCh:
		_ = sc.CloseRead(ctx) // Needed to handle pings/pongs.
		return sc
	case <-ctx.Done():
		t.Fatal("timed out waiting for server websocket accept")
		return nil
	}
}

// probeRecorder is a simple wrapper around a channel used to record probe results.
type probeRecorder struct {
	T testing.TB
	C chan ProbeResult
}

func (r *probeRecorder) record(_ context.Context, result ProbeResult) {
	select {
	case r.C <- result:
	default:
		r.T.Errorf("probeRecorder.C is full, dropping result %s", result)
	}
}

func TestWSWatcher(t *testing.T) {
	t.Parallel()

	t.Run("ServerSideClose", func(t *testing.T) {
		t.Parallel()
		ctx := testutil.Context(t, testutil.WaitShort)

		sink := testutil.NewFakeSink(t)
		logger := sink.Logger()
		mClock := quartz.NewMock(t)
		rec := &probeRecorder{T: t, C: make(chan ProbeResult, 1)}

		trap := mClock.Trap().NewTicker("WSWatcher")
		defer trap.Close()

		serverConn := websocketPair(ctx, t)

		w := &WSWatcher{rec: rec.record, clk: mClock, interval: time.Second}
		watchCtx := w.Watch(ctx, logger, serverConn)

		// Wait for the ticker to be created, then release.
		trap.MustWait(ctx).MustRelease(ctx)

		// Close the server-side connection before the tick fires.
		// The next ping will get a close/net.ErrClosed error.
		_ = serverConn.Close(websocket.StatusGoingAway, "simulated teardown")

		// Advance clock to trigger the tick.
		mClock.Advance(time.Second).MustWait(ctx)

		// The watch context should be canceled after probe failure.
		select {
		case <-watchCtx.Done():
		case <-ctx.Done():
			t.Fatal("timed out waiting for watch context to be canceled")
		}

		gotRes := testutil.RequireReceive(ctx, t, rec.C)
		assert.Equal(t, ProbePeerClosed, gotRes, "expected ProbePeerClosed result")

		// A closed connection is a normal shutdown condition. The
		// error should be logged at Debug, not Error.
		errorEntries := sink.Entries(func(e slog.SinkEntry) bool { return e.Level == slog.LevelError })
		assert.Empty(t, errorEntries,
			"closed connection should not produce error-level logs, got: %+v", errorEntries)
		debugEntries := sink.Entries(func(e slog.SinkEntry) bool { return e.Level == slog.LevelDebug })
		assert.NotEmpty(t, debugEntries,
			"expected a debug-level log entry for the closed connection")
	})

	t.Run("ContextCanceled", func(t *testing.T) {
		t.Parallel()
		ctx := testutil.Context(t, testutil.WaitShort)

		sink := testutil.NewFakeSink(t)
		logger := sink.Logger()
		mClock := quartz.NewMock(t)
		rec := &probeRecorder{T: t, C: make(chan ProbeResult, 1)}

		trap := mClock.Trap().NewTicker("WSWatcher")
		defer trap.Close()

		serverCtx, serverCancel := context.WithCancel(ctx)
		serverConn := websocketPair(ctx, t)

		w := &WSWatcher{rec: rec.record, clk: mClock, interval: time.Second}
		watchCtx := w.Watch(serverCtx, logger, serverConn)

		trap.MustWait(ctx).MustRelease(ctx)

		// Cancel the parent context. The watcher should exit via
		// the <-ctx.Done() branch without closing the conn.
		serverCancel()

		select {
		case <-watchCtx.Done():
		case <-ctx.Done():
			t.Fatal("timed out waiting for watch context to be canceled")
		}

		errorEntries := sink.Entries(func(e slog.SinkEntry) bool { return e.Level == slog.LevelError })
		assert.Empty(t, errorEntries,
			"context cancellation should not produce error-level logs, got: %+v", errorEntries)
		assert.Empty(t, rec.C, "expected no probes when context is canceled before tick")
	})

	t.Run("PingSucceeds", func(t *testing.T) {
		t.Parallel()
		ctx, cancel := context.WithCancel(testutil.Context(t, testutil.WaitShort))
		defer cancel()

		sink := testutil.NewFakeSink(t)
		logger := sink.Logger()
		mClock := quartz.NewMock(t)
		rec := &probeRecorder{T: t, C: make(chan ProbeResult, 3)}

		trap := mClock.Trap().NewTicker("WSWatcher")
		defer trap.Close()

		serverConn := websocketPair(ctx, t)

		w := &WSWatcher{rec: rec.record, clk: mClock, interval: time.Second}
		watchCtx := w.Watch(ctx, logger, serverConn)
		t.Cleanup(func() {
			<-watchCtx.Done()
		})

		trap.MustWait(ctx).MustRelease(ctx)

		// Fire several ticks; pings should succeed each time.
		for i := range 3 {
			mClock.Advance(time.Second).MustWait(ctx)
			gotRes := testutil.RequireReceive(ctx, t, rec.C)
			assert.Equal(t, ProbeOK, gotRes, "expected probe result to be ProbeOK at tick %d", i+1)
		}

		// No logs should be emitted during normal operation.
		errorEntries := sink.Entries(func(e slog.SinkEntry) bool { return e.Level == slog.LevelError })
		assert.Empty(t, errorEntries,
			"successful pings should not produce error-level logs, got: %+v", errorEntries)
		debugEntries := sink.Entries(func(e slog.SinkEntry) bool { return e.Level == slog.LevelDebug })
		assert.Empty(t, debugEntries,
			"successful pings should not produce debug-level logs, got: %+v", debugEntries)
	})

	t.Run("RecordsPrometheusCounter", func(t *testing.T) {
		t.Parallel()

		// Use a real prometheus registry to verify end-to-end metric recording.
		registry := prometheus.NewRegistry()
		probes := prometheus.NewCounterVec(prometheus.CounterOpts{
			Namespace: "coderd",
			Subsystem: "api",
			Name:      "websocket_probes_total",
			Help:      "test",
		}, []string{"path", "result"})
		registry.MustRegister(probes)

		recorder := func(ctx context.Context, r ProbeResult) {
			probes.WithLabelValues("/test/path", string(r)).Inc()
		}

		sink := testutil.NewFakeSink(t)
		logger := sink.Logger()
		mClock := quartz.NewMock(t)

		trap := mClock.Trap().NewTicker("WSWatcher")
		defer trap.Close()

		ctx, cancel := context.WithCancel(testutil.Context(t, testutil.WaitShort))
		defer cancel()
		serverConn := websocketPair(ctx, t)

		w := &WSWatcher{rec: recorder, clk: mClock, interval: time.Second}
		watchCtx := w.Watch(ctx, logger, serverConn)
		t.Cleanup(func() {
			<-watchCtx.Done()
		})

		trap.MustWait(ctx).MustRelease(ctx)
		mClock.Advance(time.Second).MustWait(ctx)

		testutil.Eventually(ctx, t, func(context.Context) bool {
			select {
			case <-watchCtx.Done():
				t.Fatal("watch context should not be canceled when pings succeed")
			default:
			}
			metrics, err := registry.Gather()
			require.NoError(t, err)
			return testutil.PromCounterHasValue(t, metrics, 1,
				"coderd_api_websocket_probes_total", "/test/path", "ok")
		}, testutil.IntervalFast, "probe counter not incremented")
	})

	t.Run("ProbeTimeout", func(t *testing.T) {
		t.Parallel()
		ctx := testutil.Context(t, testutil.WaitShort)

		mClock := quartz.NewMock(t)
		trap := mClock.Trap().NewTicker("WSWatcher")
		defer trap.Close()

		sink := testutil.NewFakeSink(t)
		logger := sink.Logger()
		rec := &probeRecorder{T: t, C: make(chan ProbeResult, 1)}

		pingCh := make(chan struct{})
		closeCodeCh := make(chan websocket.StatusCode, 1)
		fConn := &fakePingCloser{
			pingFn: func(context.Context) error {
				t.Log("ping")
				close(pingCh)
				// Determinism tradeoff: by returning DeadlineExceeded directly
				// we lose coverage of the WithTimeout path in probe().
				return context.DeadlineExceeded
			},
			closeFn: func(code websocket.StatusCode, _ string) error {
				closeCodeCh <- code
				return nil
			},
		}

		w := &WSWatcher{rec: rec.record, clk: mClock, interval: time.Second}
		watchCtx := w.Watch(ctx, logger, fConn)

		trap.MustWait(ctx).MustRelease(ctx)
		mClock.Advance(time.Second).MustWait(ctx)

		_, _ = testutil.SoftTryReceive(ctx, t, pingCh)

		select {
		case <-watchCtx.Done():
		case <-ctx.Done():
			t.Fatal("timed out waiting for watch context to be canceled")
		}

		gotRes := testutil.RequireReceive(ctx, t, rec.C)
		assert.Equal(t, ProbeTimeout, gotRes, "expected ProbeTimeout result")
		gotCode := testutil.RequireReceive(ctx, t, closeCodeCh)
		assert.Equal(t, websocket.StatusGoingAway, gotCode, "expected StatusGoingAway code")

		// Timeout is an expected condition, should be Debug not Error.
		errorEntries := sink.Entries(func(e slog.SinkEntry) bool { return e.Level == slog.LevelError })
		assert.Empty(t, errorEntries,
			"probe timeout should not produce error-level logs, got: %+v", errorEntries)
	})

	t.Run("ProbeError", func(t *testing.T) {
		t.Parallel()

		sink := testutil.NewFakeSink(t)
		logger := sink.Logger()
		mClock := quartz.NewMock(t)
		trap := mClock.Trap().NewTicker("WSWatcher")
		defer trap.Close()

		rec := &probeRecorder{T: t, C: make(chan ProbeResult, 1)}
		closeCodeCh := make(chan websocket.StatusCode, 1)

		fConn := &fakePingCloser{
			pingFn: func(context.Context) error {
				return assert.AnError
			},
			closeFn: func(code websocket.StatusCode, _ string) error {
				t.Log("close error", code)
				closeCodeCh <- code
				return nil
			},
		}

		w := &WSWatcher{rec: rec.record, clk: mClock, interval: time.Second}

		ctx, cancel := context.WithCancel(testutil.Context(t, testutil.WaitShort))
		defer cancel()
		watchCtx := w.Watch(ctx, logger, fConn)
		t.Cleanup(func() {
			<-watchCtx.Done()
		})

		trap.MustWait(ctx).MustRelease(ctx)
		mClock.Advance(time.Second).MustWait(ctx)

		gotRes := testutil.RequireReceive(ctx, t, rec.C)
		assert.Equal(t, ProbeError, gotRes, "expected ProbeError result")

		gotCode := testutil.RequireReceive(ctx, t, closeCodeCh)
		assert.Equal(t, websocket.StatusGoingAway, gotCode)

		// ProbeError should log at Error level (unlike other failures).
		errorEntries := sink.Entries(func(e slog.SinkEntry) bool {
			return e.Level == slog.LevelError
		})
		assert.NotEmpty(t, errorEntries, "ProbeError should produce error-level log")
	})
}

// fakePingCloser is a test double for the pingCloser interface.
type fakePingCloser struct {
	pingFn  func(context.Context) error
	closeFn func(websocket.StatusCode, string) error
}

func (f *fakePingCloser) Ping(ctx context.Context) error {
	return f.pingFn(ctx)
}

func (f *fakePingCloser) Close(code websocket.StatusCode, reason string) error {
	return f.closeFn(code, reason)
}