From 41eb204936e9303f9ea90a930c8d90f39c5b8380 Mon Sep 17 00:00:00 2001
From: Damilola Edwards <damilolaedwards@gmail.com>
Date: Thu, 23 Jul 2026 16:27:40 +0100
Subject: [PATCH] Make storage upload atomic

Upload opened the destination directly with O_TRUNC and streamed the
request body straight into it. The source is a live sandbox HTTP
stream that can break mid-transfer, which is a normal occurrence, not
an edge case, but a mid-copy failure had already truncated whatever
was at that key. A re-upload to the same filename, which happens any
time a script or session step reruns and writes the same output name
again, could destroy a prior good file and replace it with a partial
one while the caller was told the upload failed and nothing happened.

Write to a temp file in the same directory instead and rename it over
the destination only once the copy and close both succeed. Any
failure along the way removes the temp file and leaves the existing
key untouched.
---
 pkg/storage/service.go      | 27 +++++++++++++----
 pkg/storage/service_test.go | 59 +++++++++++++++++++++++++++++++++++++
 2 files changed, 81 insertions(+), 5 deletions(-)

diff --git a/pkg/storage/service.go b/pkg/storage/service.go
index 2de3b95d..7d4298c9 100644
--- a/pkg/storage/service.go
+++ b/pkg/storage/service.go
@@ -90,20 +90,37 @@ func (s *service) Upload(scope Scope, name string, body io.Reader) (UploadResult
 
 	filePath := filepath.Join(dir, filepath.FromSlash(rel))
 
-	f, err := s.fs.OpenFile(filePath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o644)
+	// Write to a temp file in the same directory and rename it over the
+	// destination only once the copy fully succeeds, so a failed or
+	// interrupted upload (a normal occurrence: the source is a live sandbox
+	// HTTP stream that can break mid-transfer) can never truncate or corrupt
+	// whatever was already stored at this key. Same directory keeps the
+	// rename on one filesystem, which is what makes it atomic.
+	tmp, err := afero.TempFile(s.fs, dir, ".upload-*")
 	if err != nil {
-		return UploadResult{}, fmt.Errorf("creating file: %w", err)
+		return UploadResult{}, fmt.Errorf("creating temp file: %w", err)
 	}
+	tmpPath := tmp.Name()
+
+	if _, err := io.Copy(tmp, body); err != nil {
+		_ = tmp.Close()
+		_ = s.fs.Remove(tmpPath)
 
-	if _, err := io.Copy(f, body); err != nil {
-		_ = f.Close()
 		return UploadResult{}, fmt.Errorf("writing file: %w", err)
 	}
 
-	if err := f.Close(); err != nil {
+	if err := tmp.Close(); err != nil {
+		_ = s.fs.Remove(tmpPath)
+
 		return UploadResult{}, fmt.Errorf("closing file: %w", err)
 	}
 
+	if err := s.fs.Rename(tmpPath, filePath); err != nil {
+		_ = s.fs.Remove(tmpPath)
+
+		return UploadResult{}, fmt.Errorf("finalizing file: %w", err)
+	}
+
 	return UploadResult{
 		Key:  rel,
 		URL:  s.fileURL(scope, rel),
diff --git a/pkg/storage/service_test.go b/pkg/storage/service_test.go
index 19b80c26..3365097b 100644
--- a/pkg/storage/service_test.go
+++ b/pkg/storage/service_test.go
@@ -233,3 +233,62 @@ func TestServeFilePathTraversal(t *testing.T) {
 	svc.ServeFile(w, r, "../../etc/passwd")
 	assert.Equal(t, http.StatusNotFound, w.Code)
 }
+
+// breaksAfter is a reader that returns n bytes and then a fixed error,
+// simulating a source stream (a live sandbox HTTP body, in production) that
+// breaks partway through.
+type breaksAfter struct {
+	data []byte
+	err  error
+}
+
+func (r *breaksAfter) Read(p []byte) (int, error) {
+	if len(r.data) == 0 {
+		return 0, r.err
+	}
+
+	n := copy(p, r.data)
+	r.data = r.data[n:]
+
+	return n, nil
+}
+
+func TestUploadFailureLeavesPriorFileUntouched(t *testing.T) {
+	t.Parallel()
+
+	svc, fs := newTestService()
+
+	_, err := svc.Upload(Scope{ExecutionID: "exec-1"}, "output.txt", bytes.NewBufferString("the original good contents"))
+	require.NoError(t, err)
+
+	_, err = svc.Upload(Scope{ExecutionID: "exec-1"}, "output.txt", &breaksAfter{
+		data: []byte("partial"),
+		err:  assert.AnError,
+	})
+	require.Error(t, err)
+
+	contents, err := afero.ReadFile(fs, "/data/exec-1/output.txt")
+	require.NoError(t, err)
+	assert.Equal(t, "the original good contents", string(contents),
+		"a failed re-upload to the same key must not disturb the existing file")
+
+	files, err := svc.List(Scope{ExecutionID: "exec-1"}, "")
+	require.NoError(t, err)
+	require.Len(t, files, 1, "no leftover temp file should remain after a failed upload")
+}
+
+func TestUploadFailureOnNewKeyLeavesNoFile(t *testing.T) {
+	t.Parallel()
+
+	svc, _ := newTestService()
+
+	_, err := svc.Upload(Scope{ExecutionID: "exec-1"}, "output.txt", &breaksAfter{
+		data: []byte("partial"),
+		err:  assert.AnError,
+	})
+	require.Error(t, err)
+
+	files, err := svc.List(Scope{ExecutionID: "exec-1"}, "")
+	require.NoError(t, err)
+	assert.Empty(t, files, "a failed upload to a brand-new key should leave nothing behind")
+}