diff --git a/README.md b/README.md index c8d6fdb7..2cef5e1e 100644 --- a/README.md +++ b/README.md @@ -178,6 +178,8 @@ See [docs/metrics.md](docs/metrics.md) for complete metrics reference. TAG supports transparent proxy mode (default) with local SigV4 validation and per-bucket authorization caching, as well as signing mode with local credential stores. +For requests routed upstream through the signing forwarder, header-authenticated `STREAMING-AWS4-HMAC-SHA256-PAYLOAD` uploads are verified before their decoded bodies are forwarded from temporary storage; concurrent stages account for at most half of the temporary-filesystem space available when the first stage is admitted. This guarantee does not cover tiered-mode PUTs stored directly in the local cache. Presigned URL validation retains its existing `UNSIGNED-PAYLOAD` behavior. + See [docs/security.md](docs/security.md) for authentication, access control, and security architecture. ## Contributing diff --git a/auth/validator.go b/auth/validator.go index d4a3b383..e5103260 100644 --- a/auth/validator.go +++ b/auth/validator.go @@ -45,7 +45,8 @@ func NewRequestValidator(keyProvider KeyProvider) *RequestValidator { // ValidateRequest validates the AWS SigV4 signature of an incoming request. // This requires the X-Amz-Content-Sha256 header to be present (all AWS SDKs set this). -// This enables zero-copy streaming of the request body. +// It validates the request header without reading r.Body; callers choose the body handling +// appropriate to the payload marker. // Returns the access key if validation succeeds, or an error if it fails. func (v *RequestValidator) ValidateRequest(r *http.Request) (string, error) { // Require X-Amz-Content-Sha256 header - all AWS SDKs set this diff --git a/docs/s3-compatibility.md b/docs/s3-compatibility.md index 4ec6c94d..1526d718 100644 --- a/docs/s3-compatibility.md +++ b/docs/s3-compatibility.md @@ -1,6 +1,6 @@ # S3 API Compatibility -TAG implements all the commonly used bucket and object APIs from the Amazon S3 REST API, acting as a caching proxy between S3 clients and Tigris object storage. By default, TAG operates in transparent proxy mode where client requests are forwarded as-is with proxy headers. In signing mode, requests are validated, re-signed, and forwarded to the Tigris upstream endpoint. +TAG implements all the commonly used bucket and object APIs from the Amazon S3 REST API, acting as a caching proxy between S3 clients and Tigris object storage. By default, TAG operates in transparent proxy mode where client requests are forwarded as-is with proxy headers. Requests that signing mode routes upstream are validated, re-signed, and forwarded to the configured endpoint. ## Supported Operations @@ -63,7 +63,13 @@ TAG supports AWS chunked transfer encoding (streaming SigV4), used by tools like In **transparent proxy mode** (default), chunked uploads are forwarded as-is since the original signatures remain valid (the Host header is not changed). -In **signing mode**, when a client sends a PUT with `X-Amz-Content-Sha256: STREAMING-AWS4-HMAC-SHA256-PAYLOAD`, the body is wrapped in an AWS-proprietary chunked format where each chunk includes a signature chained from the request-level seed signature. Because TAG re-signs requests with the upstream Tigris hostname (changing the `Host` header), these per-chunk signatures become invalid. TAG handles this by decoding the chunked body on-the-fly — stripping the chunk framing and signatures — and forwarding the raw payload as `UNSIGNED-PAYLOAD`. The decoded content length is read from the `X-Amz-Decoded-Content-Length` header. This is a streaming operation with no buffering of the full object. +On the **signing-mode upstream-forwarding path**, a header-authenticated PUT with `X-Amz-Content-Sha256: STREAMING-AWS4-HMAC-SHA256-PAYLOAD` carries chunk signatures chained from the request-level signature. TAG validates every chunk and the terminal signature before forwarding any payload. Tiered-mode PUTs stored directly in the local cache bypass this forwarding path and are not covered by this guarantee. + +Because the upstream request has a different `Host` header, TAG re-signs the decoded bytes as `UNSIGNED-PAYLOAD`. It stages those bytes in a temporary file first. Memory use is bounded, but temporary storage proportional to the upload is required and upstream dispatch waits until verification completes. + +Before copying data, TAG reserves the declared decoded length. If no nonnegative decoded length is available, it reserves each next chunk. The process-wide pool admits reservations up to half of the temporary-filesystem space available when the first stage is admitted. The other half is left unclaimed by signed-stream staging. This accounting is process-local; other users can still consume available space. Insufficient capacity rejects the request before the unreserved chunk is written. Reservations remain held until the staged file is removed. If staging aborts before the request body is fully consumed, TAG closes the HTTP/1 connection instead of reusing it. A nonnegative `X-Amz-Decoded-Content-Length` must match the staged body. + +Presigned URL validation continues to use `UNSIGNED-PAYLOAD`; it does not establish the forwarded-path chunk-signature guarantee described above. ## Caching Behavior diff --git a/docs/security.md b/docs/security.md index 53e533ba..9eec672a 100644 --- a/docs/security.md +++ b/docs/security.md @@ -19,13 +19,14 @@ Tigris independently validates both the client's SigV4 signature and TAG's proxy ### Signing Mode -In signing mode, TAG terminates the client signature and re-issues the upstream one itself: +For requests routed upstream in signing mode, TAG terminates the client signature and re-issues the upstream one itself: 1. The client signs its request with its own SigV4 credentials. -2. TAG looks up the secret for the client's access key in its **local credential store** and cryptographically validates the incoming signature. -3. TAG re-signs the (possibly transformed) request for the upstream endpoint using standard AWS SigV4 with **the same access key and secret**, then streams it upstream. +2. TAG looks up the secret for the client's access key in its **local credential store** and cryptographically validates the incoming request signature. +3. On this upstream-forwarding path, for header-authenticated requests with `STREAMING-AWS4-HMAC-SHA256-PAYLOAD`, TAG reserves temporary storage, verifies each chained chunk signature and the terminal signature, checks a declared decoded length against the staged bytes, and stages the payload before forwarding. +4. TAG re-signs the (possibly transformed) request for the upstream endpoint using standard AWS SigV4 with **the same access key and secret**, then streams it upstream. -TAG re-signs rather than forwarding the original signature because it may transform the request — for example decoding AWS chunked transfer encoding to `UNSIGNED-PAYLOAD` — which would otherwise invalidate the client's signature. The upstream sees the **same identity** as the client; this is not identity translation. +TAG re-signs rather than forwarding the original signature because it may transform the request — for example decoding AWS chunked transfer encoding to `UNSIGNED-PAYLOAD` — which would otherwise invalidate the client's signature. These upstream-forwarded header-authenticated streaming uploads require temporary storage proportional to the payload and are not dispatched upstream until the full chain has been verified. Concurrent signed-stream stages reserve at most half of the temporary-filesystem space available when the first stage is admitted; other filesystem users are outside this process-local accounting. Tiered-mode PUTs stored directly in the local cache bypass this upstream-forwarding path and are not covered by this chunk-signature guarantee. Presigned URL validation continues to use `UNSIGNED-PAYLOAD` and does not establish this guarantee. The upstream sees the **same identity** as the client; this is not identity translation. Because TAG must know the secret for every access key it serves, those credentials must be present in its local credential store. In production the store is populated only from `AWS_ACCESS_KEY_ID` / `AWS_SECRET_ACCESS_KEY` (a single pair), so **clients must authenticate with those same credentials**, and TAG re-signs upstream with them. If the store is empty, TAG rejects all requests. diff --git a/handlers/response_tracker_test.go b/handlers/response_tracker_test.go new file mode 100644 index 00000000..252e0d55 --- /dev/null +++ b/handlers/response_tracker_test.go @@ -0,0 +1,30 @@ +package handlers + +import ( + "net/http" + "net/http/httptest" + "testing" + "time" +) + +type responseTrackerDeadlineWriter struct { + http.ResponseWriter + deadline time.Time +} + +func (w *responseTrackerDeadlineWriter) SetReadDeadline(deadline time.Time) error { + w.deadline = deadline + return nil +} + +func TestResponseTrackerUnwrapsReadDeadline(t *testing.T) { + writer := &responseTrackerDeadlineWriter{ResponseWriter: httptest.NewRecorder()} + wrapped := &responseTracker{ResponseWriter: writer} + deadline := time.Now() + if err := http.NewResponseController(wrapped).SetReadDeadline(deadline); err != nil { + t.Fatalf("set read deadline through responseTracker: %v", err) + } + if !writer.deadline.Equal(deadline) { + t.Fatalf("underlying read deadline = %v, want %v", writer.deadline, deadline) + } +} diff --git a/handlers/server.go b/handlers/server.go index 8c2e8436..afb65d91 100644 --- a/handlers/server.go +++ b/handlers/server.go @@ -288,6 +288,11 @@ type responseTracker struct { committed bool } +// Unwrap lets http.ResponseController reach connection controls through this wrapper. +func (rt *responseTracker) Unwrap() http.ResponseWriter { + return rt.ResponseWriter +} + func (rt *responseTracker) WriteHeader(code int) { rt.committed = true rt.ResponseWriter.WriteHeader(code) diff --git a/proxy/chunked_reader.go b/proxy/chunked_reader.go index fc6ddf45..26d31165 100644 --- a/proxy/chunked_reader.go +++ b/proxy/chunked_reader.go @@ -54,9 +54,9 @@ func IsStreamingPayload(bodyHash string) bool { // // 0[;chunk-signature=]\r\n // -// The reader strips the framing and returns only the raw chunk data. -// Chunk signatures and trailing checksums are not validated (the -// request-level signature was already verified by the auth validator). +// The reader strips the framing and returns only the raw chunk data. It does +// not validate chunk signatures or trailing checksums. The signing forwarder +// verifies and stages the standard HMAC streaming-payload marker separately. type awsChunkedReader struct { reader *bufio.Reader remaining int @@ -191,8 +191,10 @@ func (r *awsChunkedReader) readTrailingCRLF() error { // decoded content length, and chunked=true. Otherwise returns the original values // unchanged with chunked=false. // -// Supports both signed (STREAMING-AWS4-HMAC-SHA256-PAYLOAD) and unsigned -// (STREAMING-UNSIGNED-PAYLOAD-TRAILER) chunked formats. +// Supports the framing of both signed (STREAMING-AWS4-HMAC-SHA256-PAYLOAD) and +// unsigned (STREAMING-UNSIGNED-PAYLOAD-TRAILER) chunked formats. This low-level +// decoder does not authenticate signed chunks; the signing forwarder handles the +// standard HMAC streaming-payload marker before forwarding. func decodeChunkedIfNeeded(r *http.Request) (body io.ReadCloser, bodyHash string, contentLength int64, chunked bool) { bodyHash = r.Header.Get("X-Amz-Content-Sha256") diff --git a/proxy/forwarder.go b/proxy/forwarder.go index 076f79fa..85118b89 100644 --- a/proxy/forwarder.go +++ b/proxy/forwarder.go @@ -500,6 +500,17 @@ func NewForwarder(credStore *auth.CredentialStore, tigrisEndpoint, region string } } +// headersForForwardedSigning excludes headers that are removed after decoding, so +// the generated Authorization header matches the request sent to the upstream. +func headersForForwardedSigning(headers http.Header, chunked bool) http.Header { + if !chunked { + return headers + } + forwarded := headers.Clone() + forwarded.Del("X-Amz-Decoded-Content-Length") + return forwarded +} + // prepareForwardedRequest sets Content-Length on the forwarded request and strips // AWS chunked-encoding headers when the request body was decoded from chunked format. func prepareForwardedRequest(fwdReq *http.Request, contentLength int64, chunked bool) { diff --git a/proxy/forwarder_signed_stream_test.go b/proxy/forwarder_signed_stream_test.go new file mode 100644 index 00000000..1d66bcba --- /dev/null +++ b/proxy/forwarder_signed_stream_test.go @@ -0,0 +1,521 @@ +package proxy + +import ( + "bytes" + "context" + "crypto/hmac" + "crypto/sha256" + "encoding/hex" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + + "github.com/tigrisdata/tag/auth" +) + +const ( + signedStreamTestAccessKey = "AKIAIOSFODNN7EXAMPLE" + signedStreamTestSecretKey = "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" +) + +type signedStreamObservation struct { + body []byte + contentHash string + contentLength int64 + contentCoding string + decodedLength string +} + +type signedStreamSeed struct { + endpoint string + path string + headers http.Header + host string + seedSignature string + signingKey []byte + signingTime string + credentialScope string +} + +type signedStreamWire struct { + wire []byte + data []byte + dataSignatureStart []int + dataSignatureExtension []int + terminalSignatureStart int + terminalSignatureExt int +} + +// TestSigningForwarderSignedStreamIntegrity checks the public signing-mode +// forwarder against independently generated SigV4 chunk chains. The local +// upstream only records requests; it does not authenticate or reject them. +func TestSigningForwarderSignedStreamIntegrity(t *testing.T) { + var ( + dispatched atomic.Int32 + mu sync.Mutex + seen []signedStreamObservation + ) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + dispatched.Add(1) + body, err := io.ReadAll(r.Body) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + mu.Lock() + seen = append(seen, signedStreamObservation{ + body: body, + contentHash: r.Header.Get("X-Amz-Content-Sha256"), + contentLength: r.ContentLength, + contentCoding: r.Header.Get("Content-Encoding"), + decodedLength: r.Header.Get("X-Amz-Decoded-Content-Length"), + }) + mu.Unlock() + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + credentials := auth.NewCredentialStore() + credentials.AddCredential(signedStreamTestAccessKey, signedStreamTestSecretKey) + forwarder := NewForwarder(credentials, upstream.URL, "us-east-1", 1, nil, nil) + // Keep the integrity assertion independent of the host's available temp space. + useSignedStreamStageBudget(t, newSignedStreamStageBudget(func(string) (signedStreamStageSpace, error) { + return signedStreamStageSpace{availableBytes: 1 << 30, blockSize: 1}, nil + })) + chunks := [][]byte{[]byte("first chunk "), []byte("second chunk")} + decodedLength := len(chunks[0]) + len(chunks[1]) + seed := newSignedStreamSeed(t, upstream.URL, StreamingPayloadHash, decodedLength, true) + validWire := makeSignedStreamWire(t, seed, chunks...) + + methods := []struct { + name string + call func(*http.Request) (int, error) + }{ + { + name: "Forward", + call: func(r *http.Request) (int, error) { + w := httptest.NewRecorder() + err := forwarder.Forward(r.Context(), w, r) + return w.Code, err + }, + }, + { + name: "ForwardWithCapture", + call: func(r *http.Request) (int, error) { + capture, err := forwarder.ForwardWithCapture(r.Context(), httptest.NewRecorder(), r) + if capture == nil { + return 0, err + } + return capture.StatusCode, err + }, + }, + { + name: "ForwardTeeingBody", + call: func(r *http.Request) (int, error) { + tee, ok := forwarder.(interface { + ForwardTeeingBody(context.Context, http.ResponseWriter, *http.Request, io.Writer) (int, http.Header, string, string, error) + }) + if !ok { + return 0, fmt.Errorf("signing forwarder does not implement ForwardTeeingBody") + } + status, _, _, _, err := tee.ForwardTeeingBody(r.Context(), httptest.NewRecorder(), r, io.Discard) + return status, err + }, + }, + } + + for _, method := range methods { + t.Run(method.name, func(t *testing.T) { + assertSignedStreamForwarded(t, method.call, seed, validWire.wire, validWire.data, &dispatched, &mu, &seen) + + invalid := []struct { + name string + wire []byte + }{ + {name: "same-length payload mutation", wire: mutateStreamPayload(t, validWire)}, + {name: "missing data signature", wire: removeStreamSignature(t, validWire, validWire.dataSignatureExtension[0], validWire.dataSignatureStart[0])}, + {name: "changed later data signature", wire: changeStreamSignature(t, validWire.wire, validWire.dataSignatureStart[1])}, + {name: "changed terminal signature", wire: changeStreamSignature(t, validWire.wire, validWire.terminalSignatureStart)}, + {name: "missing terminal signature", wire: removeStreamSignature(t, validWire, validWire.terminalSignatureExt, validWire.terminalSignatureStart)}, + } + for _, tc := range invalid { + t.Run(tc.name, func(t *testing.T) { + assertSignedStreamRejectedBeforeDispatch(t, method.call, seed, tc.wire, &dispatched) + }) + } + + emptySeed := newSignedStreamSeed(t, upstream.URL, StreamingPayloadHash, 0, true) + emptyWire := makeSignedStreamWire(t, emptySeed) + t.Run("valid zero-byte terminal signature", func(t *testing.T) { + assertSignedStreamForwarded(t, method.call, emptySeed, emptyWire.wire, nil, &dispatched, &mu, &seen) + }) + badEmptyTerminal := changeStreamSignature(t, emptyWire.wire, emptyWire.terminalSignatureStart) + t.Run("invalid zero-byte terminal signature", func(t *testing.T) { + assertSignedStreamRejectedBeforeDispatch(t, method.call, emptySeed, badEmptyTerminal, &dispatched) + }) + }) + } + + // Compatibility controls: ordinary signed payloads and the separately + // supported unsigned-trailer stream keep their existing forwarding behavior. + payload := []byte("compatibility payload") + nonStreamingHash := sha256Hex(payload) + nonStreamingSeed := newSignedStreamSeed(t, upstream.URL, nonStreamingHash, -1, false) + assertCompatibilityStreamForwarded(t, methods[0].call, nonStreamingSeed, payload, nonStreamingHash, payload, &dispatched, &mu, &seen) + + unsignedSeed := newSignedStreamSeed(t, upstream.URL, StreamingUnsignedTrailerHash, len(payload), true) + unsignedWire := makeUnsignedStreamWire(payload) + assertCompatibilityStreamForwarded(t, methods[0].call, unsignedSeed, unsignedWire, "UNSIGNED-PAYLOAD", payload, &dispatched, &mu, &seen) +} + +func TestSigningForwarderDecodedStreamOutgoingSignature(t *testing.T) { + credentials := auth.NewCredentialStore() + credentials.AddCredential(signedStreamTestAccessKey, signedStreamTestSecretKey) + upstreamValidator := auth.NewRequestValidator(credentials) + var ( + dispatched atomic.Int32 + mu sync.Mutex + seen []signedStreamObservation + ) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + dispatched.Add(1) + if _, err := upstreamValidator.ValidateRequest(r); err != nil { + http.Error(w, fmt.Sprintf("outbound SigV4 validation failed: %v", err), http.StatusForbidden) + return + } + body, err := io.ReadAll(r.Body) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + mu.Lock() + seen = append(seen, signedStreamObservation{ + body: body, + contentHash: r.Header.Get("X-Amz-Content-Sha256"), + contentLength: r.ContentLength, + contentCoding: r.Header.Get("Content-Encoding"), + decodedLength: r.Header.Get("X-Amz-Decoded-Content-Length"), + }) + mu.Unlock() + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + forwarder := NewForwarder(credentials, upstream.URL, "us-east-1", 1, nil, nil) + useSignedStreamStageBudget(t, newSignedStreamStageBudget(func(string) (signedStreamStageSpace, error) { + return signedStreamStageSpace{availableBytes: 1 << 20, blockSize: 1}, nil + })) + payload := []byte("payload signed for the decoded upstream request") + seed := newSignedStreamSeed(t, upstream.URL, StreamingPayloadHash, len(payload), true) + stream := makeSignedStreamWire(t, seed, payload) + methods := []struct { + name string + call func(*http.Request) (int, error) + }{ + { + name: "Forward", + call: func(r *http.Request) (int, error) { + w := httptest.NewRecorder() + err := forwarder.Forward(r.Context(), w, r) + return w.Code, err + }, + }, + { + name: "ForwardWithCapture", + call: func(r *http.Request) (int, error) { + capture, err := forwarder.ForwardWithCapture(r.Context(), httptest.NewRecorder(), r) + if capture == nil { + return 0, err + } + return capture.StatusCode, err + }, + }, + { + name: "ForwardTeeingBody", + call: func(r *http.Request) (int, error) { + tee, ok := forwarder.(interface { + ForwardTeeingBody(context.Context, http.ResponseWriter, *http.Request, io.Writer) (int, http.Header, string, string, error) + }) + if !ok { + return 0, fmt.Errorf("signing forwarder does not implement ForwardTeeingBody") + } + status, _, _, _, err := tee.ForwardTeeingBody(r.Context(), httptest.NewRecorder(), r, io.Discard) + return status, err + }, + }, + } + for _, method := range methods { + t.Run(method.name, func(t *testing.T) { + before := dispatched.Load() + seenBefore := observationCount(&mu, &seen) + status, err := method.call(seed.request(t, stream.wire)) + if err != nil { + t.Fatalf("forward decoded signed stream: %v", err) + } + if status != http.StatusOK || dispatched.Load() != before+1 { + t.Fatalf("forwarded status=%d dispatches=%d; want one SigV4-accepted request", status, dispatched.Load()-before) + } + got, ok := observationAt(&mu, &seen, seenBefore) + if !ok { + t.Fatalf("upstream dispatch %d produced no accepted request observation", seenBefore) + } + if !bytes.Equal(got.body, payload) || got.contentHash != "UNSIGNED-PAYLOAD" || got.contentLength != int64(len(payload)) || got.contentCoding != "" || got.decodedLength != "" { + t.Fatalf("upstream accepted a different decoded request: %+v want bytes=%q", got, payload) + } + }) + } +} + +func newSignedStreamSeed(t *testing.T, endpoint, bodyHash string, decodedLength int, chunked bool) signedStreamSeed { + t.Helper() + + path := "/bucket/object" + headers := make(http.Header) + if chunked { + headers.Set("Content-Encoding", "aws-chunked") + if decodedLength >= 0 { + headers.Set("X-Amz-Decoded-Content-Length", fmt.Sprint(decodedLength)) + } + } + signed, err := auth.NewRequestSigner(endpoint, "us-east-1").SignRequest( + t.Context(), http.MethodPut, path, nil, bodyHash, + signedStreamTestAccessKey, signedStreamTestSecretKey, headers, + ) + if err != nil { + t.Fatalf("sign request header: %v", err) + } + info, err := auth.ParseAuthInfo(signed) + if err != nil { + t.Fatalf("parse generated request signature: %v", err) + } + requestTime, err := auth.ParseHTTPDate(signed.Header.Get("X-Amz-Date")) + if err != nil { + t.Fatalf("parse generated request date: %v", err) + } + + return signedStreamSeed{ + endpoint: endpoint, + path: path, + headers: signed.Header.Clone(), + host: signed.Host, + seedSignature: info.Signature, + signingKey: deriveSignedStreamTestKey(signedStreamTestSecretKey, info.Date, info.Region), + signingTime: requestTime.UTC().Format(auth.TimeFormat), + credentialScope: info.Date + "/" + info.Region + "/s3/aws4_request", + } +} + +func (seed signedStreamSeed) request(t *testing.T, wire []byte) *http.Request { + t.Helper() + req, err := http.NewRequestWithContext(t.Context(), http.MethodPut, seed.endpoint+seed.path, bytes.NewReader(wire)) + if err != nil { + t.Fatalf("create signed-stream request: %v", err) + } + req.Header = seed.headers.Clone() + req.Host = seed.host + req.ContentLength = int64(len(wire)) + return req +} + +func makeSignedStreamWire(t *testing.T, seed signedStreamSeed, chunks ...[]byte) signedStreamWire { + t.Helper() + + var wire signedStreamWire + var encoded bytes.Buffer + wire.dataSignatureStart = make([]int, 0, len(chunks)) + wire.dataSignatureExtension = make([]int, 0, len(chunks)) + previous := seed.seedSignature + for _, chunk := range chunks { + if len(chunk) == 0 { + t.Fatal("test chunk must be non-empty") + } + signature := signedStreamTestChunkSignature(seed.signingKey, seed.signingTime, seed.credentialScope, previous, chunk) + fmt.Fprintf(&encoded, "%x", len(chunk)) + wire.dataSignatureExtension = append(wire.dataSignatureExtension, encoded.Len()) + encoded.WriteString(";chunk-signature=") + wire.dataSignatureStart = append(wire.dataSignatureStart, encoded.Len()) + encoded.WriteString(signature) + encoded.WriteString("\r\n") + wire.data = append(wire.data, chunk...) + encoded.Write(chunk) + encoded.WriteString("\r\n") + previous = signature + } + + terminal := signedStreamTestChunkSignature(seed.signingKey, seed.signingTime, seed.credentialScope, previous, nil) + encoded.WriteString("0") + wire.terminalSignatureExt = encoded.Len() + encoded.WriteString(";chunk-signature=") + wire.terminalSignatureStart = encoded.Len() + encoded.WriteString(terminal) + encoded.WriteString("\r\n\r\n") + wire.wire = bytes.Clone(encoded.Bytes()) + return wire +} + +func signedStreamTestChunkSignature(signingKey []byte, signingTime, scope, previous string, chunk []byte) string { + stringToSign := strings.Join([]string{ + "AWS4-HMAC-SHA256-PAYLOAD", + signingTime, + scope, + previous, + sha256Hex(nil), + sha256Hex(chunk), + }, "\n") + mac := hmac.New(sha256.New, signingKey) + _, _ = mac.Write([]byte(stringToSign)) + return hex.EncodeToString(mac.Sum(nil)) +} + +func deriveSignedStreamTestKey(secret, date, region string) []byte { + kDate := signedStreamTestHMAC([]byte("AWS4"+secret), []byte(date)) + kRegion := signedStreamTestHMAC(kDate, []byte(region)) + kService := signedStreamTestHMAC(kRegion, []byte("s3")) + return signedStreamTestHMAC(kService, []byte("aws4_request")) +} + +func signedStreamTestHMAC(key, message []byte) []byte { + mac := hmac.New(sha256.New, key) + _, _ = mac.Write(message) + return mac.Sum(nil) +} + +func sha256Hex(data []byte) string { + sum := sha256.Sum256(data) + return hex.EncodeToString(sum[:]) +} + +func mutateStreamPayload(t *testing.T, stream signedStreamWire) []byte { + t.Helper() + if len(stream.dataSignatureStart) < 2 { + t.Fatal("test stream needs two data chunks") + } + wire := bytes.Clone(stream.wire) + dataStart := bytes.Index(wire, []byte("second chunk")) + if dataStart < 0 { + t.Fatal("test stream payload not found") + } + wire[dataStart] ^= 1 + return wire +} + +func changeStreamSignature(t *testing.T, wire []byte, signatureStart int) []byte { + t.Helper() + changed := bytes.Clone(wire) + if signatureStart < 0 || signatureStart+64 > len(changed) { + t.Fatalf("invalid test signature offset %d", signatureStart) + } + if changed[signatureStart] == '0' { + changed[signatureStart] = '1' + } else { + changed[signatureStart] = '0' + } + return changed +} + +func removeStreamSignature(t *testing.T, stream signedStreamWire, extensionStart, signatureStart int) []byte { + t.Helper() + end := signatureStart + 64 + if extensionStart < 0 || signatureStart < extensionStart || end > len(stream.wire) { + t.Fatalf("invalid test signature span %d..%d", extensionStart, end) + } + changed := make([]byte, 0, len(stream.wire)-(end-extensionStart)) + changed = append(changed, stream.wire[:extensionStart]...) + changed = append(changed, stream.wire[end:]...) + return changed +} + +func assertSignedStreamForwarded(t *testing.T, call func(*http.Request) (int, error), seed signedStreamSeed, wire, wantBody []byte, dispatched *atomic.Int32, mu *sync.Mutex, seen *[]signedStreamObservation) { + t.Helper() + before := dispatched.Load() + seenBefore := observationCount(mu, seen) + status, err := call(seed.request(t, wire)) + if err != nil { + if authErr, ok := IsAuthError(err); ok && authErr.Code == ErrCodeSignatureMismatch { + signedStreamViolation(t, "valid signed stream was rejected as an authentication failure") + } else { + t.Fatalf("valid stream setup or forwarding failed: %v", err) + } + } + if status != http.StatusOK || dispatched.Load() != before+1 { + signedStreamViolation(t, "valid stream was not dispatched successfully: status=%d dispatches=%d", status, dispatched.Load()-before) + return + } + got, ok := observationAt(mu, seen, seenBefore) + if !ok { + t.Fatalf("upstream dispatch %d produced no complete observation", seenBefore) + } + if !bytes.Equal(got.body, wantBody) || got.contentHash != "UNSIGNED-PAYLOAD" || got.contentCoding != "" || got.decodedLength != "" || got.contentLength != int64(len(wantBody)) { + signedStreamViolation(t, "valid stream was not re-signed as the exact decoded payload: observation=%+v want bytes=%q", got, wantBody) + } +} + +func assertSignedStreamRejectedBeforeDispatch(t *testing.T, call func(*http.Request) (int, error), seed signedStreamSeed, wire []byte, dispatched *atomic.Int32) { + t.Helper() + before := dispatched.Load() + _, err := call(seed.request(t, wire)) + after := dispatched.Load() + authErr, isAuthErr := IsAuthError(err) + if err != nil && !isAuthErr && after == before { + t.Fatalf("invalid signed stream was not conclusively rejected by payload authentication: %v", err) + } + if err == nil || !isAuthErr || authErr.Code != ErrCodeSignatureMismatch || after != before { + signedStreamViolation(t, "invalid signed stream was not rejected as a signature mismatch before upstream dispatch: err=%v auth=%v dispatches=%d", err, isAuthErr, after-before) + } +} + +func assertCompatibilityStreamForwarded(t *testing.T, call func(*http.Request) (int, error), seed signedStreamSeed, wire []byte, outboundHash string, payload []byte, dispatched *atomic.Int32, mu *sync.Mutex, seen *[]signedStreamObservation) { + t.Helper() + before := dispatched.Load() + seenBefore := observationCount(mu, seen) + status, err := call(seed.request(t, wire)) + if err != nil { + t.Fatalf("compatibility control failed: %v", err) + } + if status != http.StatusOK { + t.Fatalf("compatibility control status = %d, want %d", status, http.StatusOK) + } + if dispatched.Load() != before+1 { + t.Fatalf("compatibility control dispatched %d times, want exactly once", dispatched.Load()-before) + } + got, ok := observationAt(mu, seen, seenBefore) + if !ok { + t.Fatalf("compatibility control dispatch %d produced no complete observation", seenBefore) + } + if !bytes.Equal(got.body, payload) || got.contentHash != outboundHash { + t.Fatalf("compatibility control changed body or payload hash: got body=%q hash=%q, want body=%q hash=%q", got.body, got.contentHash, payload, outboundHash) + } +} + +func observationCount(mu *sync.Mutex, seen *[]signedStreamObservation) int { + mu.Lock() + defer mu.Unlock() + return len(*seen) +} + +func observationAt(mu *sync.Mutex, seen *[]signedStreamObservation, index int) (signedStreamObservation, bool) { + mu.Lock() + defer mu.Unlock() + if index < 0 || index >= len(*seen) { + return signedStreamObservation{}, false + } + return (*seen)[index], true +} + +func signedStreamViolation(t *testing.T, format string, args ...any) { + t.Helper() + t.Errorf(format, args...) +} + +func makeUnsignedStreamWire(payload []byte) []byte { + var wire bytes.Buffer + fmt.Fprintf(&wire, "%x\r\n", len(payload)) + wire.Write(payload) + wire.WriteString("\r\n0\r\n\r\n") + return wire.Bytes() +} diff --git a/proxy/forwarder_signing.go b/proxy/forwarder_signing.go index 9f802d90..7676cde7 100644 --- a/proxy/forwarder_signing.go +++ b/proxy/forwarder_signing.go @@ -22,13 +22,10 @@ type signingForwarder struct { // Forward forwards a request to Tigris and writes the response to the client. // Validates the incoming request signature, re-signs with upstream credentials, -// and streams the response back. If the request uses AWS chunked transfer encoding -// (streaming SigV4), the body is decoded on-the-fly and forwarded as UNSIGNED-PAYLOAD. +// and streams the response back. Signed AWS streaming bodies are authenticated and +// staged before their decoded bytes are forwarded as UNSIGNED-PAYLOAD. func (f *signingForwarder) Forward(ctx context.Context, w http.ResponseWriter, r *http.Request) error { - // Decode AWS chunked encoding if present, otherwise pass through unchanged - body, bodyHash, contentLength, chunked := decodeChunkedIfNeeded(r) - - // Validate incoming request signature + // Validate the request header before consuming any body data. accessKey, err := f.validator.ValidateRequest(r) if err != nil { log.Warn().Err(err).Str("path", r.URL.Path).Msg("Request signature validation failed") @@ -41,18 +38,25 @@ func (f *signingForwarder) Forward(ctx context.Context, w http.ResponseWriter, r return mapAuthError(err) } + body, bodyHash, contentLength, chunked, staged, err := f.decodeIncomingBody(ctx, w, r, accessKey) + if err != nil { + return err + } + defer func() { closeStagedSignedChunkedBody(staged) }() + // Build the path with query string path := r.URL.Path if r.URL.RawQuery != "" { path = path + "?" + r.URL.RawQuery } - // Create signed request (passes body hash, streams body directly) - fwdReq, err := f.signer.SignRequest(ctx, r.Method, path, body, bodyHash, accessKey, secretKey, r.Header) + // Create signed request. Verified AWS streams use the staged decoded body. + fwdReq, err := f.signer.SignRequest(ctx, r.Method, path, body, bodyHash, accessKey, secretKey, headersForForwardedSigning(r.Header, chunked)) if err != nil { return err } prepareForwardedRequest(fwdReq, contentLength, chunked) + handoffStagedSignedChunkedBody(&staged, fwdReq) return f.executeAndStream(w, fwdReq, contentLength, nil) } @@ -63,15 +67,14 @@ func (f *signingForwarder) Forward(ctx context.Context, w http.ResponseWriter, r // status code and a clone of the response headers (for the ETag) so the caller can build // cache metadata for the just-written object. // -// Only the signing forwarder implements this: it decodes any AWS chunked encoding, so it -// sees the assembled object bytes. The transparent forwarder preserves the client's opaque -// (possibly chunked) body and signature, so it can't tee cleanly and falls back to -// warm-on-write. The `tee` writer must never return an error — that would truncate the -// upstream stream via io.TeeReader — so callers pass a capped, non-erroring buffer and +// Only the signing forwarder implements this: it decodes AWS chunked encoding, so it +// sees the assembled object bytes. Signed HMAC streams are verified and staged first. +// The transparent forwarder preserves the client's opaque (possibly chunked) body +// and signature, so it can't tee cleanly and falls back to warm-on-write. The `tee` +// writer must never return an error because that would truncate the upstream stream +// via io.TeeReader; callers pass a capped, non-erroring buffer and // check overflow out of band. func (f *signingForwarder) ForwardTeeingBody(ctx context.Context, w http.ResponseWriter, r *http.Request, tee io.Writer) (int, http.Header, string, string, error) { - body, bodyHash, contentLength, chunked := decodeChunkedIfNeeded(r) - accessKey, err := f.validator.ValidateRequest(r) if err != nil { log.Warn().Err(err).Str("path", r.URL.Path).Msg("Request signature validation failed") @@ -82,18 +85,28 @@ func (f *signingForwarder) ForwardTeeingBody(ctx context.Context, w http.Respons return 0, nil, "", "", mapAuthError(err) } + body, bodyHash, contentLength, chunked, staged, err := f.decodeIncomingBody(ctx, w, r, accessKey) + if err != nil { + return 0, nil, "", "", err + } + defer func() { closeStagedSignedChunkedBody(staged) }() + path := r.URL.Path if r.URL.RawQuery != "" { path = path + "?" + r.URL.RawQuery } // Tee the decoded body into the caller's buffer as it is streamed upstream. - teedBody := io.TeeReader(body, tee) - fwdReq, err := f.signer.SignRequest(ctx, r.Method, path, teedBody, bodyHash, accessKey, secretKey, r.Header) + var teedBody io.Reader = io.TeeReader(body, tee) + if staged != nil { + teedBody = &stagedTeeReadCloser{Reader: teedBody, Closer: staged} + } + fwdReq, err := f.signer.SignRequest(ctx, r.Method, path, teedBody, bodyHash, accessKey, secretKey, headersForForwardedSigning(r.Header, chunked)) if err != nil { return 0, nil, "", "", err } prepareForwardedRequest(fwdReq, contentLength, chunked) + handoffStagedSignedChunkedBody(&staged, fwdReq) // Return the validated credentials so the caller can HEAD/warm without re-validating. status, headers, err := f.executeAndStreamReturningMeta(w, fwdReq, contentLength, nil) @@ -104,10 +117,7 @@ func (f *signingForwarder) ForwardTeeingBody(ctx context.Context, w http.Respons // Validates and re-signs like Forward, but also captures the response body // for caching while streaming to the client. func (f *signingForwarder) ForwardWithCapture(ctx context.Context, w http.ResponseWriter, r *http.Request) (*ResponseCapture, error) { - // Decode AWS chunked encoding if present, otherwise pass through unchanged - body, bodyHash, contentLength, chunked := decodeChunkedIfNeeded(r) - - // Validate incoming request signature + // Validate the request header before consuming any body data. accessKey, err := f.validator.ValidateRequest(r) if err != nil { log.Warn().Err(err).Str("path", r.URL.Path).Msg("Request signature validation failed") @@ -120,18 +130,25 @@ func (f *signingForwarder) ForwardWithCapture(ctx context.Context, w http.Respon return nil, mapAuthError(err) } + body, bodyHash, contentLength, chunked, staged, err := f.decodeIncomingBody(ctx, w, r, accessKey) + if err != nil { + return nil, err + } + defer func() { closeStagedSignedChunkedBody(staged) }() + // Build the path with query string path := r.URL.Path if r.URL.RawQuery != "" { path = path + "?" + r.URL.RawQuery } - // Create signed request (passes body hash, streams body directly) - fwdReq, err := f.signer.SignRequest(ctx, r.Method, path, body, bodyHash, accessKey, secretKey, r.Header) + // Create signed request. Verified AWS streams use the staged decoded body. + fwdReq, err := f.signer.SignRequest(ctx, r.Method, path, body, bodyHash, accessKey, secretKey, headersForForwardedSigning(r.Header, chunked)) if err != nil { return nil, err } prepareForwardedRequest(fwdReq, contentLength, chunked) + handoffStagedSignedChunkedBody(&staged, fwdReq) return f.executeAndCapture(w, fwdReq, contentLength, nil) } @@ -157,7 +174,7 @@ func (f *signingForwarder) ValidateAndGetCredentials(r *http.Request) (AuthResul // Returns the raw response for streaming. Caller is responsible for closing the response body. // If the request uses AWS chunked transfer encoding, the body is decoded on-the-fly. func (f *signingForwarder) DoRequestWithCreds(ctx context.Context, r *http.Request, accessKey, secretKey string) (*http.Response, error) { - // Decode AWS chunked encoding if present, otherwise pass through unchanged + // Decode AWS chunked encoding if present, otherwise pass through unchanged. body, bodyHash, contentLength, chunked := decodeChunkedIfNeeded(r) path := r.URL.Path diff --git a/proxy/service.go b/proxy/service.go index eb9e9958..538817ab 100644 --- a/proxy/service.go +++ b/proxy/service.go @@ -669,6 +669,11 @@ type statusRecorder struct { status int } +// Unwrap lets http.ResponseController reach connection controls through this wrapper. +func (rec *statusRecorder) Unwrap() http.ResponseWriter { + return rec.ResponseWriter +} + func (rec *statusRecorder) WriteHeader(code int) { rec.status = code rec.ResponseWriter.WriteHeader(code) diff --git a/proxy/signed_chunked_body.go b/proxy/signed_chunked_body.go new file mode 100644 index 00000000..d3eb5516 --- /dev/null +++ b/proxy/signed_chunked_body.go @@ -0,0 +1,688 @@ +package proxy + +import ( + "bufio" + "context" + "crypto/hmac" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "io" + "math" + "net/http" + "os" + "strconv" + "strings" + "sync" + "syscall" + "time" + + "github.com/rs/zerolog/log" + "github.com/tigrisdata/tag/auth" +) + +const signedChunkStringToSignPrefix = "AWS4-HMAC-SHA256-PAYLOAD" + +var ( + errSignedStreamStagingCapacity = errors.New("signed AWS stream staging capacity is unavailable") + errSignedStreamLengthMismatch = errors.New("signed AWS stream decoded length differs from its declaration") +) + +// signedStreamStageSpace describes writable capacity on the filesystem that +// holds staged request bodies. +type signedStreamStageSpace struct { + availableBytes int64 + blockSize int64 +} + +type signedStreamStageSpaceReader func(string) (signedStreamStageSpace, error) + +// signedStreamStageBudget reserves at most half of the temporary filesystem's +// currently available space for all concurrent signed-stream stages, leaving the +// other half unclaimed by this pool. Reservations stay held until the staged file +// is removed, including while the HTTP transport owns it. +type signedStreamStageBudget struct { + mu sync.Mutex + spaceReader signedStreamStageSpaceReader + dir string + limit int64 + headroom int64 + reserved int64 + unallocated int64 +} + +func newSignedStreamStageBudget(spaceReader signedStreamStageSpaceReader) *signedStreamStageBudget { + if spaceReader == nil { + spaceReader = readSignedStreamStageSpace + } + return &signedStreamStageBudget{spaceReader: spaceReader} +} + +var defaultSignedStreamStageBudget = newSignedStreamStageBudget(nil) + +func readSignedStreamStageSpace(dir string) (signedStreamStageSpace, error) { + var stat syscall.Statfs_t + if err := syscall.Statfs(dir, &stat); err != nil { + return signedStreamStageSpace{}, err + } + blockSize := int64(stat.Bsize) + if blockSize <= 0 { + return signedStreamStageSpace{}, fmt.Errorf("temporary filesystem reported invalid block size %d", blockSize) + } + availableBlocks := stat.Bavail + if availableBlocks > uint64(math.MaxInt64/blockSize) { + return signedStreamStageSpace{availableBytes: math.MaxInt64, blockSize: blockSize}, nil + } + return signedStreamStageSpace{ + availableBytes: int64(availableBlocks) * blockSize, + blockSize: blockSize, + }, nil +} + +func (b *signedStreamStageBudget) reserve(dir string, space signedStreamStageSpace, amount int64) error { + if amount < 0 { + return fmt.Errorf("%w: negative reservation %d", errSignedStreamStagingCapacity, amount) + } + if amount == 0 { + return b.checkLocked(dir, space) + } + if b.reserved == 0 { + b.dir = dir + b.limit = space.availableBytes / 2 + b.headroom = space.availableBytes - b.limit + } else if b.dir != dir { + return fmt.Errorf("%w: temporary directory changed while a stage is active", errSignedStreamStagingCapacity) + } + if b.limit <= 0 || b.reserved > b.limit || amount > b.limit-b.reserved { + return errSignedStreamStagingCapacity + } + if err := b.checkLocked(dir, space); err != nil { + return err + } + if amount > space.availableBytes-b.headroom-b.unallocated { + return errSignedStreamStagingCapacity + } + b.reserved += amount + b.unallocated += amount + return nil +} + +func (b *signedStreamStageBudget) checkLocked(dir string, space signedStreamStageSpace) error { + if b.reserved == 0 { + return nil + } + if b.dir != dir || space.availableBytes < b.headroom || b.unallocated > space.availableBytes-b.headroom { + return errSignedStreamStagingCapacity + } + return nil +} + +func (b *signedStreamStageBudget) release(reserved, unallocated int64) { + if reserved == 0 { + return + } + b.mu.Lock() + defer b.mu.Unlock() + b.reserved -= reserved + b.unallocated -= unallocated + if b.reserved == 0 { + b.dir = "" + b.limit = 0 + b.headroom = 0 + b.unallocated = 0 + } +} + +// signedStreamStageReservation tracks the rounded disk space reserved for one +// staged body and the portion already allocated by writes to its file. +type signedStreamStageReservation struct { + mu sync.Mutex + budget *signedStreamStageBudget + dir string + blockSize int64 + reservedBytes int64 + allocated int64 + written int64 +} + +func (r *signedStreamStageReservation) reservePayload(total int64) error { + if total < 0 { + return fmt.Errorf("%w: negative decoded length %d", errSignedStreamStagingCapacity, total) + } + if total == 0 { + return nil + } + + r.mu.Lock() + defer r.mu.Unlock() + r.budget.mu.Lock() + defer r.budget.mu.Unlock() + space, err := r.budget.spaceReader(r.dir) + if err != nil { + return fmt.Errorf("%w: inspect temporary filesystem: %w", errSignedStreamStagingCapacity, err) + } + if space.availableBytes < 0 || space.blockSize <= 0 { + return fmt.Errorf("%w: invalid temporary filesystem capacity", errSignedStreamStagingCapacity) + } + if r.blockSize == 0 { + r.blockSize = space.blockSize + } + required, err := roundSignedStreamStageBytes(total, r.blockSize) + if err != nil { + return err + } + if required < r.reservedBytes { + return fmt.Errorf("signed-stream reservation shrank from %d to %d bytes", r.reservedBytes, required) + } + if err := r.budget.reserve(r.dir, space, required-r.reservedBytes); err != nil { + return err + } + r.reservedBytes = required + return nil +} + +func (r *signedStreamStageReservation) checkAvailable() error { + r.mu.Lock() + defer r.mu.Unlock() + if r.reservedBytes == 0 { + return nil + } + r.budget.mu.Lock() + defer r.budget.mu.Unlock() + space, err := r.budget.spaceReader(r.dir) + if err != nil { + return fmt.Errorf("%w: inspect temporary filesystem: %w", errSignedStreamStagingCapacity, err) + } + return r.budget.checkLocked(r.dir, space) +} + +// write serializes I/O and accounting for one staged file, but holds the shared +// budget lock only while checking and updating reservations. +func (r *signedStreamStageReservation) write(file io.Writer, p []byte) (int, error) { + if len(p) == 0 { + return 0, nil + } + r.mu.Lock() + defer r.mu.Unlock() + + writeLength := int64(len(p)) + if writeLength > math.MaxInt64-r.written { + return 0, errors.New("signed-stream written length overflows int64") + } + maxAllocated, err := roundSignedStreamStageBytes(r.written+writeLength, r.blockSize) + if err != nil { + return 0, err + } + if maxAllocated > r.reservedBytes { + return 0, fmt.Errorf("signed-stream write would use %d bytes beyond %d-byte reservation", maxAllocated, r.reservedBytes) + } + + reservationDelta := maxAllocated - r.allocated + r.budget.mu.Lock() + if reservationDelta > r.budget.unallocated { + r.budget.mu.Unlock() + return 0, fmt.Errorf("signed-stream write would exceed unallocated staging reservations") + } + // Treat the maximum allocation as occupied before releasing the shared lock. + // A partial write restores the unused portion after its actual size is known. + r.budget.unallocated -= reservationDelta + r.budget.mu.Unlock() + + n, writeErr := file.Write(p) + if n < 0 || n > len(p) { + allocated := maxAllocated + r.written += writeLength + r.allocated = allocated + return n, fmt.Errorf("signed-stream writer returned invalid count %d for %d-byte input", n, len(p)) + } + if n <= 0 { + r.budget.mu.Lock() + r.budget.unallocated += reservationDelta + r.budget.mu.Unlock() + return n, writeErr + } + written := int64(n) + allocated, accountingErr := roundSignedStreamStageBytes(r.written+written, r.blockSize) + if accountingErr != nil { + // Charge the conservative maximum if accounting unexpectedly overflows + // after bytes were written, so cleanup cannot release their disk budget. + allocated = maxAllocated + } + allocatedDelta := allocated - r.allocated + if allocatedDelta < 0 || allocatedDelta > reservationDelta { + allocated = maxAllocated + allocatedDelta = reservationDelta + accountingErr = errors.New("signed-stream write exceeded its reserved staging budget") + } + + r.budget.mu.Lock() + r.budget.unallocated += reservationDelta - allocatedDelta + r.written += written + r.allocated = allocated + r.budget.mu.Unlock() + + if accountingErr != nil { + return n, accountingErr + } + return n, writeErr +} + +func (r *signedStreamStageReservation) release(fileRemoved bool) { + r.mu.Lock() + defer r.mu.Unlock() + r.releaseLocked(fileRemoved) +} + +func (r *signedStreamStageReservation) releaseLocked(fileRemoved bool) { + unallocated := r.reservedBytes - r.allocated + releasable := unallocated + if fileRemoved { + releasable += r.allocated + } + r.budget.release(releasable, unallocated) + r.reservedBytes = 0 + r.allocated = 0 + r.written = 0 +} + +func roundSignedStreamStageBytes(size, blockSize int64) (int64, error) { + if size < 0 || blockSize <= 0 { + return 0, fmt.Errorf("invalid signed-stream size %d or block size %d", size, blockSize) + } + if size == 0 { + return 0, nil + } + remainder := size % blockSize + if remainder == 0 { + return size, nil + } + padding := blockSize - remainder + if size > math.MaxInt64-padding { + return 0, errors.New("signed-stream staging size overflows int64") + } + return size + padding, nil +} + +// stagedSignedChunkedBody owns the verified payload file until the outbound +// request has finished using it. Close is idempotent because the transport and +// the local cleanup path may both close the request body. +type stagedSignedChunkedBody struct { + *os.File + path string + reservation *signedStreamStageReservation + closeOnce sync.Once + closeErr error +} + +func (b *stagedSignedChunkedBody) Write(p []byte) (int, error) { + return b.reservation.write(b.File, p) +} + +func (b *stagedSignedChunkedBody) Close() error { + b.closeOnce.Do(func() { + b.reservation.mu.Lock() + defer b.reservation.mu.Unlock() + closeErr := b.File.Close() + removeErr := os.Remove(b.path) + removed := removeErr == nil || errors.Is(removeErr, os.ErrNotExist) + if errors.Is(removeErr, os.ErrNotExist) { + removeErr = nil + } + b.reservation.releaseLocked(removed) + b.closeErr = errors.Join(closeErr, removeErr) + }) + return b.closeErr +} + +// closeStagedSignedChunkedBody runs as the forwarding method's deferred owner. +// A cleanup error cannot undo a request already sent upstream, so report it +// rather than silently leaving sensitive request bytes on disk. +func closeStagedSignedChunkedBody(body *stagedSignedChunkedBody) { + if body == nil { + return + } + if err := body.Close(); err != nil { + log.Error().Err(err).Msg("Failed to clean up verified AWS streaming payload") + } +} + +// stagedTeeReadCloser preserves the staging file's Close ownership while the +// outbound transport reads through the write-through tee. +type stagedTeeReadCloser struct { + io.Reader + io.Closer +} + +// handoffStagedSignedChunkedBody transfers a non-empty staged body to the +// outbound request. A zero-length request is rewritten to http.NoBody, so its +// staging file has no transport owner and is closed here instead. +func handoffStagedSignedChunkedBody(body **stagedSignedChunkedBody, request *http.Request) { + if *body == nil { + return + } + if request.Body == http.NoBody { + closeStagedSignedChunkedBody(*body) + } + *body = nil +} + +// cleanupSignedStreamIngress closes the consumed body and interrupts any incomplete +// read before returning control to net/http. A read deadline wakes its HTTP/1 body +// reader without calling Close while that reader holds its lock. +func cleanupSignedStreamIngress(ctx context.Context, w http.ResponseWriter, r *http.Request, body io.ReadCloser) func(bool) { + if body == nil { + return func(bool) {} + } + var ( + closeOnce sync.Once + interruptOnce sync.Once + ) + closeBody := func() { + closeOnce.Do(func() { _ = body.Close() }) + } + interruptRead := func() { + interruptOnce.Do(func() { + if w != nil && r != nil && r.ProtoMajor == 1 { + w.Header().Set("Connection", "close") + } + if w != nil { + if err := http.NewResponseController(w).SetReadDeadline(time.Now()); err == nil { + return + } + } + closeBody() + }) + } + callbackDone := make(chan struct{}) + stop := context.AfterFunc(ctx, func() { + defer close(callbackDone) + interruptRead() + }) + return func(bodyComplete bool) { + if !bodyComplete { + interruptRead() + } + // Keep cancellation armed while Close runs: HTTP/1 Close can itself drain + // the remainder of an incomplete body and block on the same request read. + closeBody() + if !stop() { + <-callbackDone + } + } +} + +// stageSignedAWSChunkedBody authenticates every SigV4 chunk and the terminal +// chunk before returning a seeked, cleanup-owned file containing only payload +// bytes. It reserves staging space before copying chunk data and does not expose +// the body to the upstream request until the complete chain has been verified. +func stageSignedAWSChunkedBody(ctx context.Context, source io.Reader, signingKey []byte, signingTime, credentialScope, seedSignature string, declaredLength int64, budget *signedStreamStageBudget) (*stagedSignedChunkedBody, error) { + if err := ctx.Err(); err != nil { + return nil, fmt.Errorf("read signed AWS stream: %w", err) + } + if budget == nil { + budget = defaultSignedStreamStageBudget + } + dir := os.TempDir() + reservation := &signedStreamStageReservation{budget: budget, dir: dir} + if declaredLength >= 0 { + if err := reservation.reservePayload(declaredLength); err != nil { + return nil, fmt.Errorf("reserve temporary storage for signed AWS stream: %w", err) + } + } + file, err := os.CreateTemp(dir, "tag-signed-stream-*") + if err != nil { + reservation.release(true) + return nil, fmt.Errorf("create temporary file for signed AWS stream: %w", err) + } + staged := &stagedSignedChunkedBody{File: file, path: file.Name(), reservation: reservation} + reader := bufio.NewReaderSize(source, awsChunkedReaderBufSize) + previousSignature := seedSignature + decodedLength := int64(0) + + for { + if err := ctx.Err(); err != nil { + return nil, cleanupFailedSignedStage(staged, fmt.Errorf("read signed AWS stream: %w", err)) + } + line, err := readSignedChunkLine(reader) + if err != nil { + return nil, cleanupFailedSignedStage(staged, fmt.Errorf("read signed AWS chunk header: %w", signedStreamReadError(ctx, err))) + } + chunkLength, signature, err := parseSignedChunkHeader(line) + if err != nil { + return nil, cleanupFailedSignedStage(staged, err) + } + + if chunkLength == 0 { + expected := signedChunkSignature(signingKey, signingTime, credentialScope, previousSignature, sha256HexBytes(nil)) + if !hmac.Equal([]byte(signature), []byte(expected)) { + return nil, cleanupFailedSignedStage(staged, auth.ErrSignatureMismatch) + } + if err := readSignedChunkTerminator(reader); err != nil { + return nil, cleanupFailedSignedStage(staged, fmt.Errorf("read signed AWS terminal chunk terminator: %w", signedStreamReadError(ctx, err))) + } + if _, err := reader.ReadByte(); err == nil { + return nil, cleanupFailedSignedStage(staged, errors.New("signed AWS stream has data after its terminal chunk")) + } else if !errors.Is(err, io.EOF) { + return nil, cleanupFailedSignedStage(staged, fmt.Errorf("finish signed AWS stream: %w", signedStreamReadError(ctx, err))) + } + if err := ctx.Err(); err != nil { + return nil, cleanupFailedSignedStage(staged, fmt.Errorf("finish signed AWS stream: %w", err)) + } + if declaredLength >= 0 && decodedLength != declaredLength { + return nil, cleanupFailedSignedStage(staged, fmt.Errorf("%w: received %d bytes, declared %d", errSignedStreamLengthMismatch, decodedLength, declaredLength)) + } + break + } + + if chunkLength > math.MaxInt64-decodedLength { + return nil, cleanupFailedSignedStage(staged, errors.New("signed AWS stream payload length overflows int64")) + } + nextLength := decodedLength + chunkLength + if declaredLength < 0 || nextLength > declaredLength { + if err := reservation.reservePayload(nextLength); err != nil { + return nil, cleanupFailedSignedStage(staged, fmt.Errorf("reserve temporary storage for signed AWS stream: %w", err)) + } + } else if err := reservation.checkAvailable(); err != nil { + return nil, cleanupFailedSignedStage(staged, fmt.Errorf("reserve temporary storage for signed AWS stream: %w", err)) + } + + chunkHash := sha256.New() + written, copyErr := io.CopyN(io.MultiWriter(staged, chunkHash), reader, chunkLength) + if copyErr != nil { + return nil, cleanupFailedSignedStage(staged, fmt.Errorf("stage signed AWS chunk data: %w", signedStreamReadError(ctx, copyErr))) + } + if written != chunkLength { + return nil, cleanupFailedSignedStage(staged, io.ErrUnexpectedEOF) + } + if err := ctx.Err(); err != nil { + return nil, cleanupFailedSignedStage(staged, fmt.Errorf("stage signed AWS stream: %w", err)) + } + decodedLength += written + + expected := signedChunkSignature(signingKey, signingTime, credentialScope, previousSignature, hex.EncodeToString(chunkHash.Sum(nil))) + if !hmac.Equal([]byte(signature), []byte(expected)) { + return nil, cleanupFailedSignedStage(staged, auth.ErrSignatureMismatch) + } + previousSignature = signature + if err := readSignedChunkTerminator(reader); err != nil { + return nil, cleanupFailedSignedStage(staged, fmt.Errorf("read signed AWS chunk terminator: %w", signedStreamReadError(ctx, err))) + } + } + + if err := reservation.checkAvailable(); err != nil { + return nil, cleanupFailedSignedStage(staged, fmt.Errorf("reserve temporary storage for signed AWS stream: %w", err)) + } + if _, err := file.Seek(0, io.SeekStart); err != nil { + return nil, cleanupFailedSignedStage(staged, fmt.Errorf("rewind verified AWS stream: %w", err)) + } + return staged, nil +} + +func cleanupFailedSignedStage(staged *stagedSignedChunkedBody, cause error) error { + if cleanupErr := staged.Close(); cleanupErr != nil { + return errors.Join(cause, fmt.Errorf("clean up signed AWS stream staging file: %w", cleanupErr)) + } + return cause +} + +func signedStreamReadError(ctx context.Context, err error) error { + if ctxErr := ctx.Err(); ctxErr != nil { + return ctxErr + } + return err +} + +func readSignedChunkLine(reader *bufio.Reader) (string, error) { + line := make([]byte, 0, 64) + for { + b, err := reader.ReadByte() + if err != nil { + return "", err + } + if b == '\n' { + if len(line) == 0 || line[len(line)-1] != '\r' { + return "", errors.New("signed AWS chunk header is not CRLF-terminated") + } + return string(line[:len(line)-1]), nil + } + line = append(line, b) + if len(line) > maxChunkHeaderLen { + return "", fmt.Errorf("signed AWS chunk header exceeds maximum length (%d bytes)", maxChunkHeaderLen) + } + } +} + +func parseSignedChunkHeader(line string) (int64, string, error) { + sizeText, extensions, hasExtensions := strings.Cut(line, ";") + size, err := strconv.ParseUint(strings.TrimSpace(sizeText), 16, 64) + if err != nil { + return 0, "", fmt.Errorf("parsing signed AWS chunk size %q: %w", sizeText, err) + } + if size > math.MaxInt64 { + return 0, "", fmt.Errorf("signed AWS chunk size %d exceeds int64", size) + } + + var signature string + haveSignature := false + if hasExtensions { + for _, extension := range strings.Split(extensions, ";") { + name, value, hasValue := strings.Cut(extension, "=") + if name != "chunk-signature" { + continue + } + if haveSignature || !hasValue || len(value) != sha256.Size*2 { + return 0, "", auth.ErrSignatureMismatch + } + if _, err := hex.DecodeString(value); err != nil || strings.ToLower(value) != value { + return 0, "", auth.ErrSignatureMismatch + } + signature = value + haveSignature = true + } + } + if !haveSignature { + return 0, "", auth.ErrSignatureMismatch + } + return int64(size), signature, nil +} + +func readSignedChunkTerminator(reader *bufio.Reader) error { + var terminator [2]byte + if _, err := io.ReadFull(reader, terminator[:]); err != nil { + return err + } + if terminator != [2]byte{'\r', '\n'} { + return fmt.Errorf("invalid signed AWS chunk terminator %q", terminator) + } + return nil +} + +func signedChunkSignature(signingKey []byte, signingTime, credentialScope, previousSignature, chunkHash string) string { + stringToSign := strings.Join([]string{ + signedChunkStringToSignPrefix, + signingTime, + credentialScope, + previousSignature, + sha256HexBytes(nil), + chunkHash, + }, "\n") + mac := hmac.New(sha256.New, signingKey) + _, _ = mac.Write([]byte(stringToSign)) + return hex.EncodeToString(mac.Sum(nil)) +} + +func sha256HexBytes(data []byte) string { + sum := sha256.Sum256(data) + return hex.EncodeToString(sum[:]) +} + +// decodeIncomingBody authenticates the only signed HMAC marker repaired by this +// path. Other supported streaming markers keep the existing decoder behavior. +func (f *signingForwarder) decodeIncomingBody(ctx context.Context, w http.ResponseWriter, r *http.Request, accessKey string) (io.ReadCloser, string, int64, bool, *stagedSignedChunkedBody, error) { + if r.Header.Get("X-Amz-Content-Sha256") != StreamingPayloadHash { + body, bodyHash, contentLength, chunked := decodeChunkedIfNeeded(r) + return body, bodyHash, contentLength, chunked, nil, nil + } + + info, err := auth.ParseAuthInfo(r) + if err != nil { + return nil, "", 0, false, nil, mapAuthError(err) + } + if info.IsPresigned { + // Presigned validation is based on UNSIGNED-PAYLOAD and does not establish + // the header-signature seed required by the chunk HMAC chain. + body, bodyHash, contentLength, chunked := decodeChunkedIfNeeded(r) + return body, bodyHash, contentLength, chunked, nil, nil + } + + // Staging consumes the request body before forwarding. Abort incomplete ingress + // reads on cancellation or validation failure, then close the body exactly once. + bodyComplete := false + cleanupBody := cleanupSignedStreamIngress(ctx, w, r, r.Body) + defer func() { cleanupBody(bodyComplete) }() + + dateText := r.Header.Get("X-Amz-Date") + if dateText == "" { + dateText = r.Header.Get("Date") + } + requestTime, err := auth.ParseHTTPDate(dateText) + if err != nil { + return nil, "", 0, false, nil, mapAuthError(auth.ErrInvalidDate) + } + date := info.Date + if date == "" { + date = requestTime.UTC().Format("20060102") + } + signingKey, err := f.credStore.GetSigningKey(accessKey, date, info.Region) + if err != nil { + return nil, "", 0, false, nil, mapAuthError(err) + } + + decodedLength := int64(-1) + if value := r.Header.Get("X-Amz-Decoded-Content-Length"); value != "" { + if parsed, parseErr := strconv.ParseInt(value, 10, 64); parseErr == nil { + decodedLength = parsed + } + } + + staged, err := stageSignedAWSChunkedBody( + ctx, + r.Body, + signingKey, + requestTime.UTC().Format(auth.TimeFormat), + date+"/"+info.Region+"/s3/aws4_request", + info.Signature, + decodedLength, + defaultSignedStreamStageBudget, + ) + if err != nil { + if errors.Is(err, auth.ErrSignatureMismatch) { + return nil, "", 0, false, nil, mapAuthError(err) + } + return nil, "", 0, false, nil, fmt.Errorf("validate signed AWS streaming payload: %w", err) + } + + bodyComplete = true + return staged, "UNSIGNED-PAYLOAD", decodedLength, true, staged, nil +} diff --git a/proxy/signed_chunked_body_test.go b/proxy/signed_chunked_body_test.go new file mode 100644 index 00000000..484fa0c1 --- /dev/null +++ b/proxy/signed_chunked_body_test.go @@ -0,0 +1,1095 @@ +package proxy + +import ( + "bytes" + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/tigrisdata/tag/auth" +) + +func useSignedStreamStageBudget(t *testing.T, budget *signedStreamStageBudget) { + t.Helper() + previous := defaultSignedStreamStageBudget + defaultSignedStreamStageBudget = budget + t.Cleanup(func() { defaultSignedStreamStageBudget = previous }) +} + +func TestStatusRecorderUnwrapsReadDeadline(t *testing.T) { + writer := &readDeadlineResponseWriter{ResponseWriter: httptest.NewRecorder()} + wrapped := &statusRecorder{ResponseWriter: writer} + deadline := time.Now() + if err := http.NewResponseController(wrapped).SetReadDeadline(deadline); err != nil { + t.Fatalf("set read deadline through statusRecorder: %v", err) + } + if !writer.deadline.Equal(deadline) { + t.Fatalf("underlying read deadline = %v, want %v", writer.deadline, deadline) + } +} + +type readDeadlineResponseWriter struct { + http.ResponseWriter + deadline time.Time +} + +func (w *readDeadlineResponseWriter) SetReadDeadline(deadline time.Time) error { + w.deadline = deadline + return nil +} + +type blockingReadDeadlineWriter struct { + http.ResponseWriter + body io.Closer + started chan struct{} + release chan struct{} + start sync.Once + unblock sync.Once +} + +func (w *blockingReadDeadlineWriter) SetReadDeadline(time.Time) error { + w.start.Do(func() { close(w.started) }) + _ = w.body.Close() + <-w.release + return nil +} + +func (w *blockingReadDeadlineWriter) releaseDeadline() { + w.unblock.Do(func() { close(w.release) }) +} + +func TestSigningForwarderWaitsForCancellationCallbackBeforeReturning(t *testing.T) { + tempDir := t.TempDir() + t.Setenv("TMPDIR", tempDir) + budget := newSignedStreamStageBudget(func(string) (signedStreamStageSpace, error) { + return signedStreamStageSpace{availableBytes: 1 << 20, blockSize: 1}, nil + }) + credentials := auth.NewCredentialStore() + credentials.AddCredential(signedStreamTestAccessKey, signedStreamTestSecretKey) + forwarder := NewForwarder(credentials, "http://stage.test", "us-east-1", 1, nil, nil).(*signingForwarder) + useSignedStreamStageBudget(t, budget) + body := &blockingSignedStreamBody{readStarted: make(chan struct{}), closed: make(chan struct{})} + writer := &blockingReadDeadlineWriter{ + ResponseWriter: httptest.NewRecorder(), + body: body, + started: make(chan struct{}), + release: make(chan struct{}), + } + ctx, cancel := context.WithCancel(t.Context()) + defer func() { + cancel() + _ = body.Close() + writer.releaseDeadline() + }() + payload := []byte("staging before canceled read") + seed := newSignedStreamSeed(t, "http://stage.test", StreamingPayloadHash, len(payload), true) + stream := makeSignedStreamWire(t, seed, payload) + request := seed.request(t, stream.wire).WithContext(ctx) + request.Body = body + done := make(chan error, 1) + go func() { done <- forwarder.Forward(ctx, writer, request) }() + select { + case <-body.readStarted: + case <-time.After(5 * time.Second): + t.Fatal("forwarder did not begin the ingress body read") + } + budget.mu.Lock() + reservedBeforeCancel := budget.reserved + budget.mu.Unlock() + if reservedBeforeCancel == 0 { + t.Fatal("staging reservation was not held before the body read") + } + cancel() + select { + case <-writer.started: + case <-time.After(5 * time.Second): + t.Fatal("cancellation did not invoke the read-deadline callback") + } + select { + case err := <-done: + t.Fatalf("forwarding returned before the cancellation callback completed: %v", err) + default: + } + writer.releaseDeadline() + select { + case err := <-done: + if !errors.Is(err, context.Canceled) { + t.Fatalf("canceled forwarding error = %v, want context.Canceled", err) + } + case <-time.After(5 * time.Second): + t.Fatal("forwarding did not return after the cancellation callback completed") + } + budget.mu.Lock() + reservedAfterCancel := budget.reserved + budget.mu.Unlock() + if reservedAfterCancel != 0 { + t.Fatalf("canceled forwarding retained %d reserved bytes", reservedAfterCancel) + } + assertStageDirectoryEmpty(t, tempDir) +} + +func TestSigningForwarderStagingCleanup(t *testing.T) { + var dispatched atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + dispatched.Add(1) + if _, err := io.Copy(io.Discard, r.Body); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + credentials := auth.NewCredentialStore() + credentials.AddCredential(signedStreamTestAccessKey, signedStreamTestSecretKey) + forwarder := NewForwarder(credentials, upstream.URL, "us-east-1", 1, nil, nil) + payload := []byte("staged signed payload") + seed := newSignedStreamSeed(t, upstream.URL, StreamingPayloadHash, len(payload), true) + stream := makeSignedStreamWire(t, seed, payload) + + t.Run("success removes stage file", func(t *testing.T) { + tempDir := t.TempDir() + t.Setenv("TMPDIR", tempDir) + before := dispatched.Load() + if err := forwarder.Forward(t.Context(), httptest.NewRecorder(), seed.request(t, stream.wire)); err != nil { + t.Fatalf("forward valid stream: %v", err) + } + if dispatched.Load() != before+1 { + t.Fatalf("valid stream dispatch count = %d, want 1", dispatched.Load()-before) + } + assertStageDirectoryEmpty(t, tempDir) + }) + + t.Run("zero-byte stream removes its unused stage file", func(t *testing.T) { + tempDir := t.TempDir() + t.Setenv("TMPDIR", tempDir) + emptySeed := newSignedStreamSeed(t, upstream.URL, StreamingPayloadHash, 0, true) + emptyStream := makeSignedStreamWire(t, emptySeed) + before := dispatched.Load() + if err := forwarder.Forward(t.Context(), httptest.NewRecorder(), emptySeed.request(t, emptyStream.wire)); err != nil { + t.Fatalf("forward valid empty stream: %v", err) + } + if dispatched.Load() != before+1 { + t.Fatalf("empty stream dispatch count = %d, want 1", dispatched.Load()-before) + } + assertStageDirectoryEmpty(t, tempDir) + }) + + t.Run("invalid signature removes stage file", func(t *testing.T) { + tempDir := t.TempDir() + t.Setenv("TMPDIR", tempDir) + badWire := changeStreamSignature(t, stream.wire, stream.terminalSignatureStart) + before := dispatched.Load() + err := forwarder.Forward(t.Context(), httptest.NewRecorder(), seed.request(t, badWire)) + if authErr, ok := IsAuthError(err); !ok || authErr.Code != ErrCodeSignatureMismatch { + t.Fatalf("invalid terminal signature error = %v, want signature-mismatch AuthError", err) + } + if dispatched.Load() != before { + t.Fatalf("invalid terminal signature dispatched upstream %d times", dispatched.Load()-before) + } + assertStageDirectoryEmpty(t, tempDir) + }) + + t.Run("cancellation interrupts a blocked ingress read and removes stage file", func(t *testing.T) { + tempDir := t.TempDir() + t.Setenv("TMPDIR", tempDir) + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + body := &blockingSignedStreamBody{readStarted: make(chan struct{}), closed: make(chan struct{})} + defer body.Close() + req := seed.request(t, stream.wire).WithContext(ctx) + req.Body = body + before := dispatched.Load() + done := make(chan error, 1) + go func() { + done <- forwarder.Forward(ctx, httptest.NewRecorder(), req) + }() + select { + case <-body.readStarted: + case <-time.After(5 * time.Second): + t.Fatal("forwarder did not begin reading the signed stream") + } + cancel() + var err error + select { + case err = <-done: + case <-time.After(5 * time.Second): + t.Fatal("cancellation did not interrupt the signed-stream read") + } + if !errors.Is(err, context.Canceled) { + t.Fatalf("canceled staging error = %v, want context.Canceled", err) + } + if dispatched.Load() != before { + t.Fatalf("canceled stream dispatched upstream %d times", dispatched.Load()-before) + } + assertStageDirectoryEmpty(t, tempDir) + }) + + t.Run("temporary storage creation failure does not dispatch", func(t *testing.T) { + parent := t.TempDir() + blockedPath := filepath.Join(parent, "not-a-directory") + if err := os.WriteFile(blockedPath, []byte("sentinel"), 0o600); err != nil { + t.Fatalf("create blocked temp path: %v", err) + } + t.Setenv("TMPDIR", blockedPath) + before := dispatched.Load() + err := forwarder.Forward(t.Context(), httptest.NewRecorder(), seed.request(t, stream.wire)) + if err == nil { + t.Fatal("forward succeeded without a usable staging directory") + } + if dispatched.Load() != before { + t.Fatalf("staging creation failure dispatched upstream %d times", dispatched.Load()-before) + } + if data, readErr := os.ReadFile(blockedPath); readErr != nil || string(data) != "sentinel" { + t.Fatalf("blocked temp path changed: data=%q err=%v", data, readErr) + } + }) +} + +type blockingSignedStreamBody struct { + readStarted chan struct{} + closed chan struct{} + readOnce sync.Once + closeOnce sync.Once +} + +func (b *blockingSignedStreamBody) Read([]byte) (int, error) { + b.readOnce.Do(func() { close(b.readStarted) }) + <-b.closed + return 0, io.ErrClosedPipe +} + +func (b *blockingSignedStreamBody) Close() error { + b.closeOnce.Do(func() { close(b.closed) }) + return nil +} + +func TestSigningForwarderTransfersStagedBodyOwnership(t *testing.T) { + payload := []byte("transport-owned staged payload") + for _, route := range []string{"Forward", "ForwardWithCapture", "ForwardTeeingBody"} { + t.Run(route, func(t *testing.T) { + tempDir := t.TempDir() + t.Setenv("TMPDIR", tempDir) + upstream := httptest.NewServer(http.NotFoundHandler()) + defer upstream.Close() + + credentials := auth.NewCredentialStore() + credentials.AddCredential(signedStreamTestAccessKey, signedStreamTestSecretKey) + forwarder := NewForwarder(credentials, upstream.URL, "us-east-1", 1, nil, nil).(*signingForwarder) + transport := &delayedSignedBodyTransport{ + readStarted: make(chan struct{}), + release: make(chan struct{}), + completed: make(chan stagedBodyReadResult, 1), + } + defer func() { + select { + case <-transport.release: + default: + close(transport.release) + } + }() + forwarder.httpClient.Transport = transport + + seed := newSignedStreamSeed(t, upstream.URL, StreamingPayloadHash, len(payload), true) + stream := makeSignedStreamWire(t, seed, payload) + request := seed.request(t, stream.wire) + var err error + switch route { + case "Forward": + err = forwarder.Forward(t.Context(), httptest.NewRecorder(), request) + case "ForwardWithCapture": + _, err = forwarder.ForwardWithCapture(t.Context(), httptest.NewRecorder(), request) + case "ForwardTeeingBody": + var tee bytes.Buffer + _, _, _, _, err = forwarder.ForwardTeeingBody(t.Context(), httptest.NewRecorder(), request, &tee) + defer func() { + if err == nil && !bytes.Equal(tee.Bytes(), payload) { + t.Errorf("tee body = %q, want %q", tee.Bytes(), payload) + } + }() + } + if err != nil { + t.Fatalf("forward signed stream: %v", err) + } + select { + case <-transport.readStarted: + case <-time.After(5 * time.Second): + t.Fatal("round trip did not begin") + } + select { + case result := <-transport.completed: + t.Fatalf("transport consumed the body before release: %q, %v", result.body, result.readErr) + default: + } + + close(transport.release) + select { + case result := <-transport.completed: + if result.readErr != nil { + t.Fatalf("transport read staged body after forwarding returned: %v", result.readErr) + } + if result.closeErr != nil { + t.Fatalf("transport closed staged body: %v", result.closeErr) + } + if !bytes.Equal(result.body, payload) { + t.Fatalf("transport body = %q, want %q", result.body, payload) + } + case <-time.After(5 * time.Second): + t.Fatal("transport did not complete staged-body read") + } + assertStageDirectoryEmpty(t, tempDir) + }) + } +} + +func TestSigningForwarderRejectsSignedStreamLengthMismatch(t *testing.T) { + tempDir := t.TempDir() + t.Setenv("TMPDIR", tempDir) + var dispatched atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + dispatched.Add(1) + _, _ = io.Copy(io.Discard, r.Body) + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + credentials := auth.NewCredentialStore() + credentials.AddCredential(signedStreamTestAccessKey, signedStreamTestSecretKey) + forwarder := NewForwarder(credentials, upstream.URL, "us-east-1", 1, nil, nil).(*signingForwarder) + useSignedStreamStageBudget(t, newSignedStreamStageBudget(func(string) (signedStreamStageSpace, error) { + return signedStreamStageSpace{availableBytes: 1 << 30, blockSize: 1}, nil + })) + + methods := []struct { + name string + call func(*http.Request) error + }{ + { + name: "Forward", + call: func(r *http.Request) error { + return forwarder.Forward(r.Context(), httptest.NewRecorder(), r) + }, + }, + { + name: "ForwardWithCapture", + call: func(r *http.Request) error { + _, err := forwarder.ForwardWithCapture(r.Context(), httptest.NewRecorder(), r) + return err + }, + }, + { + name: "ForwardTeeingBody", + call: func(r *http.Request) error { + _, _, _, _, err := forwarder.ForwardTeeingBody(r.Context(), httptest.NewRecorder(), r, io.Discard) + return err + }, + }, + } + mismatches := []struct { + name string + declaredLength int + payload []byte + }{ + {name: "declared zero with nonempty payload", declaredLength: 0, payload: []byte("not empty")}, + {name: "declared length shorter than payload", declaredLength: 3, payload: []byte("longer")}, + {name: "declared length longer than payload", declaredLength: 8, payload: []byte("short")}, + } + for _, method := range methods { + t.Run(method.name, func(t *testing.T) { + for _, mismatch := range mismatches { + t.Run(mismatch.name, func(t *testing.T) { + seed := newSignedStreamSeed(t, upstream.URL, StreamingPayloadHash, mismatch.declaredLength, true) + stream := makeSignedStreamWire(t, seed, mismatch.payload) + before := dispatched.Load() + err := method.call(seed.request(t, stream.wire)) + if !errors.Is(err, errSignedStreamLengthMismatch) || dispatched.Load() != before { + t.Errorf("length mismatch err=%v, dispatches=%d; want length rejection before dispatch", err, dispatched.Load()-before) + } + }) + } + }) + } + assertStageDirectoryEmpty(t, tempDir) +} + +func TestSigningForwarderForwardsUnknownLengthSignedStream(t *testing.T) { + tempDir := t.TempDir() + t.Setenv("TMPDIR", tempDir) + var ( + dispatched atomic.Int32 + mu sync.Mutex + seen []signedStreamObservation + ) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + dispatched.Add(1) + body, err := io.ReadAll(r.Body) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + mu.Lock() + seen = append(seen, signedStreamObservation{ + body: body, + contentHash: r.Header.Get("X-Amz-Content-Sha256"), + contentLength: r.ContentLength, + decodedLength: r.Header.Get("X-Amz-Decoded-Content-Length"), + }) + mu.Unlock() + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + credentials := auth.NewCredentialStore() + credentials.AddCredential(signedStreamTestAccessKey, signedStreamTestSecretKey) + forwarder := NewForwarder(credentials, upstream.URL, "us-east-1", 1, nil, nil).(*signingForwarder) + useSignedStreamStageBudget(t, newSignedStreamStageBudget(func(string) (signedStreamStageSpace, error) { + return signedStreamStageSpace{availableBytes: 1 << 20, blockSize: 1}, nil + })) + payload := []byte("unknown decoded length") + seed := newSignedStreamSeed(t, upstream.URL, StreamingPayloadHash, -1, true) + stream := makeSignedStreamWire(t, seed, payload) + response := httptest.NewRecorder() + if err := forwarder.Forward(t.Context(), response, seed.request(t, stream.wire)); err != nil { + t.Fatalf("forward signed stream without decoded-length header: %v", err) + } + if response.Code != http.StatusOK || dispatched.Load() != 1 { + t.Fatalf("unknown-length response=%d dispatches=%d, want one successful dispatch", response.Code, dispatched.Load()) + } + got, ok := observationAt(&mu, &seen, 0) + if !ok || !bytes.Equal(got.body, payload) || got.contentHash != "UNSIGNED-PAYLOAD" || got.contentLength != -1 || got.decodedLength != "" { + t.Fatalf("unknown-length forwarded observation=%+v want exact bytes=%q and unknown length", got, payload) + } + assertStageDirectoryEmpty(t, tempDir) +} + +func TestSigningForwarderPreservesLargeChunkSegmentation(t *testing.T) { + tempDir := t.TempDir() + t.Setenv("TMPDIR", tempDir) + var ( + dispatched atomic.Int32 + mu sync.Mutex + seen []signedStreamObservation + ) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + dispatched.Add(1) + body, err := io.ReadAll(r.Body) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + mu.Lock() + seen = append(seen, signedStreamObservation{ + body: body, + contentHash: r.Header.Get("X-Amz-Content-Sha256"), + contentLength: r.ContentLength, + contentCoding: r.Header.Get("Content-Encoding"), + decodedLength: r.Header.Get("X-Amz-Decoded-Content-Length"), + }) + mu.Unlock() + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + credentials := auth.NewCredentialStore() + credentials.AddCredential(signedStreamTestAccessKey, signedStreamTestSecretKey) + forwarder := NewForwarder(credentials, upstream.URL, "us-east-1", 1, nil, nil) + useSignedStreamStageBudget(t, newSignedStreamStageBudget(func(string) (signedStreamStageSpace, error) { + return signedStreamStageSpace{availableBytes: 1 << 20, blockSize: 1}, nil + })) + payload := bytes.Repeat([]byte("x"), 256<<10) + seed := newSignedStreamSeed(t, upstream.URL, StreamingPayloadHash, len(payload), true) + t.Run("one large chunk", func(t *testing.T) { + stream := makeSignedStreamWire(t, seed, payload) + assertSignedStreamForwarded(t, func(r *http.Request) (int, error) { + w := httptest.NewRecorder() + err := forwarder.Forward(r.Context(), w, r) + return w.Code, err + }, seed, stream.wire, payload, &dispatched, &mu, &seen) + assertStageDirectoryEmpty(t, tempDir) + }) + t.Run("same payload split into smaller chunks", func(t *testing.T) { + chunks := [][]byte{payload[:64<<10], payload[64<<10 : 128<<10], payload[128<<10 : 192<<10], payload[192<<10:]} + stream := makeSignedStreamWire(t, seed, chunks...) + assertSignedStreamForwarded(t, func(r *http.Request) (int, error) { + w := httptest.NewRecorder() + err := forwarder.Forward(r.Context(), w, r) + return w.Code, err + }, seed, stream.wire, payload, &dispatched, &mu, &seen) + assertStageDirectoryEmpty(t, tempDir) + }) +} + +func TestSignedStreamStageBudgetSerializesConcurrentReservations(t *testing.T) { + budget := newSignedStreamStageBudget(func(string) (signedStreamStageSpace, error) { + return signedStreamStageSpace{availableBytes: 100, blockSize: 1}, nil + }) + dir := t.TempDir() + reservations := []*signedStreamStageReservation{ + {budget: budget, dir: dir}, + {budget: budget, dir: dir}, + } + start := make(chan struct{}) + ready := sync.WaitGroup{} + ready.Add(len(reservations)) + type result struct { + reservation *signedStreamStageReservation + err error + } + results := make(chan result, len(reservations)) + for _, reservation := range reservations { + go func(reservation *signedStreamStageReservation) { + ready.Done() + <-start + results <- result{reservation: reservation, err: reservation.reservePayload(40)} + }(reservation) + } + ready.Wait() + close(start) + + var accepted *signedStreamStageReservation + for range reservations { + result := <-results + if result.err == nil { + if accepted != nil { + t.Fatal("both concurrent requests reserved 40 bytes from a 50-byte pool") + } + accepted = result.reservation + } else if !errors.Is(result.err, errSignedStreamStagingCapacity) { + t.Fatalf("concurrent reservation error = %v, want capacity error", result.err) + } + } + if accepted == nil { + t.Fatal("neither request acquired the available staging reservation") + } + accepted.release(true) +} + +func TestSigningForwarderReservesKnownLengthBeforeReading(t *testing.T) { + tempDir := t.TempDir() + t.Setenv("TMPDIR", tempDir) + var dispatched atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + dispatched.Add(1) + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + budget := newSignedStreamStageBudget(func(string) (signedStreamStageSpace, error) { + return signedStreamStageSpace{availableBytes: 8, blockSize: 1}, nil + }) + credentials := auth.NewCredentialStore() + credentials.AddCredential(signedStreamTestAccessKey, signedStreamTestSecretKey) + forwarder := NewForwarder(credentials, upstream.URL, "us-east-1", 1, nil, nil).(*signingForwarder) + useSignedStreamStageBudget(t, budget) + + payload := []byte("12345") + seed := newSignedStreamSeed(t, upstream.URL, StreamingPayloadHash, len(payload), true) + stream := makeSignedStreamWire(t, seed, payload) + body := &readCountingSignedStreamBody{reader: bytes.NewReader(stream.wire)} + request := seed.request(t, stream.wire) + request.Body = body + err := forwarder.Forward(t.Context(), httptest.NewRecorder(), request) + if !errors.Is(err, errSignedStreamStagingCapacity) { + t.Fatalf("forward without enough staging reservation = %v, want capacity error", err) + } + if body.reads.Load() != 0 { + t.Fatalf("request body read %d times before its full decoded length was reserved", body.reads.Load()) + } + if dispatched.Load() != 0 { + t.Fatalf("capacity rejection dispatched upstream %d times", dispatched.Load()) + } + assertStageDirectoryEmpty(t, tempDir) +} + +func TestSigningForwarderReservesUnknownChunkBeforeCopy(t *testing.T) { + tempDir := t.TempDir() + t.Setenv("TMPDIR", tempDir) + var dispatched atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + dispatched.Add(1) + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + budget := newSignedStreamStageBudget(func(string) (signedStreamStageSpace, error) { + return signedStreamStageSpace{availableBytes: 8, blockSize: 1}, nil + }) + credentials := auth.NewCredentialStore() + credentials.AddCredential(signedStreamTestAccessKey, signedStreamTestSecretKey) + forwarder := NewForwarder(credentials, upstream.URL, "us-east-1", 1, nil, nil).(*signingForwarder) + useSignedStreamStageBudget(t, budget) + + payload := []byte("12345") + seed := newSignedStreamSeed(t, upstream.URL, StreamingPayloadHash, -1, true) + stream := makeSignedStreamWire(t, seed, payload) + headerEnd := bytes.Index(stream.wire, []byte("\r\n")) + len("\r\n") + body := &splitSignedStreamBody{first: stream.wire[:headerEnd], rest: stream.wire[headerEnd:]} + request := seed.request(t, stream.wire) + request.Body = body + err := forwarder.Forward(t.Context(), httptest.NewRecorder(), request) + if !errors.Is(err, errSignedStreamStagingCapacity) { + t.Fatalf("forward unknown-length chunk without a reservation = %v, want capacity error", err) + } + if body.reads.Load() != 1 { + t.Fatalf("unknown-length staging read request body %d times, want only the chunk header before reservation", body.reads.Load()) + } + if dispatched.Load() != 0 { + t.Fatalf("capacity rejection dispatched upstream %d times", dispatched.Load()) + } + assertStageDirectoryEmpty(t, tempDir) +} + +type splitSignedStreamBody struct { + first []byte + rest []byte + reads atomic.Int32 +} + +func (b *splitSignedStreamBody) Read(p []byte) (int, error) { + switch b.reads.Add(1) { + case 1: + return copy(p, b.first), nil + case 2: + return copy(p, b.rest), nil + default: + return 0, io.EOF + } +} + +func (*splitSignedStreamBody) Close() error { return nil } + +func TestSigningForwarderAccountsForPartialChunkWritesDuringConcurrentAdmission(t *testing.T) { + tempDir := t.TempDir() + t.Setenv("TMPDIR", tempDir) + var ( + dispatched atomic.Int32 + mu sync.Mutex + bodies [][]byte + ) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + dispatched.Add(1) + body, err := io.ReadAll(r.Body) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + mu.Lock() + bodies = append(bodies, body) + mu.Unlock() + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + budget := newSignedStreamStageBudget(func(dir string) (signedStreamStageSpace, error) { + entries, err := os.ReadDir(dir) + if err != nil { + return signedStreamStageSpace{}, err + } + var used int64 + for _, entry := range entries { + info, err := entry.Info() + if err != nil { + return signedStreamStageSpace{}, err + } + used += info.Size() + } + return signedStreamStageSpace{availableBytes: 100 - used, blockSize: 1}, nil + }) + credentials := auth.NewCredentialStore() + credentials.AddCredential(signedStreamTestAccessKey, signedStreamTestSecretKey) + forwarder := NewForwarder(credentials, upstream.URL, "us-east-1", 1, nil, nil).(*signingForwarder) + useSignedStreamStageBudget(t, budget) + + firstPayload := bytes.Repeat([]byte("a"), 40) + firstSeed := newSignedStreamSeed(t, upstream.URL, StreamingPayloadHash, len(firstPayload), true) + firstWire := makeSignedStreamWire(t, firstSeed, firstPayload) + headerEnd := bytes.Index(firstWire.wire, []byte("\r\n")) + len("\r\n") + dataEnd := headerEnd + 30 + firstBody := &blockedSignedStreamBody{ + header: firstWire.wire[:headerEnd], + partial: firstWire.wire[headerEnd:dataEnd], + rest: firstWire.wire[dataEnd:], + blocked: make(chan struct{}), + release: make(chan struct{}), + } + defer firstBody.unblock() + firstRequest := firstSeed.request(t, firstWire.wire) + firstRequest.Body = firstBody + firstDone := make(chan error, 1) + go func() { + firstDone <- forwarder.Forward(t.Context(), httptest.NewRecorder(), firstRequest) + }() + select { + case <-firstBody.blocked: + case <-time.After(5 * time.Second): + t.Fatal("first upload did not block after writing its partial chunk") + } + if dispatched.Load() != 0 { + t.Fatalf("incomplete first stream dispatched upstream %d times", dispatched.Load()) + } + + secondPayload := bytes.Repeat([]byte("b"), 10) + secondSeed := newSignedStreamSeed(t, upstream.URL, StreamingPayloadHash, len(secondPayload), true) + secondWire := makeSignedStreamWire(t, secondSeed, secondPayload) + if err := forwarder.Forward(t.Context(), httptest.NewRecorder(), secondSeed.request(t, secondWire.wire)); err != nil { + t.Fatalf("admit second upload while first is ingress-blocked: %v", err) + } + if dispatched.Load() != 1 { + t.Fatalf("second upload dispatch count = %d, want 1", dispatched.Load()) + } + + firstBody.unblock() + select { + case err := <-firstDone: + if err != nil { + t.Fatalf("finish first staged upload: %v", err) + } + case <-time.After(5 * time.Second): + t.Fatal("first upload did not finish after ingress release") + } + if dispatched.Load() != 2 { + t.Fatalf("both admitted upload dispatches = %d, want 2", dispatched.Load()) + } + mu.Lock() + defer mu.Unlock() + if len(bodies) != 2 || !bytes.Equal(bodies[0], secondPayload) || !bytes.Equal(bodies[1], firstPayload) { + t.Fatalf("upstream bodies = %q, want second %q then first %q", bodies, secondPayload, firstPayload) + } + assertStageDirectoryEmpty(t, tempDir) +} + +func TestSignedStreamStageWriteDoesNotHoldSharedBudgetLock(t *testing.T) { + dir := t.TempDir() + budget := newSignedStreamStageBudget(func(string) (signedStreamStageSpace, error) { + return signedStreamStageSpace{availableBytes: 100, blockSize: 1}, nil + }) + first := &signedStreamStageReservation{budget: budget, dir: dir} + if err := first.reservePayload(40); err != nil { + t.Fatalf("reserve first upload staging: %v", err) + } + + writer := &blockedSignedStreamStageWriter{started: make(chan struct{}), unblocked: make(chan struct{})} + firstDone := make(chan error, 1) + go func() { + payload := []byte("first payload") + n, err := first.write(writer, payload) + if err == nil && n != len(payload) { + err = errors.New("first staged write was short") + } + firstDone <- err + }() + select { + case <-writer.started: + case <-time.After(5 * time.Second): + writer.unblock() + t.Fatal("first upload did not enter its blocked filesystem write") + } + + overBudget := &signedStreamStageReservation{budget: budget, dir: dir} + if err := overBudget.reservePayload(11); !errors.Is(err, errSignedStreamStagingCapacity) { + t.Fatalf("reserve beyond the shared in-flight staging cap = %v, want capacity error", err) + } + + second := &signedStreamStageReservation{budget: budget, dir: dir} + secondDone := make(chan error, 1) + go func() { + if err := second.reservePayload(10); err != nil { + secondDone <- err + return + } + payload := []byte("second") + n, err := second.write(io.Discard, payload) + if err == nil && n != len(payload) { + err = errors.New("second staged write was short") + } + second.release(true) + secondDone <- err + }() + + blocked := false + var secondErr error + select { + case secondErr = <-secondDone: + case <-time.After(5 * time.Second): + blocked = true + } + writer.unblock() + if err := <-firstDone; err != nil { + t.Fatalf("finish first staged write: %v", err) + } + if blocked { + secondErr = <-secondDone + if secondErr != nil { + t.Fatalf("complete second reservation after unblocking first write: %v", secondErr) + } + t.Fatal("second upload waited for the first upload's filesystem write") + } + if secondErr != nil { + t.Fatalf("reserve, write, and release second upload during first write: %v", secondErr) + } + + first.release(true) + budget.mu.Lock() + reserved, unallocated := budget.reserved, budget.unallocated + budget.mu.Unlock() + if reserved != 0 || unallocated != 0 { + t.Fatalf("staging budget after both uploads: reserved=%d unallocated=%d, want 0", reserved, unallocated) + } +} + +func TestSignedStreamStageWriteAccountsForPartialFailures(t *testing.T) { + writeErr := io.ErrShortWrite + for _, tc := range []struct { + name string + written int + writeErr error + wantAllocated int64 + wantUnallocated int64 + }{ + {name: "partial write", written: 5, writeErr: writeErr, wantAllocated: 8, wantUnallocated: 8}, + {name: "no bytes written", written: 0, writeErr: writeErr, wantAllocated: 0, wantUnallocated: 16}, + } { + t.Run(tc.name, func(t *testing.T) { + budget := newSignedStreamStageBudget(func(string) (signedStreamStageSpace, error) { + return signedStreamStageSpace{availableBytes: 128, blockSize: 4}, nil + }) + reservation := &signedStreamStageReservation{budget: budget, dir: t.TempDir()} + if err := reservation.reservePayload(16); err != nil { + t.Fatalf("reserve staged body: %v", err) + } + + payload := bytes.Repeat([]byte("x"), 8) + writer := stagedWriteResult{written: tc.written, err: tc.writeErr} + n, err := reservation.write(writer, payload) + if n != tc.written || !errors.Is(err, tc.writeErr) { + t.Fatalf("staged write = (%d, %v), want (%d, %v)", n, err, tc.written, tc.writeErr) + } + + budget.mu.Lock() + reserved, unallocated := budget.reserved, budget.unallocated + budget.mu.Unlock() + if reserved != 16 || unallocated != tc.wantUnallocated || reservation.allocated != tc.wantAllocated { + t.Fatalf("partial-write accounting: reserved=%d unallocated=%d allocated=%d", reserved, unallocated, reservation.allocated) + } + reservation.release(true) + budget.mu.Lock() + reserved, unallocated = budget.reserved, budget.unallocated + budget.mu.Unlock() + if reserved != 0 || unallocated != 0 { + t.Fatalf("budget after cleanup: reserved=%d unallocated=%d, want 0", reserved, unallocated) + } + }) + } +} + +type stagedWriteResult struct { + written int + err error +} + +func (w stagedWriteResult) Write([]byte) (int, error) { return w.written, w.err } + +type blockedSignedStreamStageWriter struct { + started chan struct{} + unblocked chan struct{} + startOnce sync.Once + stopOnce sync.Once +} + +func (w *blockedSignedStreamStageWriter) Write(p []byte) (int, error) { + w.startOnce.Do(func() { close(w.started) }) + <-w.unblocked + return len(p), nil +} + +func (w *blockedSignedStreamStageWriter) unblock() { + w.stopOnce.Do(func() { close(w.unblocked) }) +} + +type blockedSignedStreamBody struct { + header []byte + partial []byte + rest []byte + reads int + blocked chan struct{} + release chan struct{} + once sync.Once +} + +func (b *blockedSignedStreamBody) Read(p []byte) (int, error) { + switch b.reads { + case 0: + b.reads++ + return copy(p, b.header), nil + case 1: + b.reads++ + return copy(p, b.partial), nil + case 2: + b.reads++ + close(b.blocked) + <-b.release + return copy(p, b.rest), nil + default: + return 0, io.EOF + } +} + +func (b *blockedSignedStreamBody) unblock() { + b.once.Do(func() { close(b.release) }) +} + +func (b *blockedSignedStreamBody) Close() error { + b.unblock() + return nil +} + +func TestSigningForwarderStagingReservationLivesUntilTransportClose(t *testing.T) { + tempDir := t.TempDir() + t.Setenv("TMPDIR", tempDir) + budget := newSignedStreamStageBudget(func(string) (signedStreamStageSpace, error) { + return signedStreamStageSpace{availableBytes: 100, blockSize: 1}, nil + }) + credentials := auth.NewCredentialStore() + credentials.AddCredential(signedStreamTestAccessKey, signedStreamTestSecretKey) + forwarder := NewForwarder(credentials, "http://stage.test", "us-east-1", 1, nil, nil).(*signingForwarder) + useSignedStreamStageBudget(t, budget) + transport := &holdingSignedBodyTransport{} + forwarder.httpClient.Transport = transport + + firstPayload := bytes.Repeat([]byte("a"), 40) + firstSeed := newSignedStreamSeed(t, "http://stage.test", StreamingPayloadHash, len(firstPayload), true) + firstWire := makeSignedStreamWire(t, firstSeed, firstPayload) + if err := forwarder.Forward(t.Context(), httptest.NewRecorder(), firstSeed.request(t, firstWire.wire)); err != nil { + t.Fatalf("forward first staged request: %v", err) + } + if transport.requestCount() != 1 { + t.Fatalf("first request dispatch count = %d, want 1", transport.requestCount()) + } + + secondPayload := bytes.Repeat([]byte("b"), 11) + secondSeed := newSignedStreamSeed(t, "http://stage.test", StreamingPayloadHash, len(secondPayload), true) + secondWire := makeSignedStreamWire(t, secondSeed, secondPayload) + secondRequest := secondSeed.request(t, secondWire.wire) + err := forwarder.Forward(t.Context(), httptest.NewRecorder(), secondRequest) + if !errors.Is(err, errSignedStreamStagingCapacity) { + t.Fatalf("concurrent request without staging reservation = %v, want capacity error", err) + } + if transport.requestCount() != 1 { + t.Fatalf("capacity-rejected request reached transport; dispatch count = %d", transport.requestCount()) + } + + firstRequest := transport.requestAt(0) + got, err := io.ReadAll(firstRequest.Body) + if err != nil || !bytes.Equal(got, firstPayload) { + t.Fatalf("transport-owned first body = %q, err=%v, want %q", got, err, firstPayload) + } + if err := firstRequest.Body.Close(); err != nil { + t.Fatalf("close first transport-owned body: %v", err) + } + assertStageDirectoryEmpty(t, tempDir) + + if err := forwarder.Forward(t.Context(), httptest.NewRecorder(), secondSeed.request(t, secondWire.wire)); err != nil { + t.Fatalf("forward second request after reservation release: %v", err) + } + if transport.requestCount() != 2 { + t.Fatalf("retry dispatch count = %d, want 2", transport.requestCount()) + } + secondForwarded := transport.requestAt(1) + got, err = io.ReadAll(secondForwarded.Body) + if err != nil || !bytes.Equal(got, secondPayload) { + t.Fatalf("transport-owned second body = %q, err=%v, want %q", got, err, secondPayload) + } + if err := secondForwarded.Body.Close(); err != nil { + t.Fatalf("close second transport-owned body: %v", err) + } + assertStageDirectoryEmpty(t, tempDir) +} + +type readCountingSignedStreamBody struct { + reader io.Reader + reads atomic.Int32 +} + +func (b *readCountingSignedStreamBody) Read(p []byte) (int, error) { + b.reads.Add(1) + return b.reader.Read(p) +} + +func (*readCountingSignedStreamBody) Close() error { return nil } + +type holdingSignedBodyTransport struct { + mu sync.Mutex + requests []*http.Request +} + +func (t *holdingSignedBodyTransport) RoundTrip(request *http.Request) (*http.Response, error) { + t.mu.Lock() + t.requests = append(t.requests, request) + t.mu.Unlock() + return &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: http.NoBody, + Request: request, + }, nil +} + +func (t *holdingSignedBodyTransport) requestCount() int { + t.mu.Lock() + defer t.mu.Unlock() + return len(t.requests) +} + +func (t *holdingSignedBodyTransport) requestAt(index int) *http.Request { + t.mu.Lock() + defer t.mu.Unlock() + return t.requests[index] +} + +type stagedBodyReadResult struct { + body []byte + readErr error + closeErr error +} + +type delayedSignedBodyTransport struct { + readStarted chan struct{} + release chan struct{} + completed chan stagedBodyReadResult +} + +func (t *delayedSignedBodyTransport) RoundTrip(request *http.Request) (*http.Response, error) { + body := request.Body + close(t.readStarted) + go func() { + <-t.release + data, readErr := io.ReadAll(body) + closeErr := body.Close() + t.completed <- stagedBodyReadResult{body: data, readErr: readErr, closeErr: closeErr} + }() + return &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: http.NoBody, + Request: request, + }, nil +} + +func assertStageDirectoryEmpty(t *testing.T, dir string) { + t.Helper() + entries, err := os.ReadDir(dir) + if err != nil { + t.Fatalf("read staging directory: %v", err) + } + if len(entries) != 0 { + t.Fatalf("staging directory retains files after request: %v", entries) + } +} diff --git a/proxy/signed_stream_http_cancel_test.go b/proxy/signed_stream_http_cancel_test.go new file mode 100644 index 00000000..d46c8582 --- /dev/null +++ b/proxy/signed_stream_http_cancel_test.go @@ -0,0 +1,417 @@ +package proxy + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "os" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/tigrisdata/tag/auth" +) + +type deadlineGateConnContextKey struct{} + +type signedStreamCancelResult struct { + err error + connectionHeader string + stageFiles int + reservedBytes int64 + readDirErr error +} + +type deadlineGateConn struct { + net.Conn + armed atomic.Bool + bodyReads atomic.Int32 + blockedRead chan struct{} + deadline chan struct{} + deadlineOne sync.Once +} + +func (c *deadlineGateConn) Read(p []byte) (int, error) { + if c.armed.Load() && c.bodyReads.Add(1) == 2 { + close(c.blockedRead) + <-c.deadline + return 0, os.ErrDeadlineExceeded + } + return c.Conn.Read(p) +} + +func (c *deadlineGateConn) unblockRead() { + c.deadlineOne.Do(func() { close(c.deadline) }) +} + +func (c *deadlineGateConn) SetReadDeadline(deadline time.Time) error { + if c.armed.Load() && !deadline.IsZero() && !deadline.After(time.Now()) { + c.unblockRead() + } + return c.Conn.SetReadDeadline(deadline) +} + +type deadlineGateListener struct { + net.Listener + accepted chan *deadlineGateConn +} + +func (l *deadlineGateListener) Accept() (net.Conn, error) { + conn, err := l.Listener.Accept() + if err != nil { + return nil, err + } + gated := &deadlineGateConn{ + Conn: conn, + blockedRead: make(chan struct{}), + deadline: make(chan struct{}), + } + l.accepted <- gated + return gated, nil +} + +func TestSigningForwarderHMACFailureInterruptsRealHTTP1BodyClose(t *testing.T) { + var upstreamCalls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamCalls.Add(1) + _, _ = io.Copy(io.Discard, r.Body) + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + tempDir := t.TempDir() + t.Setenv("TMPDIR", tempDir) + budget := newSignedStreamStageBudget(func(string) (signedStreamStageSpace, error) { + return signedStreamStageSpace{availableBytes: 1 << 20, blockSize: 1}, nil + }) + credentials := auth.NewCredentialStore() + credentials.AddCredential(signedStreamTestAccessKey, signedStreamTestSecretKey) + forwarder := NewForwarder(credentials, upstream.URL, "us-east-1", 1, nil, nil).(*signingForwarder) + useSignedStreamStageBudget(t, budget) + + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen for ingress: %v", err) + } + gatedListener := &deadlineGateListener{Listener: listener, accepted: make(chan *deadlineGateConn, 1)} + handlerStarted := make(chan struct{}) + handlerDone := make(chan struct{}) + releaseHandler := make(chan struct{}) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(releaseHandler) }) } + type result struct { + err error + connectionHeader string + stageFiles int + reservedBytes int64 + readDirErr error + } + resultCh := make(chan result, 1) + server := &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gatedConn := r.Context().Value(deadlineGateConnContextKey{}).(*deadlineGateConn) + gatedConn.armed.Store(true) + close(handlerStarted) + err := forwarder.Forward(r.Context(), w, r) + entries, readDirErr := os.ReadDir(tempDir) + budget.mu.Lock() + reserved := budget.reserved + budget.mu.Unlock() + resultCh <- result{err: err, connectionHeader: w.Header().Get("Connection"), stageFiles: len(entries), reservedBytes: reserved, readDirErr: readDirErr} + <-releaseHandler + if err != nil { + w.WriteHeader(http.StatusBadRequest) + } + close(handlerDone) + }), ConnContext: func(ctx context.Context, conn net.Conn) context.Context { + return context.WithValue(ctx, deadlineGateConnContextKey{}, conn.(*deadlineGateConn)) + }} + serveDone := make(chan error, 1) + go func() { serveDone <- server.Serve(gatedListener) }() + var conn net.Conn + var gatedConn *deadlineGateConn + defer func() { + release() + if gatedConn != nil { + gatedConn.unblockRead() + } + if conn != nil { + _ = conn.Close() + } + _ = server.Close() + _ = gatedListener.Close() + <-serveDone + }() + conn, err = net.Dial("tcp", listener.Addr().String()) + if err != nil { + t.Fatalf("dial ingress: %v", err) + } + gatedConn = <-gatedListener.accepted + + payload := []byte("0123456789") + seed := newSignedStreamSeed(t, "http://"+listener.Addr().String(), StreamingPayloadHash, len(payload), true) + stream := makeSignedStreamWire(t, seed, payload) + badWire := changeStreamSignature(t, stream.wire, stream.dataSignatureStart[0]) + headerEnd := bytes.Index(badWire, []byte("\r\n")) + len("\r\n") + bodyEnd := headerEnd + len(payload) + len("\r\n") + var request bytes.Buffer + if _, err := fmt.Fprintf(&request, "PUT %s HTTP/1.1\r\nHost: %s\r\n", seed.path, seed.host); err != nil { + t.Fatalf("write request line: %v", err) + } + for name, values := range seed.headers { + if name == "Host" { + continue + } + for _, value := range values { + if _, err := fmt.Fprintf(&request, "%s: %s\r\n", name, value); err != nil { + t.Fatalf("write request header: %v", err) + } + } + } + if _, err := fmt.Fprintf(&request, "Content-Length: %d\r\nConnection: keep-alive\r\n\r\n", len(badWire)); err != nil { + t.Fatalf("write request framing: %v", err) + } + if _, err := conn.Write(request.Bytes()); err != nil { + t.Fatalf("write request headers: %v", err) + } + select { + case <-handlerStarted: + case <-time.After(5 * time.Second): + t.Fatal("HTTP server did not start request handling") + } + if _, err := conn.Write(badWire[:bodyEnd]); err != nil { + t.Fatalf("write signed chunk with bad HMAC: %v", err) + } + select { + case <-gatedConn.blockedRead: + case <-time.After(5 * time.Second): + t.Fatal("body Close did not reach the unread remainder after HMAC rejection") + } + select { + case got := <-resultCh: + if authErr, ok := IsAuthError(got.err); !ok || authErr.Code != ErrCodeSignatureMismatch { + t.Fatalf("bad chunk signature error = %v, want signature mismatch", got.err) + } + if got.connectionHeader != "close" { + t.Fatalf("rejected HTTP/1 response Connection = %q, want close", got.connectionHeader) + } + if got.readDirErr != nil || got.stageFiles != 0 || got.reservedBytes != 0 { + t.Fatalf("rejected stream retained staging resources: files=%d reserved=%d err=%v", got.stageFiles, got.reservedBytes, got.readDirErr) + } + case <-time.After(5 * time.Second): + t.Fatal("bad HMAC rejection blocked while closing the incomplete request body") + } + if upstreamCalls.Load() != 0 { + t.Fatalf("bad HMAC stream dispatched upstream %d times", upstreamCalls.Load()) + } + release() + _ = conn.Close() + select { + case <-handlerDone: + case <-time.After(5 * time.Second): + t.Fatal("HTTP/1 handler did not finish after rejecting the stream") + } +} + +func TestSigningForwarderCancellationInterruptsRealHTTP1Body(t *testing.T) { + var upstreamCalls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamCalls.Add(1) + _, _ = io.Copy(io.Discard, r.Body) + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + methods := []struct { + name string + call func(*signingForwarder, context.Context, http.ResponseWriter, *http.Request) error + }{ + { + name: "Forward", + call: func(f *signingForwarder, ctx context.Context, w http.ResponseWriter, r *http.Request) error { + return f.Forward(ctx, w, r) + }, + }, + { + name: "ForwardWithCapture", + call: func(f *signingForwarder, ctx context.Context, w http.ResponseWriter, r *http.Request) error { + _, err := f.ForwardWithCapture(ctx, w, r) + return err + }, + }, + { + name: "ForwardTeeingBody", + call: func(f *signingForwarder, ctx context.Context, w http.ResponseWriter, r *http.Request) error { + _, _, _, _, err := f.ForwardTeeingBody(ctx, w, r, io.Discard) + return err + }, + }, + } + for _, method := range methods { + method := method + t.Run(method.name, func(t *testing.T) { + tempDir := t.TempDir() + t.Setenv("TMPDIR", tempDir) + budget := newSignedStreamStageBudget(func(string) (signedStreamStageSpace, error) { + return signedStreamStageSpace{availableBytes: 1 << 20, blockSize: 1}, nil + }) + credentials := auth.NewCredentialStore() + credentials.AddCredential(signedStreamTestAccessKey, signedStreamTestSecretKey) + forwarder := NewForwarder(credentials, upstream.URL, "us-east-1", 1, nil, nil).(*signingForwarder) + useSignedStreamStageBudget(t, budget) + + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen for ingress: %v", err) + } + gatedListener := &deadlineGateListener{Listener: listener, accepted: make(chan *deadlineGateConn, 1)} + cancelRequest := make(chan context.CancelFunc, 1) + handlerStarted := make(chan struct{}) + handlerDone := make(chan struct{}) + releaseHandler := make(chan struct{}) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(releaseHandler) }) } + resultCh := make(chan signedStreamCancelResult, 1) + server := &http.Server{ + Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gatedConn := r.Context().Value(deadlineGateConnContextKey{}).(*deadlineGateConn) + gatedConn.armed.Store(true) + close(handlerStarted) + ctx, cancel := context.WithCancel(r.Context()) + cancelRequest <- cancel + err := method.call(forwarder, ctx, w, r) + entries, readDirErr := os.ReadDir(tempDir) + budget.mu.Lock() + reserved := budget.reserved + budget.mu.Unlock() + resultCh <- signedStreamCancelResult{ + err: err, + connectionHeader: w.Header().Get("Connection"), + stageFiles: len(entries), + reservedBytes: reserved, + readDirErr: readDirErr, + } + <-releaseHandler + if err != nil { + w.WriteHeader(http.StatusBadRequest) + } + close(handlerDone) + }), + ConnContext: func(ctx context.Context, conn net.Conn) context.Context { + return context.WithValue(ctx, deadlineGateConnContextKey{}, conn.(*deadlineGateConn)) + }, + } + serveDone := make(chan error, 1) + go func() { serveDone <- server.Serve(gatedListener) }() + var ( + conn net.Conn + gatedConn *deadlineGateConn + ) + defer func() { + release() + if gatedConn != nil { + gatedConn.unblockRead() + } + if conn != nil { + _ = conn.Close() + } + _ = server.Close() + _ = gatedListener.Close() + <-serveDone + }() + conn, err = net.Dial("tcp", listener.Addr().String()) + if err != nil { + t.Fatalf("dial ingress: %v", err) + } + gatedConn = <-gatedListener.accepted + + payload := []byte("0123456789") + seed := newSignedStreamSeed(t, "http://"+listener.Addr().String(), StreamingPayloadHash, len(payload), true) + stream := makeSignedStreamWire(t, seed, payload) + headerEnd := bytes.Index(stream.wire, []byte("\r\n")) + len("\r\n") + if headerEnd < len("\r\n") || headerEnd+2 >= len(stream.wire) { + t.Fatalf("invalid signed-stream test wire length %d", len(stream.wire)) + } + + var request bytes.Buffer + if _, err := fmt.Fprintf(&request, "PUT %s HTTP/1.1\r\nHost: %s\r\n", seed.path, seed.host); err != nil { + t.Fatalf("write request line: %v", err) + } + for name, values := range seed.headers { + if name == "Host" { + continue + } + for _, value := range values { + if _, err := fmt.Fprintf(&request, "%s: %s\r\n", name, value); err != nil { + t.Fatalf("write request header: %v", err) + } + } + } + if _, err := fmt.Fprintf(&request, "Content-Length: %d\r\nConnection: keep-alive\r\n\r\n", len(stream.wire)); err != nil { + t.Fatalf("write request framing: %v", err) + } + if _, err := conn.Write(request.Bytes()); err != nil { + t.Fatalf("write request headers: %v", err) + } + select { + case <-handlerStarted: + case <-time.After(5 * time.Second): + t.Fatal("HTTP server did not start request handling") + } + if _, err := conn.Write(stream.wire[:headerEnd+2]); err != nil { + t.Fatalf("write partial signed stream: %v", err) + } + cancel := <-cancelRequest + select { + case <-gatedConn.blockedRead: + case <-time.After(5 * time.Second): + t.Fatal("HTTP/1 request body did not block inside its second read") + } + entries, err := os.ReadDir(tempDir) + if err != nil { + t.Fatalf("read active staging directory: %v", err) + } + if len(entries) == 0 { + t.Fatal("staging file was not owned while the ingress body was blocked") + } + budget.mu.Lock() + reservedBeforeCancel := budget.reserved + budget.mu.Unlock() + if reservedBeforeCancel == 0 { + t.Fatal("staging reservation was not held while the ingress body was blocked") + } + + cancel() + var result signedStreamCancelResult + select { + case result = <-resultCh: + case <-time.After(5 * time.Second): + t.Fatal("canceled forwarding remained blocked until the client closed ingress") + } + if !errors.Is(result.err, context.Canceled) { + t.Fatalf("canceled forwarding error = %v, want context.Canceled", result.err) + } + if result.connectionHeader != "close" { + t.Fatalf("canceled HTTP/1 response Connection = %q, want close", result.connectionHeader) + } + if result.readDirErr != nil || result.stageFiles != 0 || result.reservedBytes != 0 { + t.Fatalf("cancellation retained staging resources: files=%d reserved=%d err=%v", result.stageFiles, result.reservedBytes, result.readDirErr) + } + if upstreamCalls.Load() != 0 { + t.Fatalf("canceled signed stream dispatched upstream %d times", upstreamCalls.Load()) + } + release() + _ = conn.Close() + select { + case <-handlerDone: + case <-time.After(5 * time.Second): + t.Fatal("HTTP/1 handler did not finish after ingress cleanup") + } + }) + } +}