From ff2b45789b97bfc1675f0f6b0cfc3a395d86cf1c Mon Sep 17 00:00:00 2001 From: Dmitri Dolguikh Date: Wed, 9 Sep 2026 14:45:37 +0200 Subject: [PATCH] ws to grpc connection adapter Signed-off-by: Dmitri Dolguikh --- util/wsproxy/server/proxy.go | 156 ++++++-------------- util/wsproxy/server/ws_conn_adapter.go | 146 ++++++++++++++++++ util/wsproxy/server/ws_conn_adapter_test.go | 152 +++++++++++++++++++ 3 files changed, 343 insertions(+), 111 deletions(-) create mode 100644 util/wsproxy/server/ws_conn_adapter.go create mode 100644 util/wsproxy/server/ws_conn_adapter_test.go diff --git a/util/wsproxy/server/proxy.go b/util/wsproxy/server/proxy.go index ffb622200..bccaf5b11 100644 --- a/util/wsproxy/server/proxy.go +++ b/util/wsproxy/server/proxy.go @@ -1,11 +1,7 @@ package server import ( - "context" - "io" - "net" "net/http" - "sync" "time" "github.com/coder/websocket" @@ -53,14 +49,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 *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 +74,50 @@ 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 } + serverConn := &wsConnAdapter{ + ctx: ctx, + conn: wsConn, + metrics: ph.metrics, + clientAddr: r.RemoteAddr, + } defer func() { - _ = wsConn.Close(websocket.StatusNormalClosure, "") - }() - - clientConn, serverConn := net.Pipe() - defer func() { - _ = clientConn.Close() _ = serverConn.Close() }() + ph.conn = 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{ + // MaxConcurrentStreams: 20, + // IdleTimeout: 3 * time.Second, + }).ServeConn(serverConn, &http2.ServeConnOpts{ + Context: ctx, + Handler: ph.handler, + BaseConfig: &http.Server{ + // // this is disabled in http2/server.go in "processHeaders" + // // after headers have been read + ReadHeaderTimeout: 2 * time.Second, + ReadTimeout: 5 * time.Second, + // // IdleTimeout: 10 * time.Second, + }, + }) - p.proxyData(ctx, wsConn, clientConn, r.RemoteAddr) + log.Debugf("WebSocket proxy closing: %s -> gRPC handler", r.RemoteAddr) + + // 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() +type timeoutResettingHandler struct { + wrappedHandler http.Handler } -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)) - } - } +func (th *timeoutResettingHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { + // ctrl := http.NewResponseController(w) + // ctrl.SetReadDeadline(time.Time{}) + th.wrappedHandler.ServeHTTP(w, r) } diff --git a/util/wsproxy/server/ws_conn_adapter.go b/util/wsproxy/server/ws_conn_adapter.go new file mode 100644 index 000000000..d0b7ef127 --- /dev/null +++ b/util/wsproxy/server/ws_conn_adapter.go @@ -0,0 +1,146 @@ +package server + +import ( + "bytes" + "context" + "net" + "time" + + "github.com/coder/websocket" + log "github.com/sirupsen/logrus" + "golang.org/x/net/http2" + "golang.org/x/net/http2/hpack" +) + +type wsConnAdapter struct { + prefix string + ctx context.Context + conn *websocket.Conn + metrics MetricsRecorder + clientAddr string + closed bool + bufferedRead []byte + frameBuffer *bytes.Buffer + framer *http2.Framer + frameDecoder *hpack.Decoder +} + +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) WithFrameSnooper() { + ws.frameBuffer = bytes.NewBuffer(make([]byte, 0, 512)) + ws.framer = http2.NewFramer(nil, ws.frameBuffer) + ws.frameDecoder = hpack.NewDecoder(0, nil) +} + +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) + + f, err := ws.frameDecoder.ReadFrame() + if err != nil { + return hs.wrappedConn.Write(b) + } + + 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 = 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 { + time.AfterFunc(time.Until(t), ws.onReadTimeout) + 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 +} + +func (ws *wsConnAdapter) onReadTimeout() { + ws.Close() +} 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..4cb37920c --- /dev/null +++ b/util/wsproxy/server/ws_conn_adapter_test.go @@ -0,0 +1,152 @@ +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 TestXxx(t *testing.T) { + var cases = []struct { + description string + casenum int + frameHandlerFunc func(b []byte) (n int, err error) + }{ + // {"client-side ws connection is closed", 0, nil}, + // {"server-side ws connection is closed", 1, nil}, + // {"client-side context is cancelled", 2, nil}, + // {"server-side context is cancelled", 3, nil}, + {"client slow to start a stream", 4, func(b []byte) (n int, err error) { return len(b), nil }}, + } + + for _, c := range cases { + t.Run(c.description, func(t *testing.T) { + 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() + w.Write([]byte("echo: " + string(buf))) + })) + + handler, ok := proxy.Handler().(*proxyHandler) + assert.True(t, ok) + + protocols := new(http.Protocols) + protocols.SetHTTP1(true) + protocols.SetUnencryptedHTTP2(true) + httpServer := http.Server{ + Handler: handler, + IdleTimeout: 3 * time.Second, + // Handler: h2c.NewHandler(handler, &http2.Server{ + // IdleTimeout: 500 * time.Millisecond, + // }), + } + go httpServer.Serve(l) + + clientconn, _, err := websocket.Dial(context.Background(), "http://whatever", &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()) + 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: clientCtx, + conn: clientconn, + }, frameHandlerFunc: c.frameHandlerFunc}, nil + }, + }} + + resp, err := h2client.Post("http://whatever", "text/html", strings.NewReader("g'day")) + assert.NoError(t, err) + + body, err := io.ReadAll(resp.Body) + + 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.Close() + case 2: + cancel() + case 3: + resp.Body.Close() + h2client.CloseIdleConnections() + } + + assert.EventuallyWithT(t, func(c *assert.CollectT) { + assert.True(c, handler.conn.IsClosed()) + }, 5*time.Second, 100*time.Millisecond) + }) + } +} + +type h2ConnectionSnooper struct { + wrappedConn net.Conn + frameHandlerFunc func(b []byte) (n int, err error) +} + +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 (f.Header().Type == http2.FrameData || f.Header().Type == http2.FrameHeaders) && hs.frameHandlerFunc != nil { + return hs.frameHandlerFunc(b) + } + + 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) +}