Make storage upload atomic

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),

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")
+}