fix: track how the request body ended before closing the connection

The body stream is now wrapped in a `bodyStreamTracker` before the handler touches it, which remembers the stream's first terminal result rather than asking fasthttp a second, unsafe question. `io.EOF` means the body was read out in full and the connection is still in sync; no terminal result means the handler stopped partway, so the leftovers are drained the way a
`Content-Length` body already was; a framing error means nothing decodable is left and the connection cannot carry another request.

`fasthttp.Request.SetBodyStream` cannot install the wrapper, as it releases the current `*requestStream` back to its pool, so `requestBodyStream` is now the accessor every body reader takes the stream from.

Broken framing no longer gives up on draining either. The connection is closed either way, so a bounded read off the raw socket costs nothing and lets the client finish its write and read the S3 error instead of a reset.

internal/httpctx/context_keys.go

@@ -32,6 +32,7 @@ const (
 	ContextKeyParsedAcl        ContextKey = "parsed-acl"
 	ContextKeySkipResBodyLog   ContextKey = "skip-res-body-log"
 	ContextKeyBodyReader       ContextKey = "body-reader"
+	ContextKeyBodyStream       ContextKey = "body-stream"
 	ContextKeySkip             ContextKey = "__skip"
 	ContextKeyStack            ContextKey = "stack"
 	ContextKeyBucketOwner      ContextKey = "bucket-owner"

s3api/middlewares/body-reader.go

@@ -88,7 +88,7 @@ var _ ChecksumReader = &MockChecksumReader{}
 func wrapBodyReader(ctx fiber.Ctx, wr func(io.Reader) io.Reader) {
 	rdr, ok := utils.ContextKeyBodyReader.Get(ctx).(io.Reader)
 	if !ok {
-		rdr = ctx.Request().BodyStream()
+		rdr = requestBodyStream(ctx)
 		// Override the body reader with an empty reader to prevent panics
 		// in case of unexpected or malformed HTTP requests.
 		if rdr == nil {

s3api/middlewares/drain-request-body.go

@@ -22,6 +22,7 @@ import (
 
 	"github.com/gofiber/fiber/v3"
 	"github.com/versity/versitygw/debuglogger"
+	"github.com/versity/versitygw/s3api/utils"
 )
 
 const (
@@ -57,36 +58,41 @@ var (
 //
 // Draining first lets the client finish its write and read the real error. It
 // also keeps keep-alive connections in sync: leftover body bytes would otherwise
-// be parsed as the start of the next request.
+// be parsed as the start of the next request, which on a connection shared by an
+// upstream proxy mixes requests across tenants.
+//
+// It wraps the body in a bodyStreamTracker on the way in, which is what tells a
+// body the handler finished from one it abandoned. Every reader in the chain
+// must therefore take the body from requestBodyStream.
 //
 // Register it before every route so it wraps all of them. It runs before the
 // fiber ErrorHandler, which fiber invokes after the handler chain returns, so
 // that handler must not reset the response header the drain may have written to.
 func DrainRequestBody() fiber.Handler {
 	return func(ctx fiber.Ctx) error {
+		var body *bodyStreamTracker
+		if stream := ctx.Request().BodyStream(); stream != nil {
+			body = &bodyStreamTracker{reader: stream}
+			utils.ContextKeyBodyStream.Set(ctx, body)
+		}
+
 		// deferred so a panic unwinding through the chain still drains
-		defer drainRequestBody(ctx)
+		defer drainRequestBody(ctx, body)
+
 		return ctx.Next()
 	}
 }
 
-func drainRequestBody(ctx fiber.Ctx) {
-	stream := ctx.Request().BodyStream()
-	if stream == nil {
-		// The body was either absent or already buffered in full by fasthttp.
+func drainRequestBody(ctx fiber.Ctx, body *bodyStreamTracker) {
+	if body == nil || ctx.Request().BodyStream() == nil {
+		// The body was either absent, or buffered in full and released by
+		// fasthttp behind the chain's back, the way ctx.Body() does.
 		return
 	}
 
-	// fasthttp's requestStream reports EOF idempotently for a Content-Length
-	// body, but not for a chunked one: past the terminating chunk it goes back
-	// to the socket for another chunk header that will never come. Reading a
-	// chunked body the handler already finished would block until the deadline
-	// and hold the response back with it, so do not drain chunked framing. The
-	// stream may still have unread bytes, though, so the connection must not be
-	// reused for another request.
-	cLength := ctx.Request().Header.ContentLength()
-	if cLength < 0 {
-		ctx.Response().Header.SetConnectionClose()
+	if errors.Is(body.end, io.EOF) {
+		// The handler read the body to its end: the socket holds nothing more
+		// of it and the connection is already in sync for the next request.
 		return
 	}
 
@@ -94,8 +100,26 @@ func drainRequestBody(ctx fiber.Ctx) {
 	if conn != nil {
 		defer conn.SetReadDeadline(time.Time{})
 	}
+
+	// The leftovers are read back through the tracker, so a drain that reaches
+	// the end of the body records it the same way the handler's reads would.
+	src := io.Reader(body)
+	if body.end != nil {
+		// The framing broke before the body ended, so what is still queued
+		// cannot be told apart from the start of the next request and the
+		// connection cannot carry one. Draining is still worth attempting: the
+		// connection is going away either way, and absorbing the client's
+		// in-flight write is what lets it read the S3 error instead of an RST.
+		// Only the raw socket can absorb it once the framing is gone.
+		ctx.Response().Header.SetConnectionClose()
+		if conn == nil {
+			return
+		}
+		src = conn
+	}
+
 	reader := &drainReader{
-		reader:   stream,
+		reader:   src,
 		conn:     conn,
 		deadline: time.Now().Add(drainTotalTimeout),
 	}
@@ -125,6 +149,51 @@ func drainRequestBody(ctx fiber.Ctx) {
 	ctx.Response().Header.SetConnectionClose()
 }
 
+// requestBodyStream returns the reader the request body must be read through.
+// Reading ctx.Request().BodyStream() directly instead hides those reads from
+// DrainRequestBody, which then cannot tell a finished body from an abandoned
+// one and falls back to closing the connection.
+//
+// It yields the raw stream where DrainRequestBody is not registered, and nil
+// when the request carries no streamed body.
+func requestBodyStream(ctx fiber.Ctx) io.Reader {
+	if body, ok := utils.ContextKeyBodyStream.Get(ctx).(*bodyStreamTracker); ok {
+		return body
+	}
+
+	return ctx.Request().BodyStream()
+}
+
+// bodyStreamTracker remembers how the request body ended, so the drain can tell
+// a body the handler read to its end from one it stopped partway through.
+//
+// fasthttp cannot be asked a second time. Its requestStream reports EOF
+// idempotently for a Content-Length body, but not for a chunked one: past the
+// terminating chunk it goes back to the socket for another chunk header that
+// will never come, so a read of a body the handler already finished blocks
+// until the deadline and holds the response back with it. The first terminal
+// result is recorded here and replayed instead.
+type bodyStreamTracker struct {
+	reader io.Reader
+	// end is the first error the stream ended on: nil while it still has more
+	// to give, io.EOF once it was read out in full, and the framing error if it
+	// broke before that.
+	end error
+}
+
+func (t *bodyStreamTracker) Read(p []byte) (int, error) {
+	if t.end != nil {
+		return 0, t.end
+	}
+
+	n, err := t.reader.Read(p)
+	if err != nil {
+		t.end = err
+	}
+
+	return n, err
+}
+
 // drainReader refreshes the connection's read deadline before every read, so a
 // client that keeps sending is never cut off mid-drain while one that goes quiet
 // is dropped after drainIdleTimeout. deadline caps the whole drain regardless.

s3api/middlewares/drain-request-body_test.go

@@ -86,15 +86,7 @@ func slowClientFinishesWriting(t *testing.T, disableKeepalive bool) {
 	}
 
 	// a fully drained body leaves the connection in sync for the next request
-	if _, err := conn.Write(putHeaders(addr, 0)); err != nil {
-		t.Fatalf("write second request: %v", err)
-	}
-	resp2, err := http.ReadResponse(br, nil)
-	if err != nil {
-		t.Fatalf("read second response: %v", err)
-	}
-	defer resp2.Body.Close()
-	if resp2.StatusCode != http.StatusBadRequest {
+	if resp2 := reuseConnection(t, conn, br, addr, resp); resp2.StatusCode != http.StatusBadRequest {
 		t.Fatalf("expected status %v on the reused connection, got %v", http.StatusBadRequest, resp2.StatusCode)
 	}
 }
@@ -133,9 +125,7 @@ func TestDrainRequestBody_closesConnectionWhenUnreadBodyExceedsTheLimit(t *testi
 }
 
 func TestDrainRequestBody_givesUpOnAClientThatStopsSending(t *testing.T) {
-	idle, total := drainIdleTimeout, drainTotalTimeout
-	drainIdleTimeout, drainTotalTimeout = 100*time.Millisecond, 250*time.Millisecond
-	t.Cleanup(func() { drainIdleTimeout, drainTotalTimeout = idle, total })
+	shortenDrainTimeouts(t)
 
 	addr := startEarlyResponder(t, true)
 
@@ -173,7 +163,8 @@ func TestDrainRequestBody_givesUpOnAClientThatStopsSending(t *testing.T) {
 
 // A chunked body the handler read to its end must not be read again: fasthttp's
 // requestStream goes back to the socket for another chunk header past the
-// terminating chunk, which would hold the response back until the deadline.
+// terminating chunk, which would hold the response back until the deadline. It
+// also has nothing left to desync the connection, so it keeps keep-alive.
 func TestDrainRequestBody_doesNotStallAChunkedBodyTheHandlerFinished(t *testing.T) {
 	addr := startFullReader(t)
 
@@ -186,13 +177,13 @@ func TestDrainRequestBody_doesNotStallAChunkedBodyTheHandlerFinished(t *testing.
 		t.Fatalf("set deadline: %v", err)
 	}
 
-	body := "PUT /object HTTP/1.1\r\nHost: " + addr + "\r\nTransfer-Encoding: chunked\r\n\r\n5\r\nhello\r\n0\r\n\r\n"
-	if _, err := conn.Write([]byte(body)); err != nil {
+	if _, err := conn.Write(chunkedRequest(addr, []byte("hello"))); err != nil {
 		t.Fatalf("write request: %v", err)
 	}
 
 	start := time.Now()
-	resp, err := http.ReadResponse(bufio.NewReader(conn), nil)
+	br := bufio.NewReader(conn)
+	resp, err := http.ReadResponse(br, nil)
 	if err != nil {
 		t.Fatalf("read response: %v", err)
 	}
@@ -204,9 +195,18 @@ func TestDrainRequestBody_doesNotStallAChunkedBodyTheHandlerFinished(t *testing.
 	if elapsed := time.Since(start); elapsed > drainIdleTimeout {
 		t.Fatalf("the response was held back for %v: the drain re-read a finished chunked body", elapsed)
 	}
+	if resp.Close {
+		t.Fatal("expected the connection to stay usable after a chunked body the handler finished")
+	}
+
+	if resp2 := reuseConnection(t, conn, br, addr, resp); resp2.StatusCode != http.StatusOK {
+		t.Fatalf("expected status %v on the reused connection, got %v", http.StatusOK, resp2.StatusCode)
+	}
 }
 
-func TestDrainRequestBody_closesConnectionForUnreadChunkedBody(t *testing.T) {
+// A chunked body the handler abandoned is drained like any other. Only the read
+// past its end is unsafe, and the drain never gets there on a body it finished.
+func TestDrainRequestBody_drainsAnUnreadChunkedBody(t *testing.T) {
 	addr := startEarlyResponder(t, false)
 
 	conn, err := net.Dial("tcp", addr)
@@ -218,16 +218,96 @@ func TestDrainRequestBody_closesConnectionForUnreadChunkedBody(t *testing.T) {
 		t.Fatalf("set deadline: %v", err)
 	}
 
-	// Leave one decoded byte unread after the handler's initial read. The
-	// middleware cannot safely probe for EOF on a chunked request.
+	// leave one decoded byte unread after the handler's initial read
 	chunk := bytes.Repeat([]byte("a"), handlerReadBytes+1)
-	body := fmt.Appendf(nil, "PUT /object HTTP/1.1\r\nHost: %s\r\nTransfer-Encoding: chunked\r\n\r\n%x\r\n", addr, len(chunk))
-	body = append(body, chunk...)
-	body = append(body, []byte("\r\n0\r\n\r\n")...)
-	if _, err := conn.Write(body); err != nil {
+	if _, err := conn.Write(chunkedRequest(addr, chunk)); err != nil {
 		t.Fatalf("write request: %v", err)
 	}
 
+	br := bufio.NewReader(conn)
+	resp, err := http.ReadResponse(br, nil)
+	if err != nil {
+		t.Fatalf("read response: %v", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusBadRequest {
+		t.Fatalf("expected status %v, got %v", http.StatusBadRequest, resp.StatusCode)
+	}
+	if resp.Close {
+		t.Fatal("expected the connection to stay usable after the chunked body was drained")
+	}
+
+	if resp2 := reuseConnection(t, conn, br, addr, resp); resp2.StatusCode != http.StatusBadRequest {
+		t.Fatalf("expected status %v on the reused connection, got %v", http.StatusBadRequest, resp2.StatusCode)
+	}
+}
+
+func TestDrainRequestBody_closesConnectionWhenUnreadChunkedBodyExceedsTheLimit(t *testing.T) {
+	addr := startEarlyResponder(t, false)
+	chunk := bytes.Repeat([]byte("a"), int(maxDrainBytes)*4)
+
+	conn, err := net.Dial("tcp", addr)
+	if err != nil {
+		t.Fatalf("dial: %v", err)
+	}
+	defer conn.Close()
+	if err := conn.SetDeadline(time.Now().Add(30 * time.Second)); err != nil {
+		t.Fatalf("set deadline: %v", err)
+	}
+
+	// the write is expected to fail once the server gives up draining
+	go conn.Write(chunkedRequest(addr, chunk))
+
+	resp, err := http.ReadResponse(bufio.NewReader(conn), nil)
+	if err != nil {
+		t.Fatalf("read response: %v", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusBadRequest {
+		t.Fatalf("expected status %v, got %v", http.StatusBadRequest, resp.StatusCode)
+	}
+	if !resp.Close {
+		t.Fatal("expected 'Connection: close', the undrained chunk bytes would desync the next request")
+	}
+}
+
+// Broken chunk framing leaves nothing that can be decoded as body bytes, so the
+// connection cannot be reused. The drain still absorbs what the client is
+// writing, off the socket itself, so it reaches the S3 error instead of an RST.
+func TestDrainRequestBody_drainsAndClosesOnBrokenChunkedFraming(t *testing.T) {
+	shortenDrainTimeouts(t)
+
+	addr := startEarlyResponder(t, false)
+	// small enough to be absorbed in full, big enough that it cannot sit in the
+	// socket buffers while the server decides to close
+	garbage := bytes.Repeat([]byte("a"), 128<<10)
+
+	conn, err := net.Dial("tcp", addr)
+	if err != nil {
+		t.Fatalf("dial: %v", err)
+	}
+	defer conn.Close()
+	if err := conn.SetDeadline(time.Now().Add(30 * time.Second)); err != nil {
+		t.Fatalf("set deadline: %v", err)
+	}
+
+	// a chunk size that is not a hex number: the handler's first read fails
+	head := fmt.Appendf(nil, "PUT /object HTTP/1.1\r\nHost: %s\r\nTransfer-Encoding: chunked\r\n\r\nzz\r\n", addr)
+	if _, err := conn.Write(head); err != nil {
+		t.Fatalf("write request: %v", err)
+	}
+
+	// dribble the rest out, the way a client keeps uploading after the gateway
+	// has already given up on the request
+	for off := 0; off < len(garbage); off += 4 << 10 {
+		if _, err := conn.Write(garbage[off:min(off+(4<<10), len(garbage))]); err != nil {
+			t.Fatalf("the server dropped the connection with %v of %v bytes sent: %v", off, len(garbage), err)
+		}
+		time.Sleep(time.Millisecond)
+	}
+
 	resp, err := http.ReadResponse(bufio.NewReader(conn), nil)
 	if err != nil {
 		t.Fatalf("read response: %v", err)
@@ -238,7 +318,7 @@ func TestDrainRequestBody_closesConnectionForUnreadChunkedBody(t *testing.T) {
 		t.Fatalf("expected status %v, got %v", http.StatusBadRequest, resp.StatusCode)
 	}
 	if !resp.Close {
-		t.Fatal("expected 'Connection: close' for an unread chunked body")
+		t.Fatal("expected 'Connection: close', broken chunk framing cannot be resynchronized")
 	}
 }
 
@@ -291,7 +371,7 @@ func startEarlyResponder(t *testing.T, disableKeepalive bool) string {
 	})
 	app.Use("*", DrainRequestBody())
 	app.Put("/object", func(ctx fiber.Ctx) error {
-		if body := ctx.Request().BodyStream(); body != nil {
+		if body := requestBodyStream(ctx); body != nil {
 			// consume a little of it, the way the chunk reader parses a chunk
 			// header before rejecting the upload
 			io.CopyN(io.Discard, body, handlerReadBytes) //nolint:errcheck
@@ -324,10 +404,10 @@ func putHeaders(addr string, contentLength int) []byte {
 func startFullReader(t *testing.T) string {
 	t.Helper()
 
-	app := fiber.New(fiber.Config{StreamRequestBody: true, DisableKeepalive: true})
+	app := fiber.New(fiber.Config{StreamRequestBody: true})
 	app.Use("*", DrainRequestBody())
 	app.Put("/object", func(ctx fiber.Ctx) error {
-		if body := ctx.Request().BodyStream(); body != nil {
+		if body := requestBodyStream(ctx); body != nil {
 			if _, err := io.Copy(io.Discard, body); err != nil {
 				return err
 			}
@@ -337,3 +417,43 @@ func startFullReader(t *testing.T) string {
 
 	return listen(t, app)
 }
+
+// chunkedRequest builds a PUT that carries body as a single chunk, the framing
+// a client uses when it cannot announce a Content-Length up front.
+func chunkedRequest(addr string, body []byte) []byte {
+	req := fmt.Appendf(nil, "PUT /object HTTP/1.1\r\nHost: %s\r\nTransfer-Encoding: chunked\r\n\r\n%x\r\n", addr, len(body))
+	req = append(req, body...)
+
+	return append(req, "\r\n0\r\n\r\n"...)
+}
+
+// reuseConnection sends a second, bodiless request on the same connection once
+// prev is off the wire. It is answered only while the connection is still in
+// sync with the client.
+func reuseConnection(t *testing.T, conn net.Conn, br *bufio.Reader, addr string, prev *http.Response) *http.Response {
+	t.Helper()
+
+	if _, err := io.Copy(io.Discard, prev.Body); err != nil {
+		t.Fatalf("read response body: %v", err)
+	}
+	if _, err := conn.Write(putHeaders(addr, 0)); err != nil {
+		t.Fatalf("write second request: %v", err)
+	}
+	resp, err := http.ReadResponse(br, nil)
+	if err != nil {
+		t.Fatalf("read second response: %v", err)
+	}
+	t.Cleanup(func() { resp.Body.Close() }) //nolint:errcheck
+
+	return resp
+}
+
+// shortenDrainTimeouts keeps a test that waits the drain out from taking the
+// production timeouts to finish.
+func shortenDrainTimeouts(t *testing.T) {
+	t.Helper()
+
+	idle, total := drainIdleTimeout, drainTotalTimeout
+	drainIdleTimeout, drainTotalTimeout = 100*time.Millisecond, 250*time.Millisecond
+	t.Cleanup(func() { drainIdleTimeout, drainTotalTimeout = idle, total })
+}

s3api/middlewares/object-post-auth.go

@@ -72,7 +72,7 @@ func AuthorizePostObject(root RootUserConfig, iam auth.IAMService, region string
 			return s3err.GetAPIError(s3err.ErrMalformedPOSTRequest)
 		}
 
-		bodyRdr := ctx.Request().BodyStream()
+		bodyRdr := requestBodyStream(ctx)
 		if bodyRdr == nil {
 			bodyRdr = bytes.NewReader(ctx.BodyRaw())
 		}

s3api/middlewares/public-bucket.go

@@ -98,7 +98,7 @@ func AuthorizePublicBucketAccess(be backend.Backend, s3action string, policyPerm
 				return err
 			} else if utils.IsUnsignedPaylod(payloadHash) {
 				// for UNSIGNED-PAYLOD simply store the body reader in context locals
-				utils.ContextKeyBodyReader.Set(ctx, ctx.Request().BodyStream())
+				utils.ContextKeyBodyReader.Set(ctx, requestBodyStream(ctx))
 				return nil
 			} else {
 				// stack a hash reader to calculated the payload sha256 hash

s3api/utils/context-keys.go

@@ -32,6 +32,7 @@ const (
 	ContextKeyParsedAcl        = httpctx.ContextKeyParsedAcl
 	ContextKeySkipResBodyLog   = httpctx.ContextKeySkipResBodyLog
 	ContextKeyBodyReader       = httpctx.ContextKeyBodyReader
+	ContextKeyBodyStream       = httpctx.ContextKeyBodyStream
 	ContextKeySkip             = httpctx.ContextKeySkip
 	ContextKeyStack            = httpctx.ContextKeyStack
 	ContextKeyBucketOwner      = httpctx.ContextKeyBucketOwner