diff --git a/util/wsproxy/server/proxy.go b/util/wsproxy/server/proxy.go index ffb622200..0618beb91 100644 --- a/util/wsproxy/server/proxy.go +++ b/util/wsproxy/server/proxy.go @@ -1,12 +1,8 @@ package server import ( - "context" - "io" - "net" "net/http" - "sync" - "time" + "sync/atomic" "github.com/coder/websocket" log "github.com/sirupsen/logrus" @@ -15,11 +11,6 @@ import ( "github.com/netbirdio/netbird/util/wsproxy" ) -const ( - bufferSize = 32 * 1024 - ioTimeout = 5 * time.Second -) - // Config contains the configuration for the WebSocket proxy. type Config struct { Handler http.Handler @@ -53,14 +44,23 @@ func New(handler http.Handler, opts ...Option) *Proxy { // Handler returns an http.Handler that proxies WebSocket connections to the local gRPC server. func (p *Proxy) Handler() http.Handler { - return http.HandlerFunc(p.handleWebSocket) + return &proxyHandler{ + metrics: p.config.MetricsRecorder, + handler: p.config.Handler, + } } -func (p *Proxy) handleWebSocket(w http.ResponseWriter, r *http.Request) { +type proxyHandler struct { + metrics MetricsRecorder + handler http.Handler + conn atomic.Pointer[wsConnAdapter] +} + +func (ph *proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { ctx := r.Context() - p.metrics.RecordConnection(ctx) - defer p.metrics.RecordDisconnection(ctx) + ph.metrics.RecordConnection(ctx) + defer ph.metrics.RecordDisconnection(ctx) log.Debugf("WebSocket proxy handling connection from %s, forwarding to internal gRPC handler", r.RemoteAddr) acceptOptions := &websocket.AcceptOptions{ @@ -69,121 +69,44 @@ func (p *Proxy) handleWebSocket(w http.ResponseWriter, r *http.Request) { wsConn, err := websocket.Accept(w, r, acceptOptions) if err != nil { - p.metrics.RecordError(ctx, "websocket_accept_failed") + ph.metrics.RecordError(ctx, "websocket_accept_failed") log.Errorf("WebSocket upgrade failed from %s: %v", r.RemoteAddr, err) return } - defer func() { - _ = wsConn.Close(websocket.StatusNormalClosure, "") - }() + serverConn := (&wsConnAdapter{ + ctx: ctx, + conn: wsConn, + metrics: ph.metrics, + clientAddr: r.RemoteAddr, + }) - clientConn, serverConn := net.Pipe() defer func() { - _ = clientConn.Close() _ = serverConn.Close() }() + ph.conn.Store(serverConn) // used in tests only + log.Debugf("WebSocket proxy established: %s -> gRPC handler", r.RemoteAddr) - go func() { - (&http2.Server{}).ServeConn(serverConn, &http2.ServeConnOpts{ - Context: ctx, - Handler: p.config.Handler, - }) - }() + (&http2.Server{ + // TODO (dmitri) we should limit the number of concurrent streams per connection (peer) + // and idle timeouts + // MaxConcurrentStreams: 20, + // IdleTimeout: 10 * time.Second, + }).ServeConn(serverConn, &http2.ServeConnOpts{ + Context: ctx, + Handler: ph.handler, + BaseConfig: &http.Server{ + // b/c we are wrapping a ws connection, read and write connection deadlines normally set + // via ReadTimeout and WriteTimeout http.Server fields aren't available to us. The ws + // library doesn't expose connection deadline timer config, and we ignore these calls in "wsConnAdapter". + // + // Another issue is that Server.ServeConn() call bypasses setting of connection deadlines altogether, + // ReadTimeout and Writetimeout set here would only apply to h2 streams, i.e. after a HEADERS frame + // arrival and processing, turning ReadTimeout into a request body read deadline, and WriteTimeout into + // a response deadline (the latter not useful for streaming requests). + }, + }) - p.proxyData(ctx, wsConn, clientConn, r.RemoteAddr) -} - -func (p *Proxy) proxyData(ctx context.Context, wsConn *websocket.Conn, pipeConn net.Conn, clientAddr string) { - proxyCtx, cancel := context.WithCancel(ctx) - defer cancel() - - var wg sync.WaitGroup - wg.Add(2) - - go p.wsToPipe(proxyCtx, cancel, &wg, wsConn, pipeConn, clientAddr) - go p.pipeToWS(proxyCtx, cancel, &wg, wsConn, pipeConn, clientAddr) - - wg.Wait() -} - -func (p *Proxy) wsToPipe(ctx context.Context, cancel context.CancelFunc, wg *sync.WaitGroup, wsConn *websocket.Conn, pipeConn net.Conn, clientAddr string) { - defer wg.Done() - defer cancel() - - for { - msgType, data, err := wsConn.Read(ctx) - if err != nil { - switch { - case ctx.Err() != nil: - log.Debugf("WebSocket from %s terminating due to context cancellation", clientAddr) - case websocket.CloseStatus(err) != -1: - log.Debugf("WebSocket from %s disconnected", clientAddr) - default: - p.metrics.RecordError(ctx, "websocket_read_error") - log.Debugf("WebSocket read error from %s: %v", clientAddr, err) - } - return - } - - if msgType != websocket.MessageBinary { - log.Warnf("Unexpected WebSocket message type from %s: %v", clientAddr, msgType) - continue - } - - if ctx.Err() != nil { - log.Tracef("wsToPipe goroutine terminating due to context cancellation before pipe write") - return - } - - if err := pipeConn.SetWriteDeadline(time.Now().Add(ioTimeout)); err != nil { - log.Debugf("Failed to set pipe write deadline: %v", err) - } - - n, err := pipeConn.Write(data) - if err != nil { - p.metrics.RecordError(ctx, "pipe_write_error") - log.Warnf("Pipe write error for %s: %v", clientAddr, err) - return - } - - p.metrics.RecordBytesTransferred(ctx, "ws_to_grpc", int64(n)) - } -} - -func (p *Proxy) pipeToWS(ctx context.Context, cancel context.CancelFunc, wg *sync.WaitGroup, wsConn *websocket.Conn, pipeConn net.Conn, clientAddr string) { - defer wg.Done() - defer cancel() - - buf := make([]byte, bufferSize) - for { - n, err := pipeConn.Read(buf) - if err != nil { - if ctx.Err() != nil { - log.Tracef("pipeToWS goroutine terminating due to context cancellation") - return - } - - if err != io.EOF { - log.Debugf("Pipe read error for %s: %v", clientAddr, err) - } - return - } - - if ctx.Err() != nil { - log.Tracef("pipeToWS goroutine terminating due to context cancellation before WebSocket write") - return - } - - if n > 0 { - if err := wsConn.Write(ctx, websocket.MessageBinary, buf[:n]); err != nil { - p.metrics.RecordError(ctx, "websocket_write_error") - log.Warnf("WebSocket write error for %s: %v", clientAddr, err) - return - } - - p.metrics.RecordBytesTransferred(ctx, "grpc_to_ws", int64(n)) - } - } + log.Debugf("WebSocket proxy closing: %s -> gRPC handler", r.RemoteAddr) } diff --git a/util/wsproxy/server/ws_conn_adapter.go b/util/wsproxy/server/ws_conn_adapter.go new file mode 100644 index 000000000..eb29ab0cb --- /dev/null +++ b/util/wsproxy/server/ws_conn_adapter.go @@ -0,0 +1,126 @@ +package server + +import ( + "context" + "net" + "sync/atomic" + "time" + + "github.com/coder/websocket" + log "github.com/sirupsen/logrus" +) + +type wsConnAdapter struct { + prefix string + ctx context.Context + conn *websocket.Conn + metrics MetricsRecorder + clientAddr string + closed atomic.Bool + bufferedRead []byte +} + +var _ net.Conn = &wsConnAdapter{} + +type wsAddr struct{ prefix string } + +func (wa wsAddr) Network() string { return wa.prefix + "ws-proxy" } +func (wa wsAddr) String() string { return wa.prefix + "ws-proxy" } + +func (ws *wsConnAdapter) Read(b []byte) (int, error) { + if len(ws.bufferedRead) > 0 { + return ws.readFromBuffer(b) + } + + msgType, data, err := ws.conn.Read(ws.ctx) + if err != nil { + switch { + case ws.ctx.Err() != nil: + log.Debugf("WebSocket from %s terminating due to context cancellation", ws.clientAddr) + case websocket.CloseStatus(err) != -1: + log.Debugf("WebSocket from %s disconnected", ws.clientAddr) + default: + ws.recordError(ws.ctx, "websocket_read_error") + log.Debugf("WebSocket read error from %s: %v", ws.clientAddr, err) + } + return copy(b, data), err + } + if msgType != websocket.MessageBinary { + log.Warnf("Unexpected WebSocket message type from %s: %v", ws.clientAddr, msgType) + return 0, nil + } + + ws.bufferedRead = data + return ws.readFromBuffer(b) +} + +func (ws *wsConnAdapter) readFromBuffer(b []byte) (int, error) { + n := copy(b, ws.bufferedRead) + + ws.recordBytesTransferred(ws.ctx, "ws_to_grpc", n) + if n == len(ws.bufferedRead) { + ws.bufferedRead = nil + return n, nil + } else { + ws.bufferedRead = ws.bufferedRead[n:] + } + return n, nil +} + +func (ws *wsConnAdapter) Write(b []byte) (int, error) { + maybeErr := ws.ctx.Err() + + n := len(b) + if n == 0 { + return n, maybeErr + } + if maybeErr != nil { + return 0, maybeErr + } + if err := ws.conn.Write(ws.ctx, websocket.MessageBinary, b[:n]); err != nil { + ws.recordError(ws.ctx, "websocket_write_error") + log.Warnf("WebSocket write error for %s: %v", ws.clientAddr, err) + return 0, err // we don't know how many bytes have been written + } + + ws.recordBytesTransferred(ws.ctx, "grpc_to_ws", n) + return n, nil +} + +func (ws *wsConnAdapter) Close() error { + ws.closed.Store(true) + return ws.conn.Close(websocket.StatusNormalClosure, "") +} + +func (ws *wsConnAdapter) LocalAddr() net.Addr { return wsAddr{ws.prefix} } +func (ws *wsConnAdapter) RemoteAddr() net.Addr { return wsAddr{ws.prefix} } + +func (ws *wsConnAdapter) SetDeadline(t time.Time) error { + return nil +} + +func (ws *wsConnAdapter) SetReadDeadline(t time.Time) error { + return nil +} + +func (ws *wsConnAdapter) SetWriteDeadline(t time.Time) error { + return nil +} + +func (ws *wsConnAdapter) recordError(ctx context.Context, errorType string) { + if ws.metrics == nil { + return + } + ws.metrics.RecordError(ctx, errorType) +} + +func (ws *wsConnAdapter) recordBytesTransferred(ctx context.Context, direction string, bytes int) { + if ws.metrics == nil { + return + } + ws.metrics.RecordBytesTransferred(ctx, direction, int64(bytes)) +} + +func (ws *wsConnAdapter) IsClosed() bool { + return ws.closed.Load() +} diff --git a/util/wsproxy/server/ws_conn_adapter_test.go b/util/wsproxy/server/ws_conn_adapter_test.go new file mode 100644 index 000000000..5369b2362 --- /dev/null +++ b/util/wsproxy/server/ws_conn_adapter_test.go @@ -0,0 +1,204 @@ +package server + +import ( + "bytes" + "context" + "crypto/tls" + "io" + "math/rand/v2" + "net" + "net/http" + "os" + "path/filepath" + "strconv" + "strings" + "testing" + "time" + + "github.com/coder/websocket" + "github.com/stretchr/testify/assert" + "golang.org/x/net/http2" + "golang.org/x/net/http2/hpack" +) + +func TestAdapterHandlingConnectionClosures(t *testing.T) { + var cases = []struct { + description string + casenum int + }{ + {"client-side ws connection is closed", 0}, + {"server-side ws connection is closed", 1}, + {"client-side context is cancelled", 2}, + {"server-side context is cancelled", 3}, + } + + for _, c := range cases { + t.Run(c.description, func(t *testing.T) { + serversock := filepath.Join("/tmp", "http-server-"+strconv.FormatInt(rand.Int64(), 10)+".sock") + t.Cleanup(func() { os.Remove(serversock) }) + + l, err := net.Listen("unix", serversock) + assert.NoError(t, err) + + proxy := New(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + buf, _ := io.ReadAll(r.Body) + defer r.Body.Close() + w.Write([]byte("echo: " + string(buf))) //nolint:errcheck + })) + + handler, ok := proxy.Handler().(*proxyHandler) + assert.True(t, ok) + + protocols := new(http.Protocols) + protocols.SetHTTP1(true) + protocols.SetUnencryptedHTTP2(true) + httpServer := http.Server{ + Handler: handler, + } + go httpServer.Serve(l) //nolint:errcheck + t.Cleanup(func() { httpServer.Close() }) + + clientconn, _, err := websocket.Dial(context.Background(), "http://whatever", //nolint:bodyclose + &websocket.DialOptions{HTTPClient: &http.Client{ + Transport: &http.Transport{ + DialContext: func(_ context.Context, _, _ string) (net.Conn, error) { + return net.Dial("unix", serversock) + }, + }}}) + assert.NoError(t, err) + + clientCtx, cancel := context.WithCancel(context.Background()) //nolint:govet + h2client := &http.Client{ + Transport: &http2.Transport{ + AllowHTTP: true, + DialTLSContext: func(_ context.Context, _, _ string, _ *tls.Config) (net.Conn, error) { + return &wsConnAdapter{ + prefix: "test-client", + ctx: clientCtx, + conn: clientconn, + }, nil + }, + }} + + resp, err := h2client.Post("http://whatever", "text/html", strings.NewReader("g'day")) + assert.NoError(t, err) + + body, err := io.ReadAll(resp.Body) + defer resp.Body.Close() + + assert.NoError(t, err) + assert.Equal(t, "echo: g'day", string(body)) + + switch c.casenum { + case 0: + clientconn.Close(websocket.StatusNormalClosure, "") + case 1: + handler.conn.Load().Close() + case 2: + cancel() + case 3: + resp.Body.Close() + h2client.CloseIdleConnections() + } + + assert.EventuallyWithT(t, func(c *assert.CollectT) { + assert.True(c, handler.conn.Load().IsClosed()) + }, 3*time.Second, 100*time.Millisecond) + }) //nolint:govet + } +} + +func TestAdapterHandlingHttpConnection_NoHeadersSent(t *testing.T) { + t.Skip("currently disabled as it requires idle timeout to be set") + + serversock := filepath.Join("/tmp", "http-server-"+strconv.FormatInt(rand.Int64(), 10)+".sock") + defer os.Remove(serversock) + + l, err := net.Listen("unix", serversock) + assert.NoError(t, err) + + proxy := New(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + buf, _ := io.ReadAll(r.Body) + defer r.Body.Close() //nolint:errcheck + w.Write([]byte("echo: " + string(buf))) //nolint:errcheck + })) + + handler, ok := proxy.Handler().(*proxyHandler) + assert.True(t, ok) + + protocols := new(http.Protocols) + protocols.SetHTTP1(true) + protocols.SetUnencryptedHTTP2(true) + httpServer := http.Server{ + Handler: handler, + } + go httpServer.Serve(l) //nolint:errcheck + + clientconn, _, err := websocket.Dial(context.Background(), "http://whatever", //nolint:bodyclose + &websocket.DialOptions{HTTPClient: &http.Client{ + Transport: &http.Transport{ + DialContext: func(_ context.Context, _, _ string) (net.Conn, error) { + return net.Dial("unix", serversock) + }, + }}}) + assert.NoError(t, err) + + h2client := &http.Client{ + Transport: &http2.Transport{ + AllowHTTP: true, + DialTLSContext: func(_ context.Context, _, _ string, _ *tls.Config) (net.Conn, error) { + return &h2ConnectionSnooper{wrappedConn: &wsConnAdapter{ + prefix: "test-client", + ctx: context.Background(), + conn: clientconn, + }, shouldDropFrame: func(f http2.FrameType) bool { return f == http2.FrameHeaders || f == http2.FrameData }}, nil + }, + }} + + _, err = h2client.Post("http://whatever", "text/html", strings.NewReader("g'day")) + assert.Error(t, err) + + assert.EventuallyWithT(t, func(c *assert.CollectT) { + assert.True(c, handler.conn.Load().IsClosed()) + }, 3*time.Second, 100*time.Millisecond) +} + +type h2ConnectionSnooper struct { + wrappedConn net.Conn + shouldDropFrame func(f http2.FrameType) bool +} + +func (hs *h2ConnectionSnooper) Read(b []byte) (n int, err error) { + return hs.wrappedConn.Read(b) +} + +func (hs *h2ConnectionSnooper) Write(b []byte) (n int, err error) { + fr := http2.NewFramer(nil, bytes.NewReader(b)) + fr.ReadMetaHeaders = hpack.NewDecoder(0, nil) + f, err := fr.ReadFrame() + if err != nil { + return hs.wrappedConn.Write(b) + } + + if hs.shouldDropFrame != nil && hs.shouldDropFrame(f.Header().Type) { + return len(b), nil + } + + return hs.wrappedConn.Write(b) +} + +func (hs *h2ConnectionSnooper) Close() error { return hs.wrappedConn.Close() } + +func (hs *h2ConnectionSnooper) LocalAddr() net.Addr { return hs.wrappedConn.LocalAddr() } + +func (hs *h2ConnectionSnooper) RemoteAddr() net.Addr { return hs.wrappedConn.RemoteAddr() } + +func (hs *h2ConnectionSnooper) SetDeadline(t time.Time) error { return hs.wrappedConn.SetDeadline(t) } + +func (hs *h2ConnectionSnooper) SetReadDeadline(t time.Time) error { + return hs.wrappedConn.SetReadDeadline(t) +} + +func (hs *h2ConnectionSnooper) SetWriteDeadline(t time.Time) error { + return hs.wrappedConn.SetWriteDeadline(t) +}