Merge pull request #2345 from versity/sis/drain-request-body

fix: drain unread request bodies before closing the connection
2026-09-22 16:04:15 +00:00 · 2026-09-02 10:38:09 -07:00
parent ee179f0462 ca8d434f8a
commit d2b53687ef
3 changed files with 476 additions and 0 deletions
@@ -0,0 +1,168 @@
+// Copyright 2026 Versity Software
+// This file is licensed under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package middlewares
+
+import (
+	"errors"
+	"io"
+	"net"
+	"time"
+
+	"github.com/gofiber/fiber/v3"
+	"github.com/versity/versitygw/debuglogger"
+)
+
+const (
+	// maxDrainBytes caps how much unread request body is read and discarded
+	// after the response has been decided. It matches net/http's
+	// maxPostHandlerReadBytes: enough to cover request bodies that are rejected
+	// early (bad chunk framing, auth failures, missing buckets), small enough
+	// that a rejected multi-gigabyte upload is not streamed through the gateway
+	// just to be thrown away. A body with more than this still unread is left
+	// alone and the connection is closed, so its client can still see a reset.
+	maxDrainBytes int64 = 256 << 10
+)
+
+// The drain is bounded twice: idle time, so a client that stops sending is cut
+// loose quickly, and total time, so a client that trickles cannot hold a worker
+// for long.
+var (
+	drainIdleTimeout  = time.Second
+	drainTotalTimeout = 5 * time.Second
+)
+
+// DrainRequestBody reads and discards whatever is left of the request body once
+// the rest of the handler chain is done with it.
+//
+// The gateway can decide a response long before the client has finished sending
+// the body: an invalid chunk size is detected a few kilobytes into an aws-chunked
+// upload, a signature check fails before any payload is read, and so on. fasthttp
+// streams request bodies (StreamRequestBody) and does not drain what the handler
+// left behind, so the connection is closed with unread bytes still queued in the
+// socket. The kernel answers the client's in-flight writes with an RST, and the
+// client reports "connection reset by peer" instead of the S3 error the gateway
+// took the trouble to produce.
+//
+// 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.
+//
+// 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 {
+		// deferred so a panic unwinding through the chain still drains
+		defer drainRequestBody(ctx)
+		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.
+		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 only Content-Length framing is
+	// drained. Nothing is lost for the aws-chunked uploads this exists for --
+	// STREAMING-* payloads carry a Content-Length.
+	cLength := ctx.Request().Header.ContentLength()
+	if cLength <= 0 {
+		return
+	}
+
+	conn := requestConn(ctx)
+	if conn != nil {
+		defer conn.SetReadDeadline(time.Time{})
+	}
+	reader := &drainReader{
+		reader:   stream,
+		conn:     conn,
+		deadline: time.Now().Add(drainTotalTimeout),
+	}
+
+	n, err := io.CopyN(io.Discard, reader, maxDrainBytes)
+	if err == nil {
+		// Filled the cap exactly. One more read tells a body that happened to
+		// end there from one with more still to come.
+		err = reader.atEOF()
+	}
+	if errors.Is(err, io.EOF) {
+		if n > 0 {
+			debuglogger.Logf("discarded %v unread request body bytes before responding", n)
+		}
+		return
+	}
+
+	if err != nil {
+		debuglogger.Logf("failed to discard the unread request body after %v bytes: %v", n, err)
+	} else {
+		debuglogger.Logf("unread request body exceeds the %v byte drain limit: %v bytes discarded", maxDrainBytes, n)
+	}
+
+	// The body was not consumed to its end, so the bytes still in flight would
+	// be parsed as the start of the next request on a keep-alive connection.
+	// Tell fasthttp to close it instead.
+	ctx.Response().Header.SetConnectionClose()
+}
+
+// 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.
+type drainReader struct {
+	reader   io.Reader
+	conn     net.Conn
+	deadline time.Time
+}
+
+func (dr *drainReader) Read(p []byte) (int, error) {
+	if dr.conn != nil {
+		next := time.Now().Add(drainIdleTimeout)
+		if next.After(dr.deadline) {
+			next = dr.deadline
+		}
+		if err := dr.conn.SetReadDeadline(next); err != nil {
+			return 0, err
+		}
+	}
+
+	return dr.reader.Read(p)
+}
+
+// atEOF reports io.EOF when the body ends exactly at the drain limit.
+func (dr *drainReader) atEOF() error {
+	var b [1]byte
+	n, err := dr.Read(b[:])
+	if n == 0 && err == nil {
+		return nil
+	}
+
+	return err
+}
+
+func requestConn(ctx fiber.Ctx) net.Conn {
+	rctx := ctx.RequestCtx()
+	if rctx == nil {
+		return nil
+	}
+
+	return rctx.Conn()
+}
@@ -0,0 +1,303 @@
+// Copyright 2026 Versity Software
+// This file is licensed under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package middlewares
+
+import (
+	"bufio"
+	"bytes"
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"testing"
+	"time"
+
+	"github.com/gofiber/fiber/v3"
+)
+
+func TestDrainRequestBody_slowClientFinishesWritingAndReadsTheResponse(t *testing.T) {
+	// keep-alive is off by default in the gateway and on with --keep-alive
+	for _, disableKeepalive := range []bool{true, false} {
+		t.Run(fmt.Sprintf("disableKeepalive=%v", disableKeepalive), func(t *testing.T) {
+			slowClientFinishesWriting(t, disableKeepalive)
+		})
+	}
+}
+
+func slowClientFinishesWriting(t *testing.T, disableKeepalive bool) {
+	t.Helper()
+
+	addr := startEarlyResponder(t, disableKeepalive)
+	// small enough to be drained in full, big enough that it cannot sit in the
+	// socket buffers while the server decides to close
+	body := 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)
+	}
+
+	if _, err := conn.Write(putHeaders(addr, len(body))); err != nil {
+		t.Fatalf("write headers: %v", err)
+	}
+
+	// dribble the body out, so the response is decided well before the last byte
+	for off := 0; off < len(body); off += 4 << 10 {
+		if _, err := conn.Write(body[off:min(off+(4<<10), len(body))]); err != nil {
+			t.Fatalf("the server dropped the connection with %v of %v body bytes sent: %v", off, len(body), err)
+		}
+		time.Sleep(time.Millisecond)
+	}
+
+	br := bufio.NewReader(conn)
+	resp, err := http.ReadResponse(br, nil)
+	if err != nil {
+		t.Fatalf("read response: %v", err)
+	}
+	defer resp.Body.Close()
+	if _, err := io.Copy(io.Discard, resp.Body); err != nil {
+		t.Fatalf("read response body: %v", err)
+	}
+
+	if resp.StatusCode != http.StatusBadRequest {
+		t.Fatalf("expected status %v, got %v", http.StatusBadRequest, resp.StatusCode)
+	}
+	if disableKeepalive {
+		return
+	}
+	if resp.Close {
+		t.Fatal("expected the connection to stay usable after the body was drained")
+	}
+
+	// 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 {
+		t.Fatalf("expected status %v on the reused connection, got %v", http.StatusBadRequest, resp2.StatusCode)
+	}
+}
+
+func TestDrainRequestBody_closesConnectionWhenUnreadBodyExceedsTheLimit(t *testing.T) {
+	addr := startEarlyResponder(t, false)
+	body := 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)
+	}
+
+	if _, err := conn.Write(putHeaders(addr, len(body))); err != nil {
+		t.Fatalf("write headers: %v", err)
+	}
+	// the write is expected to fail once the server gives up draining
+	go conn.Write(body)
+
+	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 body bytes would desync the next request")
+	}
+}
+
+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 })
+
+	addr := startEarlyResponder(t, true)
+
+	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)
+	}
+
+	// announce a body, send only what the handler consumes, then go quiet
+	if _, err := conn.Write(putHeaders(addr, 64<<10)); err != nil {
+		t.Fatalf("write headers: %v", err)
+	}
+	if _, err := conn.Write(bytes.Repeat([]byte("a"), handlerReadBytes)); err != nil {
+		t.Fatalf("write body: %v", err)
+	}
+
+	start := time.Now()
+	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 elapsed := time.Since(start); elapsed > 5*time.Second {
+		t.Fatalf("the drain held the response for %v, the timeout should have cut it short", elapsed)
+	}
+}
+
+// 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.
+func TestDrainRequestBody_doesNotStallAChunkedBodyTheHandlerFinished(t *testing.T) {
+	addr := startFullReader(t)
+
+	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)
+	}
+
+	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 {
+		t.Fatalf("write request: %v", err)
+	}
+
+	start := time.Now()
+	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.StatusOK {
+		t.Fatalf("expected status %v, got %v", http.StatusOK, resp.StatusCode)
+	}
+	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)
+	}
+}
+
+// The same, for a Content-Length body: fasthttp reports EOF idempotently there,
+// so it is drained, but a handler that already finished it must not be delayed.
+func TestDrainRequestBody_doesNotStallABodyTheHandlerFinished(t *testing.T) {
+	addr := startFullReader(t)
+
+	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)
+	}
+
+	if _, err := conn.Write(append(putHeaders(addr, 5), "hello"...)); err != nil {
+		t.Fatalf("write request: %v", err)
+	}
+
+	start := time.Now()
+	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.StatusOK {
+		t.Fatalf("expected status %v, got %v", http.StatusOK, resp.StatusCode)
+	}
+	if elapsed := time.Since(start); elapsed > drainIdleTimeout {
+		t.Fatalf("the response was held back for %v on a body the handler had finished", elapsed)
+	}
+}
+
+// fasthttp pre-reads up to 8KB of a declared body before it hands the request to
+// the handler, so a test body has to be at least that big to reach the drain.
+const handlerReadBytes = 8 << 10
+
+// startEarlyResponder serves a fiber app that reads only the head of the request
+// body and then answers, the way the gateway rejects an upload on a bad chunk
+// header or a failed authorization long before the client is done sending.
+func startEarlyResponder(t *testing.T, disableKeepalive bool) string {
+	t.Helper()
+
+	app := fiber.New(fiber.Config{
+		StreamRequestBody: true,
+		DisableKeepalive:  disableKeepalive,
+	})
+	app.Use("*", DrainRequestBody())
+	app.Put("/object", func(ctx fiber.Ctx) error {
+		if body := ctx.Request().BodyStream(); 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
+		}
+		return ctx.Status(http.StatusBadRequest).SendString("rejected")
+	})
+
+	return listen(t, app)
+}
+
+func listen(t *testing.T, app *fiber.App) string {
+	t.Helper()
+
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("listen: %v", err)
+	}
+	go app.Listener(ln, fiber.ListenConfig{DisableStartupMessage: true}) //nolint:errcheck
+	t.Cleanup(func() { app.Shutdown() })                                 //nolint:errcheck
+
+	return ln.Addr().String()
+}
+
+func putHeaders(addr string, contentLength int) []byte {
+	return fmt.Appendf(nil, "PUT /object HTTP/1.1\r\nHost: %s\r\nContent-Length: %d\r\n\r\n", addr, contentLength)
+}
+
+// startFullReader serves a fiber app whose handler consumes the whole request
+// body, so the drain has nothing left to do.
+func startFullReader(t *testing.T) string {
+	t.Helper()
+
+	app := fiber.New(fiber.Config{StreamRequestBody: true, DisableKeepalive: true})
+	app.Use("*", DrainRequestBody())
+	app.Put("/object", func(ctx fiber.Ctx) error {
+		if body := ctx.Request().BodyStream(); body != nil {
+			if _, err := io.Copy(io.Discard, body); err != nil {
+				return err
+			}
+		}
+		return ctx.SendStatus(http.StatusOK)
+	})
+
+	return listen(t, app)
+}
@@ -126,6 +126,11 @@ func New(
 			StackTraceHandler: stackTraceHandler,
 		}))

+	// initialize the request body drainer. it goes right after the panic
+	// recovery and before every route, so it wraps all of them and a panic in
+	// the drain itself is still recovered
+	app.Use("*", middlewares.DrainRequestBody())
+
 	// Logging middlewares
 	if !server.quiet {
 		app.Use("*", logger.New(logger.Config{