ws to grpc connection adapter

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>
2026-09-13 18:29:07 +02:00 · 2026-09-09 14:45:37 +02:00
parent 066af82c3e
commit ff2b45789b
3 changed files with 343 additions and 111 deletions
@@ -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)
 }
@@ -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()
+}
@@ -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)
+}