diff --git a/util/wsproxy/server/proxy.go b/util/wsproxy/server/proxy.go index 512d42ac9..9d4d9807c 100644 --- a/util/wsproxy/server/proxy.go +++ b/util/wsproxy/server/proxy.go @@ -51,17 +51,15 @@ 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 &proxyHandler{ - metrics: p.config.MetricsRecorder, - handler: p.config.Handler, - headersReadTimeout: 10 * time.Second, + metrics: p.config.MetricsRecorder, + handler: p.config.Handler, } } type proxyHandler struct { - metrics MetricsRecorder - handler http.Handler - conn atomic.Pointer[wsConnAdapter] - headersReadTimeout time.Duration + metrics MetricsRecorder + handler http.Handler + conn atomic.Pointer[wsConnAdapter] } func (ph *proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { @@ -86,7 +84,7 @@ func (ph *proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { conn: wsConn, metrics: ph.metrics, clientAddr: r.RemoteAddr, - }).WithFrameSnooper(ph.headersReadTimeout) + }) defer func() { _ = serverConn.Close() @@ -100,7 +98,7 @@ func (ph *proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { // TODO (dmitri) we should limit the number of concurrent streams per connection (peer) // and idle timeouts // MaxConcurrentStreams: 20, - // IdleTimeout: 60 * time.Second, + // IdleTimeout: 10 * time.Second, }).ServeConn(serverConn, &http2.ServeConnOpts{ Context: ctx, Handler: ph.handler, diff --git a/util/wsproxy/server/ws_conn_adapter.go b/util/wsproxy/server/ws_conn_adapter.go index 4f1441c0e..eb29ab0cb 100644 --- a/util/wsproxy/server/ws_conn_adapter.go +++ b/util/wsproxy/server/ws_conn_adapter.go @@ -1,29 +1,23 @@ package server import ( - "bytes" "context" - "io" "net" + "sync/atomic" "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 - headerReadDeadlineTimer *time.Timer + prefix string + ctx context.Context + conn *websocket.Conn + metrics MetricsRecorder + clientAddr string + closed atomic.Bool + bufferedRead []byte } var _ net.Conn = &wsConnAdapter{} @@ -33,18 +27,6 @@ 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(d time.Duration) *wsConnAdapter { - if d == 0 { - return ws - } - - ws.frameBuffer = bytes.NewBuffer(make([]byte, 0, 512)) - ws.framer = http2.NewFramer(nil, ws.frameBuffer) - ws.framer.ReadMetaHeaders = hpack.NewDecoder(0, nil) - ws.headerReadDeadlineTimer = time.AfterFunc(d, ws.onReadTimeout) - return ws -} - func (ws *wsConnAdapter) Read(b []byte) (int, error) { if len(ws.bufferedRead) > 0 { return ws.readFromBuffer(b) @@ -75,16 +57,6 @@ func (ws *wsConnAdapter) Read(b []byte) (int, error) { func (ws *wsConnAdapter) readFromBuffer(b []byte) (int, error) { n := copy(b, ws.bufferedRead) - // check if we started receiving data, stop the header read timeout timer - if ws.isFramerActive() { - _, _ = ws.frameBuffer.Write(b) // we don't care about the number of bytes copied and no errors are returned from Write - if frame, err := ws.framer.ReadFrame(); err != nil && frame != nil && frame.Header().Type == http2.FrameData { - ws.headerReadDeadlineTimer.Stop() - ws.cleanupFramer() - } - } - - io.Pipe() ws.recordBytesTransferred(ws.ctx, "ws_to_grpc", n) if n == len(ws.bufferedRead) { ws.bufferedRead = nil @@ -116,7 +88,7 @@ func (ws *wsConnAdapter) Write(b []byte) (int, error) { } func (ws *wsConnAdapter) Close() error { - ws.closed = true + ws.closed.Store(true) return ws.conn.Close(websocket.StatusNormalClosure, "") } @@ -150,19 +122,5 @@ func (ws *wsConnAdapter) recordBytesTransferred(ctx context.Context, direction s } func (ws *wsConnAdapter) IsClosed() bool { - return ws.closed -} - -func (ws *wsConnAdapter) onReadTimeout() { - ws.Close() -} - -func (ws *wsConnAdapter) isFramerActive() bool { - return ws.framer != nil && ws.headerReadDeadlineTimer != nil -} - -func (ws *wsConnAdapter) cleanupFramer() { - ws.frameBuffer = nil - ws.framer = nil - ws.headerReadDeadlineTimer = nil + return ws.closed.Load() } diff --git a/util/wsproxy/server/ws_conn_adapter_test.go b/util/wsproxy/server/ws_conn_adapter_test.go index b701f59cd..d672bd92c 100644 --- a/util/wsproxy/server/ws_conn_adapter_test.go +++ b/util/wsproxy/server/ws_conn_adapter_test.go @@ -106,6 +106,8 @@ func TestAdapterHandlingConnectionClosures(t *testing.T) { } 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) @@ -120,7 +122,6 @@ func TestAdapterHandlingHttpConnection_NoHeadersSent(t *testing.T) { handler, ok := proxy.Handler().(*proxyHandler) assert.True(t, ok) - handler.headersReadTimeout = 1 * time.Second protocols := new(http.Protocols) protocols.SetHTTP1(true) @@ -146,7 +147,7 @@ func TestAdapterHandlingHttpConnection_NoHeadersSent(t *testing.T) { prefix: "test-client", ctx: context.Background(), conn: clientconn, - }, shouldDropFrame: func(f http2.FrameType) bool { return f == http2.FrameHeaders }}, nil + }, shouldDropFrame: func(f http2.FrameType) bool { return f == http2.FrameHeaders || f == http2.FrameData }}, nil }, }}