feat(worker): Implement proxy handler websocket instrumentation

Add metrics on read/written bytes and active worker websocket
connections.

internal/servers/worker/handler.go

@@ -39,7 +39,7 @@ func (w *Worker) handler(props HandlerProperties) (http.Handler, error) {
 	if err != nil {
 		return nil, fmt.Errorf("%s: %w", op, err)
 	}
-	mux.Handle("/v1/proxy", h)
+	mux.Handle("/v1/proxy", metric.InstrumentWebsocketWrapper(h))
 
 	genericWrappedHandler := w.wrapGenericHandler(mux, props)
 	metricHandler := metric.InstrumentHttpHandler(genericWrappedHandler)

internal/servers/worker/internal/metric/proxy_ws.go

@@ -0,0 +1,163 @@
+package metric
+
+import (
+	"bufio"
+	"bytes"
+	"io"
+	"net"
+	"net/http"
+	"sync"
+
+	"github.com/hashicorp/boundary/globals"
+	"github.com/prometheus/client_golang/prometheus"
+)
+
+const (
+	proxyWebsocketSubsystem = "worker_proxy_websocket"
+)
+
+var (
+	activeConnCount = prometheus.NewGauge(
+		prometheus.GaugeOpts{
+			Namespace: globals.MetricNamespace,
+			Subsystem: proxyWebsocketSubsystem,
+			Name:      "active_connections",
+			Help:      "Count of open websocket proxy connections (to Boundary workers).",
+		},
+	)
+
+	bytesReceived = prometheus.NewCounter(
+		prometheus.CounterOpts{
+			Namespace: globals.MetricNamespace,
+			Subsystem: proxyWebsocketSubsystem,
+			Name:      "received_bytes_total",
+			Help:      "Count of received bytes for Worker proxy websocket connections.",
+		},
+	)
+
+	bytesSent = prometheus.NewCounter(
+		prometheus.CounterOpts{
+			Namespace: globals.MetricNamespace,
+			Subsystem: proxyWebsocketSubsystem,
+			Name:      "sent_bytes_total",
+			Help:      "Count of sent bytes for Worker proxy websocket connections.",
+		},
+	)
+)
+
+// InitializeWebsocketCollectors registers the websocket collectors onto `r`.
+// It panics upon the first registration that causes an error.
+func InitializeWebsocketCollectors(r prometheus.Registerer) {
+	if r == nil {
+		return
+	}
+	r.MustRegister(activeConnCount, bytesReceived, bytesSent)
+}
+
+// InstrumentWebsocketWrapper expects an http.Handler where websockets are used
+// and wraps it to instrument that websocket with metrics (bytes received, sent, etc).
+func InstrumentWebsocketWrapper(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		nextRw := rw
+		if rwh, ok := rw.(http.Hijacker); ok {
+			nextRw = wrapProxyHandler(&hijackerWrapper{ResponseWriter: rw, Hijacker: rwh}, rw)
+		}
+		next.ServeHTTP(nextRw, req)
+	})
+}
+
+// wrapProxyHandler wraps the incoming `http.ResponseWriter` in `hijackerWrapper` such
+// that we don't lose access to the other possible implementations that could be present
+// in the original object.
+//
+// See: https://medium.com/@cep21/interface-wrapping-method-erasure-c523b3549912
+func wrapProxyHandler(with *hijackerWrapper, wrap http.ResponseWriter) http.ResponseWriter {
+	const op = "metric.wrapProxyHandler"
+	flusher, _ := wrap.(http.Flusher)
+	pusher, _ := wrap.(http.Pusher)
+
+	// We don't check for http.Hijacker because hijackerWrapper is always an http.Hijacker.
+	switch {
+	case flusher == nil && pusher == nil:
+		return with
+	case flusher != nil && pusher == nil:
+		return struct {
+			*hijackerWrapper
+			http.Flusher
+		}{with, flusher}
+	case flusher == nil && pusher != nil:
+		return struct {
+			*hijackerWrapper
+			http.Pusher
+		}{with, pusher}
+	default:
+		return struct {
+			*hijackerWrapper
+			http.Flusher
+			http.Pusher
+		}{with, flusher, pusher}
+	}
+}
+
+type wsConn struct {
+	*sync.Once
+	net.Conn
+}
+
+func (wc *wsConn) Read(b []byte) (int, error) {
+	n, err := wc.Conn.Read(b)
+	if err != nil && err != io.EOF {
+		return n, err
+	}
+	bytesReceived.Add(float64(n))
+	return n, err
+}
+
+func (wc *wsConn) Write(b []byte) (int, error) {
+	n, err := wc.Conn.Write(b)
+	if err != nil {
+		return n, err
+	}
+	bytesSent.Add(float64(n))
+	return n, nil
+}
+
+func (wc *wsConn) Close() error {
+	wc.Once.Do(func() { activeConnCount.Dec() })
+	return wc.Conn.Close()
+}
+
+type hijackerWrapper struct {
+	// The `Once` objects constrain incrementing/decrementing active connections.
+	// Only a precautionary measure as:
+	// 1. http.(Hijacker).Hijack() cannot be called more than once (2nd call is an error).
+	// 2. websocket.(*Conn).Close() cannot be called more than once (2nd call is no-op).
+	inc sync.Once
+	dec sync.Once
+
+	// `ServeHTTP` takes an `http.ResponseWriter`, so `metricsWrapper`
+	// needs to implement it. Other than that, `ResponseWriter` is unused.
+	http.ResponseWriter
+
+	http.Hijacker
+}
+
+func (hw *hijackerWrapper) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	c, rw, err := hw.Hijacker.Hijack()
+	if err != nil {
+		return nil, nil, err
+	}
+	c = &wsConn{Conn: c, Once: &hw.dec}
+
+	// Our underlying `Hijack` call returns a `bufio.ReadWriter` object with the
+	// `Reader` and `Writer` set to the `net.Conn` object *before* we wrap it.
+	// The websocket library uses this buffer to Read and Write. This means
+	// our Conn implementation would never be called if we didn't replace the
+	// `Reader` and `Writer`.
+	buffered, _ := rw.Reader.Peek(rw.Reader.Buffered())
+	rw.Reader.Reset(io.MultiReader(bytes.NewReader(buffered), c))
+	rw = bufio.NewReadWriter(rw.Reader, bufio.NewWriter(c))
+
+	hw.inc.Do(func() { activeConnCount.Inc() })
+	return c, rw, err
+}

internal/servers/worker/internal/metric/proxy_ws_test.go

@@ -0,0 +1,464 @@
+package metric
+
+import (
+	"bufio"
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/stretchr/testify/require"
+)
+
+func TestInitializeWebsocketCollectors(t *testing.T) {
+	// Assert that calling the function with a nil Registerer is a valid operation.
+	require.NotPanics(t, func() { InitializeWebsocketCollectors(nil) })
+
+	// Try with an actual Registerer.
+	r := prometheus.NewRegistry()
+	require.NotPanics(t, func() { InitializeWebsocketCollectors(r) })
+
+	// Expect that some metrics exist.
+	f, err := r.Gather()
+	require.NoError(t, err)
+	require.Greater(t, len(f), 0)
+}
+
+func TestWrapProxyHandler(t *testing.T) {
+	type rw struct {
+		http.ResponseWriter
+	}
+	type flusher struct {
+		http.ResponseWriter
+		http.Flusher
+	}
+	type pusher struct {
+		http.ResponseWriter
+		http.Pusher
+	}
+	type flusherPusher struct {
+		http.ResponseWriter
+		http.Flusher
+		http.Pusher
+	}
+
+	tests := []struct {
+		name        string
+		wrap        http.ResponseWriter
+		wantFlusher bool
+		wantPusher  bool
+	}{
+		{
+			name:        "response writer only",
+			wrap:        rw{},
+			wantFlusher: false,
+			wantPusher:  false,
+		},
+		{
+			name:        "flusher only",
+			wrap:        flusher{},
+			wantFlusher: true,
+			wantPusher:  false,
+		},
+		{
+			name:        "pusher only",
+			wrap:        pusher{},
+			wantFlusher: false,
+			wantPusher:  true,
+		},
+		{
+			name:        "flusher and pusher",
+			wrap:        flusherPusher{},
+			wantFlusher: true,
+			wantPusher:  true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			wrapped := wrapProxyHandler(&hijackerWrapper{}, tt.wrap)
+			require.NotNil(t, wrapped)
+
+			// The wrapped interface should always be a Hijacker and ResponseWriter.
+			require.Implements(t, (*http.ResponseWriter)(nil), wrapped)
+			require.Implements(t, (*http.Hijacker)(nil), wrapped)
+
+			if tt.wantFlusher {
+				require.Implements(t, (*http.Flusher)(nil), wrapped)
+			}
+			if tt.wantPusher {
+				require.Implements(t, (*http.Pusher)(nil), wrapped)
+			}
+		})
+	}
+}
+
+func TestHandlerWrapper_WithHijacker(t *testing.T) {
+	handlerCalled := false
+	h := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) { // Mock for the proxy handler.
+		handlerCalled = true
+
+		require.NotNil(t, req)
+		require.NotNil(t, rw)
+
+		// We only wrap the response writer in metrics code if it implements Hijacker.
+		require.IsType(t, &hijackerWrapper{}, rw)
+
+		mw := rw.(*hijackerWrapper)
+		require.NotNil(t, mw.Hijacker)
+		require.NotNil(t, mw.ResponseWriter)
+	})
+
+	rw := &testHttpHijacker{ResponseWriter: httptest.NewRecorder(), t: t}
+	req := httptest.NewRequest("GET", "/v1/proxy", nil)
+	wrapped := InstrumentWebsocketWrapper(h)
+	wrapped.ServeHTTP(rw, req)
+
+	require.True(t, handlerCalled)
+}
+
+func TestHandlerWrapper_WithNoHijacker(t *testing.T) {
+	handlerCalled := false
+	h := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) { // Mock for the proxy handler.
+		handlerCalled = true
+
+		require.NotNil(t, req)
+		require.NotNil(t, rw)
+		require.IsType(t, &httptest.ResponseRecorder{}, rw) // The websocket metrics code must not have wrapped `rw`
+	})
+
+	rw := httptest.NewRecorder()
+	req := httptest.NewRequest("GET", "/v1/proxy", nil)
+	wrapped := InstrumentWebsocketWrapper(h)
+	wrapped.ServeHTTP(rw, req)
+
+	require.True(t, handlerCalled)
+}
+
+func TestHijack(t *testing.T) {
+	ogConnCount := activeConnCount
+	ogBytesReceived := bytesReceived
+	ogBytesSent := bytesSent
+	t.Cleanup(func() {
+		activeConnCount = ogConnCount
+		bytesReceived = ogBytesReceived
+		bytesSent = ogBytesSent
+	})
+
+	tpg := &testPrometheusGauge{t: t}
+	activeConnCount = tpg
+
+	tpcBytesRecv := &testPrometheusCounter{t: t}
+	bytesReceived = tpcBytesRecv
+
+	tpcBytesSent := &testPrometheusCounter{t: t}
+	bytesSent = tpcBytesSent
+
+	underlyingConn := &testConn{t: t, readEOF: true, retWriteLen: true}
+	mw := &hijackerWrapper{Hijacker: &testHttpHijacker{t: t, retConn: underlyingConn}}
+
+	c, brw, err := mw.Hijack()
+	require.NoError(t, err)
+	require.NotNil(t, c)
+	require.NotNil(t, brw)
+
+	require.Equal(t, 1, tpg.incCalledN)
+	require.Equal(t, 0, tpg.decCalledN)
+	require.IsType(t, &wsConn{}, c)
+
+	wsConn := c.(*wsConn)
+	require.Equal(t, underlyingConn, wsConn.Conn) // Underlying conn must be our testConn.
+	require.Equal(t, &mw.dec, wsConn.Once)        // wsConn.Once must be set to a ptr to `dec`.
+
+	// Assert that the buffer Reader and Writer are backed up
+	// by the underlying Conn.
+	_, err = io.ReadAll(brw.Reader)
+	require.NoError(t, err)
+	require.True(t, underlyingConn.readCalled)
+	require.True(t, tpcBytesRecv.addCalled)
+
+	_, err = io.WriteString(brw.Writer, "test")
+	require.NoError(t, err)
+	err = brw.Flush()
+	require.NoError(t, err)
+	require.True(t, underlyingConn.writeCalled)
+	require.True(t, tpcBytesSent.addCalled)
+}
+
+func TestHijack_Error(t *testing.T) {
+	ogConnCount := activeConnCount
+	t.Cleanup(func() { activeConnCount = ogConnCount })
+
+	tpg := &testPrometheusGauge{t: t}
+	activeConnCount = tpg
+
+	mw := &hijackerWrapper{
+		Hijacker: &testHttpHijacker{
+			t:      t,
+			retErr: fmt.Errorf("oops!"), // Trigger underlying `Hijack` error
+		},
+	}
+
+	c, brw, err := mw.Hijack()
+	require.Nil(t, c)
+	require.Nil(t, brw)
+	require.EqualError(t, err, "oops!")
+
+	require.Zero(t, tpg.incCalledN)
+	require.Zero(t, tpg.decCalledN)
+}
+
+func TestHijack_MultipleCalls(t *testing.T) {
+	ogConnCount := activeConnCount
+	t.Cleanup(func() { activeConnCount = ogConnCount })
+
+	tpg := &testPrometheusGauge{t: t}
+	activeConnCount = tpg
+
+	thh := &testHttpHijacker{t: t, retConn: &testConn{t: t}}
+	mw := &hijackerWrapper{Hijacker: thh}
+
+	n := 10
+	for i := 0; i < n; i++ {
+		mw.Hijack()
+	}
+	require.Equal(t, n, thh.hijackCalledN)
+	require.Equal(t, 1, tpg.incCalledN)
+	require.Zero(t, tpg.decCalledN)
+}
+
+func TestWsConnRead(t *testing.T) {
+	ogBytesReceived := bytesReceived
+	t.Cleanup(func() { bytesReceived = ogBytesReceived })
+
+	tpc := &testPrometheusCounter{t: t}
+	bytesReceived = tpc
+
+	tc := &testConn{t: t}
+	wc := &wsConn{Once: &sync.Once{}, Conn: tc}
+
+	n, err := wc.Read([]byte{})
+	require.NoError(t, err)
+	require.True(t, tc.readCalled)
+	require.True(t, tpc.addCalled)
+	require.Equal(t, 1, n)
+}
+
+func TestWsConnRead_Error(t *testing.T) {
+	ogBytesReceived := bytesReceived
+	t.Cleanup(func() { bytesReceived = ogBytesReceived })
+
+	tpc := &testPrometheusCounter{t: t}
+	bytesReceived = tpc
+
+	tc := &testConn{t: t, readErr: fmt.Errorf("oops!")} // Trigger underlying `Read` error.
+	wc := &wsConn{Once: &sync.Once{}, Conn: tc}
+
+	n, err := wc.Read([]byte{})
+	require.EqualError(t, err, "oops!")
+	require.True(t, tc.readCalled)
+	require.False(t, tpc.addCalled)
+	require.Zero(t, n)
+}
+
+func TestWsConnWrite(t *testing.T) {
+	ogBytesSent := bytesSent
+	t.Cleanup(func() { bytesSent = ogBytesSent })
+
+	tpc := &testPrometheusCounter{t: t}
+	bytesSent = tpc
+
+	tc := &testConn{t: t}
+	wc := &wsConn{Once: &sync.Once{}, Conn: tc}
+
+	n, err := wc.Write([]byte{})
+	require.NoError(t, err)
+	require.True(t, tc.writeCalled)
+	require.True(t, tpc.addCalled)
+	require.Equal(t, 1, n)
+}
+
+func TestWsConnWrite_Error(t *testing.T) {
+	ogBytesSent := bytesSent
+	t.Cleanup(func() { bytesSent = ogBytesSent })
+
+	tpc := &testPrometheusCounter{t: t}
+	bytesSent = tpc
+
+	tc := &testConn{t: t, writeErr: fmt.Errorf("oops!")} // Trigger underlying `Write` error.
+	wc := &wsConn{Once: &sync.Once{}, Conn: tc}
+
+	n, err := wc.Write([]byte{})
+	require.EqualError(t, err, "oops!")
+	require.True(t, tc.writeCalled)
+	require.False(t, tpc.addCalled)
+	require.Equal(t, 0, n)
+}
+
+func TestWsConnClose(t *testing.T) {
+	ogConnCount := activeConnCount
+	t.Cleanup(func() { activeConnCount = ogConnCount })
+
+	tpg := &testPrometheusGauge{t: t}
+	activeConnCount = tpg
+
+	tc := &testConn{t: t}
+	wsConn := &wsConn{Once: &sync.Once{}, Conn: tc}
+
+	err := wsConn.Close()
+	require.NoError(t, err)
+	require.Equal(t, 1, tc.closeCalledN)
+	require.Equal(t, 1, tpg.decCalledN)
+	require.Equal(t, 0, tpg.incCalledN)
+}
+
+func TestWsConnClose_Error(t *testing.T) {
+	ogConnCount := activeConnCount
+	t.Cleanup(func() { activeConnCount = ogConnCount })
+
+	tpg := &testPrometheusGauge{t: t}
+	activeConnCount = tpg
+
+	wsConn := &wsConn{
+		Once: &sync.Once{},
+		Conn: &testConn{
+			t:        t,
+			closeErr: fmt.Errorf("oops!"), // Trigger underlying `Close` error.
+		},
+	}
+
+	err := wsConn.Close()
+	require.EqualError(t, err, "oops!")
+}
+
+func TestWsConnClose_MultipleCalls(t *testing.T) {
+	ogConnCount := activeConnCount
+	t.Cleanup(func() { activeConnCount = ogConnCount })
+
+	tpg := &testPrometheusGauge{t: t}
+	activeConnCount = tpg
+
+	tc := &testConn{t: t}
+	wsConn := &wsConn{Once: &sync.Once{}, Conn: tc}
+
+	calledCount := 0
+	for i := 0; i < 10; i++ {
+		wsConn.Close()
+		calledCount++
+	}
+
+	require.Equal(t, calledCount, tc.closeCalledN)
+	require.Equal(t, 1, tpg.decCalledN) // We should only decrement the metric once
+	require.Zero(t, tpg.incCalledN)
+}
+
+type testHttpHijacker struct {
+	// Mocks the underlying http.Hijacker in metricsWrapper.
+	http.ResponseWriter
+	t             *testing.T
+	retConn       net.Conn
+	retErr        error
+	hijackCalledN int
+}
+
+func (h *testHttpHijacker) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	h.hijackCalledN++
+
+	if h.retErr != nil {
+		return nil, nil, h.retErr
+	}
+	if h.retConn == nil {
+		return nil, nil, nil
+	}
+
+	brw := bufio.NewReadWriter(bufio.NewReader(h.retConn), bufio.NewWriter(h.retConn))
+	return h.retConn, brw, nil
+}
+
+type testConn struct {
+	// Mocks the underlying net.Conn in wsConn.
+	readCalled bool
+	readEOF    bool
+	readErr    error
+
+	writeCalled bool
+	retWriteLen bool
+	writeErr    error
+
+	closeCalledN int
+	closeErr     error
+	t            *testing.T
+}
+
+func (tc *testConn) Read(b []byte) (n int, err error) {
+	tc.readCalled = true
+	if tc.readErr != nil {
+		return 0, tc.readErr
+	}
+	if tc.readEOF { // Simulate an actual read and tx one byte
+		b[0] = 1
+		return 1, io.EOF
+	}
+
+	return 1, nil
+}
+
+func (tc *testConn) Write(b []byte) (n int, err error) {
+	tc.writeCalled = true
+	if tc.writeErr != nil {
+		return 0, tc.writeErr
+	}
+	if tc.retWriteLen {
+		return len(b), nil
+	}
+	return 1, nil
+}
+
+func (tc *testConn) Close() error {
+	tc.closeCalledN++
+	if tc.closeErr != nil {
+		return tc.closeErr
+	}
+	return nil
+}
+func (tc *testConn) LocalAddr() net.Addr                { return nil } // Unused
+func (tc *testConn) RemoteAddr() net.Addr               { return nil } // Unused
+func (tc *testConn) SetDeadline(t time.Time) error      { return nil } // Unused
+func (tc *testConn) SetReadDeadline(t time.Time) error  { return nil } // Unused
+func (tc *testConn) SetWriteDeadline(t time.Time) error { return nil } // Unused
+
+type testPrometheusCounter struct {
+	// Mocks a Prometheus Counter as used in Proxy Websocket metrics.
+	prometheus.Metric
+	prometheus.Collector
+
+	addCalled bool
+	t         *testing.T
+}
+
+func (tpc *testPrometheusCounter) Inc()        { tpc.t.Fatal("testPrometheusCounter Inc() called") }
+func (tpc *testPrometheusCounter) Add(float64) { tpc.addCalled = true }
+
+type testPrometheusGauge struct {
+	// Mocks a Prometheus Gauge as used in Proxy Websocket metrics.
+	prometheus.Metric
+	prometheus.Collector
+
+	incCalledN int
+	decCalledN int
+	t          *testing.T
+}
+
+func (tpg *testPrometheusGauge) Set(float64) { tpg.t.Fatal("testPrometheusGauge Set() called") }
+func (tpg *testPrometheusGauge) Inc()        { tpg.incCalledN++ }
+func (tpg *testPrometheusGauge) Dec()        { tpg.decCalledN++ }
+func (tpg *testPrometheusGauge) Add(float64) { tpg.t.Fatal("testPrometheusGauge Add() called") }
+func (tpg *testPrometheusGauge) Sub(float64) { tpg.t.Fatal("testPrometheusGauge Sub() called") }
+func (tpg *testPrometheusGauge) SetToCurrentTime() {
+	tpg.t.Fatal("testPrometheusGauge SetToCurrentTime() called")
+}

internal/servers/worker/worker.go

@@ -67,6 +67,7 @@ type Worker struct {
 
 func New(conf *Config) (*Worker, error) {
 	metric.InitializeHttpCollectors(conf.PrometheusRegisterer)
+	metric.InitializeWebsocketCollectors(conf.PrometheusRegisterer)
 	w := &Worker{
 		conf:                  conf,
 		logger:                conf.Logger.Named("worker"),