support for timeouts on reading h2 stream headers

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

util/wsproxy/server/proxy.go

@@ -56,9 +56,10 @@ func (p *Proxy) Handler() http.Handler {
 }
 
 type proxyHandler struct {
-	metrics MetricsRecorder
-	handler http.Handler
-	conn    *wsConnAdapter
+	metrics            MetricsRecorder
+	handler            http.Handler
+	conn               *wsConnAdapter
+	headersReadTimeout time.Duration
 }
 
 func (ph *proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
@@ -78,12 +79,13 @@ func (ph *proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		log.Errorf("WebSocket upgrade failed from %s: %v", r.RemoteAddr, err)
 		return
 	}
-	serverConn := &wsConnAdapter{
+	serverConn := (&wsConnAdapter{
 		ctx:        ctx,
 		conn:       wsConn,
 		metrics:    ph.metrics,
 		clientAddr: r.RemoteAddr,
-	}
+	}).WithFrameSnooper(ph.headersReadTimeout)
+
 	defer func() {
 		_ = serverConn.Close()
 	}()
@@ -94,16 +96,13 @@ func (ph *proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 	(&http2.Server{
 		// MaxConcurrentStreams: 20,
-		// IdleTimeout: 3 * time.Second,
+		// IdleTimeout: 60 * time.Second,
 	}).ServeConn(serverConn, &http2.ServeConnOpts{
-		Context: ctx,
-		Handler: ph.handler,
+		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,
+			// we don't set read/write timeouts here,
+			// as they interfere with streaming grpc calls
 		},
 	})
 

util/wsproxy/server/ws_conn_adapter.go

@@ -13,16 +13,16 @@ import (
 )
 
 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
+	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
 }
 
 var _ net.Conn = &wsConnAdapter{}
@@ -32,10 +32,16 @@ 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() {
+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.frameDecoder = hpack.NewDecoder(0, nil)
+	ws.framer.ReadMetaHeaders = hpack.NewDecoder(0, nil)
+	ws.headerReadDeadlineTimer = time.AfterFunc(d, ws.onReadTimeout)
+	return ws
 }
 
 func (ws *wsConnAdapter) Read(b []byte) (int, error) {
@@ -68,9 +74,12 @@ func (ws *wsConnAdapter) Read(b []byte) (int, error) {
 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)
+	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()
+		}
 	}
 
 	ws.recordBytesTransferred(ws.ctx, "ws_to_grpc", n)
@@ -116,9 +125,9 @@ func (ws *wsConnAdapter) SetDeadline(t time.Time) error {
 }
 
 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
 }
@@ -144,3 +153,13 @@ func (ws *wsConnAdapter) IsClosed() bool {
 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
+}

util/wsproxy/server/ws_conn_adapter_test.go

@@ -21,17 +21,15 @@ import (
 	"golang.org/x/net/http2/hpack"
 )
 
-func TestXxx(t *testing.T) {
+func TestAdapterHandlingConnectionClosures(t *testing.T) {
 	var cases = []struct {
-		description      string
-		casenum          int
-		frameHandlerFunc func(b []byte) (n int, err error)
+		description string
+		casenum     int
 	}{
-		// {"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 }},
+		{"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 {
@@ -55,11 +53,7 @@ func TestXxx(t *testing.T) {
 			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,
-				// }),
+				Handler: handler,
 			}
 			go httpServer.Serve(l)
 
@@ -76,11 +70,11 @@ func TestXxx(t *testing.T) {
 				Transport: &http2.Transport{
 					AllowHTTP: true,
 					DialTLSContext: func(_ context.Context, _, _ string, _ *tls.Config) (net.Conn, error) {
-						return &h2ConnectionSnooper{wrappedConn: &wsConnAdapter{
+						return &wsConnAdapter{
 							prefix: "test-client",
 							ctx:    clientCtx,
 							conn:   clientconn,
-						}, frameHandlerFunc: c.frameHandlerFunc}, nil
+						}, nil
 					},
 				}}
 
@@ -106,14 +100,83 @@ func TestXxx(t *testing.T) {
 
 			assert.EventuallyWithT(t, func(c *assert.CollectT) {
 				assert.True(c, handler.conn.IsClosed())
-			}, 5*time.Second, 100*time.Millisecond)
+			}, 3*time.Second, 100*time.Millisecond)
+		})
+	}
+}
+
+func TestAdapterHandlingBuggyHttpConnection(t *testing.T) {
+	var cases = []struct {
+		description              string
+		shouldDropFrameCondition func(f http2.FrameType) bool
+		responseError            string
+	}{
+		{"client slow to start a stream (nothing past settings frame)",
+			func(f http2.FrameType) bool { return f == http2.FrameData || f == http2.FrameHeaders },
+			"failed to get reader"},
+		{"client slow to start a stream (only headers received)", func(f http2.FrameType) bool { return f == http2.FrameData },
+			"failed to get reader"},
+	}
+
+	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)
+			handler.headersReadTimeout = 1 * time.Second
+
+			protocols := new(http.Protocols)
+			protocols.SetHTTP1(true)
+			protocols.SetUnencryptedHTTP2(true)
+			httpServer := http.Server{
+				Handler: handler,
+			}
+			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)
+
+			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: c.shouldDropFrameCondition}, nil
+					},
+				}}
+
+			_, err = h2client.Post("http://whatever", "text/html", strings.NewReader("g'day"))
+			assert.ErrorContains(t, err, c.responseError)
+
+			assert.EventuallyWithT(t, func(c *assert.CollectT) {
+				assert.True(c, handler.conn.IsClosed())
+			}, 3*time.Second, 100*time.Millisecond)
 		})
 	}
 }
 
 type h2ConnectionSnooper struct {
-	wrappedConn      net.Conn
-	frameHandlerFunc func(b []byte) (n int, err error)
+	wrappedConn     net.Conn
+	shouldDropFrame func(f http2.FrameType) bool
 }
 
 func (hs *h2ConnectionSnooper) Read(b []byte) (n int, err error) {
@@ -128,8 +191,8 @@ func (hs *h2ConnectionSnooper) Write(b []byte) (n int, err error) {
 		return hs.wrappedConn.Write(b)
 	}
 
-	if (f.Header().Type == http2.FrameData || f.Header().Type == http2.FrameHeaders) && hs.frameHandlerFunc != nil {
-		return hs.frameHandlerFunc(b)
+	if hs.shouldDropFrame != nil && hs.shouldDropFrame(f.Header().Type) {
+		return len(b), nil
 	}
 
 	return hs.wrappedConn.Write(b)