diff --git a/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go b/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go index 4049ef2e70..7eca7406e8 100644 --- a/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go +++ b/packages/orchestrator/pkg/sandbox/block/streaming_chunk.go @@ -290,13 +290,18 @@ func (c *Chunker) progressiveFetch(ctx context.Context, s *fetchSession, mmapSli var closeErr error res.stats, closeErr = reader.Close(context.WithoutCancel(ctx)) + // A compressed frame's CRC is only verified once its footer is consumed + // on Close, so a Close error is a verification failure that must fail + // the fetch (don't cache or release a corrupt frame). A read error, if + // any, takes precedence. + if err == nil { + err = closeErr + } + ct := ft.CompressionType() attrs := storage.OKAttrs(c.objType, source, ct) - switch { - case err != nil: + if err != nil { attrs = storage.ErrAttrs(c.objType, source, ct, err) - case closeErr != nil: - attrs = storage.ErrAttrs(c.objType, source, ct, closeErr) } var readDur time.Duration @@ -318,9 +323,14 @@ func (c *Chunker) progressiveFetch(ctx context.Context, s *fetchSession, mmapSli n, readErr := io.ReadFull(reader, mmapSlice[totalRead:readEnd]) totalRead += int64(n) - if n > 0 { + if n > 0 && !ft.IsCompressed() { // Dirty marking is deferred to runFetch after the full chunk is fetched. // With coarse dirty granularity, marking here would expose partially-written data. + // + // Compressed chunks are a single frame whose CRC is only verified + // once the footer is consumed on Close. Releasing waiters here + // would hand out bytes that a later CRC failure proves corrupt, so + // their release is deferred to setDone after verification. s.advance(totalRead) } diff --git a/packages/orchestrator/pkg/sandbox/block/streaming_chunk_test.go b/packages/orchestrator/pkg/sandbox/block/streaming_chunk_test.go index fcd09e48dd..2e6d207298 100644 --- a/packages/orchestrator/pkg/sandbox/block/streaming_chunk_test.go +++ b/packages/orchestrator/pkg/sandbox/block/streaming_chunk_test.go @@ -50,10 +50,12 @@ func makeTestData(size int) []byte { // fakeSeekable implements storage.Seekable backed by in-memory data. // When ctrl is non-nil, reads are gated through its channels for concurrency tests. type fakeSeekable struct { - data []byte - failAfter int64 // >0: truncate reads at this offset; 0 = disabled - fetchCount atomic.Int64 - ctrl *testControl // nil = ungated immediate reads + data []byte + failAfter int64 // >0: truncate reads at this offset; 0 = disabled + corrupt bool // enable corruptByte injection + corruptByte int64 // absolute offset to XOR 0xFF in served bytes (requires corrupt=true) + fetchCount atomic.Int64 + ctrl *testControl // nil = ungated immediate reads } var _ storage.Seekable = (*fakeSeekable)(nil) @@ -126,7 +128,13 @@ func (s *fakeSeekable) OpenRangeReader(_ context.Context, offsetU int64, length end = min(end, s.failAfter) } - r := io.Reader(bytes.NewReader(s.data[fetchOff:end])) + served := s.data[fetchOff:end] + if s.corrupt && s.corruptByte >= fetchOff && s.corruptByte < end { + served = bytes.Clone(served) + served[s.corruptByte-fetchOff] ^= 0xFF + } + + r := io.Reader(bytes.NewReader(served)) if frameTable.IsCompressed() { dec, err := storage.NewDecompressReader(storage.NewRangeReader(io.NopCloser(r)), frameTable.CompressionType(), storage.UnknownSource, storage.UnknownSeekableObjectType) @@ -645,3 +653,36 @@ func (r *controlledReader) Close(context.Context) (*storage.ReadStats, error) { return nil, nil } + +// TestChunker_CorruptCompressedFrameNotServedOrCached verifies the integrity +// gap fix: a compressed frame whose CRC only fails at the footer (content +// decodes to plausible bytes) must not be released to waiters or marked +// cached. Without the fix, an exact-size read never triggers the codec's +// footer CRC, the Close error was ignored, and the frame was advanced to +// waiters and cached. Corrupts the final byte of the first compressed frame. +func TestChunker_CorruptCompressedFrameNotServedOrCached(t *testing.T) { + t.Parallel() + + data := makeTestData(testFileSize) + ft, file := makeCompressedTestData(t, data) + + // Corrupt the last byte of the first frame's compressed range (footer). + r, err := ft.LocateCompressed(0) + require.NoError(t, err) + file.corrupt = true + file.corruptByte = r.Offset + int64(r.Length) - 1 + + chunker := newTestChunker(t, int64(len(data))) + defer chunker.Close() + + _, err = chunker.Slice(t.Context(), 0, testBlockSize, file, ft) + require.Error(t, err, "corrupt compressed frame must surface an error, not serve unverified bytes") + require.False(t, chunker.IsCached(t.Context(), 0, testBlockSize), + "corrupt compressed frame must not be marked cached") + + // A later read of a clean frame still works (chunker remains usable). + lastOff := int64(testFileSize) - testBlockSize + slice, err := chunker.Slice(t.Context(), lastOff, testBlockSize, file, ft) + require.NoError(t, err) + require.Equal(t, data[lastOff:], slice) +} diff --git a/packages/shared/pkg/storage/compress_decode.go b/packages/shared/pkg/storage/compress_decode.go index 3c5bf76b35..c3c997fbc2 100644 --- a/packages/shared/pkg/storage/compress_decode.go +++ b/packages/shared/pkg/storage/compress_decode.go @@ -117,6 +117,17 @@ func (r *decompressReader) Read(p []byte) (int, error) { } func (r *decompressReader) Close(ctx context.Context) (*ReadStats, error) { + // Drain any remaining decoded bytes so the codec consumes the frame footer + // and surfaces a CRC/truncation error. A caller reading only the exact + // uncompressed frame size never pulls the footer through, so zstd would + // otherwise report success on a footer-corrupted or truncated frame. + // Bounded: the source is a single frame. + if r.readErr == nil { + if _, err := io.Copy(io.Discard, r.meteredOut); err != nil && !errors.Is(err, io.EOF) { + r.readErr = err + } + } + r.releaseCodec() stats := &ReadStats{ @@ -130,5 +141,9 @@ func (r *decompressReader) Close(ctx context.Context) (*ReadStats, error) { _, innerErr := r.inner.Close(ctx) + if r.readErr != nil { + return stats, r.readErr + } + return stats, innerErr } diff --git a/packages/shared/pkg/storage/compress_decode_test.go b/packages/shared/pkg/storage/compress_decode_test.go new file mode 100644 index 0000000000..683f0c8c8f --- /dev/null +++ b/packages/shared/pkg/storage/compress_decode_test.go @@ -0,0 +1,52 @@ +package storage + +import ( + "bytes" + "crypto/sha256" + "io" + "testing" + + "github.com/stretchr/testify/require" +) + +// TestDecompressReaderVerifiesCRCOnClose ensures a footer-corrupt frame is not +// reported as success when the caller reads exactly the uncompressed size. +// zstd only verifies the frame checksum once the footer is consumed, which an +// exact-size read never triggers, so Close must drain-verify and surface the +// error. +func TestDecompressReaderVerifiesCRCOnClose(t *testing.T) { + t.Parallel() + + const frameU = 512 * 1024 + data := generateSemiRandomData(frameU) + + up := &memPartUploader{} + fullFT, _, err := compressStream(t.Context(), bytes.NewReader(data), defaultCfg(CompressionZstd, 1, frameU), up, 1, nil) + require.NoError(t, err) + require.Equal(t, 1, fullFT.Table().NumFrames()) + blob := up.Assemble() + + // Corrupt the frame footer (last byte) so only the trailing checksum, + // verified after the last content byte, is wrong. + corrupt := bytes.Clone(blob) + corrupt[len(corrupt)-1] ^= 0xFF + + dec, err := NewDecompressReader(bytesRangeReader(corrupt), CompressionZstd, SourceAWS, SeekableObjectType(0)) + require.NoError(t, err) + + // Read EXACTLY the uncompressed frame size — no read past EOF. + buf := make([]byte, frameU) + _, _ = io.ReadFull(dec, buf) + _, closeErr := dec.Close(t.Context()) + require.Error(t, closeErr, "exact-size read of a footer-corrupt frame must fail on Close") + + // A clean frame still round-trips and closes without error. + dec2, err := NewDecompressReader(bytesRangeReader(blob), CompressionZstd, SourceAWS, SeekableObjectType(0)) + require.NoError(t, err) + got := make([]byte, frameU) + _, err = io.ReadFull(dec2, got) + require.NoError(t, err) + _, closeErr = dec2.Close(t.Context()) + require.NoError(t, closeErr) + require.Equal(t, sha256.Sum256(data), sha256.Sum256(got)) +} diff --git a/packages/shared/pkg/storage/compress_frame_table.go b/packages/shared/pkg/storage/compress_frame_table.go index 196f44d681..62cc2bd1e1 100644 --- a/packages/shared/pkg/storage/compress_frame_table.go +++ b/packages/shared/pkg/storage/compress_frame_table.go @@ -236,6 +236,14 @@ func (ft *FrameTable) LocateCompressed(offset int64) (Range, error) { return Range{}, err } + // The compressed range covers a whole frame; a mid-frame uncompressed + // offset would fetch and decode that frame from its start, silently + // returning data for the wrong position. Callers must frame-align via + // LocateUncompressed, so reject anything else rather than mis-serve. + if offset != e.StartU { + return Range{}, fmt.Errorf("offset %d is not frame-aligned (frame starts at %d); align via LocateUncompressed", offset, e.StartU) + } + return Range{Offset: e.StartC, Length: int(e.SizeC)}, nil } diff --git a/packages/shared/pkg/storage/compress_frame_table_test.go b/packages/shared/pkg/storage/compress_frame_table_test.go index 9927d612c4..c9333d6839 100644 --- a/packages/shared/pkg/storage/compress_frame_table_test.go +++ b/packages/shared/pkg/storage/compress_frame_table_test.go @@ -66,6 +66,16 @@ func TestLocate(t *testing.T) { require.Error(t, err) }) + t.Run("mid-frame offset errors", func(t *testing.T) { + t.Parallel() + // A mid-frame uncompressed offset would fetch and decode the whole + // containing frame from its start, silently returning data for the + // wrong position. It must be rejected so callers frame-align first. + _, err := ft.LocateCompressed((1 << 20) / 2) + require.Error(t, err) + require.Contains(t, err.Error(), "frame-aligned") + }) + t.Run("nil table errors", func(t *testing.T) { t.Parallel() _, err := (*FrameTable)(nil).LocateCompressed(0) diff --git a/packages/shared/pkg/storage/gcp_multipart.go b/packages/shared/pkg/storage/gcp_multipart.go index def9d037c8..24e029f290 100644 --- a/packages/shared/pkg/storage/gcp_multipart.go +++ b/packages/shared/pkg/storage/gcp_multipart.go @@ -15,6 +15,7 @@ import ( "net/http" "os" "slices" + "strings" "sync" "time" @@ -55,6 +56,7 @@ func createRetryableClient(ctx context.Context, config RetryConfig) *retryableht client.RetryMax = config.MaxAttempts - 1 // go-retryablehttp counts retries, not total attempts client.RetryWaitMin = config.InitialBackoff client.RetryWaitMax = config.MaxBackoff + client.CheckRetry = retryOnCompleteError // Custom backoff function with full jitter to avoid thundering herd client.Backoff = func(start, maxBackoff time.Duration, attemptNum int, _ *http.Response) time.Duration { @@ -95,6 +97,40 @@ func createRetryableClient(ctx context.Context, config RetryConfig) *retryableht return client } +// retryOnCompleteError extends the default retry policy so a transient +// CompleteMultipartUpload failure is retried like a 5xx. S3 and the GCS XML +// API can return HTTP 200 with an body (e.g. InternalError) for a +// commit that didn't happen; the default policy only looks at the status code, +// so without this such a response fails on the first attempt and ignores the +// configured retry budget. Only the complete request (URL query "uploadId=…", +// small XML body) is inspected, so buffering the body to peek + restore it is +// cheap. +func retryOnCompleteError(ctx context.Context, resp *http.Response, err error) (bool, error) { + if retry, rerr := retryablehttp.DefaultRetryPolicy(ctx, resp, err); retry || rerr != nil { + return retry, rerr + } + + if resp == nil || resp.StatusCode != http.StatusOK || resp.Request == nil || + !strings.HasPrefix(resp.Request.URL.RawQuery, "uploadId=") || resp.Body == nil { + return false, nil + } + + body, readErr := io.ReadAll(resp.Body) + resp.Body.Close() + // Restore the body so the caller (or a retry) can read it. + resp.Body = io.NopCloser(bytes.NewReader(body)) + + // Retry when the body couldn't be fully read (transient truncation/reset + // after the headers) or when it carries a transient (e.g. + // InternalError). Crucially, a read failure must not fall through as + // no-retry: that would hand completeUpload a clean, buffered partial body + // and let it commit on a truncated response. + var apiErr completeMultipartError + retry := readErr != nil || (xml.Unmarshal(body, &apiErr) == nil && apiErr.Code != "") + + return retry, nil //nolint:nilerr // decision is carried by the bool; returning readErr would abort the retry loop instead of retrying +} + // zapLogger adapts zap.Logger to retryablehttp.LeveledLogger interface var _ retryablehttp.LeveledLogger = &leveledLogger{} @@ -134,6 +170,15 @@ type Part struct { ETag string `xml:"ETag"` } +// completeMultipartError matches the S3/GCS XML error payload that can arrive +// with an HTTP 200 status on CompleteMultipartUpload. Unmarshal only succeeds +// (Code populated) when the response root is . +type completeMultipartError struct { + XMLName xml.Name `xml:"Error"` + Code string `xml:"Code"` + Message string `xml:"Message"` +} + type MultipartUploader struct { bucketName string objectName string @@ -257,13 +302,20 @@ func (m *MultipartUploader) uploadPart(ctx context.Context, uploadID string, par url := fmt.Sprintf("%s/%s?partNumber=%d&uploadId=%s", m.baseURL, m.objectName, partNumber, uploadID) - req, err := retryablehttp.NewRequestWithContext(ctx, "PUT", url, bytes.NewReader(data)) + // A non-nil zero-length body counts as "unknown length" to net/http and + // is sent chunked, which S3-compatible XML backends reject with 411. Pass + // no body for empty parts so Content-Length: 0 is sent instead. + var body any + if len(data) > 0 { + body = bytes.NewReader(data) + } + + req, err := retryablehttp.NewRequestWithContext(ctx, "PUT", url, body) if err != nil { return "", err } req.Header.Set("Authorization", "Bearer "+m.token) - req.Header.Set("Content-Length", fmt.Sprintf("%d", len(data))) sum := md5.Sum(data) //nolint:gosec // GCS multipart uses Content-MD5 for transport integrity. req.Header.Set("Content-MD5", base64.StdEncoding.EncodeToString(sum[:])) @@ -297,18 +349,26 @@ func (m *MultipartUploader) uploadPartSlices(ctx context.Context, uploadID strin url := fmt.Sprintf("%s/%s?partNumber=%d&uploadId=%s", m.baseURL, m.objectName, partNumber, uploadID) - // Use a ReaderFunc so the retryable client can replay the body on retries - bodyFn := func() (io.Reader, error) { - return newMultiSliceReader(slices), nil + // Use a ReaderFunc so the retryable client can replay the body on + // retries. The multiSliceReader's Len makes retryablehttp set the + // request's ContentLength, so parts are sent with an explicit + // Content-Length rather than chunked transfer encoding (which GCS + // tolerates but S3-compatible XML backends reject with 411). Empty parts + // send no body at all for the same reason: a zero-length reader still + // counts as "unknown length" to net/http and would be chunked. + var body any + if totalLen > 0 { + body = retryablehttp.ReaderFunc(func() (io.Reader, error) { + return newMultiSliceReader(slices), nil + }) } - req, err := retryablehttp.NewRequestWithContext(ctx, "PUT", url, retryablehttp.ReaderFunc(bodyFn)) + req, err := retryablehttp.NewRequestWithContext(ctx, "PUT", url, body) if err != nil { return "", err } req.Header.Set("Authorization", "Bearer "+m.token) - req.Header.Set("Content-Length", fmt.Sprintf("%d", totalLen)) h := md5.New() //nolint:gosec // GCS multipart uses Content-MD5 for transport integrity. for _, s := range slices { _, _ = h.Write(s) @@ -365,12 +425,24 @@ func (m *MultipartUploader) completeUpload(ctx context.Context, uploadID string, } defer resp.Body.Close() - if resp.StatusCode != http.StatusOK { - body, _ := io.ReadAll(resp.Body) + body, readErr := io.ReadAll(resp.Body) + if readErr != nil { + return fmt.Errorf("read complete upload response (status %d): %w", resp.StatusCode, readErr) + } + if resp.StatusCode != http.StatusOK { return fmt.Errorf("failed to complete upload (status %d): %s", resp.StatusCode, string(body)) } + // S3 and the GCS XML API can return HTTP 200 with an body when the + // commit fails server-side (documented for internal errors/timeouts). + // Treating that as success would record a frame table for an object that + // was never committed, so reject any response whose root is an Error. + var apiErr completeMultipartError + if xml.Unmarshal(body, &apiErr) == nil && apiErr.Code != "" { + return fmt.Errorf("failed to complete upload (status %d, code %s): %s", resp.StatusCode, apiErr.Code, apiErr.Message) + } + return nil } diff --git a/packages/shared/pkg/storage/gcp_multipart_test.go b/packages/shared/pkg/storage/gcp_multipart_test.go index 5964a4ecbc..5866abc830 100644 --- a/packages/shared/pkg/storage/gcp_multipart_test.go +++ b/packages/shared/pkg/storage/gcp_multipart_test.go @@ -1098,3 +1098,99 @@ func TestRetryableClient_ActualRetryBehavior(t *testing.T) { totalTime := time.Since(startTime) t.Logf("Total time: %v, Retry delays: %v", totalTime, retryDelays) } + +// TestGCPCompleteRejects200WithErrorBody verifies CompleteMultipartUpload +// treats an HTTP 200 carrying an body as a failure. S3-dialect +// servers (including the GCS XML API) can return 200 with an error payload on +// internal timeouts; accepting it would record a frame table for an object +// that was never committed. +// fastRetryConfig retries a few times with negligible backoff so retry-path +// tests stay fast. +func fastRetryConfig() RetryConfig { + return RetryConfig{MaxAttempts: 3, InitialBackoff: time.Millisecond, MaxBackoff: 2 * time.Millisecond, BackoffMultiplier: 2} +} + +func newCompleteErrorUploader(t *testing.T, completeAttempts *atomic.Int32, complete http.HandlerFunc) *MultipartUploader { + t.Helper() + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.Method == http.MethodPost && r.URL.RawQuery == "uploads": + w.Write([]byte(`id-1`)) + case r.Method == http.MethodPut && r.URL.Query().Get("partNumber") != "": + io.Copy(io.Discard, r.Body) + w.Header().Set("ETag", `"e1"`) + case r.Method == http.MethodPost && r.URL.Query().Get("uploadId") != "": + completeAttempts.Add(1) + complete(w, r) + default: + t.Errorf("unexpected request %s %s", r.Method, r.URL) + } + })) + t.Cleanup(server.Close) + + return &MultipartUploader{ + bucketName: "b", objectName: "o", token: "t", + client: createRetryableClient(t.Context(), fastRetryConfig()), + baseURL: server.URL + "/b", + } +} + +func TestGCPCompleteRejects200WithErrorBody(t *testing.T) { + t.Parallel() + + // Persistent 200-with-: retried up to the budget, then reported as + // a failed commit (never treated as success). + var attempts atomic.Int32 + up := newCompleteErrorUploader(t, &attempts, func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) // 200 OK, but the body is an error — commit did NOT happen. + w.Write([]byte(`InternalErrorplease try again`)) + }) + require.NoError(t, up.Start(t.Context())) + require.NoError(t, up.UploadPart(t.Context(), 1, []byte("data"))) + + err := up.Complete(t.Context()) + require.Error(t, err, "200-with--body must be treated as a failed commit, never as success") + require.Equal(t, int32(3), attempts.Load(), "a transient 200-error must be retried up to the configured budget") +} + +func TestGCPCompleteRetriesTransient200Error(t *testing.T) { + t.Parallel() + + // First complete attempt returns a 200-with-; the retry succeeds. + var attempts atomic.Int32 + up := newCompleteErrorUploader(t, &attempts, func(w http.ResponseWriter, _ *http.Request) { + if attempts.Load() == 1 { + w.WriteHeader(http.StatusOK) + w.Write([]byte(`InternalErrorplease try again`)) + + return + } + w.Write([]byte(`"final"`)) + }) + require.NoError(t, up.Start(t.Context())) + require.NoError(t, up.UploadPart(t.Context(), 1, []byte("data"))) + + require.NoError(t, up.Complete(t.Context()), "a transient 200-error should be retried to success") + require.Equal(t, int32(2), attempts.Load()) +} + +func TestGCPCompleteRetriesTruncatedBody(t *testing.T) { + t.Parallel() + + // 200 headers, but the body is truncated (declared Content-Length exceeds + // what's written), so reading it fails after the headers. This must be + // retried and ultimately fail — never masked as a clean commit. + var attempts atomic.Int32 + up := newCompleteErrorUploader(t, &attempts, func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Length", "4096") + w.WriteHeader(http.StatusOK) + w.Write([]byte(" short -> skip. + c.writeFrameBack(t.Context(), framePath, 0, len(compressed)+100, SourceFS, CompressionLZ4, captured) c.wg.Wait() _, err = os.Stat(framePath) require.True(t, os.IsNotExist(err), "mismatched frame should not be cached") }) } + +// TestCorruptFetchPoisonsCacheRecovers guards the compressed cache against +// poisoning: a fetch that returns corrupt bytes at the right length passes the +// size-only writeback guard, so it must not be cached (and a cache hit that +// fails to decode must be evicted). After the upstream heals the read +// recovers, rather than failing forever. +func TestCorruptFetchPoisonsCacheRecovers(t *testing.T) { + t.Parallel() + + data := generateSemiRandomData(1 * megabyte) // single frame + up := &memPartUploader{} + fullFT, _, err := compressStream(t.Context(), bytes.NewReader(data), defaultCfg(CompressionZstd, 2, 2*megabyte), up, 2, nil) + require.NoError(t, err) + blob := up.Assemble() + + corrupt := bytes.Clone(blob) + corrupt[len(corrupt)/2] ^= 0xFF + + var serveCorrupt atomic.Bool + serveCorrupt.Store(true) + + inner := NewMockSeekable(t) + inner.EXPECT().OpenRangeReader(mock.Anything, mock.Anything, mock.Anything, mock.Anything). + RunAndReturn(func(_ context.Context, off, length int64, _ *FrameTable) (RangeReader, Source, error) { + b := blob + if serveCorrupt.Load() { + b = corrupt + } + + return bytesRangeReader(b[off : off+length]), SourceAWS, nil + }).Maybe() + + c := cachedSeekable{path: t.TempDir(), inner: inner, tracer: noopTracer, chunkSize: 1024} + ft := fullFT.Table() + + // First read: corrupt upstream -> decode error; the corrupt frame must + // not be cached. + rr, _, err := c.OpenRangeReader(t.Context(), 0, 0, ft) + require.NoError(t, err) + _, readErr := io.Copy(io.Discard, rr) + rr.Close(t.Context()) + require.Error(t, readErr) + c.wg.Wait() + + // Upstream heals; the read must now succeed (refetch, not a poisoned hit). + serveCorrupt.Store(false) + rr2, _, err := c.OpenRangeReader(t.Context(), 0, 0, ft) + require.NoError(t, err) + var got bytes.Buffer + _, readErr2 := got.ReadFrom(rr2) + rr2.Close(t.Context()) + c.wg.Wait() + require.NoError(t, readErr2, "read must recover after upstream heals (cache not poisoned)") + require.Equal(t, data, got.Bytes()) +} + +// TestTruncatedFetchNotCached verifies a short fetch is not cached and the +// next read recovers once the upstream heals. +func TestTruncatedFetchNotCached(t *testing.T) { + t.Parallel() + + data := generateSemiRandomData(1 * megabyte) + up := &memPartUploader{} + fullFT, _, err := compressStream(t.Context(), bytes.NewReader(data), defaultCfg(CompressionZstd, 2, 2*megabyte), up, 2, nil) + require.NoError(t, err) + blob := up.Assemble() + + var serveTruncated atomic.Bool + serveTruncated.Store(true) + + inner := NewMockSeekable(t) + inner.EXPECT().OpenRangeReader(mock.Anything, mock.Anything, mock.Anything, mock.Anything). + RunAndReturn(func(_ context.Context, off, length int64, _ *FrameTable) (RangeReader, Source, error) { + b := blob[off : off+length] + if serveTruncated.Load() { + b = b[:len(b)/2] + } + + return bytesRangeReader(b), SourceAWS, nil + }).Maybe() + + c := cachedSeekable{path: t.TempDir(), inner: inner, tracer: noopTracer, chunkSize: 1024} + ft := fullFT.Table() + + rr, _, err := c.OpenRangeReader(t.Context(), 0, 0, ft) + require.NoError(t, err) + _, readErr := io.Copy(io.Discard, rr) + rr.Close(t.Context()) + require.Error(t, readErr) + c.wg.Wait() + + frameFile := makeFrameFilename(c.path, Range{Offset: 0, Length: int(ft.CompressedSize())}) + _, statErr := os.Stat(frameFile) + require.True(t, os.IsNotExist(statErr), "truncated frame must not be cached") + + serveTruncated.Store(false) + rr2, _, err := c.OpenRangeReader(t.Context(), 0, 0, ft) + require.NoError(t, err) + var got bytes.Buffer + _, readErr2 := got.ReadFrom(rr2) + rr2.Close(t.Context()) + c.wg.Wait() + require.NoError(t, readErr2, "read must recover after upstream heals") + require.Equal(t, data, got.Bytes()) +} diff --git a/packages/shared/pkg/storage/storage_cache_seekable_compressed.go b/packages/shared/pkg/storage/storage_cache_seekable_compressed.go index 965a7b44e5..693deb0ba8 100644 --- a/packages/shared/pkg/storage/storage_cache_seekable_compressed.go +++ b/packages/shared/pkg/storage/storage_cache_seekable_compressed.go @@ -46,7 +46,16 @@ func (c *cachedSeekable) openReaderCompressed(ctx context.Context, offsetU int64 } RecordReadOpen(ctx, time.Since(start), c.objType, SourceNFS, ct, err) if err == nil { - return dec, SourceNFS, nil + // A cached frame that decodes cleanly returns here. NewDecompressReader's + // Close drains and CRC-verifies the frame, so a non-nil Close error means + // the cached bytes no longer decode (bit rot / torn write that still has + // the right size, which the size check above cannot catch) — evict it so + // the next read refetches instead of failing forever. + return &closeHookReader{RangeReader: dec, onClose: func(_ context.Context, err error) { + if err != nil { + _ = os.Remove(path) + } + }}, SourceNFS, nil } // Cache miss: fetch raw compressed bytes via OpenRangeReader(nil frameTable). @@ -55,10 +64,16 @@ func (c *cachedSeekable) openReaderCompressed(ctx context.Context, offsetU int64 return nil, innerSource, fmt.Errorf("raw fetch at C=%d: %w", rng.Offset, err) } + // The captureReader tees the raw compressed bytes into `captured`; write + // them back only if the frame decoded cleanly (Close error == nil), so a + // corrupt-but-right-sized frame is never cached. + var captured []byte + capturing := !skipCacheWriteback(ctx) frameReader := raw - if !skipCacheWriteback(ctx) { - frameReader = newCaptureReader(raw, rng.Length, true, - c.compressedFrameWriteback(path, offsetU, rng.Length, innerSource, ct)) + if capturing { + frameReader = newCaptureReader(raw, rng.Length, true, func(_ context.Context, frame []byte) { + captured = frame + }) } dec, err = NewDecompressReader(frameReader, ct, innerSource, c.objType) @@ -68,36 +83,59 @@ func (c *cachedSeekable) openReaderCompressed(ctx context.Context, offsetU int64 return nil, innerSource, fmt.Errorf("create decompressor: %w", err) } - return dec, innerSource, nil + return &closeHookReader{RangeReader: dec, onClose: func(ctx context.Context, err error) { + if err != nil || !capturing { + return + } + c.writeFrameBack(ctx, path, offsetU, rng.Length, innerSource, ct, captured) + }}, innerSource, nil } -// compressedFrameWriteback returns a captureReader callback that -// persists the captured frame to the NFS cache in a detached goroutine. -// Best-effort: a short capture is logged and skipped — the caller already -// got valid decompressed bytes. -func (c *cachedSeekable) compressedFrameWriteback(framePath string, offset int64, expectedSize int, src Source, codec CompressionType) func(context.Context, []byte) { - return func(ctx context.Context, frame []byte) { - if !isCompleteRead(len(frame), expectedSize, nil) { - logger.L().Warn(ctx, "compressed frame cache writeback short, skipping", - zap.Int("got", len(frame)), zap.Int("expected", expectedSize), zap.String("path", framePath)) +// closeHookReader runs onClose exactly once when the wrapped reader is closed, +// passing its Close error. The compressed cache relies on NewDecompressReader's +// Close draining and CRC-verifying the frame, so that error is the decode +// verdict: on a miss, write the frame back only when it's nil; on a hit, evict +// when it isn't. Reads pass straight through the embedded RangeReader. +type closeHookReader struct { + RangeReader - return - } + onClose func(ctx context.Context, err error) +} + +func (r *closeHookReader) Close(ctx context.Context) (*ReadStats, error) { + stats, err := r.RangeReader.Close(ctx) + if r.onClose != nil { + r.onClose(ctx, err) + r.onClose = nil + } - c.goCtx(ctx, func(ctx context.Context) { - ctx, span := c.tracer.Start(ctx, "write compressed frame back to cache") - defer span.End() + return stats, err +} - start := time.Now() - err := c.writeToCache(ctx, offset, framePath, frame) - recordWriteback(ctx, time.Since(start), int64(len(frame)), c.objType, src, codec, TriggerRead, err) +// writeFrameBack persists a fully-read compressed frame to the NFS cache in a +// detached goroutine. Best-effort: a short frame is logged and skipped — the +// caller already has valid decompressed bytes. +func (c *cachedSeekable) writeFrameBack(ctx context.Context, framePath string, offset int64, expectedSize int, src Source, codec CompressionType, frame []byte) { + if !isCompleteRead(len(frame), expectedSize, nil) { + logger.L().Warn(ctx, "compressed frame cache writeback short, skipping", + zap.Int("got", len(frame)), zap.Int("expected", expectedSize), zap.String("path", framePath)) - if err != nil && !errors.Is(err, lock.ErrLockAlreadyHeld) { - recordError(span, err) - logger.L().Warn(ctx, "failed to write frame back to cache", zap.Error(err)) - } - }) + return } + + c.goCtx(ctx, func(ctx context.Context) { + ctx, span := c.tracer.Start(ctx, "write compressed frame back to cache") + defer span.End() + + start := time.Now() + err := c.writeToCache(ctx, offset, framePath, frame) + recordWriteback(ctx, time.Since(start), int64(len(frame)), c.objType, src, codec, TriggerRead, err) + + if err != nil && !errors.Is(err, lock.ErrLockAlreadyHeld) { + recordError(span, err) + logger.L().Warn(ctx, "failed to write frame back to cache", zap.Error(err)) + } + }) } // makeFrameFilename returns the NFS cache path for a compressed frame.