diff --git a/internal/servers/worker/handler.go b/internal/servers/worker/handler.go index 74ab26c4f..aea2c15ce 100644 --- a/internal/servers/worker/handler.go +++ b/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) diff --git a/internal/servers/worker/internal/metric/proxy_ws.go b/internal/servers/worker/internal/metric/proxy_ws.go new file mode 100644 index 000000000..76e1bd205 --- /dev/null +++ b/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 +} diff --git a/internal/servers/worker/internal/metric/proxy_ws_test.go b/internal/servers/worker/internal/metric/proxy_ws_test.go new file mode 100644 index 000000000..0ebfd490f --- /dev/null +++ b/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") +} diff --git a/internal/servers/worker/worker.go b/internal/servers/worker/worker.go index a6f54aef5..ad403ef0a 100644 --- a/internal/servers/worker/worker.go +++ b/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"),