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