diff --git a/s3api/middlewares/drain-request-body.go b/s3api/middlewares/drain-request-body.go new file mode 100644 index 00000000..25c6b02e --- /dev/null +++ b/s3api/middlewares/drain-request-body.go @@ -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() +} diff --git a/s3api/middlewares/drain-request-body_test.go b/s3api/middlewares/drain-request-body_test.go new file mode 100644 index 00000000..293217e1 --- /dev/null +++ b/s3api/middlewares/drain-request-body_test.go @@ -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) +} diff --git a/s3api/server.go b/s3api/server.go index e8563f3b..fdbd30fd 100644 --- a/s3api/server.go +++ b/s3api/server.go @@ -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{