From 9bba6f8e35c3aa03cfea6901205ad544278d9228 Mon Sep 17 00:00:00 2001
From: zili-simular <zili@simular.ai>
Date: Thu, 24 Sep 2026 18:32:06 +0800
Subject: [PATCH] fix: replay OCI blob bodies on retry

---
 oci/oci.go                 | 17 +++++++++++------
 oci/oci_test.go            | 22 ++++++++++++++++++++++
 snapshot/roundtrip_test.go |  2 +-
 snapshot/snapshot.go       |  2 +-
 snapshot/snapshot_test.go  |  2 +-
 5 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/oci/oci.go b/oci/oci.go
index f46b67a..59130ee 100644
--- a/oci/oci.go
+++ b/oci/oci.go
@@ -103,7 +103,7 @@ func (r *OCIRegistry) HasManifest(ctx context.Context, repo, tag string) (bool,
 	return true, nil
 }
 
-func (r *OCIRegistry) PutBlob(ctx context.Context, repo, digest string, body io.Reader, size int64) error {
+func (r *OCIRegistry) PutBlob(ctx context.Context, repo, digest string, body io.ReadSeeker, size int64) error {
 	repoRef, err := name.NewRepository(r.base + "/" + repo)
 	if err != nil {
 		return fmt.Errorf("parse repo %s: %w", repo, err)
@@ -182,16 +182,21 @@ func isNotFound(err error) bool {
 	return ok && terr.StatusCode == http.StatusNotFound
 }
 
-// streamLayer is a single-use v1.Layer over a body of known digest and size, so PutBlob streams without buffering.
+// streamLayer is a replayable v1.Layer over a body of known digest and size, so PutBlob streams without buffering.
 type streamLayer struct {
 	hash v1.Hash
 	size int64
-	body io.Reader
+	body io.ReadSeeker
 }
 
-func (l *streamLayer) Digest() (v1.Hash, error)             { return l.hash, nil }
-func (l *streamLayer) Size() (int64, error)                 { return l.size, nil }
-func (l *streamLayer) Compressed() (io.ReadCloser, error)   { return io.NopCloser(l.body), nil }
+func (l *streamLayer) Digest() (v1.Hash, error) { return l.hash, nil }
+func (l *streamLayer) Size() (int64, error)     { return l.size, nil }
+func (l *streamLayer) Compressed() (io.ReadCloser, error) {
+	if _, err := l.body.Seek(0, io.SeekStart); err != nil {
+		return nil, fmt.Errorf("rewind blob: %w", err)
+	}
+	return io.NopCloser(l.body), nil
+}
 func (l *streamLayer) MediaType() (types.MediaType, error)  { return types.OCILayer, nil }
 func (l *streamLayer) DiffID() (v1.Hash, error)             { return v1.Hash{}, errBlobUncompressed }
 func (l *streamLayer) Uncompressed() (io.ReadCloser, error) { return nil, errBlobUncompressed }
diff --git a/oci/oci_test.go b/oci/oci_test.go
index 8f7322d..f73571f 100644
--- a/oci/oci_test.go
+++ b/oci/oci_test.go
@@ -36,6 +36,28 @@ func TestBulkTransportUsesHTTP1(t *testing.T) {
 	}
 }
 
+func TestStreamLayerCompressedReplaysBody(t *testing.T) {
+	want := []byte("complete replayable blob")
+	layer := &streamLayer{body: bytes.NewReader(want)}
+
+	for call := 1; call <= 2; call++ {
+		rc, err := layer.Compressed()
+		if err != nil {
+			t.Fatalf("Compressed call %d: %v", call, err)
+		}
+		got, err := io.ReadAll(rc)
+		if err != nil {
+			t.Fatalf("read Compressed call %d: %v", call, err)
+		}
+		if err := rc.Close(); err != nil {
+			t.Fatalf("close Compressed call %d: %v", call, err)
+		}
+		if !bytes.Equal(got, want) {
+			t.Fatalf("Compressed call %d = %q, want %q", call, got, want)
+		}
+	}
+}
+
 func TestOCIRegistryRoundTrip(t *testing.T) {
 	srv := httptest.NewServer(registry.New())
 	t.Cleanup(srv.Close)
diff --git a/snapshot/roundtrip_test.go b/snapshot/roundtrip_test.go
index 7775074..fb837f7 100644
--- a/snapshot/roundtrip_test.go
+++ b/snapshot/roundtrip_test.go
@@ -872,7 +872,7 @@ type countingUploader struct {
 	putBlobCalls atomic.Int64
 }
 
-func (c *countingUploader) PutBlob(ctx context.Context, name, digest string, body io.Reader, size int64) error {
+func (c *countingUploader) PutBlob(ctx context.Context, name, digest string, body io.ReadSeeker, size int64) error {
 	c.putBlobCalls.Add(1)
 	return c.fakeUploader.PutBlob(ctx, name, digest, body, size)
 }
diff --git a/snapshot/snapshot.go b/snapshot/snapshot.go
index fddb731..4edf35d 100644
--- a/snapshot/snapshot.go
+++ b/snapshot/snapshot.go
@@ -34,7 +34,7 @@ type ProgressFunc = func(string)
 // Uploader abstracts OCI blob and manifest uploads.
 type Uploader interface {
 	HasBlob(ctx context.Context, name, digest string) (bool, error)
-	PutBlob(ctx context.Context, name, digest string, body io.Reader, size int64) error
+	PutBlob(ctx context.Context, name, digest string, body io.ReadSeeker, size int64) error
 	PutManifest(ctx context.Context, name, tag string, data []byte, contentType string) error
 }
 
diff --git a/snapshot/snapshot_test.go b/snapshot/snapshot_test.go
index 1452c66..3d456e0 100644
--- a/snapshot/snapshot_test.go
+++ b/snapshot/snapshot_test.go
@@ -467,7 +467,7 @@ func (f *fakeUploader) HasBlob(_ context.Context, _, digest string) (bool, error
 	return ok, nil
 }
 
-func (f *fakeUploader) PutBlob(_ context.Context, _, digest string, body io.Reader, _ int64) error {
+func (f *fakeUploader) PutBlob(_ context.Context, _, digest string, body io.ReadSeeker, _ int64) error {
 	data, err := io.ReadAll(body)
 	if err != nil {
 		return err