refactor(snapshot): extract blob spool helper from PushBlobStreaming

pkg/snapshot/registry.go

@@ -89,37 +89,22 @@ func PushBlobStreaming(
 		return "", 0, fmt.Errorf("parse repository %q: %w", repository, err)
 	}
 
-	f, err := os.CreateTemp("", "devsy-snapshot-blob-*")
-	if err != nil {
-		return "", 0, fmt.Errorf("create temp file for blob spool: %w", err)
-	}
-	path := f.Name()
-	defer func() { _ = os.Remove(path) }()
-
-	size, err := io.Copy(f, r)
-	closeErr := f.Close()
-	if err != nil {
-		return "", 0, fmt.Errorf("spool blob to temp file: %w", err)
-	}
-	if closeErr != nil {
-		return "", 0, fmt.Errorf("close spooled blob temp file: %w", closeErr)
-	}
-
-	digest, err := digestOfFile(path)
+	blob, err := spoolBlob(r)
 	if err != nil {
 		return "", 0, err
 	}
+	defer blob.remove()
 
 	opts, err := remoteOptions(ctx)
 	if err != nil {
 		return "", 0, err
 	}
 
 	layer := &fileLayer{
-		path:      path,
+		path:      blob.path,
 		mediaType: types.MediaType(mediaType),
-		digest:    digest,
-		size:      size,
+		digest:    blob.digest,
+		size:      blob.size,
 	}
 	if err := remote.WriteLayer(repo, layer, opts...); err != nil {
 		return "", 0, fmt.Errorf(
@@ -128,7 +113,7 @@ func PushBlobStreaming(
 			image.SanitizeRegistryError(err),
 		)
 	}
-	return digest.String(), size, nil
+	return blob.digest.String(), blob.size, nil
 }
 
 func digestOfFile(path string) (v1.Hash, error) {
@@ -145,6 +130,56 @@ func digestOfFile(path string) (v1.Hash, error) {
 	return digest, nil
 }
 
+// spooledBlob is a blob staged in a temp file. The holder owns path and must
+// call remove once the blob is no longer needed.
+type spooledBlob struct {
+	path   string
+	digest v1.Hash
+	size   int64
+}
+
+func (b spooledBlob) remove() {
+	_ = os.Remove(b.path)
+}
+
+func spoolBlob(r io.Reader) (spooledBlob, error) {
+	f, err := os.CreateTemp("", "devsy-snapshot-blob-*")
+	if err != nil {
+		return spooledBlob{}, fmt.Errorf("create temp file for blob spool: %w", err)
+	}
+	blob := spooledBlob{path: f.Name()}
+	handedOff := false
+	defer func() {
+		if !handedOff {
+			// Close before remove: Windows cannot delete an open file. Also runs on panic or runtime.Goexit from r.
+			_ = f.Close()
+			blob.remove()
+		}
+	}()
+
+	blob.size, err = copyToSpool(f, r)
+	if err == nil {
+		blob.digest, err = digestOfFile(blob.path)
+	}
+	if err != nil {
+		return spooledBlob{}, err
+	}
+	handedOff = true
+	return blob, nil
+}
+
+func copyToSpool(f *os.File, r io.Reader) (int64, error) {
+	size, err := io.Copy(f, r)
+	closeErr := f.Close()
+	if err != nil {
+		return 0, fmt.Errorf("spool blob to temp file: %w", err)
+	}
+	if closeErr != nil {
+		return 0, fmt.Errorf("close spooled blob temp file: %w", closeErr)
+	}
+	return size, nil
+}
+
 // fileLayer is a v1.Layer backed by a file on disk holding the blob's exact
 // bytes verbatim (no gzip envelope), mirroring static.NewLayer's semantics
 // (Digest == DiffID, Compressed == Uncompressed) but without requiring the

pkg/snapshot/registry_test.go

@@ -5,13 +5,17 @@ import (
 	"context"
 	"crypto/sha256"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"os"
+	"runtime"
 	"strconv"
 	"strings"
 	"testing"
+	"testing/iotest"
 	"time"
 
 	"github.com/google/go-containerregistry/pkg/name"
@@ -101,6 +105,120 @@ func TestPushBlobStreaming_MatchesPushBlobDigest(t *testing.T) {
 	require.Equal(t, bufferedSize, streamedSize)
 }
 
+func spoolFilesIn(t *testing.T, dir string) []string {
+	t.Helper()
+	entries, err := os.ReadDir(dir)
+	require.NoError(t, err)
+	names := make([]string, 0, len(entries))
+	for _, entry := range entries {
+		names = append(names, entry.Name())
+	}
+	return names
+}
+
+func isolateTempDir(t *testing.T) string {
+	t.Helper()
+	dir := t.TempDir()
+	t.Setenv("TMPDIR", dir)
+	t.Setenv("TMP", dir)
+	t.Setenv("TEMP", dir)
+	return dir
+}
+
+func TestPushBlobStreaming_RemovesSpoolWhenReadFails(t *testing.T) {
+	spoolDir := isolateTempDir(t)
+	host := newTestRegistry(t)
+	readErr := errors.New("source disconnected")
+	src := io.MultiReader(strings.NewReader("partial"), iotest.ErrReader(readErr))
+
+	_, _, err := PushBlobStreaming(
+		context.Background(),
+		host+"/acme/snapshots",
+		VolumesMediaType,
+		src,
+	)
+
+	require.ErrorIs(t, err, readErr)
+	require.Empty(t, spoolFilesIn(t, spoolDir))
+}
+
+func TestPushBlobStreaming_RemovesSpoolWhenUploadFails(t *testing.T) {
+	spoolDir := isolateTempDir(t)
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusForbidden)
+	}))
+	t.Cleanup(srv.Close)
+	host := strings.TrimPrefix(srv.URL, "http://")
+
+	_, _, err := PushBlobStreaming(
+		context.Background(),
+		host+"/acme/snapshots",
+		VolumesMediaType,
+		strings.NewReader("upload me"),
+	)
+
+	require.ErrorContains(t, err, "push blob to")
+	require.Empty(t, spoolFilesIn(t, spoolDir))
+}
+
+func TestPushBlobStreaming_RejectsInvalidRepositoryBeforeSpooling(t *testing.T) {
+	spoolDir := isolateTempDir(t)
+	src := strings.NewReader("never read")
+
+	_, _, err := PushBlobStreaming(context.Background(), "UPPERCASE/Repo", VolumesMediaType, src)
+
+	require.ErrorContains(t, err, "parse repository")
+	require.Equal(t, len("never read"), src.Len())
+	require.Empty(t, spoolFilesIn(t, spoolDir))
+}
+
+// panicReader and goexitReader simulate a source that aborts mid-copy.
+type panicReader struct{}
+
+func (panicReader) Read([]byte) (int, error) { panic("source exploded") }
+
+type goexitReader struct{}
+
+func (goexitReader) Read([]byte) (int, error) {
+	runtime.Goexit()
+	return 0, nil
+}
+
+func TestPushBlobStreaming_RemovesSpoolWhenReaderPanics(t *testing.T) {
+	spoolDir := isolateTempDir(t)
+	host := newTestRegistry(t)
+
+	require.PanicsWithValue(t, "source exploded", func() {
+		_, _, _ = PushBlobStreaming(
+			context.Background(),
+			host+"/acme/snapshots",
+			VolumesMediaType,
+			panicReader{},
+		)
+	})
+
+	require.Empty(t, spoolFilesIn(t, spoolDir))
+}
+
+func TestPushBlobStreaming_RemovesSpoolWhenReaderGoexits(t *testing.T) {
+	spoolDir := isolateTempDir(t)
+	host := newTestRegistry(t)
+	done := make(chan struct{})
+
+	go func() {
+		defer close(done)
+		_, _, _ = PushBlobStreaming(
+			context.Background(),
+			host+"/acme/snapshots",
+			VolumesMediaType,
+			goexitReader{},
+		)
+	}()
+	<-done
+
+	require.Empty(t, spoolFilesIn(t, spoolDir))
+}
+
 func TestPushPullManifest_RoundTrip(t *testing.T) {
 	host := newTestRegistry(t)
 	ctx := context.Background()