diff --git a/cmd/api/src/api/registration/registration.go b/cmd/api/src/api/registration/registration.go index d7ff1f490b94..b6e08026585f 100644 --- a/cmd/api/src/api/registration/registration.go +++ b/cmd/api/src/api/registration/registration.go @@ -32,6 +32,7 @@ import ( "github.com/specterops/bloodhound/cmd/api/src/model/appcfg" "github.com/specterops/bloodhound/cmd/api/src/queries" "github.com/specterops/bloodhound/cmd/api/src/services/dogtags" + "github.com/specterops/bloodhound/cmd/api/src/services/storage" "github.com/specterops/bloodhound/cmd/api/src/services/upload" "github.com/specterops/bloodhound/packages/go/cache" "github.com/specterops/dawgs/graph" @@ -68,6 +69,7 @@ func RegisterFossRoutes( authenticator api.Authenticator, authorizer auth.Authorizer, ingestSchema upload.IngestSchema, + fileServiceResolver storage.FileServiceResolver, dogtagsService dogtags.Service, openGraphSchemaService v2.OpenGraphSchemaService, ) { @@ -88,7 +90,6 @@ func RegisterFossRoutes( // Static asset handling for the UI. This route intentionally sits outside the default API rate limiter // because a single page load can request many static HTML, JavaScript, CSS, and media assets. routerInst.PathPrefix(api.UserInterfacePath, static.AssetHandler) - - var resources = v2.NewResources(rdms, graphDB, cfg, apiCache, graphQuery, collectorManifests, authorizer, authenticator, ingestSchema, dogtagsService, openGraphSchemaService) + var resources = v2.NewResources(rdms, graphDB, cfg, apiCache, graphQuery, collectorManifests, authorizer, authenticator, ingestSchema, fileServiceResolver, dogtagsService, openGraphSchemaService) NewV2API(resources, routerInst) } diff --git a/cmd/api/src/api/tools/ingest.go b/cmd/api/src/api/tools/ingest.go index 9cb7700db532..08b87c97c976 100644 --- a/cmd/api/src/api/tools/ingest.go +++ b/cmd/api/src/api/tools/ingest.go @@ -23,30 +23,30 @@ import ( "io" "log/slog" "net/http" - "os" - "path/filepath" + "path" + "strings" "sync" "time" "github.com/specterops/bloodhound/cmd/api/src/api" - "github.com/specterops/bloodhound/cmd/api/src/config" "github.com/specterops/bloodhound/cmd/api/src/database/types" "github.com/specterops/bloodhound/cmd/api/src/model/appcfg" "github.com/specterops/bloodhound/packages/go/bhlog/attr" "github.com/specterops/bloodhound/packages/go/headers" + "github.com/specterops/bloodhound/packages/go/storage" ) type IngestControl struct { - cfg config.Configuration - retainedFileLock *sync.RWMutex - parameterService appcfg.ParameterService + retainedFileLock *sync.RWMutex + parameterService appcfg.ParameterService + retainedFileService storage.FileService } -func NewIngestControlTool(cfg config.Configuration, parameterService appcfg.ParameterService) IngestControl { +func NewIngestControlTool(parameterService appcfg.ParameterService, retainedFileService storage.FileService) IngestControl { return IngestControl{ - cfg: cfg, - retainedFileLock: &sync.RWMutex{}, - parameterService: parameterService, + retainedFileLock: &sync.RWMutex{}, + parameterService: parameterService, + retainedFileService: retainedFileService, } } @@ -69,6 +69,39 @@ func (s *IngestControl) setIngestFileRetention(ctx context.Context, parameter ap return nil } +func retainedArchiveName(filePath string) string { + filePath = path.Clean(strings.TrimPrefix(filePath, "/")) + if filePath == "." || filePath == ".." || strings.HasPrefix(filePath, "../") { + return path.Base(filePath) + } + return filePath +} + +func (s *IngestControl) writeRetainedFileToTar(ctx context.Context, tarWriter *tar.Writer, filePath string) error { + reader, fileInfo, err := s.retainedFileService.GetFile(ctx, filePath) + if err != nil { + return fmt.Errorf("opening retained file: %w", err) + } + defer reader.Close() + + tarHeader := &tar.Header{ + Name: retainedArchiveName(fileInfo.Path), + Size: fileInfo.Size, + Mode: 0o600, + ModTime: fileInfo.LastModified, + } + + if err := tarWriter.WriteHeader(tarHeader); err != nil { + return fmt.Errorf("writing tar header: %w", err) + } + + if _, err := io.Copy(tarWriter, reader); err != nil { + return fmt.Errorf("copying retained file: %w", err) + } + + return nil +} + func (s *IngestControl) FetchRetainedIngestFiles(response http.ResponseWriter, request *http.Request) { if !s.retainedFileLock.TryRLock() { // Unable to acquire a write lock at this time - notify of a conflict. User is expected to retry @@ -85,101 +118,38 @@ func (s *IngestControl) FetchRetainedIngestFiles(response http.ResponseWriter, r return } - retainedFilesDirectory := s.cfg.RetainedFilesDirectory() - - if retainedFilesDirEntries, err := os.ReadDir(retainedFilesDirectory); err != nil { - // Unable to stat the retained files directory. Log a warning and inform the user that the - // operation failed. - slog.WarnContext( - request.Context(), - "Failed reading retained files directory", - slog.String("directory", retainedFilesDirectory), - attr.Error(err), - ) + files, err := s.retainedFileService.ListFiles(request.Context(), "", storage.ListOptions{Recursive: true}) + if err != nil { + slog.WarnContext(request.Context(), "Failed listing retained files", attr.Error(err)) response.WriteHeader(http.StatusInternalServerError) - } else { - // Author the response - response.Header().Set(headers.ContentType.String(), "application/gzip") - response.Header().Set(headers.ContentEncoding.String(), "gzip") - response.Header().Set(headers.ContentDisposition.String(), fmt.Sprintf(`attachment; filename="retained-ingest-%s.tar.gz"`, time.Now().UTC().Format("20060102T150405Z"))) - response.WriteHeader(http.StatusOK) + return + } - var ( - gzipWriter = gzip.NewWriter(response) - tarWriter = tar.NewWriter(gzipWriter) - ) + response.Header().Set(headers.ContentType.String(), "application/gzip") + response.Header().Set(headers.ContentEncoding.String(), "gzip") + response.Header().Set(headers.ContentDisposition.String(), fmt.Sprintf(`attachment; filename="retained-ingest-%s.tar.gz"`, time.Now().UTC().Format("20060102T150405Z"))) + response.WriteHeader(http.StatusOK) - for _, retainedFilesDirEntry := range retainedFilesDirEntries { - retainedFilePath := filepath.Join(retainedFilesDirectory, retainedFilesDirEntry.Name()) + gzipWriter := gzip.NewWriter(response) + defer gzipWriter.Close() - if retainedFilesDirEntry.IsDir() { - // Log a warning for directory entries and skip reading it. - slog.WarnContext( - request.Context(), - "Unexpected directory in retained files directory", - slog.String("path", retainedFilePath), - slog.String("directory", retainedFilesDirectory), - ) - continue - } + tarWriter := tar.NewWriter(gzipWriter) + defer tarWriter.Close() - if fileInfo, err := retainedFilesDirEntry.Info(); err != nil { - slog.WarnContext( - request.Context(), - "Unable to stat file", - slog.String("path", retainedFilePath), - attr.Error(err), - ) - break - } else { - if tarHeader, err := tar.FileInfoHeader(fileInfo, ""); err != nil { - slog.WarnContext( - request.Context(), - "Unable to convert file info for file", - slog.String("path", retainedFilePath), - attr.Error(err), - ) - break - } else if err := tarWriter.WriteHeader(tarHeader); err != nil { - slog.WarnContext( - request.Context(), - "Failed writing tar file header to response", - slog.String("path", retainedFilePath), - attr.Error(err), - ) - break - } - - if fin, err := os.Open(retainedFilePath); err != nil { - slog.WarnContext( - request.Context(), - "Failed to open retained file", - slog.String("path", retainedFilePath), - slog.String("directory", retainedFilesDirectory), - attr.Error(err), - ) - break - } else { - // Copy and inspect the error only after closing the file. - _, copyErr := io.Copy(tarWriter, fin) - fin.Close() - - if copyErr != nil { - slog.WarnContext( - request.Context(), - "Failed writing file content from to response", - slog.String("path", retainedFilePath), - attr.Error(copyErr), - ) - break - } - } - } + for _, file := range files { + if file.IsDir { + continue } - // Attempt to flush and close the tar.gz writers - best effort - tarWriter.Close() - gzipWriter.Close() + if err := s.writeRetainedFileToTar(request.Context(), tarWriter, file.Path); err != nil { + slog.WarnContext( + request.Context(), + "Failed writing retained file to response", + slog.String("path", file.Path), + attr.Error(err), + ) + break + } } } @@ -215,6 +185,7 @@ func (s *IngestControl) DisableIngestFileRetention(response http.ResponseWriter, if err := s.setIngestFileRetention(request.Context(), appcfg.RetainIngestedFilesParameter{ Enabled: false, }); err != nil { + s.retainedFileLock.Unlock() api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusInternalServerError, err.Error(), request), response) return } @@ -223,31 +194,29 @@ func (s *IngestControl) DisableIngestFileRetention(response http.ResponseWriter, // Clear all retained files in a goroutine that releases the lock once finished. This is best effort as // it is still possible for a file to make it into the retained directory even after the parameter is set. - go func() { + cleanupCtx := context.WithoutCancel(request.Context()) + + go func(ctx context.Context) { defer s.retainedFileLock.Unlock() - retainedFilesDirectory := s.cfg.RetainedFilesDirectory() + files, err := s.retainedFileService.ListFiles(ctx, "", storage.ListOptions{Recursive: true}) + if err != nil { + slog.WarnContext(ctx, "Failed listing retained files", attr.Error(err)) + } - if retainedFilesDirEntries, err := os.ReadDir(retainedFilesDirectory); err != nil { - slog.WarnContext( - request.Context(), - "Failed reading retained files directory", - slog.String("directory", retainedFilesDirectory), - attr.Error(err), - ) - } else { - for _, retainedFilesDirEntry := range retainedFilesDirEntries { - retainedFilePath := filepath.Join(retainedFilesDirectory, retainedFilesDirEntry.Name()) - - if err := os.RemoveAll(retainedFilePath); err != nil { - slog.WarnContext( - request.Context(), - "Failed removing retained file", - slog.String("path", retainedFilePath), - attr.Error(err), - ) - } + for _, file := range files { + if file.IsDir { + continue + } + + if err := s.retainedFileService.DeleteFile(ctx, file.Path); err != nil { + slog.WarnContext( + ctx, + "Failed removing retained file", + slog.String("path", file.Path), + attr.Error(err), + ) } } - }() + }(cleanupCtx) } diff --git a/cmd/api/src/api/tools/ingest_test.go b/cmd/api/src/api/tools/ingest_test.go new file mode 100644 index 000000000000..2031dc882d7b --- /dev/null +++ b/cmd/api/src/api/tools/ingest_test.go @@ -0,0 +1,564 @@ +// Copyright 2026 Specter Ops, Inc. +// +// Licensed under the Apache License, Version 2.0 +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 + +package tools + +import ( + "archive/tar" + "bytes" + "compress/gzip" + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/specterops/bloodhound/cmd/api/src/database/types" + "github.com/specterops/bloodhound/cmd/api/src/model/appcfg" + "github.com/specterops/bloodhound/packages/go/headers" + "github.com/specterops/bloodhound/packages/go/storage" + storagemocks "github.com/specterops/bloodhound/packages/go/storage/mocks" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" +) + +type retainedIngestParameterService struct { + retain bool + setErr error +} + +func (s retainedIngestParameterService) GetAllConfigurationParameters(ctx context.Context) (appcfg.Parameters, error) { + return nil, nil +} + +func (s retainedIngestParameterService) GetConfigurationParameter(ctx context.Context, parameterKey appcfg.ParameterKey) (appcfg.Parameter, error) { + value, err := types.NewJSONBObject(appcfg.RetainIngestedFilesParameter{ + Enabled: s.retain, + }) + if err != nil { + return appcfg.Parameter{}, err + } + + return appcfg.Parameter{ + Key: parameterKey, + Value: value, + }, nil +} + +func (s retainedIngestParameterService) SetConfigurationParameter(ctx context.Context, configurationParameter appcfg.Parameter) error { + return s.setErr +} + +type retainedIngestErrorReadCloser struct { + err error +} + +func (s retainedIngestErrorReadCloser) Read([]byte) (int, error) { + return 0, s.err +} + +func (s retainedIngestErrorReadCloser) Close() error { + return nil +} + +func readRetainedTar(t *testing.T, body []byte) map[string]string { + t.Helper() + + gzipReader, err := gzip.NewReader(bytes.NewReader(body)) + require.NoError(t, err) + defer gzipReader.Close() + + var ( + files = map[string]string{} + tarReader = tar.NewReader(gzipReader) + ) + + for { + header, err := tarReader.Next() + if errors.Is(err, io.EOF) { + break + } + require.NoError(t, err) + + content, err := io.ReadAll(tarReader) + require.NoError(t, err) + files[header.Name] = string(content) + } + + return files +} + +func TestRetainedArchiveName(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + filePath string + expected string + }{ + { + name: "keeps nested logical path", + filePath: "nested/file.json", + expected: "nested/file.json", + }, + { + name: "trims leading slash", + filePath: "/nested/file.json", + expected: "nested/file.json", + }, + { + name: "keeps only base name for parent traversal", + filePath: "../outside.json", + expected: "outside.json", + }, + { + name: "keeps only base name for absolute parent traversal", + filePath: "/../outside.json", + expected: "outside.json", + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + require.Equal(t, testCase.expected, retainedArchiveName(testCase.filePath)) + }) + } +} + +func TestIngestControl_writeRetainedFileToTar(t *testing.T) { + t.Parallel() + + var ( + errGet = errors.New("get failed") + errRead = errors.New("read failed") + ) + + type expected struct { + errIs error + errContains string + files map[string]string + } + + type testData struct { + name string + filePath string + setupMock func(t *testing.T, ctx context.Context, mockFileService *storagemocks.MockFileService) + expected expected + } + + tests := []testData{ + { + name: "writes retained file to tar", + filePath: "nested/file.json", + setupMock: func(t *testing.T, ctx context.Context, mockFileService *storagemocks.MockFileService) { + mockFileService.EXPECT(). + GetFile(ctx, "nested/file.json"). + Return( + io.NopCloser(bytes.NewReader([]byte("file content"))), + storage.FileInfo{ + Path: "nested/file.json", + Size: int64(len("file content")), + LastModified: time.Now().UTC(), + }, + nil, + ) + }, + expected: expected{ + files: map[string]string{ + "nested/file.json": "file content", + }, + }, + }, + { + name: "get error returns wrapped error", + filePath: "missing.json", + setupMock: func(t *testing.T, ctx context.Context, mockFileService *storagemocks.MockFileService) { + mockFileService.EXPECT(). + GetFile(ctx, "missing.json"). + Return(nil, storage.FileInfo{}, errGet) + }, + expected: expected{ + errIs: errGet, + errContains: "opening retained file", + }, + }, + { + name: "reader error returns wrapped error", + filePath: "file.json", + setupMock: func(t *testing.T, ctx context.Context, mockFileService *storagemocks.MockFileService) { + mockFileService.EXPECT(). + GetFile(ctx, "file.json"). + Return( + retainedIngestErrorReadCloser{err: errRead}, + storage.FileInfo{ + Path: "file.json", + Size: int64(len("file content")), + }, + nil, + ) + }, + expected: expected{ + errIs: errRead, + errContains: "copying retained file", + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + var ( + ctx = context.Background() + archiveBuffer = &bytes.Buffer{} + tarWriter = tar.NewWriter(archiveBuffer) + mockFileService = storagemocks.NewMockFileService(gomock.NewController(t)) + ingestControl = NewIngestControlTool(retainedIngestParameterService{retain: true}, mockFileService) + ) + + testCase.setupMock(t, ctx, mockFileService) + + // Act + err := ingestControl.writeRetainedFileToTar(ctx, tarWriter, testCase.filePath) + + // Assert + if testCase.expected.errIs != nil { + require.ErrorIs(t, err, testCase.expected.errIs) + require.Contains(t, err.Error(), testCase.expected.errContains) + return + } + + require.NoError(t, err) + require.NoError(t, tarWriter.Close()) + + tarReader := tar.NewReader(bytes.NewReader(archiveBuffer.Bytes())) + for expectedName, expectedContent := range testCase.expected.files { + header, err := tarReader.Next() + require.NoError(t, err) + require.Equal(t, expectedName, header.Name) + + content, err := io.ReadAll(tarReader) + require.NoError(t, err) + require.Equal(t, expectedContent, string(content)) + } + }) + } +} + +func TestIngestControl_FetchRetainedIngestFiles(t *testing.T) { + t.Parallel() + + var ( + errList = errors.New("list failed") + errGet = errors.New("get failed") + ) + + type expected struct { + statusCode int + files map[string]string + } + + type testData struct { + name string + retain bool + setupMock func(t *testing.T, ctx context.Context, mockFileService *storagemocks.MockFileService) + expected expected + } + + tests := []testData{ + { + name: "retention disabled returns not found", + retain: false, + expected: expected{ + statusCode: http.StatusNotFound, + }, + }, + { + name: "list error returns internal server error", + retain: true, + setupMock: func(t *testing.T, ctx context.Context, mockFileService *storagemocks.MockFileService) { + mockFileService.EXPECT(). + ListFiles(ctx, "", storage.ListOptions{Recursive: true}). + Return(nil, errList) + }, + expected: expected{ + statusCode: http.StatusInternalServerError, + }, + }, + { + name: "writes retained files to gzip tar and skips directories", + retain: true, + setupMock: func(t *testing.T, ctx context.Context, mockFileService *storagemocks.MockFileService) { + mockFileService.EXPECT(). + ListFiles(ctx, "", storage.ListOptions{Recursive: true}). + Return([]storage.FileInfo{ + {Path: "nested", IsDir: true}, + {Path: "nested/one.json"}, + {Path: "../outside.json"}, + }, nil) + mockFileService.EXPECT(). + GetFile(ctx, "nested/one.json"). + Return( + io.NopCloser(bytes.NewReader([]byte("one"))), + storage.FileInfo{ + Path: "nested/one.json", + Size: int64(len("one")), + LastModified: time.Now().UTC(), + }, + nil, + ) + mockFileService.EXPECT(). + GetFile(ctx, "../outside.json"). + Return( + io.NopCloser(bytes.NewReader([]byte("outside"))), + storage.FileInfo{ + Path: "../outside.json", + Size: int64(len("outside")), + LastModified: time.Now().UTC(), + }, + nil, + ) + }, + expected: expected{ + statusCode: http.StatusOK, + files: map[string]string{ + "nested/one.json": "one", + "outside.json": "outside", + }, + }, + }, + { + name: "get error ends archive without failing response", + retain: true, + setupMock: func(t *testing.T, ctx context.Context, mockFileService *storagemocks.MockFileService) { + mockFileService.EXPECT(). + ListFiles(ctx, "", storage.ListOptions{Recursive: true}). + Return([]storage.FileInfo{{Path: "missing.json"}}, nil) + mockFileService.EXPECT(). + GetFile(ctx, "missing.json"). + Return(nil, storage.FileInfo{}, errGet) + }, + expected: expected{ + statusCode: http.StatusOK, + files: map[string]string{}, + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + var ( + request = httptest.NewRequest(http.MethodGet, "/api/v2/tools/ingest/retained", nil) + response = httptest.NewRecorder() + mockFileService = storagemocks.NewMockFileService(gomock.NewController(t)) + ingestControl = NewIngestControlTool(retainedIngestParameterService{retain: testCase.retain}, mockFileService) + ) + + if testCase.setupMock != nil { + testCase.setupMock(t, request.Context(), mockFileService) + } + + // Act + ingestControl.FetchRetainedIngestFiles(response, request) + + // Assert + require.Equal(t, testCase.expected.statusCode, response.Code) + if testCase.expected.statusCode != http.StatusOK { + return + } + + require.Equal(t, "application/gzip", response.Header().Get(headers.ContentType.String())) + require.Equal(t, "gzip", response.Header().Get(headers.ContentEncoding.String())) + require.Contains(t, response.Header().Get(headers.ContentDisposition.String()), `attachment; filename="retained-ingest-`) + require.Equal(t, testCase.expected.files, readRetainedTar(t, response.Body.Bytes())) + }) + } +} + +func TestIngestControl_DisableIngestFileRetention(t *testing.T) { + t.Parallel() + + t.Run("deletes retained files and skips directories", func(t *testing.T) { + t.Parallel() + + // Arrange + var ( + request = httptest.NewRequest(http.MethodPost, "/api/v2/tools/ingest/retained/disable", nil) + response = httptest.NewRecorder() + done = make(chan struct{}) + mockFileService = storagemocks.NewMockFileService(gomock.NewController(t)) + ingestControl = NewIngestControlTool(retainedIngestParameterService{}, mockFileService) + ) + + mockFileService.EXPECT(). + ListFiles(gomock.Any(), "", storage.ListOptions{Recursive: true}). + Return([]storage.FileInfo{ + {Path: "nested", IsDir: true}, + {Path: "one.json"}, + {Path: "two.json"}, + }, nil) + mockFileService.EXPECT(). + DeleteFile(gomock.Any(), "one.json"). + Return(nil) + mockFileService.EXPECT(). + DeleteFile(gomock.Any(), "two.json"). + DoAndReturn(func(context.Context, string) error { + close(done) + return nil + }) + + // Act + ingestControl.DisableIngestFileRetention(response, request) + + // Assert + require.Equal(t, http.StatusAccepted, response.Code) + require.Eventually(t, func() bool { + select { + case <-done: + return true + default: + return false + } + }, time.Second, 10*time.Millisecond) + require.Eventually(t, func() bool { + if ingestControl.retainedFileLock.TryLock() { + ingestControl.retainedFileLock.Unlock() + return true + } + return false + }, time.Second, 10*time.Millisecond) + }) + + t.Run("parameter error releases lock", func(t *testing.T) { + t.Parallel() + + // Arrange + var ( + errSet = errors.New("set failed") + request = httptest.NewRequest(http.MethodPost, "/api/v2/tools/ingest/retained/disable", nil) + response = httptest.NewRecorder() + mockFileService = storagemocks.NewMockFileService(gomock.NewController(t)) + ingestControl = NewIngestControlTool(retainedIngestParameterService{setErr: errSet}, mockFileService) + ) + + // Act + ingestControl.DisableIngestFileRetention(response, request) + + // Assert + require.Equal(t, http.StatusInternalServerError, response.Code) + require.True(t, ingestControl.retainedFileLock.TryLock()) + ingestControl.retainedFileLock.Unlock() + }) + + t.Run("list error releases lock", func(t *testing.T) { + t.Parallel() + + // Arrange + var ( + errList = errors.New("list failed") + request = httptest.NewRequest(http.MethodPost, "/api/v2/tools/ingest/retained/disable", nil) + response = httptest.NewRecorder() + done = make(chan struct{}) + mockFileService = storagemocks.NewMockFileService(gomock.NewController(t)) + ingestControl = NewIngestControlTool(retainedIngestParameterService{}, mockFileService) + ) + + mockFileService.EXPECT(). + ListFiles(gomock.Any(), "", storage.ListOptions{Recursive: true}). + DoAndReturn(func(context.Context, string, storage.ListOptions) ([]storage.FileInfo, error) { + close(done) + return nil, errList + }) + + // Act + ingestControl.DisableIngestFileRetention(response, request) + + // Assert + require.Equal(t, http.StatusAccepted, response.Code) + require.Eventually(t, func() bool { + select { + case <-done: + return true + default: + return false + } + }, time.Second, 10*time.Millisecond) + require.Eventually(t, func() bool { + if ingestControl.retainedFileLock.TryLock() { + ingestControl.retainedFileLock.Unlock() + return true + } + return false + }, time.Second, 10*time.Millisecond) + }) +} + +func TestIngestControl_LockConflict(t *testing.T) { + t.Parallel() + + t.Run("fetch returns conflict while write lock is held", func(t *testing.T) { + t.Parallel() + + // Arrange + var ( + request = httptest.NewRequest(http.MethodGet, "/api/v2/tools/ingest/retained", nil) + response = httptest.NewRecorder() + mockFileService = storagemocks.NewMockFileService(gomock.NewController(t)) + ingestControl = NewIngestControlTool(retainedIngestParameterService{retain: true}, mockFileService) + ) + + ingestControl.retainedFileLock.Lock() + defer ingestControl.retainedFileLock.Unlock() + + // Act + ingestControl.FetchRetainedIngestFiles(response, request) + + // Assert + require.Equal(t, http.StatusConflict, response.Code) + }) + + t.Run("disable returns conflict while lock is held", func(t *testing.T) { + t.Parallel() + + // Arrange + var ( + request = httptest.NewRequest(http.MethodPost, "/api/v2/tools/ingest/retained/disable", nil) + response = httptest.NewRecorder() + mockFileService = storagemocks.NewMockFileService(gomock.NewController(t)) + ingestControl = NewIngestControlTool(retainedIngestParameterService{}, mockFileService) + ) + + ingestControl.retainedFileLock.Lock() + defer ingestControl.retainedFileLock.Unlock() + + // Act + ingestControl.DisableIngestFileRetention(response, request) + + // Assert + require.Equal(t, http.StatusConflict, response.Code) + }) +} + +var _ appcfg.ParameterService = retainedIngestParameterService{} +var _ io.ReadCloser = retainedIngestErrorReadCloser{} diff --git a/cmd/api/src/api/v2/collectors.go b/cmd/api/src/api/v2/collectors.go index 7c1567a379e1..214e59971468 100644 --- a/cmd/api/src/api/v2/collectors.go +++ b/cmd/api/src/api/v2/collectors.go @@ -27,6 +27,7 @@ import ( "github.com/specterops/bloodhound/cmd/api/src/api" "github.com/specterops/bloodhound/cmd/api/src/config" "github.com/specterops/bloodhound/packages/go/bhlog/attr" + "github.com/specterops/bloodhound/packages/go/storage" ) const ( @@ -89,7 +90,10 @@ func (s *Resources) DownloadCollectorByVersion(response http.ResponseWriter, req } else if fileName, err := retrieveCollectorZipFileName(releaseTag, collectorType, s.CollectorManifests); err != nil { slog.ErrorContext(request.Context(), "Manifest doesn't exist for collector", slog.String("collector_type", collectorType)) api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusInternalServerError, api.ErrorResponseDetailsInternalServerError, request), response) - } else if data, err := s.FileService.ReadFile(filepath.Join(s.Config.CollectorsDirectory(), collectorType, fileName)); err != nil { + } else if collectorsFileService, err := s.FileServiceResolver.Resolve(storage.FileServiceCollectors); err != nil { + slog.ErrorContext(request.Context(), "Unable to resolve collectors file service", attr.Error(err)) + api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusInternalServerError, api.ErrorResponseDetailsInternalServerError, request), response) + } else if data, err := collectorsFileService.ReadFile(request.Context(), filepath.Join(collectorType, fileName)); err != nil { slog.ErrorContext(request.Context(), "Could not open collector file for download", attr.Error(err)) api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusInternalServerError, api.ErrorResponseDetailsInternalServerError, request), response) } else { @@ -122,7 +126,10 @@ func (s *Resources) DownloadCollectorChecksumByVersion(response http.ResponseWri } else if fileName, err := retrieveCollectorSHA256FileName(releaseTag, collectorType, s.CollectorManifests); err != nil { slog.ErrorContext(request.Context(), "Manifest doesn't exist for collector", slog.String("collector_type", collectorType)) api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusInternalServerError, api.ErrorResponseDetailsInternalServerError, request), response) - } else if data, err := s.FileService.ReadFile(filepath.Join(s.Config.CollectorsDirectory(), collectorType, fileName)); err != nil { + } else if collectorsFileService, err := s.FileServiceResolver.Resolve(storage.FileServiceCollectors); err != nil { + slog.ErrorContext(request.Context(), "Unable to resolve collectors file service", attr.Error(err)) + api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusInternalServerError, api.ErrorResponseDetailsInternalServerError, request), response) + } else if data, err := collectorsFileService.ReadFile(request.Context(), filepath.Join(collectorType, fileName)); err != nil { slog.ErrorContext(request.Context(), "Could not open collector file for download", attr.Error(err)) api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusInternalServerError, api.ErrorResponseDetailsInternalServerError, request), response) } else { diff --git a/cmd/api/src/api/v2/collectors_test.go b/cmd/api/src/api/v2/collectors_test.go index 9eace73ef1c9..e99916db8258 100644 --- a/cmd/api/src/api/v2/collectors_test.go +++ b/cmd/api/src/api/v2/collectors_test.go @@ -28,10 +28,12 @@ import ( "github.com/gorilla/mux" v2 "github.com/specterops/bloodhound/cmd/api/src/api/v2" "github.com/specterops/bloodhound/cmd/api/src/config" - fsmocks "github.com/specterops/bloodhound/cmd/api/src/services/fs/mocks" + storageServiceMocks "github.com/specterops/bloodhound/cmd/api/src/services/storage/mocks" "github.com/specterops/bloodhound/cmd/api/src/utils/test" "github.com/specterops/bloodhound/packages/go/headers" "github.com/specterops/bloodhound/packages/go/mediatypes" + "github.com/specterops/bloodhound/packages/go/storage" + storagemocks "github.com/specterops/bloodhound/packages/go/storage/mocks" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -40,7 +42,8 @@ import ( func TestResources_DownloadCollectorByVersion(t *testing.T) { type mock struct { - mockFS *fsmocks.MockService + mockFS *storagemocks.MockFileService + mockFileServiceResolver *storageServiceMocks.MockFileServiceResolver } type expected struct { responseBody string @@ -89,7 +92,7 @@ func TestResources_DownloadCollectorByVersion(t *testing.T) { }, }, { - name: "Error: os.ReadFile error when retrieving collector file - Internal Server Error", + name: "Error: file service read error when retrieving collector file - Internal Server Error", buildRequest: func() *http.Request { return &http.Request{ URL: &url.URL{ @@ -99,7 +102,27 @@ func TestResources_DownloadCollectorByVersion(t *testing.T) { } }, setupMocks: func(t *testing.T, mock *mock) { - mock.mockFS.EXPECT().ReadFile("azurehound/azurehound-latest.zip").Return([]byte{}, errors.New("error")) + mock.mockFileServiceResolver.EXPECT().Resolve(storage.FileServiceCollectors).Return(mock.mockFS, nil) + mock.mockFS.EXPECT().ReadFile(gomock.Any(), "azurehound/azurehound-latest.zip").Return([]byte{}, errors.New("error")) + }, + expected: expected{ + responseCode: http.StatusInternalServerError, + responseBody: `{"errors":[{"context":"","message":"an internal error has occurred that is preventing the service from servicing this request"}],"http_status":500,"request_id":"","timestamp":"0001-01-01T00:00:00Z"}`, + responseHeader: http.Header{"Content-Type": []string{"application/json"}}, + }, + }, + { + name: "Error: file service resolver error - Internal Server Error", + buildRequest: func() *http.Request { + return &http.Request{ + URL: &url.URL{ + Path: "/api/v2/collectors/azurehound/latest", + }, + Method: http.MethodGet, + } + }, + setupMocks: func(t *testing.T, mock *mock) { + mock.mockFileServiceResolver.EXPECT().Resolve(storage.FileServiceCollectors).Return(nil, errors.New("error")) }, expected: expected{ responseCode: http.StatusInternalServerError, @@ -118,7 +141,8 @@ func TestResources_DownloadCollectorByVersion(t *testing.T) { } }, setupMocks: func(t *testing.T, mock *mock) { - mock.mockFS.EXPECT().ReadFile("azurehound/azurehound-v1.0.0.zip").Return([]byte{}, nil) + mock.mockFileServiceResolver.EXPECT().Resolve(storage.FileServiceCollectors).Return(mock.mockFS, nil) + mock.mockFS.EXPECT().ReadFile(gomock.Any(), "azurehound/azurehound-v1.0.0.zip").Return([]byte{}, nil) }, expected: expected{ responseCode: http.StatusOK, @@ -136,7 +160,8 @@ func TestResources_DownloadCollectorByVersion(t *testing.T) { } }, setupMocks: func(t *testing.T, mock *mock) { - mock.mockFS.EXPECT().ReadFile("azurehound/azurehound-latest.zip").Return([]byte{}, nil) + mock.mockFileServiceResolver.EXPECT().Resolve(storage.FileServiceCollectors).Return(mock.mockFS, nil) + mock.mockFS.EXPECT().ReadFile(gomock.Any(), "azurehound/azurehound-latest.zip").Return([]byte{}, nil) }, expected: expected{ responseCode: http.StatusOK, @@ -150,7 +175,8 @@ func TestResources_DownloadCollectorByVersion(t *testing.T) { ctrl := gomock.NewController(t) mock := &mock{ - mockFS: fsmocks.NewMockService(ctrl), + mockFS: storagemocks.NewMockFileService(ctrl), + mockFileServiceResolver: storageServiceMocks.NewMockFileServiceResolver(ctrl), } testCase.setupMocks(t, mock) @@ -162,8 +188,8 @@ func TestResources_DownloadCollectorByVersion(t *testing.T) { } resources := v2.Resources{ - CollectorManifests: collectorManifests, - FileService: mock.mockFS, + CollectorManifests: collectorManifests, + FileServiceResolver: mock.mockFileServiceResolver, } response := httptest.NewRecorder() @@ -187,7 +213,8 @@ func TestResources_DownloadCollectorByVersion(t *testing.T) { func TestResources_DownloadCollectorChecksumByVersion(t *testing.T) { type mock struct { - mockFS *fsmocks.MockService + mockFS *storagemocks.MockFileService + mockFileServiceResolver *storageServiceMocks.MockFileServiceResolver } type expected struct { responseBody string @@ -236,7 +263,27 @@ func TestResources_DownloadCollectorChecksumByVersion(t *testing.T) { }, }, { - name: "Error: os.ReadFile error when retrieving collector file - Internal Server Error", + name: "Error: file service read error when retrieving collector file - Internal Server Error", + buildRequest: func() *http.Request { + return &http.Request{ + URL: &url.URL{ + Path: "/api/v2/collectors/azurehound/latest/checksum", + }, + Method: http.MethodGet, + } + }, + setupMocks: func(t *testing.T, mock *mock) { + mock.mockFileServiceResolver.EXPECT().Resolve(storage.FileServiceCollectors).Return(mock.mockFS, nil) + mock.mockFS.EXPECT().ReadFile(gomock.Any(), "azurehound/azurehound-latest.zip.sha256").Return([]byte{}, errors.New("error")) + }, + expected: expected{ + responseCode: http.StatusInternalServerError, + responseBody: `{"errors":[{"context":"","message":"an internal error has occurred that is preventing the service from servicing this request"}],"http_status":500,"request_id":"","timestamp":"0001-01-01T00:00:00Z"}`, + responseHeader: http.Header{"Content-Type": []string{"application/json"}}, + }, + }, + { + name: "Error: file service resolver error - Internal Server Error", buildRequest: func() *http.Request { return &http.Request{ URL: &url.URL{ @@ -246,7 +293,7 @@ func TestResources_DownloadCollectorChecksumByVersion(t *testing.T) { } }, setupMocks: func(t *testing.T, mock *mock) { - mock.mockFS.EXPECT().ReadFile("azurehound/azurehound-latest.zip.sha256").Return([]byte{}, errors.New("error")) + mock.mockFileServiceResolver.EXPECT().Resolve(storage.FileServiceCollectors).Return(nil, errors.New("error")) }, expected: expected{ responseCode: http.StatusInternalServerError, @@ -265,7 +312,8 @@ func TestResources_DownloadCollectorChecksumByVersion(t *testing.T) { } }, setupMocks: func(t *testing.T, mock *mock) { - mock.mockFS.EXPECT().ReadFile("azurehound/azurehound-v1.0.0.zip.sha256").Return([]byte{}, nil) + mock.mockFileServiceResolver.EXPECT().Resolve(storage.FileServiceCollectors).Return(mock.mockFS, nil) + mock.mockFS.EXPECT().ReadFile(gomock.Any(), "azurehound/azurehound-v1.0.0.zip.sha256").Return([]byte{}, nil) }, expected: expected{ responseCode: http.StatusOK, @@ -283,7 +331,8 @@ func TestResources_DownloadCollectorChecksumByVersion(t *testing.T) { } }, setupMocks: func(t *testing.T, mock *mock) { - mock.mockFS.EXPECT().ReadFile("azurehound/azurehound-latest.zip.sha256").Return([]byte{}, nil) + mock.mockFileServiceResolver.EXPECT().Resolve(storage.FileServiceCollectors).Return(mock.mockFS, nil) + mock.mockFS.EXPECT().ReadFile(gomock.Any(), "azurehound/azurehound-latest.zip.sha256").Return([]byte{}, nil) }, expected: expected{ responseCode: http.StatusOK, @@ -304,13 +353,14 @@ func TestResources_DownloadCollectorChecksumByVersion(t *testing.T) { ctrl := gomock.NewController(t) mock := &mock{ - mockFS: fsmocks.NewMockService(ctrl), + mockFS: storagemocks.NewMockFileService(ctrl), + mockFileServiceResolver: storageServiceMocks.NewMockFileServiceResolver(ctrl), } testCase.setupMocks(t, mock) resources := v2.Resources{ - CollectorManifests: collectorManifests, - FileService: mock.mockFS, + CollectorManifests: collectorManifests, + FileServiceResolver: mock.mockFileServiceResolver, } response := httptest.NewRecorder() diff --git a/cmd/api/src/api/v2/fileingest.go b/cmd/api/src/api/v2/fileingest.go index 4df9be834f4b..d4e06302eb4e 100644 --- a/cmd/api/src/api/v2/fileingest.go +++ b/cmd/api/src/api/v2/fileingest.go @@ -22,7 +22,6 @@ import ( "log/slog" "mime" "net/http" - "os" "slices" "strconv" "strings" @@ -36,6 +35,7 @@ import ( ingestModel "github.com/specterops/bloodhound/cmd/api/src/model/ingest" "github.com/specterops/bloodhound/packages/go/bhlog/measure" "github.com/specterops/bloodhound/packages/go/headers" + "github.com/specterops/bloodhound/packages/go/storage" "github.com/specterops/bloodhound/cmd/api/src/services/job" "github.com/specterops/bloodhound/cmd/api/src/services/upload" @@ -146,7 +146,9 @@ func (s Resources) ProcessIngestTask(response http.ResponseWriter, request *http api.HandleDatabaseError(request, response, err) } else if ingestJob.Status != model.JobStatusRunning { api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusBadRequest, "job must be in running status to attach files", request), response) - } else if ingestTaskParams, err := upload.SaveIngestFile(s.Config.TempDirectory(), request, validator, ingestJob.ID); errors.Is(err, upload.ErrInvalidJSON) { + } else if ingestFileService, err := s.FileServiceResolver.Resolve(storage.FileServiceIngest); err != nil { + api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusInternalServerError, "unable to resolve file service for working directories", request), response) + } else if ingestTaskParams, err := upload.SaveIngestFile(request.Context(), ingestFileService, request, validator, ingestJob.ID); errors.Is(err, upload.ErrInvalidJSON) { api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusBadRequest, fmt.Sprintf("Error saving ingest file: %v", err), request), response) } else if report, ok := err.(upload.ValidationReport); ok { var ( @@ -169,7 +171,7 @@ func (s Resources) ProcessIngestTask(response http.ResponseWriter, request *http } else if err != nil { api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("Error saving ingest file: %v", err), request), response) } else if _, err = upload.CreateIngestTask(request.Context(), s.DB, upload.IngestTaskParams{Filename: ingestTaskParams.Filename, ProvidedFileName: checkFileName(fileName, ingestTaskParams.FileType), FileType: ingestTaskParams.FileType, RequestID: requestId, JobID: int64(jobID)}); err != nil { - if removeErr := os.Remove(ingestTaskParams.Filename); removeErr != nil { + if removeErr := ingestFileService.DeleteFile(request.Context(), ingestTaskParams.Filename); removeErr != nil { slog.WarnContext(request.Context(), "Failed to clean up file after task creation error", attr.Error(removeErr)) } api.HandleDatabaseError(request, response, err) diff --git a/cmd/api/src/api/v2/fileingest_test.go b/cmd/api/src/api/v2/fileingest_test.go index 0f675ecda18c..53a09b4f8b87 100644 --- a/cmd/api/src/api/v2/fileingest_test.go +++ b/cmd/api/src/api/v2/fileingest_test.go @@ -43,7 +43,10 @@ import ( "github.com/specterops/bloodhound/cmd/api/src/database/types/null" "github.com/specterops/bloodhound/cmd/api/src/model" "github.com/specterops/bloodhound/cmd/api/src/model/ingest" + storageServiceMocks "github.com/specterops/bloodhound/cmd/api/src/services/storage/mocks" "github.com/specterops/bloodhound/packages/go/headers" + "github.com/specterops/bloodhound/packages/go/storage" + storagemocks "github.com/specterops/bloodhound/packages/go/storage/mocks" "github.com/specterops/bloodhound/cmd/api/src/utils/test" "github.com/stretchr/testify/assert" @@ -479,9 +482,19 @@ func TestResources_ListAcceptedFileUploadTypes(t *testing.T) { }) } +func newLocalTempFileService(t *testing.T) storage.FileService { + t.Helper() + ls, err := storage.NewLocalStore(t.TempDir()) + require.NoError(t, err, "failed to create local store") + return storage.NewFileService(ls) +} + func TestResources_ProcessIngestTask(t *testing.T) { type mock struct { - mockDatabase *dbmocks.MockDatabase + mockDatabase *dbmocks.MockDatabase + mockFileService *storagemocks.MockFileService + trueFileService storage.FileService + mockFileServiceResolver *storageServiceMocks.MockFileServiceResolver } type expected struct { responseBody string @@ -489,10 +502,11 @@ func TestResources_ProcessIngestTask(t *testing.T) { responseHeader http.Header } type testData struct { - name string - buildRequest func() *http.Request - setupMocks func(t *testing.T, mock *mock) - expected expected + name string + buildRequest func() *http.Request + setupMocks func(t *testing.T, mock *mock) + fileServiceOvrd storage.FileService // if non-nil, overrides the mock for this test case + expected expected } tt := []testData{ @@ -573,13 +587,41 @@ func TestResources_ProcessIngestTask(t *testing.T) { setupMocks: func(t *testing.T, mock *mock) { t.Helper() mock.mockDatabase.EXPECT().GetIngestJob(gomock.Any(), int64(1)).Return(model.IngestJob{Status: model.JobStatusRunning}, nil) + mock.mockFileServiceResolver.EXPECT().Resolve(storage.FileServiceIngest).Return(mock.trueFileService, nil) }, + // Use a real LocalStore so that actual file I/O gives the concurrent validation + // goroutine time to set validationErr before the main goroutine reads it. + fileServiceOvrd: newLocalTempFileService(t), expected: expected{ responseCode: http.StatusBadRequest, responseBody: `{"errors":[{"context":"","message":"Error saving ingest file: file is not valid json"}],"http_status":400,"request_id":"","timestamp":"0001-01-01T00:00:00Z"}`, responseHeader: http.Header{"Content-Type": []string{"application/json"}}, }, }, + { + name: "Error: file service resolver error - Internal Server Error", + buildRequest: func() *http.Request { + return &http.Request{ + URL: &url.URL{ + Path: "/api/v2/file-upload/1", + }, + Method: http.MethodPost, + Header: http.Header{ + headers.ContentType.String(): []string{"application/json"}, + }, + } + }, + setupMocks: func(t *testing.T, mock *mock) { + t.Helper() + mock.mockDatabase.EXPECT().GetIngestJob(gomock.Any(), int64(1)).Return(model.IngestJob{Status: model.JobStatusRunning}, nil) + mock.mockFileServiceResolver.EXPECT().Resolve(storage.FileServiceIngest).Return(nil, errors.New("error")) + }, + expected: expected{ + responseCode: http.StatusInternalServerError, + responseBody: `{"errors":[{"context":"","message":"unable to resolve file service for working directories"}],"http_status":500,"request_id":"","timestamp":"0001-01-01T00:00:00Z"}`, + responseHeader: http.Header{"Content-Type": []string{"application/json"}}, + }, + }, { name: "Error: error saving ingest file - Internal Server Error", buildRequest: func() *http.Request { @@ -602,7 +644,10 @@ func TestResources_ProcessIngestTask(t *testing.T) { setupMocks: func(t *testing.T, mock *mock) { t.Helper() mock.mockDatabase.EXPECT().GetIngestJob(gomock.Any(), int64(1)).Return(model.IngestJob{Status: model.JobStatusRunning}, nil) + mock.mockFileServiceResolver.EXPECT().Resolve(storage.FileServiceIngest).Return(mock.trueFileService, nil) }, + // Use a real LocalStore for the same reason as the ErrInvalidJSON case above. + fileServiceOvrd: newLocalTempFileService(t), expected: expected{ responseCode: http.StatusInternalServerError, responseBody: `{"errors":[{"context":"","message":"Error saving ingest file: no valid meta tag or data tag found"}],"http_status":500,"request_id":"","timestamp":"0001-01-01T00:00:00Z"}`, @@ -625,7 +670,14 @@ func TestResources_ProcessIngestTask(t *testing.T) { }, setupMocks: func(t *testing.T, mock *mock) { t.Helper() + tmpFileName := "tmpFileName" mock.mockDatabase.EXPECT().GetIngestJob(gomock.Any(), int64(1)).Return(model.IngestJob{Status: model.JobStatusRunning}, nil) + mock.mockFileServiceResolver.EXPECT().Resolve(storage.FileServiceIngest).Return(mock.mockFileService, nil) + mock.mockFileService.EXPECT().WriteTempFile(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, _ string, reader io.Reader, _ storage.WriteOptions) (string, error) { + _, err := io.ReadAll(reader) + return tmpFileName, err + }) + mock.mockFileService.EXPECT().DeleteFile(gomock.Any(), tmpFileName).Return(nil) mock.mockDatabase.EXPECT().CreateIngestTask(gomock.Any(), gomock.Any()).Return(model.IngestTask{}, errors.New("error")) }, expected: expected{ @@ -651,6 +703,11 @@ func TestResources_ProcessIngestTask(t *testing.T) { setupMocks: func(t *testing.T, mock *mock) { t.Helper() mock.mockDatabase.EXPECT().GetIngestJob(gomock.Any(), int64(1)).Return(model.IngestJob{Status: model.JobStatusRunning}, nil) + mock.mockFileServiceResolver.EXPECT().Resolve(storage.FileServiceIngest).Return(mock.mockFileService, nil) + mock.mockFileService.EXPECT().WriteTempFile(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, _ string, reader io.Reader, _ storage.WriteOptions) (string, error) { + _, err := io.ReadAll(reader) + return "/tmp/test", err + }) mock.mockDatabase.EXPECT().CreateIngestTask(gomock.Any(), gomock.Any()).Return(model.IngestTask{}, nil) mock.mockDatabase.EXPECT().UpdateIngestJob(gomock.Any(), gomock.Any()).Return(errors.New("error")) }, @@ -677,6 +734,11 @@ func TestResources_ProcessIngestTask(t *testing.T) { setupMocks: func(t *testing.T, mock *mock) { t.Helper() mock.mockDatabase.EXPECT().GetIngestJob(gomock.Any(), int64(1)).Return(model.IngestJob{Status: model.JobStatusRunning}, nil) + mock.mockFileServiceResolver.EXPECT().Resolve(storage.FileServiceIngest).Return(mock.mockFileService, nil) + mock.mockFileService.EXPECT().WriteTempFile(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, _ string, reader io.Reader, _ storage.WriteOptions) (string, error) { + _, err := io.ReadAll(reader) + return "/tmp/test", err + }) mock.mockDatabase.EXPECT().CreateIngestTask(gomock.Any(), gomock.Cond(func(x model.IngestTask) bool { return x.OriginalFileName == "UnknownFileName.json" })).Return(model.IngestTask{}, nil) @@ -705,6 +767,11 @@ func TestResources_ProcessIngestTask(t *testing.T) { setupMocks: func(t *testing.T, mock *mock) { t.Helper() mock.mockDatabase.EXPECT().GetIngestJob(gomock.Any(), int64(1)).Return(model.IngestJob{Status: model.JobStatusRunning}, nil) + mock.mockFileServiceResolver.EXPECT().Resolve(storage.FileServiceIngest).Return(mock.mockFileService, nil) + mock.mockFileService.EXPECT().WriteTempFile(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, _ string, reader io.Reader, _ storage.WriteOptions) (string, error) { + _, err := io.ReadAll(reader) + return "/tmp/test", err + }) mock.mockDatabase.EXPECT().CreateIngestTask(gomock.Any(), gomock.Cond(func(x model.IngestTask) bool { return x.OriginalFileName == "Testing.json" })).Return(model.IngestTask{}, nil) @@ -751,6 +818,11 @@ func TestResources_ProcessIngestTask(t *testing.T) { setupMocks: func(t *testing.T, mock *mock) { t.Helper() mock.mockDatabase.EXPECT().GetIngestJob(gomock.Any(), int64(1)).Return(model.IngestJob{Status: model.JobStatusRunning}, nil) + mock.mockFileServiceResolver.EXPECT().Resolve(storage.FileServiceIngest).Return(mock.mockFileService, nil) + mock.mockFileService.EXPECT().WriteTempFile(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, _ string, reader io.Reader, _ storage.WriteOptions) (string, error) { + _, err := io.ReadAll(reader) + return "/tmp/test", err + }) mock.mockDatabase.EXPECT().CreateIngestTask(gomock.Any(), gomock.Cond(func(x model.IngestTask) bool { return x.OriginalFileName == "Testing.zip" })).Return(model.IngestTask{}, nil) @@ -767,15 +839,22 @@ func TestResources_ProcessIngestTask(t *testing.T) { ctrl := gomock.NewController(t) mocks := &mock{ - mockDatabase: dbmocks.NewMockDatabase(ctrl), + mockDatabase: dbmocks.NewMockDatabase(ctrl), + mockFileService: storagemocks.NewMockFileService(ctrl), + mockFileServiceResolver: storageServiceMocks.NewMockFileServiceResolver(ctrl), + } + + if testCase.fileServiceOvrd != nil { + mocks.trueFileService = testCase.fileServiceOvrd } request := testCase.buildRequest() testCase.setupMocks(t, mocks) resources := v2.Resources{ - DB: mocks.mockDatabase, - Config: config.Configuration{}, + DB: mocks.mockDatabase, + Config: config.Configuration{}, + FileServiceResolver: mocks.mockFileServiceResolver, } err := os.Mkdir(resources.Config.TempDirectory(), 0755) diff --git a/cmd/api/src/api/v2/model.go b/cmd/api/src/api/v2/model.go index cca3878b94d1..eca05d084d9a 100644 --- a/cmd/api/src/api/v2/model.go +++ b/cmd/api/src/api/v2/model.go @@ -27,7 +27,7 @@ import ( "github.com/specterops/bloodhound/cmd/api/src/queries" "github.com/specterops/bloodhound/cmd/api/src/serde" "github.com/specterops/bloodhound/cmd/api/src/services/dogtags" - "github.com/specterops/bloodhound/cmd/api/src/services/fs" + "github.com/specterops/bloodhound/cmd/api/src/services/storage" "github.com/specterops/bloodhound/cmd/api/src/services/upload" "github.com/specterops/bloodhound/packages/go/cache" "github.com/specterops/dawgs/graph" @@ -115,7 +115,7 @@ type Resources struct { Authorizer auth.Authorizer Authenticator api.Authenticator IngestSchema upload.IngestSchema - FileService fs.Service + FileServiceResolver storage.FileServiceResolver OpenGraphSchemaService OpenGraphSchemaService DogTags dogtags.Service } @@ -130,6 +130,7 @@ func NewResources( authorizer auth.Authorizer, authenticator api.Authenticator, ingestSchema upload.IngestSchema, + fileServiceResolver storage.FileServiceResolver, dogtagsService dogtags.Service, openGraphSchemaService OpenGraphSchemaService, ) Resources { @@ -145,7 +146,7 @@ func NewResources( Authorizer: authorizer, Authenticator: authenticator, IngestSchema: ingestSchema, - FileService: &fs.Client{}, + FileServiceResolver: fileServiceResolver, DogTags: dogtagsService, OpenGraphSchemaService: openGraphSchemaService, } diff --git a/cmd/api/src/bootstrap/initializer.go b/cmd/api/src/bootstrap/initializer.go index abb0ba1070e6..0fee25510d88 100644 --- a/cmd/api/src/bootstrap/initializer.go +++ b/cmd/api/src/bootstrap/initializer.go @@ -18,6 +18,7 @@ package bootstrap import ( "context" + "errors" "fmt" "log/slog" @@ -33,13 +34,15 @@ type DatabaseConnections[DBType database.Database, GraphType graph.Database] str } type DatabaseConstructor[DBType database.Database, GraphType graph.Database] func(ctx context.Context, cfg config.Configuration) (DatabaseConnections[DBType, GraphType], error) -type InitializerLogic[DBType database.Database, GraphType graph.Database] func(ctx context.Context, cfg config.Configuration, databaseConnections DatabaseConnections[DBType, GraphType]) ([]daemons.Daemon, error) +type RuntimeDependenciesConstructor[DBType database.Database, GraphType graph.Database] func(ctx context.Context, cfg config.Configuration, databaseConnections DatabaseConnections[DBType, GraphType]) (RuntimeDependencies, error) +type InitializerLogic[DBType database.Database, GraphType graph.Database] func(ctx context.Context, cfg config.Configuration, databaseConnections DatabaseConnections[DBType, GraphType], dependencies RuntimeDependencies) ([]daemons.Daemon, error) type Initializer[DBType database.Database, GraphType graph.Database] struct { Configuration config.Configuration + DBConnector DatabaseConstructor[DBType, GraphType] + RuntimeDependencies RuntimeDependenciesConstructor[DBType, GraphType] PreMigrationDaemons InitializerLogic[DBType, GraphType] Entrypoint InitializerLogic[DBType, GraphType] - DBConnector DatabaseConstructor[DBType, GraphType] } func (s Initializer[DBType, GraphType]) Launch(parentCtx context.Context, handleSignals bool) error { @@ -47,6 +50,7 @@ func (s Initializer[DBType, GraphType]) Launch(parentCtx context.Context, handle ctx = parentCtx daemonManager = daemons.NewManager(DefaultServerShutdownTimeout) databaseConnections DatabaseConnections[DBType, GraphType] + runtimeDependencies RuntimeDependencies err error ) @@ -65,9 +69,17 @@ func (s Initializer[DBType, GraphType]) Launch(parentCtx context.Context, handle defer databaseConnections.RDMS.Close(ctx) defer databaseConnections.Graph.Close(ctx) + if s.RuntimeDependencies == nil { + return errors.New("runtime dependencies constructor is required") + } + + if runtimeDependencies, err = s.RuntimeDependencies(ctx, s.Configuration, databaseConnections); err != nil { + return fmt.Errorf("failed to initialize runtime dependencies: %w", err) + } + // Daemons that start prior to blocking db migration if s.PreMigrationDaemons != nil { - if daemonInstances, err := s.PreMigrationDaemons(ctx, s.Configuration, databaseConnections); err != nil { + if daemonInstances, err := s.PreMigrationDaemons(ctx, s.Configuration, databaseConnections, runtimeDependencies); err != nil { return fmt.Errorf("failed to start services: %w", err) } else { daemonManager.Start(ctx, daemonInstances...) @@ -75,7 +87,7 @@ func (s Initializer[DBType, GraphType]) Launch(parentCtx context.Context, handle } // Daemons that start after blocking db migration - if daemonInstances, err := s.Entrypoint(ctx, s.Configuration, databaseConnections); err != nil { + if daemonInstances, err := s.Entrypoint(ctx, s.Configuration, databaseConnections, runtimeDependencies); err != nil { return fmt.Errorf("failed to start services: %w", err) } else { daemonManager.Start(ctx, daemonInstances...) diff --git a/cmd/api/src/bootstrap/server_test.go b/cmd/api/src/bootstrap/server_test.go index d4ddfb966331..9251cb67eb9f 100644 --- a/cmd/api/src/bootstrap/server_test.go +++ b/cmd/api/src/bootstrap/server_test.go @@ -17,6 +17,8 @@ package bootstrap_test import ( + "os" + "path/filepath" "testing" "github.com/specterops/bloodhound/cmd/api/src/bootstrap" @@ -71,3 +73,36 @@ func TestFillAndPopulateDefaultAdminInfo(t *testing.T) { require.NotEqual(t, "", cfg.PrincipalName) } } + +func TestEnsureServerDirectoriesCreatesRequiredDirectories(t *testing.T) { + t.Parallel() + + var ( + rootDirectory = t.TempDir() + cfg = config.Configuration{ + WorkDir: filepath.Join(rootDirectory, "work"), + CollectorsBasePath: filepath.Join(rootDirectory, "collectors"), + } + ) + + require.NoError(t, bootstrap.EnsureServerDirectories(cfg)) + + for _, directory := range []string{ + cfg.WorkDir, + cfg.TempDirectory(), + cfg.ScratchDirectory(), + cfg.RetainedFilesDirectory(), + cfg.ClientLogDirectory(), + cfg.CollectorsDirectory(), + } { + requireDirectoryExists(t, directory) + } +} + +func requireDirectoryExists(t *testing.T, path string) { + t.Helper() + + info, err := os.Stat(path) + require.NoError(t, err) + require.True(t, info.IsDir()) +} diff --git a/cmd/api/src/bootstrap/util.go b/cmd/api/src/bootstrap/util.go index c19d74a0b136..de83ec60de0a 100644 --- a/cmd/api/src/bootstrap/util.go +++ b/cmd/api/src/bootstrap/util.go @@ -26,6 +26,7 @@ import ( "github.com/specterops/bloodhound/cmd/api/src/api/dbpool" "github.com/specterops/bloodhound/cmd/api/src/api/tools" "github.com/specterops/bloodhound/cmd/api/src/config" + "github.com/specterops/bloodhound/cmd/api/src/services/storage" "github.com/specterops/dawgs" "github.com/specterops/dawgs/drivers/neo4j" "github.com/specterops/dawgs/drivers/pg" @@ -33,6 +34,14 @@ import ( "github.com/specterops/dawgs/util/size" ) +// RuntimeDependencies holds values that must be created before the entrypoint starts. For instance +// IngestControl is reliant on the FileService. In order for the pre-migration toolapi to have +// access to the FileServiceRetained, the FileServiceResolver is created prior to the +// PreMigrationDaemons and the Entrypoint. This could then be passed in. +type RuntimeDependencies struct { + FileServiceResolver storage.FileServiceResolver +} + func ensureDirectory(path string) error { if _, err := os.Stat(path); err != nil { if !os.IsNotExist(err) { @@ -58,6 +67,10 @@ func EnsureServerDirectories(cfg config.Configuration) error { return err } + if err := ensureDirectory(cfg.ScratchDirectory()); err != nil { + return err + } + if err := ensureDirectory(cfg.RetainedFilesDirectory()); err != nil { return err } @@ -70,6 +83,10 @@ func EnsureServerDirectories(cfg config.Configuration) error { return err } + if err := ensureDirectory(cfg.CollectorArtifactDirectory()); err != nil { + return err + } + return nil } diff --git a/cmd/api/src/cmd/bhapi/main.go b/cmd/api/src/cmd/bhapi/main.go index 874435a740e0..301ec39c91f4 100644 --- a/cmd/api/src/cmd/bhapi/main.go +++ b/cmd/api/src/cmd/bhapi/main.go @@ -128,6 +128,7 @@ func main() { initializer := bootstrap.Initializer[*database.BloodhoundDB, *graph.DatabaseSwitch]{ Configuration: cfg, DBConnector: services.ConnectDatabases, + RuntimeDependencies: services.CreateRuntimeDependencies, PreMigrationDaemons: services.PreMigrationDaemons, Entrypoint: services.Entrypoint, } diff --git a/cmd/api/src/config/config.go b/cmd/api/src/config/config.go index 0974d561296d..a7a44e60cbc1 100644 --- a/cmd/api/src/config/config.go +++ b/cmd/api/src/config/config.go @@ -255,6 +255,10 @@ type Configuration struct { Teleport TeleportConfiguration `json:"teleport"` } +func (s Configuration) ScratchDirectory() string { + return filepath.Join(s.WorkDir, "ingest_scratch") +} + func (s Configuration) TempDirectory() string { return filepath.Join(s.WorkDir, "tmp") } @@ -271,6 +275,10 @@ func (s Configuration) CollectorsDirectory() string { return s.CollectorsBasePath } +func (s Configuration) CollectorArtifactDirectory() string { + return filepath.Join(s.WorkDir, "collector_artifacts") +} + func WriteConfigurationFile(path string, config Configuration) error { if fout, err := os.OpenFile(path, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0644); err != nil { return fmt.Errorf("failed opening configuration file %s: %w", path, err) diff --git a/cmd/api/src/daemons/api/toolapi/api.go b/cmd/api/src/daemons/api/toolapi/api.go index decfb07a80dd..a7bed6bbbac0 100644 --- a/cmd/api/src/daemons/api/toolapi/api.go +++ b/cmd/api/src/daemons/api/toolapi/api.go @@ -34,6 +34,7 @@ import ( "github.com/specterops/bloodhound/cmd/api/src/config" "github.com/specterops/bloodhound/cmd/api/src/database" "github.com/specterops/bloodhound/packages/go/bhlog/attr" + "github.com/specterops/bloodhound/packages/go/storage" "github.com/specterops/dawgs/graph" ) @@ -43,12 +44,12 @@ type Daemon struct { server *http.Server } -func NewDaemon[DBType database.Database](ctx context.Context, connections bootstrap.DatabaseConnections[DBType, *graph.DatabaseSwitch], cfg config.Configuration, graphSchema graph.Schema, extensions ...func(router *chi.Mux)) Daemon { +func NewDaemon[DBType database.Database](ctx context.Context, connections bootstrap.DatabaseConnections[DBType, *graph.DatabaseSwitch], cfg config.Configuration, graphSchema graph.Schema, retainedFileService storage.FileService, extensions ...func(router *chi.Mux)) Daemon { var ( pgMigrator = tools.NewPGMigrator(ctx, cfg, graphSchema, connections.Graph) - ingestControl = tools.NewIngestControlTool(cfg, connections.RDMS) router = chi.NewRouter() toolContainer = tools.NewToolContainer(connections.RDMS) + ingestControl = tools.NewIngestControlTool(connections.RDMS, retainedFileService) ) router.Mount("/metrics", promhttp.Handler()) diff --git a/cmd/api/src/daemons/datapipe/cleanup.go b/cmd/api/src/daemons/datapipe/cleanup.go index 08eb8feeb91c..58da815ada8f 100644 --- a/cmd/api/src/daemons/datapipe/cleanup.go +++ b/cmd/api/src/daemons/datapipe/cleanup.go @@ -20,6 +20,8 @@ package datapipe import ( "context" + "path" + "time" "log/slog" "os" @@ -28,6 +30,12 @@ import ( "sync" "github.com/specterops/bloodhound/packages/go/bhlog/attr" + "github.com/specterops/bloodhound/packages/go/storage" +) + +const ( + orphanMinimumAge = 24 * time.Hour + scratchMinimumAge = 24 * time.Hour ) // FileOperations is an interface for describing filesystem actions. This implementation exists due to deficiencies in @@ -55,112 +63,270 @@ func (s osFileOperations) RemoveAll(path string) error { // OrphanFileSweeper is a file cleanup utility that allows only one goroutine to attempt file cleanup at a time. type OrphanFileSweeper struct { - lock *sync.Mutex - fileOps FileOperations - tempDirectoryRootPath string + lock *sync.Mutex + fileOps FileOperations + tempDirectoryRootPath string + scratchDirectoryRootPath string + excludedStoragePrefixes []string } -func NewOrphanFileSweeper(fileOps FileOperations, tempDirectoryRootPath string) *OrphanFileSweeper { - return &OrphanFileSweeper{ - lock: &sync.Mutex{}, - fileOps: fileOps, - tempDirectoryRootPath: strings.TrimSuffix(tempDirectoryRootPath, string(filepath.Separator)), +func normalizeStoragePrefixes(prefixes []string) []string { + var normalizedPrefixes []string + + for _, prefix := range prefixes { + prefix = strings.TrimSpace(filepath.ToSlash(prefix)) + prefix = strings.Trim(prefix, "/") + if prefix == "" { + continue + } + + prefix = path.Clean(prefix) + if prefix == "." || prefix == ".." || strings.HasPrefix(prefix, "../") { + continue + } + + normalizedPrefixes = append(normalizedPrefixes, prefix) } + + return normalizedPrefixes } -// Clear takes a context and a list of expected file names. The function will list all directory entries within the -// configured tempDirectoryRootPath that the sweeper was instantiated with and compare the resulting list against the -// passed expected file names. The function will then call RemoveAll on each directory entry not found in the expected -// file name slice. -func (s *OrphanFileSweeper) Clear(ctx context.Context, expectedFileNames []string) { - // Only allow one background thread to run for clearing orphaned data/ - if !s.lock.TryLock() { - return +func NewOrphanFileSweeper(fileOps FileOperations, tempDirectoryRootPath string, scratchDirectoryRootPath string, excludedStoragePrefixes ...string) *OrphanFileSweeper { + return &OrphanFileSweeper{ + lock: &sync.Mutex{}, + fileOps: fileOps, + tempDirectoryRootPath: strings.TrimSuffix(tempDirectoryRootPath, string(filepath.Separator)), + scratchDirectoryRootPath: strings.TrimSuffix(scratchDirectoryRootPath, string(filepath.Separator)), + excludedStoragePrefixes: normalizeStoragePrefixes(excludedStoragePrefixes), } +} - // Release the lock once finished - defer s.lock.Unlock() - - slog.InfoContext( - ctx, - "Running OrphanFileSweeper", - slog.String("path", s.tempDirectoryRootPath), - ) - slog.DebugContext( - ctx, - "OrphanFileSweeper expected names", - slog.String("expected_file_names", strings.Join(expectedFileNames, ",")), - ) - - if dirEntries, err := s.fileOps.ReadDir(s.tempDirectoryRootPath); err != nil { +func ClearLocalIngestScratch(ctx context.Context, scratchDirectory string, minimumAge time.Duration) { + entries, err := os.ReadDir(scratchDirectory) + if err != nil { slog.ErrorContext( ctx, - "Failed reading work directory", - slog.String("path", s.tempDirectoryRootPath), + "Error reading scratch directory", + slog.String("scratch_directory", scratchDirectory), attr.Error(err), ) - } else { - numDeleted := 0 - - // Remove expected files from the deletion list - for _, expectedFileName := range expectedFileNames { - expectedDir, expectedFile := filepath.Split(expectedFileName) - if expectedDir != "" { - expectedDir = strings.TrimSuffix(expectedDir, string(filepath.Separator)) - if expectedDir != s.tempDirectoryRootPath { - slog.WarnContext( - ctx, - "Directory does not match temp directory root path for expected file, skipping", - slog.String("expected_dir", expectedDir), - slog.String("expected_file_name", expectedFileName), - slog.String("path", s.tempDirectoryRootPath), - ) - continue - } - } - for idx, dirEntry := range dirEntries { - if expectedFile == dirEntry.Name() { - slog.DebugContext( - ctx, - "Skipping expected file", - slog.String("expected_file", expectedFile), - ) - dirEntries = append(dirEntries[:idx], dirEntries[idx+1:]...) - } - } - } - - for _, orphanedDirEntry := range dirEntries { - // Check for context cancellation before each file deletion - if ctx.Err() != nil { - break - } - - slog.InfoContext( - ctx, - "Removing orphaned file", - slog.String("name", orphanedDirEntry.Name()), - ) - fullPath := filepath.Join(s.tempDirectoryRootPath, orphanedDirEntry.Name()) + return + } + + cutoff := time.Now().Add(-minimumAge) - if err := s.fileOps.RemoveAll(fullPath); err != nil { - slog.ErrorContext( - ctx, - "Failed removing orphaned file", - slog.String("full_path", fullPath), - attr.Error(err), - ) - } + for _, entry := range entries { + if ctx.Err() != nil { + return + } - numDeleted += 1 + if entry.IsDir() { + continue } - if numDeleted > 0 { - slog.InfoContext( + info, err := entry.Info() + if err != nil { + continue + } + + if info.ModTime().After(cutoff) { + continue + } + + _ = os.Remove(filepath.Join(scratchDirectory, entry.Name())) + } +} + +func (s *OrphanFileSweeper) isExcludedStoragePath(logicalPath string) bool { + logicalPath = path.Clean(strings.TrimPrefix(filepath.ToSlash(logicalPath), "/")) + + for _, excludedPrefix := range s.excludedStoragePrefixes { + if logicalPath == excludedPrefix || strings.HasPrefix(logicalPath, excludedPrefix+"/") { + return true + } + } + + return false +} + +// addExpectedLocalPath records an expected file path for the legacy local temp-directory cleanup pass. +// Expected names may be legacy absolute filesystem paths or file-service logical paths. Logical paths are +// normalized, resolved under tempDirectoryRootPath, and ignored if they are empty or would escape the temp root. +func (s *OrphanFileSweeper) addExpectedLocalPath(expectedLocalFiles map[string]struct{}, expectedFileName string) { + expectedFileName = strings.TrimSpace(expectedFileName) + if expectedFileName == "" { + return + } + + if filepath.IsAbs(expectedFileName) { + expectedLocalFiles[filepath.Clean(expectedFileName)] = struct{}{} + return + } + + logicalPath := path.Clean(filepath.ToSlash(expectedFileName)) + if logicalPath == "." || logicalPath == ".." || strings.HasPrefix(logicalPath, "../") { + return + } + + tempDirectoryRootPath := filepath.Clean(s.tempDirectoryRootPath) + localPath := filepath.Clean(filepath.Join(tempDirectoryRootPath, filepath.FromSlash(logicalPath))) + + relativePath, err := filepath.Rel(tempDirectoryRootPath, localPath) + if err != nil || relativePath == "." || relativePath == ".." || strings.HasPrefix(relativePath, ".."+string(filepath.Separator)) || filepath.IsAbs(relativePath) { + return + } + + expectedLocalFiles[localPath] = struct{}{} +} + +// addExpectedStoragePath handles saved file paths that were saved in absolute prior to this release +func (s *OrphanFileSweeper) addExpectedStoragePath(expectedFiles map[string]struct{}, expectedFileName string) { + expectedFileName = strings.TrimSpace(expectedFileName) + if expectedFileName == "" { + return + } + + if filepath.IsAbs(expectedFileName) { + relativePath, err := filepath.Rel(s.tempDirectoryRootPath, filepath.Clean(expectedFileName)) + if err != nil { + return + } + + if relativePath == "." || relativePath == ".." || strings.HasPrefix(relativePath, ".."+string(filepath.Separator)) || filepath.IsAbs(relativePath) { + return + } + + expectedFiles[path.Clean(filepath.ToSlash(relativePath))] = struct{}{} + return + } + + logicalPath := path.Clean(filepath.ToSlash(expectedFileName)) + if logicalPath == "." || logicalPath == ".." || strings.HasPrefix(logicalPath, "../") { + return + } + + expectedFiles[logicalPath] = struct{}{} +} + +func (s *OrphanFileSweeper) clearStoredIngestFiles(ctx context.Context, ingestFileService storage.FileService, expectedFileNames []string) { + expectedFiles := map[string]struct{}{} + + for _, expectedFileName := range expectedFileNames { + s.addExpectedStoragePath(expectedFiles, expectedFileName) + } + + files, err := ingestFileService.ListFiles(ctx, "", storage.ListOptions{Recursive: true}) + if err != nil { + slog.ErrorContext(ctx, "Failed listing ingest files", attr.Error(err)) + return + } + + for _, file := range files { + if ctx.Err() != nil { + return + } + if file.IsDir { + continue + } + + logicalPath := path.Clean(file.Path) + if s.isExcludedStoragePath(logicalPath) { + continue + } + + if _, ok := expectedFiles[logicalPath]; ok { + continue + } + + if !file.LastModified.IsZero() && time.Since(file.LastModified) < orphanMinimumAge { + continue + } + + if err := ingestFileService.DeleteFile(ctx, logicalPath); err != nil { + slog.WarnContext( ctx, - "Finished removing orphaned ingest files", - slog.Int("num_deleted", numDeleted), + "Failed deleting orphaned ingest file", + slog.String("path", logicalPath), + attr.Error(err), ) } } } + +func isExpectedLocalEntry(expectedLocalFiles map[string]struct{}, localPath string, isDirectory bool) bool { + localPath = filepath.Clean(localPath) + if _, ok := expectedLocalFiles[localPath]; ok { + return true + } + + if !isDirectory { + return false + } + + for expectedLocalFile := range expectedLocalFiles { + relativePath, err := filepath.Rel(localPath, expectedLocalFile) + if err != nil { + continue + } + + if relativePath != "." && relativePath != ".." && !strings.HasPrefix(relativePath, ".."+string(filepath.Separator)) && !filepath.IsAbs(relativePath) { + return true + } + } + + return false +} + +func (s *OrphanFileSweeper) clearLegacyLocalIngestFiles(ctx context.Context, expectedFileNames []string) { + expectedLocalFiles := map[string]struct{}{} + + for _, expectedFileName := range expectedFileNames { + s.addExpectedLocalPath(expectedLocalFiles, expectedFileName) + } + + entries, err := s.fileOps.ReadDir(s.tempDirectoryRootPath) + if err != nil { + slog.WarnContext(ctx, "Failed reading legacy ingest temp directory", attr.Error(err)) + return + } + + cutoff := time.Now().Add(-orphanMinimumAge) + + for _, entry := range entries { + if ctx.Err() != nil { + return + } + + fullPath := filepath.Join(s.tempDirectoryRootPath, entry.Name()) + if isExpectedLocalEntry(expectedLocalFiles, fullPath, entry.IsDir()) { + continue + } + + if entry.IsDir() { + _ = s.fileOps.RemoveAll(fullPath) + continue + } + + info, err := entry.Info() + if err != nil || info.ModTime().After(cutoff) { + continue + } + + _ = s.fileOps.RemoveAll(fullPath) + } +} + +// Clear takes a context and a list of expected file names. The function will list all directory entries within the +// configured tempDirectoryRootPath and scratchDirectoryRootPath that the sweeper was instantiated with and compare the +// resulting list against the passed expected file names. The function will then call RemoveAll on each directory entry not +// found in the expected file name slice. In addition, ingested files are also deleted using the supplied file service. +func (s *OrphanFileSweeper) Clear(ctx context.Context, ingestFileService storage.FileService, expectedFileNames []string) { + if !s.lock.TryLock() { + return + } + defer s.lock.Unlock() + + s.clearStoredIngestFiles(ctx, ingestFileService, expectedFileNames) + s.clearLegacyLocalIngestFiles(ctx, expectedFileNames) + ClearLocalIngestScratch(ctx, s.scratchDirectoryRootPath, scratchMinimumAge) +} diff --git a/cmd/api/src/daemons/datapipe/cleanup_internal_test.go b/cmd/api/src/daemons/datapipe/cleanup_internal_test.go new file mode 100644 index 000000000000..520338197f49 --- /dev/null +++ b/cmd/api/src/daemons/datapipe/cleanup_internal_test.go @@ -0,0 +1,328 @@ +// Copyright 2026 Specter Ops, Inc. +// +// Licensed under the Apache License, Version 2.0 +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 + +package datapipe + +import ( + "context" + "io/fs" + "os" + "path/filepath" + "testing" + "time" + + "github.com/specterops/bloodhound/cmd/api/src/daemons/datapipe/mocks" + "github.com/specterops/bloodhound/packages/go/storage" + storagemocks "github.com/specterops/bloodhound/packages/go/storage/mocks" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" +) + +type cleanupInternalFileInfo struct { + name string + modTime time.Time +} + +func (s cleanupInternalFileInfo) Name() string { return s.name } +func (s cleanupInternalFileInfo) Size() int64 { return 0 } +func (s cleanupInternalFileInfo) Mode() fs.FileMode { return 0 } +func (s cleanupInternalFileInfo) ModTime() time.Time { return s.modTime } +func (s cleanupInternalFileInfo) IsDir() bool { return false } +func (s cleanupInternalFileInfo) Sys() any { return nil } + +type cleanupInternalDirEntry struct { + name string + isDir bool + mode fs.FileMode + info fs.FileInfo + infoErr error +} + +func (s cleanupInternalDirEntry) Name() string { + return s.name +} + +func (s cleanupInternalDirEntry) IsDir() bool { + return s.isDir +} + +func (s cleanupInternalDirEntry) Type() fs.FileMode { + return s.mode +} + +func (s cleanupInternalDirEntry) Info() (fs.FileInfo, error) { + return s.info, s.infoErr +} + +func TestNormalizeStoragePrefixes(t *testing.T) { + t.Parallel() + + assert.Equal(t, []string{ + "retained", + "work/tmp", + "nested/prefix", + }, normalizeStoragePrefixes([]string{ + " retained/ ", + "/work/tmp", + "../escape", + "", + ".", + "nested//prefix/", + })) +} + +func TestOrphanFileSweeper_IsExcludedStoragePath(t *testing.T) { + t.Parallel() + + sweeper := NewOrphanFileSweeper(nil, t.TempDir(), t.TempDir(), "retained", "/nested/prefix/") + + testCases := []struct { + name string + logicalPath string + expected bool + }{ + { + name: "exact prefix", + logicalPath: "retained", + expected: true, + }, + { + name: "file under prefix", + logicalPath: "retained/file.json", + expected: true, + }, + { + name: "leading slash", + logicalPath: "/nested/prefix/file.json", + expected: true, + }, + { + name: "similar prefix", + logicalPath: "retained-other/file.json", + expected: false, + }, + { + name: "traversal path", + logicalPath: "../retained/file.json", + expected: false, + }, + } + + for _, testCase := range testCases { + testCase := testCase + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + assert.Equal(t, testCase.expected, sweeper.isExcludedStoragePath(testCase.logicalPath)) + }) + } +} + +func TestOrphanFileSweeper_AddExpectedStoragePath(t *testing.T) { + t.Parallel() + + tempDirectory := t.TempDir() + sweeper := NewOrphanFileSweeper(nil, tempDirectory, t.TempDir()) + + expectedFiles := map[string]struct{}{} + sweeper.addExpectedStoragePath(expectedFiles, "nested//file.json") + sweeper.addExpectedStoragePath(expectedFiles, filepath.Join(tempDirectory, "absolute.json")) + sweeper.addExpectedStoragePath(expectedFiles, filepath.Join(filepath.Dir(tempDirectory), "outside.json")) + sweeper.addExpectedStoragePath(expectedFiles, "../escape.json") + sweeper.addExpectedStoragePath(expectedFiles, ".") + sweeper.addExpectedStoragePath(expectedFiles, " ") + + require.Len(t, expectedFiles, 2) + assert.Contains(t, expectedFiles, "nested/file.json") + assert.Contains(t, expectedFiles, "absolute.json") +} + +func TestOrphanFileSweeper_AddExpectedLocalPath(t *testing.T) { + t.Parallel() + + tempDirectory := t.TempDir() + absoluteExpectedPath := filepath.Join(tempDirectory, "absolute.json") + sweeper := NewOrphanFileSweeper(nil, tempDirectory, t.TempDir()) + + expectedLocalFiles := map[string]struct{}{} + sweeper.addExpectedLocalPath(expectedLocalFiles, "nested/file.json") + sweeper.addExpectedLocalPath(expectedLocalFiles, absoluteExpectedPath) + sweeper.addExpectedLocalPath(expectedLocalFiles, "../escape.json") + sweeper.addExpectedLocalPath(expectedLocalFiles, ".") + sweeper.addExpectedLocalPath(expectedLocalFiles, " ") + + require.Len(t, expectedLocalFiles, 2) + assert.Contains(t, expectedLocalFiles, filepath.Join(tempDirectory, "nested", "file.json")) + assert.Contains(t, expectedLocalFiles, absoluteExpectedPath) +} + +func TestOrphanFileSweeper_ClearStoredIngestFiles(t *testing.T) { + t.Parallel() + + var ( + mockCtrl = gomock.NewController(t) + mockFileService = storagemocks.NewMockFileService(mockCtrl) + tempDirectory = t.TempDir() + sweeper = NewOrphanFileSweeper(nil, tempDirectory, t.TempDir(), "retained", "/excluded/") + oldTime = time.Now().Add(-25 * time.Hour) + youngTime = time.Now() + ) + + absoluteExpectedPath := filepath.Join(tempDirectory, "absolute-expected.json") + + mockFileService.EXPECT(). + ListFiles(gomock.Any(), "", storage.ListOptions{Recursive: true}). + Return([]storage.FileInfo{ + {Path: "expected.json", LastModified: oldTime}, + {Path: "absolute-expected.json", LastModified: oldTime}, + {Path: "retained/file.json", LastModified: oldTime}, + {Path: "/excluded/file.json", LastModified: oldTime}, + {Path: "young.json", LastModified: youngTime}, + {Path: "directory", LastModified: oldTime, IsDir: true}, + {Path: "old.json", LastModified: oldTime}, + }, nil) + mockFileService.EXPECT().DeleteFile(gomock.Any(), "old.json").Return(nil) + + sweeper.clearStoredIngestFiles(context.Background(), mockFileService, []string{"expected.json", absoluteExpectedPath}) +} + +func TestOrphanFileSweeper_ClearLegacyLocalIngestFiles(t *testing.T) { + t.Parallel() + + var ( + mockCtrl = gomock.NewController(t) + mockFileOps = mocks.NewMockFileOperations(mockCtrl) + tempDirectory = t.TempDir() + sweeper = NewOrphanFileSweeper(mockFileOps, tempDirectory, t.TempDir()) + oldTime = time.Now().Add(-25 * time.Hour) + youngTime = time.Now() + ) + + mockFileOps.EXPECT().ReadDir(tempDirectory).Return([]os.DirEntry{ + cleanupInternalDirEntry{ + name: "expected.json", + }, + cleanupInternalDirEntry{ + name: "young.json", + info: cleanupInternalFileInfo{ + name: "young.json", + modTime: youngTime, + }, + }, + cleanupInternalDirEntry{ + name: "nested", + isDir: true, + }, + cleanupInternalDirEntry{ + name: "old.json", + info: cleanupInternalFileInfo{ + name: "old.json", + modTime: oldTime, + }, + }, + }, nil) + mockFileOps.EXPECT().RemoveAll(filepath.Join(tempDirectory, "nested")).Return(nil) + mockFileOps.EXPECT().RemoveAll(filepath.Join(tempDirectory, "old.json")).Return(nil) + + sweeper.clearLegacyLocalIngestFiles(context.Background(), []string{"expected.json", "../old.json"}) +} + +func TestOrphanFileSweeper_ClearLegacyLocalIngestFilesRemovesOrphanDirectories(t *testing.T) { + t.Parallel() + + var ( + mockCtrl = gomock.NewController(t) + mockFileOps = mocks.NewMockFileOperations(mockCtrl) + tempDirectory = t.TempDir() + sweeper = NewOrphanFileSweeper(mockFileOps, tempDirectory, t.TempDir()) + ) + + mockFileOps.EXPECT().ReadDir(tempDirectory).Return([]os.DirEntry{ + cleanupInternalDirEntry{ + name: "orphan", + isDir: true, + }, + cleanupInternalDirEntry{ + name: "expected-parent", + isDir: true, + }, + }, nil) + mockFileOps.EXPECT().RemoveAll(filepath.Join(tempDirectory, "orphan")).Return(nil) + + sweeper.clearLegacyLocalIngestFiles(context.Background(), []string{"expected-parent/file.json"}) +} + +func TestClearLocalIngestScratch(t *testing.T) { + t.Parallel() + + t.Run("removes old files and keeps young files and directories", func(t *testing.T) { + t.Parallel() + + var ( + scratchDirectory = t.TempDir() + oldFilePath = filepath.Join(scratchDirectory, "old.tmp") + youngFilePath = filepath.Join(scratchDirectory, "young.tmp") + nestedDirectory = filepath.Join(scratchDirectory, "nested") + oldTime = time.Now().Add(-2 * time.Hour) + ) + + require.NoError(t, os.WriteFile(oldFilePath, []byte("old"), 0600)) + require.NoError(t, os.Chtimes(oldFilePath, oldTime, oldTime)) + require.NoError(t, os.WriteFile(youngFilePath, []byte("young"), 0600)) + require.NoError(t, os.Mkdir(nestedDirectory, 0700)) + + ClearLocalIngestScratch(context.Background(), scratchDirectory, time.Hour) + + requirePathNotExists(t, oldFilePath) + requirePathExists(t, youngFilePath) + requirePathExists(t, nestedDirectory) + }) + + t.Run("exits without deleting files when context is canceled", func(t *testing.T) { + t.Parallel() + + var ( + scratchDirectory = t.TempDir() + oldFilePath = filepath.Join(scratchDirectory, "old.tmp") + oldTime = time.Now().Add(-2 * time.Hour) + ) + + require.NoError(t, os.WriteFile(oldFilePath, []byte("old"), 0600)) + require.NoError(t, os.Chtimes(oldFilePath, oldTime, oldTime)) + + ctx, done := context.WithCancel(context.Background()) + done() + + ClearLocalIngestScratch(ctx, scratchDirectory, time.Hour) + + requirePathExists(t, oldFilePath) + }) +} + +func requirePathExists(t *testing.T, path string) { + t.Helper() + + _, err := os.Stat(path) + require.NoError(t, err) +} + +func requirePathNotExists(t *testing.T, path string) { + t.Helper() + + _, err := os.Stat(path) + require.ErrorIs(t, err, os.ErrNotExist) +} diff --git a/cmd/api/src/daemons/datapipe/cleanup_test.go b/cmd/api/src/daemons/datapipe/cleanup_test.go index 94f62a996d22..f5143107da48 100644 --- a/cmd/api/src/daemons/datapipe/cleanup_test.go +++ b/cmd/api/src/daemons/datapipe/cleanup_test.go @@ -23,12 +23,27 @@ import ( "path/filepath" "sync" "testing" + "time" "github.com/specterops/bloodhound/cmd/api/src/daemons/datapipe" "github.com/specterops/bloodhound/cmd/api/src/daemons/datapipe/mocks" + "github.com/specterops/bloodhound/packages/go/storage" + storagemocks "github.com/specterops/bloodhound/packages/go/storage/mocks" "go.uber.org/mock/gomock" ) +type fileInfo struct { + name string + modTime time.Time +} + +func (s fileInfo) Name() string { return s.name } +func (s fileInfo) Size() int64 { return 0 } +func (s fileInfo) Mode() fs.FileMode { return 0 } +func (s fileInfo) ModTime() time.Time { return s.modTime } +func (s fileInfo) IsDir() bool { return false } +func (s fileInfo) Sys() any { return nil } + type dirEntry struct { name string isDir bool @@ -55,14 +70,16 @@ func (s dirEntry) Info() (fs.FileInfo, error) { func TestOrphanFileSweeper_Clear(t *testing.T) { const workDir = "/fake/work/dir" + const scratchDir = "/fake/work/dir/scratch" t.Run("Allow Only One Goroutine", func(t *testing.T) { var ( - mockCtrl = gomock.NewController(t) - mockFileOps = mocks.NewMockFileOperations(mockCtrl) - sweeper = datapipe.NewOrphanFileSweeper(mockFileOps, workDir) - wgCoordination = &sync.WaitGroup{} - wgReadDir = &sync.WaitGroup{} + mockCtrl = gomock.NewController(t) + mockFileOps = mocks.NewMockFileOperations(mockCtrl) + mockFileService = storagemocks.NewMockFileService(mockCtrl) + sweeper = datapipe.NewOrphanFileSweeper(mockFileOps, workDir, scratchDir, "") + wgCoordination = &sync.WaitGroup{} + wgReadDir = &sync.WaitGroup{} ) defer mockCtrl.Finish() @@ -71,6 +88,10 @@ func TestOrphanFileSweeper_Clear(t *testing.T) { wgCoordination.Add(1) wgReadDir.Add(1) + mockFileService.EXPECT(). + ListFiles(gomock.Any(), "", storage.ListOptions{Recursive: true}). + Return(nil, nil) + mockFileOps.EXPECT().ReadDir(workDir).DoAndReturn(func(path string) ([]os.DirEntry, error) { // Release the coordination wait group wgCoordination.Done() @@ -82,13 +103,13 @@ func TestOrphanFileSweeper_Clear(t *testing.T) { }) // Launch the clear function in a goroutine. The wait group will cause this call to block - go sweeper.Clear(context.Background(), []string{}) + go sweeper.Clear(context.Background(), mockFileService, []string{}) // Wait for the go routine to reach the ReadDir function wgCoordination.Wait() // Run the clear function in the current thread context as this should exit without blocking - sweeper.Clear(context.Background(), []string{}) + sweeper.Clear(context.Background(), mockFileService, []string{}) // Release the wait group to complete the test wgReadDir.Done() @@ -96,33 +117,73 @@ func TestOrphanFileSweeper_Clear(t *testing.T) { t.Run("Clear Orphan Files", func(t *testing.T) { var ( - mockCtrl = gomock.NewController(t) - mockFileOps = mocks.NewMockFileOperations(mockCtrl) - sweeper = datapipe.NewOrphanFileSweeper(mockFileOps, workDir) + mockCtrl = gomock.NewController(t) + mockFileOps = mocks.NewMockFileOperations(mockCtrl) + mockFileService = storagemocks.NewMockFileService(mockCtrl) + sweeper = datapipe.NewOrphanFileSweeper(mockFileOps, workDir, scratchDir, "") ) defer mockCtrl.Finish() + mockFileService.EXPECT(). + ListFiles(gomock.Any(), "", storage.ListOptions{Recursive: true}). + Return(nil, nil) + mockFileOps.EXPECT().ReadDir(workDir).Return([]os.DirEntry{ dirEntry{ name: "1", + info: fileInfo{ + name: "1", + modTime: time.Now().Add(-25 * time.Hour), + }, }, }, nil) mockFileOps.EXPECT().RemoveAll(filepath.Join(workDir, "1")).Return(nil) - sweeper.Clear(context.Background(), []string{}) + sweeper.Clear(context.Background(), mockFileService, []string{}) + }) + + t.Run("FileService - Clear Orphan Files", func(t *testing.T) { + var ( + mockCtrl = gomock.NewController(t) + mockFileOps = mocks.NewMockFileOperations(mockCtrl) + mockFileService = storagemocks.NewMockFileService(mockCtrl) + sweeper = datapipe.NewOrphanFileSweeper(mockFileOps, workDir, scratchDir, "") + ) + + defer mockCtrl.Finish() + + mockFileService.EXPECT(). + ListFiles(gomock.Any(), "", storage.ListOptions{Recursive: true}). + Return([]storage.FileInfo{ + { + Path: "1", + LastModified: time.Now().Add(-25 * time.Hour), + }, + }, nil) + + mockFileService.EXPECT().DeleteFile(gomock.Any(), "1").Return(nil) + + mockFileOps.EXPECT().ReadDir(workDir).Return(nil, nil) + + sweeper.Clear(context.Background(), mockFileService, []string{}) }) t.Run("Exclude Expected Files", func(t *testing.T) { var ( - mockCtrl = gomock.NewController(t) - mockFileOps = mocks.NewMockFileOperations(mockCtrl) - sweeper = datapipe.NewOrphanFileSweeper(mockFileOps, workDir) + mockCtrl = gomock.NewController(t) + mockFileOps = mocks.NewMockFileOperations(mockCtrl) + mockFileService = storagemocks.NewMockFileService(mockCtrl) + sweeper = datapipe.NewOrphanFileSweeper(mockFileOps, workDir, scratchDir, "") ) defer mockCtrl.Finish() + mockFileService.EXPECT(). + ListFiles(gomock.Any(), "", storage.ListOptions{Recursive: true}). + Return(nil, nil) + mockFileOps.EXPECT().ReadDir(workDir).Return([]os.DirEntry{ dirEntry{ name: "1", @@ -132,18 +193,23 @@ func TestOrphanFileSweeper_Clear(t *testing.T) { // This one is a negative assertion. Because we're passing in "1" to be excluded the listed dirEntries will be // empty and therefore the sweeper MUST NOT call RemoveAll() on the FileOperations mock. If RemoveAll() is // called then this test MUST fail. - sweeper.Clear(context.Background(), []string{"1"}) + sweeper.Clear(context.Background(), mockFileService, []string{"1"}) }) t.Run("Exclude Files With Paths", func(t *testing.T) { var ( - mockCtrl = gomock.NewController(t) - mockFileOps = mocks.NewMockFileOperations(mockCtrl) - sweeper = datapipe.NewOrphanFileSweeper(mockFileOps, workDir) + mockCtrl = gomock.NewController(t) + mockFileOps = mocks.NewMockFileOperations(mockCtrl) + mockFileService = storagemocks.NewMockFileService(mockCtrl) + sweeper = datapipe.NewOrphanFileSweeper(mockFileOps, workDir, scratchDir, "") ) defer mockCtrl.Finish() + mockFileService.EXPECT(). + ListFiles(gomock.Any(), "", storage.ListOptions{Recursive: true}). + Return(nil, nil) + mockFileOps.EXPECT().ReadDir(workDir).Return([]os.DirEntry{ dirEntry{ name: "1", @@ -156,18 +222,23 @@ func TestOrphanFileSweeper_Clear(t *testing.T) { // This one is a negative assertion. Because we're passing in paths to be excluded the listed dirEntries will be // empty and therefore the sweeper MUST NOT call RemoveAll() on the FileOperations mock. If RemoveAll() is // called then this test MUST fail. - sweeper.Clear(context.Background(), []string{"1", workDir + "/2"}) + sweeper.Clear(context.Background(), mockFileService, []string{"1", workDir + "/2"}) }) t.Run("Exit on Context Cancellation", func(t *testing.T) { var ( - mockCtrl = gomock.NewController(t) - mockFileOps = mocks.NewMockFileOperations(mockCtrl) - sweeper = datapipe.NewOrphanFileSweeper(mockFileOps, workDir) + mockCtrl = gomock.NewController(t) + mockFileOps = mocks.NewMockFileOperations(mockCtrl) + mockFileService = storagemocks.NewMockFileService(mockCtrl) + sweeper = datapipe.NewOrphanFileSweeper(mockFileOps, workDir, scratchDir, "") ) defer mockCtrl.Finish() + mockFileService.EXPECT(). + ListFiles(gomock.Any(), "", storage.ListOptions{Recursive: true}). + Return(nil, nil) + mockFileOps.EXPECT().ReadDir(workDir).Return([]os.DirEntry{ dirEntry{ name: "1", @@ -179,6 +250,6 @@ func TestOrphanFileSweeper_Clear(t *testing.T) { done() // When passed in with the cancelled context the sweeper should not call os.RemoveAll("1") - sweeper.Clear(ctx, []string{}) + sweeper.Clear(ctx, mockFileService, []string{}) }) } diff --git a/cmd/api/src/daemons/datapipe/datapipe_integration_test.go b/cmd/api/src/daemons/datapipe/datapipe_integration_test.go index 89397503b7c2..4ae6b9a6fcfc 100644 --- a/cmd/api/src/daemons/datapipe/datapipe_integration_test.go +++ b/cmd/api/src/daemons/datapipe/datapipe_integration_test.go @@ -35,6 +35,7 @@ import ( "github.com/specterops/bloodhound/cmd/api/src/database" "github.com/specterops/bloodhound/cmd/api/src/migrations" "github.com/specterops/bloodhound/cmd/api/src/services/graphify" + "github.com/specterops/bloodhound/cmd/api/src/services/storage" "github.com/specterops/bloodhound/cmd/api/src/services/upload" "github.com/specterops/bloodhound/cmd/api/src/test/integration/utils" "github.com/specterops/bloodhound/packages/go/cache" @@ -113,15 +114,20 @@ func setupIntegrationTestSuite(t *testing.T, fixturesPath string) IntegrationTes cfg.WorkDir = workDir + fileServices, err := storage.NewDefaultFileServices(cfg) + require.NoError(t, err, "error creating the default file services") + fileServiceResolver, err := storage.NewFileServiceResolver(fileServices) + require.NoError(t, err, "error creating file service resolver") + cl := changelog.NewChangelog(graphDB, db, changelog.DefaultOptions()) return IntegrationTestSuite{ Context: ctx, - GraphifyService: graphify.NewGraphifyService(ctx, db, graphDB, cfg, ingestSchema, cl), + GraphifyService: graphify.NewGraphifyService(ctx, db, graphDB, cfg, ingestSchema, fileServiceResolver, cl), GraphDB: graphDB, BHDatabase: db, WorkDir: workDir, - Daemon: datapipe.NewPipeline(ctx, cfg, db, graphDB, cache.Cache{}, ingestSchema, cl), + Daemon: datapipe.NewPipeline(ctx, cfg, db, graphDB, cache.Cache{}, ingestSchema, fileServiceResolver, cl), } } diff --git a/cmd/api/src/daemons/datapipe/delete_integration_test.go b/cmd/api/src/daemons/datapipe/delete_integration_test.go index 38ac619688ee..800f9898d2a9 100644 --- a/cmd/api/src/daemons/datapipe/delete_integration_test.go +++ b/cmd/api/src/daemons/datapipe/delete_integration_test.go @@ -29,6 +29,7 @@ import ( "github.com/specterops/bloodhound/cmd/api/src/services/graphify" "github.com/specterops/bloodhound/cmd/api/src/services/graphify/endpoint" "github.com/specterops/bloodhound/packages/go/lab/generic" + "github.com/specterops/bloodhound/packages/go/storage" "github.com/specterops/dawgs/graph" "github.com/stretchr/testify/require" ) @@ -42,15 +43,18 @@ func TestDeleteAllData(t *testing.T) { testSuite = setupIntegrationTestSuite(t, fixturesPath) files = []string{ - path.Join(testSuite.WorkDir, "base.json"), + "base.json", } ) + ingestLocalStore, err := storage.NewLocalStore(testSuite.WorkDir) + require.NoError(t, err, "error creating ingest local store") + ingestFileService := storage.NewFileService(ingestLocalStore) defer teardownIntegrationTestSuite(t, &testSuite) for _, file := range files { ingestContext := graphify.NewIngestContext(ctx, graphify.WithIngestTime(time.Now()), graphify.WithEndpointResolver(endpoint.NewResolver(testSuite.GraphDB))) - fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestContext, model.IngestTask{StoredFileName: file, FileType: model.FileTypeJson}) + fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestContext, ingestFileService, model.IngestTask{StoredFileName: file, FileType: model.FileTypeJson}) require.NoError(t, err) failed := 0 @@ -70,7 +74,7 @@ func TestDeleteAllData(t *testing.T) { sourceKinds = []graph.Kind{graph.StringKind("Base"), graph.StringKind("AZBase"), graph.StringKind("GithubBase")} ) - err := datapipe.DeleteCollectedGraphData(ctx, testSuite.GraphDB, deleteRequest, sourceKinds) + err = datapipe.DeleteCollectedGraphData(ctx, testSuite.GraphDB, deleteRequest, sourceKinds) require.Nil(t, err) expected, err := generic.LoadGraphFromFile(os.DirFS(path.Join("fixtures", t.Name())), "deleteAllExpected.json") @@ -90,15 +94,18 @@ func TestDeleteAllData_Alternative(t *testing.T) { testSuite = setupIntegrationTestSuite(t, fixturesPath) files = []string{ - path.Join(testSuite.WorkDir, "base.json"), + "base.json", } ) + ingestLocalStore, err := storage.NewLocalStore(testSuite.WorkDir) + require.NoError(t, err, "error creating ingest local store") + ingestFileService := storage.NewFileService(ingestLocalStore) defer teardownIntegrationTestSuite(t, &testSuite) for _, file := range files { ingestContext := graphify.NewIngestContext(ctx, graphify.WithIngestTime(time.Now()), graphify.WithEndpointResolver(endpoint.NewResolver(testSuite.GraphDB))) - fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestContext, model.IngestTask{StoredFileName: file, FileType: model.FileTypeJson}) + fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestContext, ingestFileService, model.IngestTask{StoredFileName: file, FileType: model.FileTypeJson}) require.NoError(t, err) failed := 0 @@ -118,7 +125,7 @@ func TestDeleteAllData_Alternative(t *testing.T) { sourceKinds = []graph.Kind{graph.StringKind("Base"), graph.StringKind("AZBase"), graph.StringKind("GithubBase")} ) - err := datapipe.DeleteCollectedGraphData(ctx, testSuite.GraphDB, deleteRequest, sourceKinds) + err = datapipe.DeleteCollectedGraphData(ctx, testSuite.GraphDB, deleteRequest, sourceKinds) require.Nil(t, err) expected, err := generic.LoadGraphFromFile(os.DirFS(path.Join("fixtures", "TestDeleteAllData")), "deleteAllExpected.json") @@ -135,15 +142,18 @@ func TestDeleteSourcelessData(t *testing.T) { testSuite = setupIntegrationTestSuite(t, fixturesPath) files = []string{ - path.Join(testSuite.WorkDir, "base.json"), + "base.json", } ) + ingestLocalStore, err := storage.NewLocalStore(testSuite.WorkDir) + require.NoError(t, err, "error creating ingest local store") + ingestFileService := storage.NewFileService(ingestLocalStore) defer teardownIntegrationTestSuite(t, &testSuite) for _, file := range files { ingestContext := graphify.NewIngestContext(ctx, graphify.WithIngestTime(time.Now()), graphify.WithEndpointResolver(endpoint.NewResolver(testSuite.GraphDB))) - fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestContext, model.IngestTask{StoredFileName: file, FileType: model.FileTypeJson}) + fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestContext, ingestFileService, model.IngestTask{StoredFileName: file, FileType: model.FileTypeJson}) require.NoError(t, err) failed := 0 @@ -162,7 +172,7 @@ func TestDeleteSourcelessData(t *testing.T) { sourceKinds = []graph.Kind{graph.StringKind("Base"), graph.StringKind("AZBase"), graph.StringKind("GithubBase")} ) - err := datapipe.DeleteCollectedGraphData(ctx, testSuite.GraphDB, deleteRequest, sourceKinds) + err = datapipe.DeleteCollectedGraphData(ctx, testSuite.GraphDB, deleteRequest, sourceKinds) require.Nil(t, err) expected, err := generic.LoadGraphFromFile(os.DirFS(path.Join("fixtures", t.Name())), "deleteSourcelessExpected.json") @@ -179,15 +189,18 @@ func TestDeleteSourceKindsData(t *testing.T) { testSuite = setupIntegrationTestSuite(t, fixturesPath) files = []string{ - path.Join(testSuite.WorkDir, "base.json"), + "base.json", } ) + ingestLocalStore, err := storage.NewLocalStore(testSuite.WorkDir) + require.NoError(t, err, "error creating ingest local store") + ingestFileService := storage.NewFileService(ingestLocalStore) defer teardownIntegrationTestSuite(t, &testSuite) for _, file := range files { ingestContext := graphify.NewIngestContext(ctx, graphify.WithIngestTime(time.Now()), graphify.WithEndpointResolver(endpoint.NewResolver(testSuite.GraphDB))) - fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestContext, model.IngestTask{StoredFileName: file, FileType: model.FileTypeJson}) + fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestContext, ingestFileService, model.IngestTask{StoredFileName: file, FileType: model.FileTypeJson}) require.NoError(t, err) failed := 0 @@ -206,7 +219,7 @@ func TestDeleteSourceKindsData(t *testing.T) { sourceKinds = []graph.Kind{graph.StringKind("Base"), graph.StringKind("AZBase"), graph.StringKind("GithubBase")} ) - err := datapipe.DeleteCollectedGraphData(ctx, testSuite.GraphDB, deleteRequest, sourceKinds) + err = datapipe.DeleteCollectedGraphData(ctx, testSuite.GraphDB, deleteRequest, sourceKinds) require.Nil(t, err) expected, err := generic.LoadGraphFromFile(os.DirFS(path.Join("fixtures", t.Name())), "deleteSourceKindsExpected.json") diff --git a/cmd/api/src/daemons/datapipe/pipeline.go b/cmd/api/src/daemons/datapipe/pipeline.go index c54af65283d6..f04ce7c92beb 100644 --- a/cmd/api/src/daemons/datapipe/pipeline.go +++ b/cmd/api/src/daemons/datapipe/pipeline.go @@ -29,6 +29,7 @@ import ( "github.com/specterops/bloodhound/cmd/api/src/model/appcfg" "github.com/specterops/bloodhound/cmd/api/src/services/graphify" "github.com/specterops/bloodhound/cmd/api/src/services/job" + storageService "github.com/specterops/bloodhound/cmd/api/src/services/storage" "github.com/specterops/bloodhound/cmd/api/src/services/upload" "github.com/specterops/bloodhound/packages/go/analysis" "github.com/specterops/bloodhound/packages/go/bhlog/attr" @@ -36,6 +37,7 @@ import ( "github.com/specterops/bloodhound/packages/go/cache" "github.com/specterops/bloodhound/packages/go/graphschema/ad" "github.com/specterops/bloodhound/packages/go/graphschema/azure" + "github.com/specterops/bloodhound/packages/go/storage" "github.com/specterops/dawgs/graph" ) @@ -48,21 +50,23 @@ type BHCEPipeline struct { cfg config.Configuration orphanedFileSweeper *OrphanFileSweeper ingestSchema upload.IngestSchema + fileServiceResolver storageService.FileServiceResolver jobService job.JobService graphifyService graphify.GraphifyService changelog *changelog.Changelog } -func NewPipeline(ctx context.Context, cfg config.Configuration, db database.Database, graphDB graph.Database, cache cache.Cache, ingestSchema upload.IngestSchema, cl *changelog.Changelog) *BHCEPipeline { +func NewPipeline(ctx context.Context, cfg config.Configuration, db database.Database, graphDB graph.Database, cache cache.Cache, ingestSchema upload.IngestSchema, fileServiceResolver storageService.FileServiceResolver, cl *changelog.Changelog) *BHCEPipeline { return &BHCEPipeline{ db: db, graphdb: graphDB, cache: cache, cfg: cfg, - orphanedFileSweeper: NewOrphanFileSweeper(NewOSFileOperations(), cfg.TempDirectory()), + orphanedFileSweeper: NewOrphanFileSweeper(NewOSFileOperations(), cfg.TempDirectory(), cfg.ScratchDirectory(), string(storage.FileServiceRetained)), ingestSchema: ingestSchema, + fileServiceResolver: fileServiceResolver, jobService: job.NewJobService(ctx, db), - graphifyService: graphify.NewGraphifyService(ctx, db, graphDB, cfg, ingestSchema, cl), + graphifyService: graphify.NewGraphifyService(ctx, db, graphDB, cfg, ingestSchema, fileServiceResolver, cl), changelog: cl, } } @@ -155,6 +159,8 @@ func filterDeletableKinds(kindsToDelete []string) []string { func (s *BHCEPipeline) PruneData(ctx context.Context) error { if ingestTasks, err := s.db.GetAllIngestTasks(ctx); err != nil { return fmt.Errorf("fetching available ingest tasks: %v", err) + } else if ingestFileService, err := s.fileServiceResolver.Resolve(storage.FileServiceIngest); err != nil { + return fmt.Errorf("error resolving ingest file service: %v", err) } else { expectedFiles := make([]string, len(ingestTasks)) @@ -162,7 +168,7 @@ func (s *BHCEPipeline) PruneData(ctx context.Context) error { expectedFiles[idx] = ingestTask.StoredFileName } - go s.orphanedFileSweeper.Clear(ctx, expectedFiles) + go s.orphanedFileSweeper.Clear(ctx, ingestFileService, expectedFiles) } return nil } diff --git a/cmd/api/src/daemons/datapipe/pipeline_integration_test.go b/cmd/api/src/daemons/datapipe/pipeline_integration_test.go index cd80757a98be..d3c9793c446f 100644 --- a/cmd/api/src/daemons/datapipe/pipeline_integration_test.go +++ b/cmd/api/src/daemons/datapipe/pipeline_integration_test.go @@ -26,6 +26,7 @@ import ( "github.com/specterops/bloodhound/cmd/api/src/model" "github.com/specterops/bloodhound/cmd/api/src/services/graphify" "github.com/specterops/bloodhound/cmd/api/src/services/graphify/endpoint" + "github.com/specterops/bloodhound/cmd/api/src/services/storage" "github.com/specterops/bloodhound/packages/go/lab/generic" "github.com/specterops/dawgs/graph" "github.com/specterops/dawgs/query" @@ -44,15 +45,18 @@ func TestDeleteData_Sourceless(t *testing.T) { testSuite = setupIntegrationTestSuite(t, fixturesPath) files = []string{ - path.Join(testSuite.WorkDir, "base.json"), + "base.json", } ) + ingestLocalStore, err := storage.NewLocalStore(testSuite.WorkDir) + require.NoError(t, err, "error creating ingest local store") + ingestFileService := storage.NewFileService(ingestLocalStore) defer teardownIntegrationTestSuite(t, &testSuite) ingestCtx := graphify.NewIngestContext(ctx, graphify.WithEndpointResolver(endpoint.NewResolver(testSuite.GraphDB))) for _, file := range files { - fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestCtx, model.IngestTask{StoredFileName: file, FileType: model.FileTypeJson}) + fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestCtx, ingestFileService, model.IngestTask{StoredFileName: file, FileType: model.FileTypeJson}) require.NoError(t, err) failed := 0 @@ -67,7 +71,7 @@ func TestDeleteData_Sourceless(t *testing.T) { } // simulate requesting deletion - err := testSuite.BHDatabase.RegisterSourceKind(ctx)(graph.StringKind("GithubBase")) + err = testSuite.BHDatabase.RegisterSourceKind(ctx)(graph.StringKind("GithubBase")) require.Nil(t, err) testSuite.BHDatabase.RequestCollectedGraphDataDeletion(ctx, model.AnalysisRequest{DeleteSourcelessGraph: true, RequestType: model.AnalysisRequestDeletion}) require.Nil(t, err) @@ -93,15 +97,18 @@ func TestDeleteData_SourceKinds(t *testing.T) { testSuite = setupIntegrationTestSuite(t, fixturesPath) files = []string{ - path.Join(testSuite.WorkDir, "base.json"), + "base.json", } ) + ingestLocalStore, err := storage.NewLocalStore(testSuite.WorkDir) + require.NoError(t, err, "error creating ingest local store") + ingestFileService := storage.NewFileService(ingestLocalStore) defer teardownIntegrationTestSuite(t, &testSuite) ingestCtx := graphify.NewIngestContext(ctx, graphify.WithEndpointResolver(endpoint.NewResolver(testSuite.GraphDB))) for _, file := range files { - fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestCtx, model.IngestTask{StoredFileName: file, FileType: model.FileTypeJson}) + fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestCtx, ingestFileService, model.IngestTask{StoredFileName: file, FileType: model.FileTypeJson}) require.NoError(t, err) failed := 0 @@ -116,7 +123,7 @@ func TestDeleteData_SourceKinds(t *testing.T) { } // simulate requesting deletion - err := testSuite.BHDatabase.RegisterSourceKind(ctx)(graph.StringKind("GithubBase")) + err = testSuite.BHDatabase.RegisterSourceKind(ctx)(graph.StringKind("GithubBase")) require.Nil(t, err) testSuite.BHDatabase.RequestCollectedGraphDataDeletion(ctx, model.AnalysisRequest{DeleteSourceKinds: []string{"GithubBase", "AZBase"}, RequestType: model.AnalysisRequestDeletion}) require.Nil(t, err) @@ -142,15 +149,18 @@ func TestDeleteData_All(t *testing.T) { testSuite = setupIntegrationTestSuite(t, fixturesPath) files = []string{ - path.Join(testSuite.WorkDir, "base.json"), + "base.json", } ) + ingestLocalStore, err := storage.NewLocalStore(testSuite.WorkDir) + require.NoError(t, err, "error creating ingest local store") + ingestFileService := storage.NewFileService(ingestLocalStore) defer teardownIntegrationTestSuite(t, &testSuite) ingestCtx := graphify.NewIngestContext(ctx, graphify.WithEndpointResolver(endpoint.NewResolver(testSuite.GraphDB))) for _, file := range files { - fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestCtx, model.IngestTask{StoredFileName: file, FileType: model.FileTypeJson}) + fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestCtx, ingestFileService, model.IngestTask{StoredFileName: file, FileType: model.FileTypeJson}) require.NoError(t, err) failed := 0 @@ -165,7 +175,7 @@ func TestDeleteData_All(t *testing.T) { } // simulate requesting deletion - err := testSuite.BHDatabase.RegisterSourceKind(ctx)(graph.StringKind("GithubBase")) + err = testSuite.BHDatabase.RegisterSourceKind(ctx)(graph.StringKind("GithubBase")) require.Nil(t, err) testSuite.BHDatabase.RequestCollectedGraphDataDeletion(ctx, model.AnalysisRequest{DeleteSourceKinds: []string{"GithubBase", "AZBase", "Base"}, RequestType: model.AnalysisRequestDeletion}) require.Nil(t, err) @@ -193,12 +203,15 @@ func TestPartialIngest(t *testing.T) { testSuite = setupIntegrationTestSuite(t, fixturesPath) ) + ingestLocalStore, err := storage.NewLocalStore(testSuite.WorkDir) + require.NoError(t, err, "error creating ingest local store") + ingestFileService := storage.NewFileService(ingestLocalStore) defer teardownIntegrationTestSuite(t, &testSuite) ingestCtx := graphify.NewIngestContext(ctx, graphify.WithEndpointResolver(endpoint.NewResolver(testSuite.GraphDB))) - fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestCtx, model.IngestTask{ - StoredFileName: path.Join(testSuite.WorkDir, "oneGoodOneInvalidRel.json"), + fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestCtx, ingestFileService, model.IngestTask{ + StoredFileName: "oneGoodOneInvalidRel.json", FileType: model.FileTypeJson, }) require.NoError(t, err) diff --git a/cmd/api/src/model/audit.go b/cmd/api/src/model/audit.go index a08355310b9c..3741bf322d5f 100644 --- a/cmd/api/src/model/audit.go +++ b/cmd/api/src/model/audit.go @@ -107,6 +107,8 @@ const ( AuditLogActionUpdateClient AuditLogAction = "UpdateClient" AuditLogActionStartClientJob AuditLogAction = "StartClientJob" + AuditLogActionUploadSupportBundle AuditLogAction = "UploadSupportBundle" + AuditLogActionImportSavedQuery AuditLogAction = "ImportSavedQueries" AuditLogActionExportSavedQuery AuditLogAction = "ExportSavedQuery" AuditLogActionExportSavedQueries AuditLogAction = "ExportSavedQueries" diff --git a/cmd/api/src/services/entrypoint.go b/cmd/api/src/services/entrypoint.go index 22983aabadee..ed4b50cf3f0d 100644 --- a/cmd/api/src/services/entrypoint.go +++ b/cmd/api/src/services/entrypoint.go @@ -41,11 +41,13 @@ import ( "github.com/specterops/bloodhound/cmd/api/src/queries" "github.com/specterops/bloodhound/cmd/api/src/services/dogtags" "github.com/specterops/bloodhound/cmd/api/src/services/opengraphschema" + storageService "github.com/specterops/bloodhound/cmd/api/src/services/storage" "github.com/specterops/bloodhound/cmd/api/src/services/upload" "github.com/specterops/bloodhound/packages/go/bhlog/attr" "github.com/specterops/bloodhound/packages/go/cache" schema "github.com/specterops/bloodhound/packages/go/graphschema" "github.com/specterops/bloodhound/packages/go/metricsregistration" + "github.com/specterops/bloodhound/packages/go/storage" "github.com/specterops/dawgs/graph" ) @@ -74,14 +76,44 @@ func ConnectDatabases(ctx context.Context, cfg config.Configuration) (bootstrap. } } +// CreateRuntimeDependencies creates the needed dependencies prior to migration. For instance, the FileService is needed for +// IngestControl which occurs prior to migration. This function can be used to make the struct to contain the services that +// are necessary for the application. +func CreateRuntimeDependencies(ctx context.Context, cfg config.Configuration, connections bootstrap.DatabaseConnections[*database.BloodhoundDB, *graph.DatabaseSwitch]) (bootstrap.RuntimeDependencies, error) { + var dependencies = bootstrap.RuntimeDependencies{} + if fileServices, err := storageService.NewDefaultFileServices(cfg); err != nil { + return dependencies, fmt.Errorf("failed to initialize file services: %w", err) + } else if fileServiceResolver, err := storageService.NewFileServiceResolver(fileServices); err != nil { + return dependencies, fmt.Errorf("failed to initialize file service resolver: %w", err) + // The FileServiceRetained is necessary for the PreMigrationDaemons where it is used in IngestControl for Cleanup. + // Checking it here ensures we have the service prior to running the application. + } else if _, err := fileServiceResolver.Resolve(storage.FileServiceRetained); err != nil { + return dependencies, fmt.Errorf("failed to resolve retained file service: %w", err) + } else { + dependencies.FileServiceResolver = fileServiceResolver + return dependencies, nil + } +} + // PreMigrationDaemons Word of caution: These daemons will be launched prior to any migration starting -func PreMigrationDaemons(ctx context.Context, cfg config.Configuration, connections bootstrap.DatabaseConnections[*database.BloodhoundDB, *graph.DatabaseSwitch]) ([]daemons.Daemon, error) { - return []daemons.Daemon{ - toolapi.NewDaemon(ctx, connections, cfg, schema.DefaultGraphSchema()), - }, nil +func PreMigrationDaemons(ctx context.Context, cfg config.Configuration, connections bootstrap.DatabaseConnections[*database.BloodhoundDB, *graph.DatabaseSwitch], dependencies bootstrap.RuntimeDependencies) ([]daemons.Daemon, error) { + if dependencies.FileServiceResolver == nil { + return nil, fmt.Errorf("file service resolver is required") + } + + if retainedFileService, err := dependencies.FileServiceResolver.Resolve(storage.FileServiceRetained); err != nil { + return nil, fmt.Errorf("error resolving FileServiceRetained: %w", err) + } else { + return []daemons.Daemon{ + toolapi.NewDaemon(ctx, connections, cfg, schema.DefaultGraphSchema(), retainedFileService), + }, nil + } } -func Entrypoint(ctx context.Context, cfg config.Configuration, connections bootstrap.DatabaseConnections[*database.BloodhoundDB, *graph.DatabaseSwitch]) ([]daemons.Daemon, error) { +func Entrypoint(ctx context.Context, cfg config.Configuration, connections bootstrap.DatabaseConnections[*database.BloodhoundDB, *graph.DatabaseSwitch], dependencies bootstrap.RuntimeDependencies) ([]daemons.Daemon, error) { + if dependencies.FileServiceResolver == nil { + return nil, fmt.Errorf("file service resolver is required") + } dogtagsService := dogtags.NewDefaultService() @@ -137,7 +169,7 @@ func Entrypoint(ctx context.Context, cfg config.Configuration, connections boots var ( cl = changelog.NewChangelog(connections.Graph, connections.RDMS, changelog.DefaultOptions()) - pipeline = datapipe.NewPipeline(ctx, cfg, connections.RDMS, connections.Graph, graphQueryCache, ingestSchema, cl) + pipeline = datapipe.NewPipeline(ctx, cfg, connections.RDMS, connections.Graph, graphQueryCache, ingestSchema, dependencies.FileServiceResolver, cl) graphQuery = queries.NewGraphQuery(connections.Graph, graphQueryCache, cfg) authorizer = auth.NewAuthorizer(connections.RDMS) datapipeDaemon = datapipe.NewDaemon(pipeline, startDelay, time.Duration(cfg.DatapipeInterval)*time.Second, connections.RDMS) @@ -147,7 +179,7 @@ func Entrypoint(ctx context.Context, cfg config.Configuration, connections boots ) registration.RegisterFossGlobalMiddleware(&routerInst, cfg, auth.NewIdentityResolver(), authenticator, connections.RDMS) - registration.RegisterFossRoutes(&routerInst, cfg, connections.RDMS, connections.Graph, graphQuery, apiCache, collectorManifests, authenticator, authorizer, ingestSchema, dogtagsService, openGraphSchemaService) + registration.RegisterFossRoutes(&routerInst, cfg, connections.RDMS, connections.Graph, graphQuery, apiCache, collectorManifests, authenticator, authorizer, ingestSchema, dependencies.FileServiceResolver, dogtagsService, openGraphSchemaService) // Set neo4j batch and flush sizes neo4jParameters := appcfg.GetNeo4jParameters(ctx, connections.RDMS) diff --git a/cmd/api/src/services/entrypoint_test.go b/cmd/api/src/services/entrypoint_test.go new file mode 100644 index 000000000000..35dddba84c3a --- /dev/null +++ b/cmd/api/src/services/entrypoint_test.go @@ -0,0 +1,119 @@ +// Copyright 2026 Specter Ops, Inc. +// +// Licensed under the Apache License, Version 2.0 +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 + +package services_test + +import ( + "context" + "path/filepath" + "testing" + + "github.com/specterops/bloodhound/cmd/api/src/bootstrap" + "github.com/specterops/bloodhound/cmd/api/src/config" + "github.com/specterops/bloodhound/cmd/api/src/database" + "github.com/specterops/bloodhound/cmd/api/src/services" + storageService "github.com/specterops/bloodhound/cmd/api/src/services/storage" + "github.com/specterops/bloodhound/packages/go/storage" + storagemocks "github.com/specterops/bloodhound/packages/go/storage/mocks" + "github.com/specterops/dawgs/graph" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" +) + +func runtimeDependencyTestConfig(t *testing.T) config.Configuration { + t.Helper() + + rootDirectory := t.TempDir() + + return config.Configuration{ + WorkDir: filepath.Join(rootDirectory, "work"), + CollectorsBasePath: filepath.Join(rootDirectory, "collectors"), + } +} + +func TestCreateRuntimeDependenciesInitializesDefaultFileServices(t *testing.T) { + t.Parallel() + + var ( + cfg = runtimeDependencyTestConfig(t) + connections = bootstrap.DatabaseConnections[*database.BloodhoundDB, *graph.DatabaseSwitch]{} + ) + + require.NoError(t, bootstrap.EnsureServerDirectories(cfg)) + + deps, err := services.CreateRuntimeDependencies(context.Background(), cfg, connections) + + require.NoError(t, err) + require.NotNil(t, deps.FileServiceResolver) + + for _, serviceName := range []storage.FileServiceName{ + storage.FileServiceWork, + storage.FileServiceIngest, + storage.FileServiceRetained, + storage.FileServiceCollectors, + } { + fileService, err := deps.FileServiceResolver.Resolve(serviceName) + require.NoError(t, err) + require.NotNil(t, fileService) + } +} + +func TestPreMigrationDaemonsResolvesRetainedFileService(t *testing.T) { + t.Parallel() + + t.Run("returns error when retained file service is missing", func(t *testing.T) { + t.Parallel() + + var ( + mockCtrl = gomock.NewController(t) + mockFileService = storagemocks.NewMockFileService(mockCtrl) + connections = bootstrap.DatabaseConnections[*database.BloodhoundDB, *graph.DatabaseSwitch]{} + ) + + fileServiceResolver, err := storageService.NewFileServiceResolver(map[storage.FileServiceName]storage.FileService{ + storage.FileServiceIngest: mockFileService, + }) + require.NoError(t, err) + + _, err = services.PreMigrationDaemons(context.Background(), config.Configuration{}, connections, bootstrap.RuntimeDependencies{ + FileServiceResolver: fileServiceResolver, + }) + + require.ErrorContains(t, err, "error resolving FileServiceRetained") + }) + + t.Run("creates daemon when retained file service resolves", func(t *testing.T) { + t.Parallel() + + var ( + mockCtrl = gomock.NewController(t) + mockRetainedService = storagemocks.NewMockFileService(mockCtrl) + connections = bootstrap.DatabaseConnections[*database.BloodhoundDB, *graph.DatabaseSwitch]{} + ) + + fileServiceResolver, err := storageService.NewFileServiceResolver(map[storage.FileServiceName]storage.FileService{ + storage.FileServiceRetained: mockRetainedService, + }) + require.NoError(t, err) + + daemonInstances, err := services.PreMigrationDaemons(context.Background(), config.Configuration{}, connections, bootstrap.RuntimeDependencies{ + FileServiceResolver: fileServiceResolver, + }) + + require.NoError(t, err) + require.Len(t, daemonInstances, 1) + }) +} diff --git a/cmd/api/src/services/fs/fs.go b/cmd/api/src/services/fs/fs.go deleted file mode 100644 index 749d47601100..000000000000 --- a/cmd/api/src/services/fs/fs.go +++ /dev/null @@ -1,38 +0,0 @@ -// Copyright 2025 Specter Ops, Inc. -// -// Licensed under the Apache License, Version 2.0 -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// -// SPDX-License-Identifier: Apache-2.0 - -package fs - -import "os" - -//go:generate go run go.uber.org/mock/mockgen -copyright_file=../../../../../LICENSE.header -destination=./mocks/fs.go -package=mocks . Service - -// Serves as a lightweight wrapper around the os package which allows for -// path management to be abstracted. -type Service interface { - CreateTemporaryDirectory(dir, pattern string) (*os.File, error) - ReadFile(name string) ([]byte, error) -} - -type Client struct{} - -func (c *Client) CreateTemporaryDirectory(dir, pattern string) (*os.File, error) { - return os.CreateTemp(dir, pattern) -} - -func (c *Client) ReadFile(name string) ([]byte, error) { - return os.ReadFile(name) -} diff --git a/cmd/api/src/services/fs/mocks/fs.go b/cmd/api/src/services/fs/mocks/fs.go deleted file mode 100644 index 5ba753d5346c..000000000000 --- a/cmd/api/src/services/fs/mocks/fs.go +++ /dev/null @@ -1,87 +0,0 @@ -// Copyright 2026 Specter Ops, Inc. -// -// Licensed under the Apache License, Version 2.0 -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// -// SPDX-License-Identifier: Apache-2.0 - -// Code generated by MockGen. DO NOT EDIT. -// Source: github.com/specterops/bloodhound/cmd/api/src/services/fs (interfaces: Service) -// -// Generated by this command: -// -// mockgen -copyright_file=../../../../../LICENSE.header -destination=./mocks/fs.go -package=mocks . Service -// - -// Package mocks is a generated GoMock package. -package mocks - -import ( - os "os" - reflect "reflect" - - gomock "go.uber.org/mock/gomock" -) - -// MockService is a mock of Service interface. -type MockService struct { - ctrl *gomock.Controller - recorder *MockServiceMockRecorder - isgomock struct{} -} - -// MockServiceMockRecorder is the mock recorder for MockService. -type MockServiceMockRecorder struct { - mock *MockService -} - -// NewMockService creates a new mock instance. -func NewMockService(ctrl *gomock.Controller) *MockService { - mock := &MockService{ctrl: ctrl} - mock.recorder = &MockServiceMockRecorder{mock} - return mock -} - -// EXPECT returns an object that allows the caller to indicate expected use. -func (m *MockService) EXPECT() *MockServiceMockRecorder { - return m.recorder -} - -// CreateTemporaryDirectory mocks base method. -func (m *MockService) CreateTemporaryDirectory(dir, pattern string) (*os.File, error) { - m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "CreateTemporaryDirectory", dir, pattern) - ret0, _ := ret[0].(*os.File) - ret1, _ := ret[1].(error) - return ret0, ret1 -} - -// CreateTemporaryDirectory indicates an expected call of CreateTemporaryDirectory. -func (mr *MockServiceMockRecorder) CreateTemporaryDirectory(dir, pattern any) *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CreateTemporaryDirectory", reflect.TypeOf((*MockService)(nil).CreateTemporaryDirectory), dir, pattern) -} - -// ReadFile mocks base method. -func (m *MockService) ReadFile(name string) ([]byte, error) { - m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "ReadFile", name) - ret0, _ := ret[0].([]byte) - ret1, _ := ret[1].(error) - return ret0, ret1 -} - -// ReadFile indicates an expected call of ReadFile. -func (mr *MockServiceMockRecorder) ReadFile(name any) *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReadFile", reflect.TypeOf((*MockService)(nil).ReadFile), name) -} diff --git a/cmd/api/src/services/graphify/graphify_integration_test.go b/cmd/api/src/services/graphify/graphify_integration_test.go index 9ae9fbe322c7..2f7620244913 100644 --- a/cmd/api/src/services/graphify/graphify_integration_test.go +++ b/cmd/api/src/services/graphify/graphify_integration_test.go @@ -34,6 +34,7 @@ import ( "github.com/specterops/bloodhound/cmd/api/src/database" "github.com/specterops/bloodhound/cmd/api/src/migrations" "github.com/specterops/bloodhound/cmd/api/src/services/graphify" + "github.com/specterops/bloodhound/cmd/api/src/services/storage" "github.com/specterops/bloodhound/cmd/api/src/services/upload" "github.com/specterops/bloodhound/cmd/api/src/test/integration/utils" "github.com/specterops/bloodhound/packages/go/graphschema" @@ -108,9 +109,14 @@ func setupIntegrationTestSuite(t *testing.T, fixturesPath string) IntegrationTes cfg.WorkDir = workDir + fileServices, err := storage.NewDefaultFileServices(cfg) + require.NoError(t, err, "error creating default file services") + fileServiceResolver, err := storage.NewFileServiceResolver(fileServices) + require.NoError(t, err, "error creating fileServiceResolver") + return IntegrationTestSuite{ Context: ctx, - GraphifyService: graphify.NewGraphifyService(ctx, db, graphDB, cfg, ingestSchema, nil), + GraphifyService: graphify.NewGraphifyService(ctx, db, graphDB, cfg, ingestSchema, fileServiceResolver, nil), GraphDB: graphDB, BHDatabase: db, WorkDir: workDir, diff --git a/cmd/api/src/services/graphify/ingest_storage.go b/cmd/api/src/services/graphify/ingest_storage.go index 614264dd15a0..e498e9ebd47e 100644 --- a/cmd/api/src/services/graphify/ingest_storage.go +++ b/cmd/api/src/services/graphify/ingest_storage.go @@ -4,7 +4,7 @@ // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // -// http://www.apache.org/licenses/LICENSE-2.0 +// http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, @@ -13,182 +13,185 @@ // limitations under the License. // // SPDX-License-Identifier: Apache-2.0 - package graphify import ( "archive/zip" "context" - "errors" + "fmt" "io" - "io/fs" "log/slog" "os" - "github.com/specterops/bloodhound/cmd/api/src/config" "github.com/specterops/bloodhound/cmd/api/src/model" "github.com/specterops/bloodhound/packages/go/bhlog/attr" - "github.com/specterops/bloodhound/packages/go/bhlog/measure" "github.com/specterops/bloodhound/packages/go/bomenc" - "github.com/specterops/bloodhound/packages/go/errorlist" + "github.com/specterops/bloodhound/packages/go/storage" + "github.com/specterops/dawgs/util" ) -type IngestFileData struct { - Name string - ParentFile string - Path string - Errors []string - UserDataErrs []string -} +func SpoolToScratch(ctx context.Context, scratchDirectory string, fileService storage.FileService, storedFileName string) (string, error) { + var ( + sourceFile io.ReadCloser + scratchFile *os.File + scratchPath string + err error + success bool + ) -// ExtractIngestFiles takes a stored ingest file path and extracts ZIP archives when necessary, returning the paths for -// files to process. ZIP archives are deleted after extraction, and extracted files are written to the configured temp -// directory with tempFilePrefix. -func ExtractIngestFiles(ctx context.Context, configuration config.Configuration, storedFileName string, providedFileName string, fileType model.FileType, tempFilePrefix string) ([]IngestFileData, error) { - if fileType == model.FileTypeJson { - // If this isn't a zip file, just return a slice with the path in it and let stuff process as normal - return []IngestFileData{ - { - Name: providedFileName, - Path: storedFileName, - Errors: []string{}, - }, - }, nil - } else if archive, err := zip.OpenReader(storedFileName); err != nil { - return []IngestFileData{}, err - } else { - var ( - errs = errorlist.NewBuilder() - fileData = make([]IngestFileData, 0) - ) + sourceFile, _, err = fileService.GetFile(ctx, storedFileName) + if err != nil { + return "", fmt.Errorf("open stored ingest file %q: %w", storedFileName, err) + } + defer sourceFile.Close() + + scratchFile, err = os.CreateTemp(scratchDirectory, "archive-*") + if err != nil { + return "", fmt.Errorf("create ingest scratch file: %w", err) + } - defer func() { - if err := archive.Close(); err != nil { - slog.ErrorContext( - ctx, - "Error closing archive", - slog.String("path", storedFileName), - attr.Error(err), - ) - } - if err := os.Remove(storedFileName); err != nil { - slog.ErrorContext( - ctx, - "Error deleting archive", - slog.String("path", storedFileName), - attr.Error(err), - ) - } - }() - - for _, file := range archive.File { - // skip directories - if file.FileInfo().IsDir() { - continue - } - - fileName, err := extractToTempFile(configuration, file, tempFilePrefix) - if err != nil { - fileData = append(fileData, IngestFileData{ - Name: file.Name, - ParentFile: providedFileName, - Errors: []string{err.Error()}, - }) - - errs.Add(err) - } else { - fileData = append(fileData, IngestFileData{ - Name: file.Name, - ParentFile: providedFileName, - Path: fileName, - }) - } + scratchPath = scratchFile.Name() + + defer func() { + if scratchFile != nil { + _ = scratchFile.Close() + } + if !success { + _ = os.Remove(scratchPath) } + }() + + if _, err = io.Copy(scratchFile, sourceFile); err != nil { + return "", fmt.Errorf("copy stored ingest file %q to scratch: %w", storedFileName, err) + } - return fileData, errs.Build() + if err = scratchFile.Close(); err != nil { + return "", fmt.Errorf("close ingest scratch file %q: %w", scratchPath, err) } + + scratchFile = nil + success = true + + return scratchPath, nil } -func extractToTempFile(configuration config.Configuration, file *zip.File, tempFilePrefix string) (string, error) { +func OpenScratchReadSeeker(ctx context.Context, scratchDirectory string, fileService storage.FileService, storedFileName string) (*os.File, string, error) { var ( - tempFile *os.File - tempFileName string - sourceFile io.ReadCloser - normalizedFile io.Reader - extractionSucceeded bool - err error + scratchPath string + scratchFile *os.File + err error ) - if tempFile, err = os.CreateTemp(configuration.TempDirectory(), tempFilePrefix); err != nil { - return "", err + if scratchPath, err = SpoolToScratch(ctx, scratchDirectory, fileService, storedFileName); err != nil { + return nil, "", err } - tempFileName = tempFile.Name() - defer func() { - if !extractionSucceeded { - tempFile.Close() - os.Remove(tempFileName) - } - }() + if scratchFile, err = os.Open(scratchPath); err != nil { + _ = os.Remove(scratchPath) + return nil, "", fmt.Errorf("error opening ingest scratch file %q: %w", scratchPath, err) + } - if sourceFile, err = file.Open(); err != nil { - return "", err + return scratchFile, scratchPath, nil +} + +func WriteArchiveFileToStorage(ctx context.Context, fileService storage.FileService, archiveFile *zip.File, prefix string) (string, error) { + var ( + sourceFile io.ReadCloser + normalizedFile io.Reader + err error + ) + + if sourceFile, err = archiveFile.Open(); err != nil { + return "", fmt.Errorf("error opening archive file %q: %w", archiveFile.Name, err) } defer sourceFile.Close() - // this creates a normalized file to feed to the copy if normalizedFile, err = bomenc.NormalizeToUTF8(sourceFile); err != nil { - return "", err - // and this is what actually copies it to disk - } else if _, err = io.Copy(tempFile, normalizedFile); err != nil { - return "", err - // try closing the temp file - } else if err = tempFile.Close(); err != nil { - return "", err - } else { - // let the deferred method above know we shouldn't delete it and return the filename - extractionSucceeded = true - return tempFileName, nil + return "", fmt.Errorf("error normalizing archive file %q to UTF8: %w", archiveFile.Name, err) + } + + extractedPath, err := fileService.WriteTempFile( + ctx, + prefix, + normalizedFile, + storage.WriteOptions{}, + ) + if err != nil { + return "", fmt.Errorf("write archive file %q to storage: %w", archiveFile.Name, err) } + + return extractedPath, nil } -func processSingleFile(ctx context.Context, fileData IngestFileData, ingestContext *IngestContext, readOpts ReadOptions) error { - defer measure.ContextLogAndMeasureWithThreshold(ctx, slog.LevelDebug, "processing single file for ingest", slog.String("filepath", fileData.Path))() +func ExtractIngestFiles(ctx context.Context, scratchDirectory string, fileService storage.FileService, storedFileName, providedFileName string, fileType model.FileType, prefix string) ([]IngestFileData, error) { + if fileType == model.FileTypeJson { + // If this isn't a zip file, just return a slice with the path in it and let stuff process as normal + return []IngestFileData{ + { + Name: providedFileName, + Path: storedFileName, + }, + }, nil + } - file, err := os.Open(fileData.Path) + // Zip Path: + scratchPath, err := SpoolToScratch(ctx, scratchDirectory, fileService, storedFileName) if err != nil { - slog.ErrorContext( - ctx, - "Error opening ingest file", - slog.String("path", fileData.Path), - attr.Error(err), - ) - return err + return []IngestFileData{ + { + Name: providedFileName, + Path: storedFileName, + Errors: []string{fmt.Sprintf("Error spooling archive to scratch: %v", err)}, + }, + }, err } + defer os.Remove(scratchPath) - defer func() { - file.Close() - - // Always remove the file after attempting to ingest it. Even if it failed - if err := os.Remove(fileData.Path); err != nil && !errors.Is(err, fs.ErrNotExist) { - slog.ErrorContext( - ctx, - "Error removing ingest file", - slog.String("path", fileData.Path), - attr.Error(err), - ) + archive, err := zip.OpenReader(scratchPath) + if err != nil { + return []IngestFileData{ + { + Name: providedFileName, + Path: storedFileName, + Errors: []string{fmt.Sprintf("Error opening archive: %v", err)}, + }, + }, err + } + defer archive.Close() + + var ( + errs = util.NewErrorCollector() + fileData = make([]IngestFileData, 0) + ) + + for _, archiveFile := range archive.File { + if archiveFile.FileInfo().IsDir() { + continue } - }() - if err := ReadFileForIngest(ingestContext, file, readOpts); err != nil { + processedFileData := IngestFileData{ + Name: archiveFile.Name, + ParentFile: providedFileName, + } + + if extractedPath, err := WriteArchiveFileToStorage(ctx, fileService, archiveFile, prefix); err != nil { + processedFileData.Errors = []string{err.Error()} + errs.Add(fmt.Errorf("error extracting file %s in archive %s: %w", archiveFile.Name, storedFileName, err)) + } else { + processedFileData.Path = extractedPath + } + + fileData = append(fileData, processedFileData) + } + + if err := fileService.DeleteFile(ctx, storedFileName); err != nil { slog.ErrorContext( ctx, - "Error reading ingest file", - slog.String("storage_path", fileData.Path), + "Error deleting archive", + slog.String("path", storedFileName), attr.Error(err), ) - return err } - return nil + return fileData, errs.Combined() } diff --git a/cmd/api/src/services/graphify/ingest_storage_test.go b/cmd/api/src/services/graphify/ingest_storage_test.go index 1a648eadf45d..b53abc79c4a2 100644 --- a/cmd/api/src/services/graphify/ingest_storage_test.go +++ b/cmd/api/src/services/graphify/ingest_storage_test.go @@ -18,108 +18,622 @@ package graphify_test import ( "archive/zip" + "bytes" "context" + "errors" + "io" "os" - "path/filepath" + "strings" "testing" - "github.com/specterops/bloodhound/cmd/api/src/config" "github.com/specterops/bloodhound/cmd/api/src/model" "github.com/specterops/bloodhound/cmd/api/src/services/graphify" + "github.com/specterops/bloodhound/packages/go/storage" + storagemocks "github.com/specterops/bloodhound/packages/go/storage/mocks" "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" ) -type zipTestFile struct { +type ingestStorageTrackingReadCloser struct { + reader io.Reader + closed bool +} + +func (s *ingestStorageTrackingReadCloser) Read(p []byte) (int, error) { + return s.reader.Read(p) +} + +func (s *ingestStorageTrackingReadCloser) Close() error { + s.closed = true + return nil +} + +type ingestStorageErrorReadCloser struct { + err error + closed bool +} + +func (s *ingestStorageErrorReadCloser) Read([]byte) (int, error) { + return 0, s.err +} + +func (s *ingestStorageErrorReadCloser) Close() error { + s.closed = true + return nil +} + +type ingestStorageZipEntry struct { name string - content []byte + content string + isDir bool } -func writeZipFile(t *testing.T, archivePath string, files []zipTestFile) { - var ( - archiveFile *os.File - zipWriter *zip.Writer - err error - ) +func buildIngestStorageZip(t *testing.T, entries ...ingestStorageZipEntry) []byte { + t.Helper() - archiveFile, err = os.Create(archivePath) - require.NoError(t, err) - defer archiveFile.Close() + var archive bytes.Buffer + archiveWriter := zip.NewWriter(&archive) - zipWriter = zip.NewWriter(archiveFile) - defer zipWriter.Close() + for _, entry := range entries { + if entry.isDir { + _, err := archiveWriter.Create(entry.name) + require.NoError(t, err) + continue + } - for _, file := range files { - fileWriter, err := zipWriter.Create(file.name) + entryWriter, err := archiveWriter.Create(entry.name) require.NoError(t, err) - _, err = fileWriter.Write(file.content) + _, err = entryWriter.Write([]byte(entry.content)) require.NoError(t, err) } - _, err = zipWriter.Create("empty-directory/") + require.NoError(t, archiveWriter.Close()) + return archive.Bytes() +} + +func openIngestStorageZipFile(t *testing.T, archiveBytes []byte, fileName string) *zip.File { + t.Helper() + + archiveReader, err := zip.NewReader(bytes.NewReader(archiveBytes), int64(len(archiveBytes))) + require.NoError(t, err) + + for _, archiveFile := range archiveReader.File { + if archiveFile.Name == fileName { + return archiveFile + } + } + + t.Fatalf("archive file %q not found", fileName) + return nil +} + +func requireScratchDirectoryEmpty(t *testing.T, scratchDirectory string) { + t.Helper() + + entries, err := os.ReadDir(scratchDirectory) + require.NoError(t, err) + require.Empty(t, entries) +} + +func requireReadFile(t *testing.T, filePath string) []byte { + t.Helper() + + data, err := os.ReadFile(filePath) require.NoError(t, err) + return data } -func TestExtractIngestFilesReturnsJSONFile(t *testing.T) { +func TestSpoolToScratch(t *testing.T) { + t.Parallel() + var ( - ctx = context.Background() - configuration = config.Configuration{WorkDir: t.TempDir()} - storedFileName = filepath.Join(configuration.WorkDir, "ingest.json") - providedFileName = "provided.json" + errGet = errors.New("get failed") + errRead = errors.New("read failed") ) - fileData, err := graphify.ExtractIngestFiles(ctx, configuration, storedFileName, providedFileName, model.FileTypeJson, "unused") - require.NoError(t, err) - require.Equal(t, []graphify.IngestFileData{ + type expected struct { + errIs error + errContains string + content string + closed bool + } + + type testData struct { + name string + storedName string + setupMock func(t *testing.T, ctx context.Context, mockFileService *storagemocks.MockFileService) io.ReadCloser + expected expected + verify func(t *testing.T, scratchDirectory string, scratchPath string) + } + + tests := []testData{ { - Name: providedFileName, - Path: storedFileName, - Errors: []string{}, + name: "copies stored file to scratch", + storedName: "stored.json", + setupMock: func(t *testing.T, ctx context.Context, mockFileService *storagemocks.MockFileService) io.ReadCloser { + readCloser := &ingestStorageTrackingReadCloser{reader: strings.NewReader("stored content")} + mockFileService.EXPECT(). + GetFile(ctx, "stored.json"). + Return(readCloser, storage.FileInfo{}, nil) + + return readCloser + }, + expected: expected{ + content: "stored content", + closed: true, + }, }, - }, fileData) + { + name: "get error returns wrapped error", + storedName: "missing.json", + setupMock: func(t *testing.T, ctx context.Context, mockFileService *storagemocks.MockFileService) io.ReadCloser { + mockFileService.EXPECT(). + GetFile(ctx, "missing.json"). + Return(nil, storage.FileInfo{}, errGet) + + return nil + }, + expected: expected{ + errIs: errGet, + errContains: `open stored ingest file "missing.json"`, + }, + verify: func(t *testing.T, scratchDirectory string, scratchPath string) { + require.Empty(t, scratchPath) + requireScratchDirectoryEmpty(t, scratchDirectory) + }, + }, + { + name: "copy error removes scratch file", + storedName: "stored.json", + setupMock: func(t *testing.T, ctx context.Context, mockFileService *storagemocks.MockFileService) io.ReadCloser { + readCloser := &ingestStorageErrorReadCloser{err: errRead} + mockFileService.EXPECT(). + GetFile(ctx, "stored.json"). + Return(readCloser, storage.FileInfo{}, nil) + + return readCloser + }, + expected: expected{ + errIs: errRead, + errContains: `copy stored ingest file "stored.json" to scratch`, + closed: true, + }, + verify: func(t *testing.T, scratchDirectory string, scratchPath string) { + require.Empty(t, scratchPath) + requireScratchDirectoryEmpty(t, scratchDirectory) + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + var ( + ctx = context.Background() + scratchDirectory = t.TempDir() + mockFileService = storagemocks.NewMockFileService(gomock.NewController(t)) + readCloser = testCase.setupMock(t, ctx, mockFileService) + ) + + // Act + scratchPath, err := graphify.SpoolToScratch(ctx, scratchDirectory, mockFileService, testCase.storedName) + + // Assert + if testCase.expected.errIs != nil { + require.ErrorIs(t, err, testCase.expected.errIs) + require.Contains(t, err.Error(), testCase.expected.errContains) + } else { + require.NoError(t, err) + require.Equal(t, testCase.expected.content, string(requireReadFile(t, scratchPath))) + } + + if readCloser != nil { + switch typedReadCloser := readCloser.(type) { + case *ingestStorageTrackingReadCloser: + require.Equal(t, testCase.expected.closed, typedReadCloser.closed) + case *ingestStorageErrorReadCloser: + require.Equal(t, testCase.expected.closed, typedReadCloser.closed) + } + } + + if testCase.verify != nil { + testCase.verify(t, scratchDirectory, scratchPath) + } + }) + } } -func TestExtractIngestFilesExpandsZIPIntoTempDirectory(t *testing.T) { +func TestOpenScratchReadSeeker(t *testing.T) { + t.Parallel() + + var errGet = errors.New("get failed") + + type expected struct { + errIs error + content string + } + + type testData struct { + name string + storedName string + setupMock func(t *testing.T, ctx context.Context, mockFileService *storagemocks.MockFileService) + expected expected + } + + tests := []testData{ + { + name: "opens copied scratch file", + storedName: "stored.json", + setupMock: func(t *testing.T, ctx context.Context, mockFileService *storagemocks.MockFileService) { + mockFileService.EXPECT(). + GetFile(ctx, "stored.json"). + Return(io.NopCloser(strings.NewReader("stored content")), storage.FileInfo{}, nil) + }, + expected: expected{ + content: "stored content", + }, + }, + { + name: "spool error returns error", + storedName: "missing.json", + setupMock: func(t *testing.T, ctx context.Context, mockFileService *storagemocks.MockFileService) { + mockFileService.EXPECT(). + GetFile(ctx, "missing.json"). + Return(nil, storage.FileInfo{}, errGet) + }, + expected: expected{ + errIs: errGet, + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + var ( + ctx = context.Background() + scratchDirectory = t.TempDir() + mockFileService = storagemocks.NewMockFileService(gomock.NewController(t)) + ) + + testCase.setupMock(t, ctx, mockFileService) + + // Act + scratchFile, scratchPath, err := graphify.OpenScratchReadSeeker(ctx, scratchDirectory, mockFileService, testCase.storedName) + + // Assert + if testCase.expected.errIs != nil { + require.ErrorIs(t, err, testCase.expected.errIs) + require.Nil(t, scratchFile) + require.Empty(t, scratchPath) + return + } + + require.NoError(t, err) + defer os.Remove(scratchPath) + defer scratchFile.Close() + + content, err := io.ReadAll(scratchFile) + require.NoError(t, err) + require.Equal(t, testCase.expected.content, string(content)) + }) + } +} + +func TestWriteArchiveFileToStorage(t *testing.T) { + t.Parallel() + + var errWrite = errors.New("write failed") + + type expected struct { + errIs error + errContains string + path string + content string + } + + type testData struct { + name string + content string + writeErr error + expected expected + } + + tests := []testData{ + { + name: "writes normalized archive file content to storage", + content: string([]byte{0xEF, 0xBB, 0xBF}) + "stored content", + expected: expected{ + path: "prefix/tmp-file", + content: "stored content", + }, + }, + { + name: "write error returns wrapped error", + content: "stored content", + writeErr: errWrite, + expected: expected{ + errIs: errWrite, + errContains: `write archive file "file.json" to storage`, + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + var ( + ctx = context.Background() + archiveBytes = buildIngestStorageZip(t, ingestStorageZipEntry{name: "file.json", content: testCase.content}) + archiveFile = openIngestStorageZipFile(t, archiveBytes, "file.json") + mockFileService = storagemocks.NewMockFileService(gomock.NewController(t)) + ) + + mockFileService.EXPECT(). + WriteTempFile(ctx, "prefix", gomock.Any(), storage.WriteOptions{}). + DoAndReturn(func(_ context.Context, _ string, reader io.Reader, _ storage.WriteOptions) (string, error) { + content, err := io.ReadAll(reader) + require.NoError(t, err) + if testCase.expected.content != "" { + require.Equal(t, testCase.expected.content, string(content)) + } + + return "prefix/tmp-file", testCase.writeErr + }) + + // Act + extractedPath, err := graphify.WriteArchiveFileToStorage(ctx, mockFileService, archiveFile, "prefix") + + // Assert + if testCase.expected.errIs != nil { + require.ErrorIs(t, err, testCase.expected.errIs) + require.Contains(t, err.Error(), testCase.expected.errContains) + require.Empty(t, extractedPath) + return + } + + require.NoError(t, err) + require.Equal(t, testCase.expected.path, extractedPath) + }) + } +} + +func TestExtractIngestFiles(t *testing.T) { + t.Parallel() + var ( - ctx = context.Background() - configuration = config.Configuration{WorkDir: t.TempDir()} - archivePath = filepath.Join(configuration.WorkDir, "archive.zip") - providedFileName = "archive.zip" - tempFilePrefix = "file_upload_job7_" - bomFileContent = append([]byte{0xEF, 0xBB, 0xBF}, []byte(`{"meta":{"type":"domains","version":5,"count":0},"data":[]}`)...) - plainFileContent = []byte(`{"meta":{"type":"users","version":5,"count":0},"data":[]}`) + errGet = errors.New("get failed") + errWrite = errors.New("write failed") ) - require.NoError(t, os.MkdirAll(configuration.TempDirectory(), 0o755)) - writeZipFile(t, archivePath, []zipTestFile{ + type expected struct { + errIs error + errContains string + fileData []graphify.IngestFileData + } + + type testData struct { + name string + fileType model.FileType + providedFileName string + setupMock func(t *testing.T, ctx context.Context, mockFileService *storagemocks.MockFileService) + expected expected + verify func(t *testing.T, scratchDirectory string) + } + + tests := []testData{ + { + name: "json file returns stored path without storage calls", + fileType: model.FileTypeJson, + providedFileName: "provided.json", + expected: expected{ + fileData: []graphify.IngestFileData{ + { + Name: "provided.json", + Path: "stored.json", + }, + }, + }, + }, + { + name: "zip file extracts files and deletes archive", + fileType: model.FileTypeZip, + setupMock: func(t *testing.T, ctx context.Context, mockFileService *storagemocks.MockFileService) { + archiveBytes := buildIngestStorageZip(t, + ingestStorageZipEntry{name: "nested/", isDir: true}, + ingestStorageZipEntry{name: "nested/one.json", content: "one"}, + ingestStorageZipEntry{name: "two.json", content: "two"}, + ) + + mockFileService.EXPECT(). + GetFile(ctx, "stored.json"). + Return(io.NopCloser(bytes.NewReader(archiveBytes)), storage.FileInfo{}, nil) + gomock.InOrder( + mockFileService.EXPECT(). + WriteTempFile(ctx, "prefix", gomock.Any(), storage.WriteOptions{}). + DoAndReturn(func(_ context.Context, _ string, reader io.Reader, _ storage.WriteOptions) (string, error) { + content, err := io.ReadAll(reader) + require.NoError(t, err) + require.Equal(t, "one", string(content)) + + return "prefix/tmp-one", nil + }), + mockFileService.EXPECT(). + WriteTempFile(ctx, "prefix", gomock.Any(), storage.WriteOptions{}). + DoAndReturn(func(_ context.Context, _ string, reader io.Reader, _ storage.WriteOptions) (string, error) { + content, err := io.ReadAll(reader) + require.NoError(t, err) + require.Equal(t, "two", string(content)) + + return "prefix/tmp-two", nil + }), + ) + mockFileService.EXPECT(). + DeleteFile(ctx, "stored.json"). + Return(nil) + }, + expected: expected{ + fileData: []graphify.IngestFileData{ + { + Name: "nested/one.json", + ParentFile: "provided.zip", + Path: "prefix/tmp-one", + }, + { + Name: "two.json", + ParentFile: "provided.zip", + Path: "prefix/tmp-two", + }, + }, + }, + verify: func(t *testing.T, scratchDirectory string) { + requireScratchDirectoryEmpty(t, scratchDirectory) + }, + }, { - name: "domains.json", - content: bomFileContent, + name: "spool error returns file data with error", + fileType: model.FileTypeZip, + setupMock: func(t *testing.T, ctx context.Context, mockFileService *storagemocks.MockFileService) { + mockFileService.EXPECT(). + GetFile(ctx, "stored.json"). + Return(nil, storage.FileInfo{}, errGet) + }, + expected: expected{ + errIs: errGet, + fileData: []graphify.IngestFileData{ + { + Name: "provided.zip", + Path: "stored.json", + Errors: []string{`Error spooling archive to scratch: open stored ingest file "stored.json": get failed`}, + }, + }, + }, + verify: func(t *testing.T, scratchDirectory string) { + requireScratchDirectoryEmpty(t, scratchDirectory) + }, }, { - name: "users.json", - content: plainFileContent, + name: "bad zip returns file data with error", + fileType: model.FileTypeZip, + setupMock: func(t *testing.T, ctx context.Context, mockFileService *storagemocks.MockFileService) { + mockFileService.EXPECT(). + GetFile(ctx, "stored.json"). + Return(io.NopCloser(strings.NewReader("not a zip")), storage.FileInfo{}, nil) + }, + expected: expected{ + errIs: zip.ErrFormat, + fileData: []graphify.IngestFileData{ + { + Name: "provided.zip", + Path: "stored.json", + Errors: []string{"Error opening archive: zip: not a valid zip file"}, + }, + }, + }, + verify: func(t *testing.T, scratchDirectory string) { + requireScratchDirectoryEmpty(t, scratchDirectory) + }, }, - }) + { + name: "archive file write error records file error and returns aggregate error", + fileType: model.FileTypeZip, + setupMock: func(t *testing.T, ctx context.Context, mockFileService *storagemocks.MockFileService) { + archiveBytes := buildIngestStorageZip(t, + ingestStorageZipEntry{name: "one.json", content: "one"}, + ingestStorageZipEntry{name: "two.json", content: "two"}, + ) - fileData, err := graphify.ExtractIngestFiles(ctx, configuration, archivePath, providedFileName, model.FileTypeZip, tempFilePrefix) - require.NoError(t, err) - require.Len(t, fileData, 2) - require.NoFileExists(t, archivePath) + mockFileService.EXPECT(). + GetFile(ctx, "stored.json"). + Return(io.NopCloser(bytes.NewReader(archiveBytes)), storage.FileInfo{}, nil) + gomock.InOrder( + mockFileService.EXPECT(). + WriteTempFile(ctx, "prefix", gomock.Any(), storage.WriteOptions{}). + DoAndReturn(func(_ context.Context, _ string, reader io.Reader, _ storage.WriteOptions) (string, error) { + _, err := io.ReadAll(reader) + require.NoError(t, err) + + return "", errWrite + }), + mockFileService.EXPECT(). + WriteTempFile(ctx, "prefix", gomock.Any(), storage.WriteOptions{}). + DoAndReturn(func(_ context.Context, _ string, reader io.Reader, _ storage.WriteOptions) (string, error) { + content, err := io.ReadAll(reader) + require.NoError(t, err) + require.Equal(t, "two", string(content)) - for _, file := range fileData { - require.Equal(t, providedFileName, file.ParentFile) - require.Empty(t, file.Errors) - require.DirExists(t, filepath.Dir(file.Path)) - require.FileExists(t, file.Path) - require.Contains(t, filepath.Base(file.Path), tempFilePrefix) + return "prefix/tmp-two", nil + }), + ) + mockFileService.EXPECT(). + DeleteFile(ctx, "stored.json"). + Return(nil) + }, + expected: expected{ + errIs: errWrite, + errContains: `error extracting file one.json in archive stored.json: write archive file "one.json" to storage: write failed`, + fileData: []graphify.IngestFileData{ + { + Name: "one.json", + ParentFile: "provided.zip", + Errors: []string{`write archive file "one.json" to storage: write failed`}, + }, + { + Name: "two.json", + ParentFile: "provided.zip", + Path: "prefix/tmp-two", + }, + }, + }, + verify: func(t *testing.T, scratchDirectory string) { + requireScratchDirectoryEmpty(t, scratchDirectory) + }, + }, } - domainsContent, err := os.ReadFile(fileData[0].Path) - require.NoError(t, err) - require.Equal(t, []byte(`{"meta":{"type":"domains","version":5,"count":0},"data":[]}`), domainsContent) + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() - usersContent, err := os.ReadFile(fileData[1].Path) - require.NoError(t, err) - require.Equal(t, plainFileContent, usersContent) + // Arrange + var ( + ctx = context.Background() + scratchDirectory = t.TempDir() + mockFileService = storagemocks.NewMockFileService(gomock.NewController(t)) + ) + + if testCase.setupMock != nil { + testCase.setupMock(t, ctx, mockFileService) + } + providedFileName := testCase.providedFileName + if providedFileName == "" { + providedFileName = "provided.zip" + } + + // Act + fileData, err := graphify.ExtractIngestFiles(ctx, scratchDirectory, mockFileService, "stored.json", providedFileName, testCase.fileType, "prefix") + + // Assert + if testCase.expected.errIs != nil { + require.ErrorIs(t, err, testCase.expected.errIs) + if testCase.expected.errContains != "" { + require.Contains(t, err.Error(), testCase.expected.errContains) + } + } else { + require.NoError(t, err) + } + require.Equal(t, testCase.expected.fileData, fileData) + + if testCase.verify != nil { + testCase.verify(t, scratchDirectory) + } + }) + } } diff --git a/cmd/api/src/services/graphify/service.go b/cmd/api/src/services/graphify/service.go index e4d23698c92c..9dba1fa6e57c 100644 --- a/cmd/api/src/services/graphify/service.go +++ b/cmd/api/src/services/graphify/service.go @@ -22,6 +22,7 @@ import ( "github.com/specterops/bloodhound/cmd/api/src/model" "github.com/specterops/bloodhound/cmd/api/src/model/appcfg" "github.com/specterops/bloodhound/cmd/api/src/services/graphify/endpoint" + "github.com/specterops/bloodhound/cmd/api/src/services/storage" "github.com/specterops/bloodhound/cmd/api/src/services/upload" "github.com/specterops/dawgs/graph" ) @@ -40,23 +41,25 @@ type GraphifyData interface { } type GraphifyService struct { - ctx context.Context - db GraphifyData - graphdb graph.Database - endpointResolver *endpoint.Resolver - cfg config.Configuration - schema upload.IngestSchema - changeManager ChangeManager + ctx context.Context + db GraphifyData + graphdb graph.Database + endpointResolver *endpoint.Resolver + cfg config.Configuration + schema upload.IngestSchema + fileServiceResolver storage.FileServiceResolver + changeManager ChangeManager } -func NewGraphifyService(ctx context.Context, db GraphifyData, graphDb graph.Database, cfg config.Configuration, schema upload.IngestSchema, changeManager ChangeManager) GraphifyService { +func NewGraphifyService(ctx context.Context, db GraphifyData, graphDb graph.Database, cfg config.Configuration, schema upload.IngestSchema, fileServiceResolver storage.FileServiceResolver, changeManager ChangeManager) GraphifyService { return GraphifyService{ - ctx: ctx, - db: db, - graphdb: graphDb, - endpointResolver: endpoint.NewResolver(graphDb), - cfg: cfg, - schema: schema, - changeManager: changeManager, + ctx: ctx, + db: db, + graphdb: graphDb, + endpointResolver: endpoint.NewResolver(graphDb), + cfg: cfg, + schema: schema, + fileServiceResolver: fileServiceResolver, + changeManager: changeManager, } } diff --git a/cmd/api/src/services/graphify/tasks.go b/cmd/api/src/services/graphify/tasks.go index b9f16092575e..e672b933a4dd 100644 --- a/cmd/api/src/services/graphify/tasks.go +++ b/cmd/api/src/services/graphify/tasks.go @@ -19,15 +19,19 @@ package graphify import ( "context" "errors" + "fmt" "io/fs" "log/slog" + "os" "time" "github.com/specterops/bloodhound/cmd/api/src/model" "github.com/specterops/bloodhound/cmd/api/src/model/appcfg" "github.com/specterops/bloodhound/cmd/api/src/services/graphify/endpoint" "github.com/specterops/bloodhound/packages/go/bhlog/attr" + "github.com/specterops/bloodhound/packages/go/bhlog/measure" "github.com/specterops/bloodhound/packages/go/errorlist" + "github.com/specterops/bloodhound/packages/go/storage" "github.com/specterops/dawgs/graph" ) @@ -44,12 +48,77 @@ func (s *GraphifyService) clearFileTask(ingestTask model.IngestTask) { } } +type IngestFileData struct { + Name string + ParentFile string + Path string + Errors []string + UserDataErrs []string +} + +func processSingleFile(ctx context.Context, fileService storage.FileService, tempDirectory string, fileData IngestFileData, ingestContext *IngestContext, readOpts ReadOptions) error { + defer measure.ContextLogAndMeasureWithThreshold(ctx, slog.LevelDebug, "processing single file for ingest", slog.String("filepath", fileData.Path))() + + file, scratchPath, err := OpenScratchReadSeeker(ctx, tempDirectory, fileService, fileData.Path) + if err != nil { + slog.ErrorContext( + ctx, + "Error opening ingest file", + slog.String("filepath", fileData.Path), + attr.Error(err), + ) + return err + } + + defer func() { + if err := file.Close(); err != nil { + slog.WarnContext( + ctx, + "Error closing ingest scratch file", + slog.String("scratch_path", scratchPath), + attr.Error(err), + ) + } + + // Always remove the file after attempting to ingest it. Even if it failed. + if err := os.Remove(scratchPath); err != nil && !errors.Is(err, fs.ErrNotExist) { + slog.WarnContext( + ctx, + "Error removing ingest scratch file", + slog.String("scratch_path", scratchPath), + attr.Error(err), + ) + } + + if err := fileService.DeleteFile(ctx, fileData.Path); err != nil { + slog.WarnContext( + ctx, + "Error removing ingest file", + slog.String("storage_path", fileData.Path), + attr.Error(err), + ) + } + }() + + if err := ReadFileForIngest(ingestContext, file, readOpts); err != nil { + slog.ErrorContext( + ctx, + "Error reading ingest file", + slog.String("storage_path", fileData.Path), + attr.Error(err), + ) + return err + } + + return nil +} + // ProcessIngestFile reads the files at the path supplied, and returns the total number of files in the // archive, the number of files that failed to ingest as JSON, and an error -func (s *GraphifyService) ProcessIngestFile(ic *IngestContext, task model.IngestTask) ([]IngestFileData, error) { +func (s *GraphifyService) ProcessIngestFile(ic *IngestContext, fileService storage.FileService, task model.IngestTask) ([]IngestFileData, error) { // Try to pre-process the file. If any of them fail, stop processing and return the error - if fileData, err := ExtractIngestFiles(s.ctx, s.cfg, task.StoredFileName, task.OriginalFileName, task.FileType, "bh"); err != nil { - return []IngestFileData{}, err + if fileData, err := ExtractIngestFiles(ic.Ctx, s.cfg.ScratchDirectory(), fileService, task.StoredFileName, task.OriginalFileName, task.FileType, fmt.Sprintf("file_upload_job_%d_", ic.JobId)); err != nil { + return fileData, err } else { errs := errorlist.NewBuilder() @@ -63,7 +132,11 @@ func (s *GraphifyService) ProcessIngestFile(ic *IngestContext, task model.Ingest RegisterSourceKind: s.RegisterSourceKind(s.ctx), } - if err := processSingleFile(ic.Ctx, data, ic, readOpts); err != nil { + if len(data.Errors) > 0 || data.Path == "" { + continue + } + + if err := processSingleFile(ic.Ctx, fileService, s.cfg.ScratchDirectory(), data, ic, readOpts); err != nil { var ( graphifyError errorlist.Error resolutionErr endpoint.ResolutionError @@ -125,6 +198,12 @@ func (s *GraphifyService) ProcessTasks(updateJob UpdateJobFunc) { return } + ingestFileService, err := s.fileServiceResolver.Resolve(storage.FileServiceIngest) + if err != nil { + slog.ErrorContext(s.ctx, "Error resolving ingest file service", attr.Error(err)) + return + } + start := time.Now() slog.InfoContext(s.ctx, "Ingest run starting", @@ -150,7 +229,7 @@ func (s *GraphifyService) ProcessTasks(updateJob UpdateJobFunc) { for _, task := range tasks { ingestCtx := s.NewIngestContext(s.ctx, time.Now().UTC(), flagChangeLogEnabled, task.JobId.ValueOrZero()) - fileData, err := s.ProcessIngestFile(ingestCtx, task) + fileData, err := s.ProcessIngestFile(ingestCtx, ingestFileService, task) switch { case errors.Is(err, fs.ErrNotExist): diff --git a/cmd/api/src/services/graphify/tasks_integration_test.go b/cmd/api/src/services/graphify/tasks_integration_test.go index 04d37a098c66..1cb03f07ab77 100644 --- a/cmd/api/src/services/graphify/tasks_integration_test.go +++ b/cmd/api/src/services/graphify/tasks_integration_test.go @@ -26,6 +26,7 @@ import ( "github.com/specterops/bloodhound/cmd/api/src/model" "github.com/specterops/bloodhound/cmd/api/src/services/graphify" "github.com/specterops/bloodhound/cmd/api/src/services/graphify/endpoint" + "github.com/specterops/bloodhound/cmd/api/src/services/storage" "github.com/specterops/bloodhound/packages/go/lab/generic" "github.com/stretchr/testify/require" ) @@ -39,22 +40,25 @@ func TestVersion5IngestJSON(t *testing.T) { testSuite = setupIntegrationTestSuite(t, fixturesPath) files = []string{ - path.Join(testSuite.WorkDir, "computers.json"), - path.Join(testSuite.WorkDir, "containers.json"), - path.Join(testSuite.WorkDir, "domains.json"), - path.Join(testSuite.WorkDir, "gpos.json"), - path.Join(testSuite.WorkDir, "groups.json"), - path.Join(testSuite.WorkDir, "ous.json"), - path.Join(testSuite.WorkDir, "sessions.json"), - path.Join(testSuite.WorkDir, "users.json"), + "computers.json", + "containers.json", + "domains.json", + "gpos.json", + "groups.json", + "ous.json", + "sessions.json", + "users.json", } ) + ingestLocalStore, err := storage.NewLocalStore(testSuite.WorkDir) + require.NoError(t, err, "error creating ingest local store") + ingestFileService := storage.NewFileService(ingestLocalStore) defer teardownIntegrationTestSuite(t, &testSuite) for _, file := range files { ingestCtx := graphify.NewIngestContext(ctx, graphify.WithEndpointResolver(endpoint.NewResolver(testSuite.GraphDB))) - fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestCtx, model.IngestTask{StoredFileName: file, FileType: model.FileTypeJson}) + fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestCtx, ingestFileService, model.IngestTask{StoredFileName: file, FileType: model.FileTypeJson}) require.NoError(t, err) failed := 0 @@ -82,15 +86,18 @@ func TestVersion5IngestZIP(t *testing.T) { testSuite = setupIntegrationTestSuite(t, fixturesPath) files = []string{ - path.Join(testSuite.WorkDir, "archive.zip"), + "archive.zip", } ) + ingestLocalStore, err := storage.NewLocalStore(testSuite.WorkDir) + require.NoError(t, err, "error creating ingest local store") + ingestFileService := storage.NewFileService(ingestLocalStore) defer teardownIntegrationTestSuite(t, &testSuite) for _, file := range files { ingestCtx := graphify.NewIngestContext(ctx, graphify.WithEndpointResolver(endpoint.NewResolver(testSuite.GraphDB))) - fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestCtx, model.IngestTask{StoredFileName: file, FileType: model.FileTypeZip}) + fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestCtx, ingestFileService, model.IngestTask{StoredFileName: file, FileType: model.FileTypeZip}) require.NoError(t, err) failed := 0 @@ -118,27 +125,30 @@ func TestVersion6ADCSJSON(t *testing.T) { testSuite = setupIntegrationTestSuite(t, fixturesPath) files = []string{ - path.Join(testSuite.WorkDir, "aiacas.json"), - path.Join(testSuite.WorkDir, "certtemplates.json"), - path.Join(testSuite.WorkDir, "computers.json"), - path.Join(testSuite.WorkDir, "containers.json"), - path.Join(testSuite.WorkDir, "domains.json"), - path.Join(testSuite.WorkDir, "enterprisecas.json"), - path.Join(testSuite.WorkDir, "gpos.json"), - path.Join(testSuite.WorkDir, "groups.json"), - path.Join(testSuite.WorkDir, "issuancepolicies.json"), - path.Join(testSuite.WorkDir, "ntauthstores.json"), - path.Join(testSuite.WorkDir, "ous.json"), - path.Join(testSuite.WorkDir, "rootcas.json"), - path.Join(testSuite.WorkDir, "users.json"), + "aiacas.json", + "certtemplates.json", + "computers.json", + "containers.json", + "domains.json", + "enterprisecas.json", + "gpos.json", + "groups.json", + "issuancepolicies.json", + "ntauthstores.json", + "ous.json", + "rootcas.json", + "users.json", } ) + ingestLocalStore, err := storage.NewLocalStore(testSuite.WorkDir) + require.NoError(t, err, "error creating ingest local store") + ingestFileService := storage.NewFileService(ingestLocalStore) defer teardownIntegrationTestSuite(t, &testSuite) for _, file := range files { ingestCtx := graphify.NewIngestContext(ctx, graphify.WithEndpointResolver(endpoint.NewResolver(testSuite.GraphDB))) - fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestCtx, model.IngestTask{StoredFileName: file, FileType: model.FileTypeJson}) + fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestCtx, ingestFileService, model.IngestTask{StoredFileName: file, FileType: model.FileTypeJson}) require.NoError(t, err) failed := 0 @@ -166,15 +176,18 @@ func TestVersion6ADCSZIP(t *testing.T) { testSuite = setupIntegrationTestSuite(t, fixturesPath) files = []string{ - path.Join(testSuite.WorkDir, "archive.zip"), + "archive.zip", } ) + ingestLocalStore, err := storage.NewLocalStore(testSuite.WorkDir) + require.NoError(t, err, "error creating ingest local store") + ingestFileService := storage.NewFileService(ingestLocalStore) defer teardownIntegrationTestSuite(t, &testSuite) for _, file := range files { ingestCtx := graphify.NewIngestContext(ctx, graphify.WithEndpointResolver(endpoint.NewResolver(testSuite.GraphDB))) - fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestCtx, model.IngestTask{StoredFileName: file, FileType: model.FileTypeZip}) + fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestCtx, ingestFileService, model.IngestTask{StoredFileName: file, FileType: model.FileTypeZip}) require.NoError(t, err) failed := 0 @@ -202,27 +215,30 @@ func TestVersion6AllJSON(t *testing.T) { testSuite = setupIntegrationTestSuite(t, fixturesPath) files = []string{ - path.Join(testSuite.WorkDir, "aiacas.json"), - path.Join(testSuite.WorkDir, "certtemplates.json"), - path.Join(testSuite.WorkDir, "computers.json"), - path.Join(testSuite.WorkDir, "containers.json"), - path.Join(testSuite.WorkDir, "domains.json"), - path.Join(testSuite.WorkDir, "enterprisecas.json"), - path.Join(testSuite.WorkDir, "gpos.json"), - path.Join(testSuite.WorkDir, "groups.json"), - path.Join(testSuite.WorkDir, "issuancepolicies.json"), - path.Join(testSuite.WorkDir, "ntauthstores.json"), - path.Join(testSuite.WorkDir, "ous.json"), - path.Join(testSuite.WorkDir, "rootcas.json"), - path.Join(testSuite.WorkDir, "users.json"), + "aiacas.json", + "certtemplates.json", + "computers.json", + "containers.json", + "domains.json", + "enterprisecas.json", + "gpos.json", + "groups.json", + "issuancepolicies.json", + "ntauthstores.json", + "ous.json", + "rootcas.json", + "users.json", } ) + ingestLocalStore, err := storage.NewLocalStore(testSuite.WorkDir) + require.NoError(t, err, "error creating ingest local store") + ingestFileService := storage.NewFileService(ingestLocalStore) defer teardownIntegrationTestSuite(t, &testSuite) for _, file := range files { ingestContext := graphify.NewIngestContext(ctx, graphify.WithEndpointResolver(endpoint.NewResolver(testSuite.GraphDB))) - fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestContext, model.IngestTask{StoredFileName: file, FileType: model.FileTypeJson}) + fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestContext, ingestFileService, model.IngestTask{StoredFileName: file, FileType: model.FileTypeJson}) require.NoError(t, err) failed := 0 @@ -250,15 +266,18 @@ func TestVersion6AllZIP(t *testing.T) { testSuite = setupIntegrationTestSuite(t, fixturesPath) files = []string{ - path.Join(testSuite.WorkDir, "archive.zip"), + "archive.zip", } ) + ingestLocalStore, err := storage.NewLocalStore(testSuite.WorkDir) + require.NoError(t, err, "error creating ingest local store") + ingestFileService := storage.NewFileService(ingestLocalStore) defer teardownIntegrationTestSuite(t, &testSuite) for _, file := range files { ingestContext := graphify.NewIngestContext(ctx, graphify.WithEndpointResolver(endpoint.NewResolver(testSuite.GraphDB))) - fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestContext, model.IngestTask{StoredFileName: file, FileType: model.FileTypeZip}) + fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestContext, ingestFileService, model.IngestTask{StoredFileName: file, FileType: model.FileTypeZip}) require.NoError(t, err) failed := 0 @@ -286,22 +305,25 @@ func TestVersion6IngestJSON(t *testing.T) { testSuite = setupIntegrationTestSuite(t, fixturesPath) files = []string{ - path.Join(testSuite.WorkDir, "computers.json"), - path.Join(testSuite.WorkDir, "containers.json"), - path.Join(testSuite.WorkDir, "domains.json"), - path.Join(testSuite.WorkDir, "gpos.json"), - path.Join(testSuite.WorkDir, "groups.json"), - path.Join(testSuite.WorkDir, "ous.json"), - path.Join(testSuite.WorkDir, "sessions.json"), - path.Join(testSuite.WorkDir, "users.json"), + "computers.json", + "containers.json", + "domains.json", + "gpos.json", + "groups.json", + "ous.json", + "sessions.json", + "users.json", } ) + ingestLocalStore, err := storage.NewLocalStore(testSuite.WorkDir) + require.NoError(t, err, "error creating ingest local store") + ingestFileService := storage.NewFileService(ingestLocalStore) defer teardownIntegrationTestSuite(t, &testSuite) for _, file := range files { ingestContext := graphify.NewIngestContext(ctx, graphify.WithEndpointResolver(endpoint.NewResolver(testSuite.GraphDB))) - fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestContext, model.IngestTask{StoredFileName: file, FileType: model.FileTypeJson}) + fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestContext, ingestFileService, model.IngestTask{StoredFileName: file, FileType: model.FileTypeJson}) require.NoError(t, err) failed := 0 @@ -329,15 +351,18 @@ func TestVersion6IngestZIP(t *testing.T) { testSuite = setupIntegrationTestSuite(t, fixturesPath) files = []string{ - path.Join(testSuite.WorkDir, "archive.zip"), + "archive.zip", } ) + ingestLocalStore, err := storage.NewLocalStore(testSuite.WorkDir) + require.NoError(t, err, "error creating ingest local store") + ingestFileService := storage.NewFileService(ingestLocalStore) defer teardownIntegrationTestSuite(t, &testSuite) for _, file := range files { ingestContext := graphify.NewIngestContext(ctx, graphify.WithEndpointResolver(endpoint.NewResolver(testSuite.GraphDB))) - fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestContext, model.IngestTask{StoredFileName: file, FileType: model.FileTypeZip}) + fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestContext, ingestFileService, model.IngestTask{StoredFileName: file, FileType: model.FileTypeZip}) require.NoError(t, err) failed := 0 diff --git a/cmd/api/src/services/storage/fileserviceresolver.go b/cmd/api/src/services/storage/fileserviceresolver.go new file mode 100644 index 000000000000..b993f2586673 --- /dev/null +++ b/cmd/api/src/services/storage/fileserviceresolver.go @@ -0,0 +1,137 @@ +// Copyright 2026 Specter Ops, Inc. +// +// Licensed under the Apache License, Version 2.0 +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 +package storage + +import ( + "errors" + "fmt" + + "github.com/specterops/bloodhound/cmd/api/src/config" + "github.com/specterops/bloodhound/packages/go/storage" +) + +//go:generate go run go.uber.org/mock/mockgen -copyright_file=../../../../../LICENSE.header -destination=./mocks/fileserviceresolver.go -package=mocks . FileServiceResolver + +// FileServiceResolver is an interface that is used to resolve the actual filestorage.FileService needed for +// a specific use case. This is ultimately map backed. +type FileServiceResolver interface { + // Resolve returns a filestorage.FileService interface if a filestorage.FileService is found with the given name. + // Otherwise, an error is returned. + Resolve(name storage.FileServiceName) (storage.FileService, error) +} + +type FileServiceMap map[storage.FileServiceName]storage.FileService + +type fileServiceResolver struct { + services FileServiceMap +} + +func NewFileServiceResolver(services FileServiceMap) (FileServiceResolver, error) { + var ( + copiedServices = make(FileServiceMap, len(services)) + ) + + for serviceName, fileService := range services { + if serviceName == "" { + return nil, errors.New("file service name is required") + } + if fileService == nil { + return nil, fmt.Errorf("file service %q is nil", serviceName) + } + + copiedServices[serviceName] = fileService + } + + return &fileServiceResolver{ + services: copiedServices, + }, nil +} + +func (s *fileServiceResolver) Resolve(name storage.FileServiceName) (storage.FileService, error) { + var ( + fileService storage.FileService + found bool + ) + + if name == "" { + return nil, fmt.Errorf("%w: empty name", storage.ErrFileServiceNotFound) + } + + fileService, found = s.services[name] + if !found { + return nil, fmt.Errorf("%w: %s", storage.ErrFileServiceNotFound, name) + } + + return fileService, nil +} + +// createLocalStore takes a location to create the storage.LocalStore, and wraps that in a +// storage.FileService. Both are returned. If there is an error in this process, nil is +// returned for both structs, and the error is returned. +func createLocalStore(location string) (*storage.LocalStore, storage.FileService, error) { + var ( + localStore *storage.LocalStore + err error + ) + + if localStore, err = storage.NewLocalStore(location); err != nil { + return nil, nil, err + } + + return localStore, storage.NewFileService(localStore), nil +} + +// closeLocalStores contains the functionality to close any storage.LocalStore that has been opened +// if there was an error. Errors from the close are joined together and returned as well. +func closeLocalStores(localStores []*storage.LocalStore) error { + var closeErr error + + for _, localStore := range localStores { + closeErr = errors.Join(closeErr, localStore.Close()) + } + + return closeErr +} + +// NewDefaultFileServices creates the file services that should be considered default with +// BloodHound. Additional FileServices can still be created prior to creating a Resolver. +func NewDefaultFileServices(cfg config.Configuration) (FileServiceMap, error) { + var ( + fileServices = make(FileServiceMap, 4) + openedStores []*storage.LocalStore + definitions = []struct { + name storage.FileServiceName + location string + }{ + {name: storage.FileServiceWork, location: cfg.WorkDir}, + {name: storage.FileServiceIngest, location: cfg.TempDirectory()}, + {name: storage.FileServiceRetained, location: cfg.RetainedFilesDirectory()}, + {name: storage.FileServiceCollectors, location: cfg.CollectorsBasePath}, + } + ) + + for _, definition := range definitions { + localStore, fileService, err := createLocalStore(definition.location) + if err != nil { + return nil, errors.Join(err, closeLocalStores(openedStores)) + } + + openedStores = append(openedStores, localStore) + fileServices[definition.name] = fileService + } + + return fileServices, nil +} diff --git a/cmd/api/src/services/storage/fileserviceresolver_test.go b/cmd/api/src/services/storage/fileserviceresolver_test.go new file mode 100644 index 000000000000..5bf2f16f5025 --- /dev/null +++ b/cmd/api/src/services/storage/fileserviceresolver_test.go @@ -0,0 +1,296 @@ +// Copyright 2026 Specter Ops, Inc. +// +// Licensed under the Apache License, Version 2.0 +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 +package storage_test + +import ( + "os" + "path/filepath" + "testing" + + "github.com/specterops/bloodhound/cmd/api/src/config" + api_storage "github.com/specterops/bloodhound/cmd/api/src/services/storage" + "github.com/specterops/bloodhound/packages/go/storage" + "github.com/specterops/bloodhound/packages/go/storage/mocks" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" +) + +func newTestStorageConfiguration(workDir string, collectorsBasePath string) config.Configuration { + return config.Configuration{ + WorkDir: workDir, + CollectorsBasePath: collectorsBasePath, + } +} + +func TestNewFileServiceResolver(t *testing.T) { + t.Parallel() + + type expected struct { + errContains string + } + + type testData struct { + name string + buildServices func(t *testing.T, controller *gomock.Controller) api_storage.FileServiceMap + expected expected + } + + tests := []testData{ + { + name: "creates resolver with services", + buildServices: func(t *testing.T, controller *gomock.Controller) api_storage.FileServiceMap { + return api_storage.FileServiceMap{ + storage.FileServiceWork: mocks.NewMockFileService(controller), + } + }, + }, + { + name: "empty service name returns error", + buildServices: func(t *testing.T, controller *gomock.Controller) api_storage.FileServiceMap { + return api_storage.FileServiceMap{ + "": mocks.NewMockFileService(controller), + } + }, + expected: expected{ + errContains: "file service name is required", + }, + }, + { + name: "nil service returns error", + buildServices: func(t *testing.T, controller *gomock.Controller) api_storage.FileServiceMap { + return api_storage.FileServiceMap{ + storage.FileServiceWork: nil, + } + }, + expected: expected{ + errContains: `file service "work" is nil`, + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + services := testCase.buildServices(t, gomock.NewController(t)) + + // Act + resolver, err := api_storage.NewFileServiceResolver(services) + + // Assert + if testCase.expected.errContains != "" { + require.Error(t, err) + require.Contains(t, err.Error(), testCase.expected.errContains) + require.Nil(t, resolver) + return + } + + require.NoError(t, err) + require.NotNil(t, resolver) + }) + } +} + +func TestFileServiceResolver_Resolve(t *testing.T) { + t.Parallel() + + type expected struct { + errIs error + service storage.FileService + } + + type testData struct { + name string + resolveName storage.FileServiceName + expected expected + } + + controller := gomock.NewController(t) + workService := mocks.NewMockFileService(controller) + services := api_storage.FileServiceMap{ + storage.FileServiceWork: workService, + } + + tests := []testData{ + { + name: "resolves service", + resolveName: storage.FileServiceWork, + expected: expected{ + service: workService, + }, + }, + { + name: "missing service returns not found", + resolveName: storage.FileServiceIngest, + expected: expected{ + errIs: storage.ErrFileServiceNotFound, + }, + }, + { + name: "empty name returns not found", + resolveName: "", + expected: expected{ + errIs: storage.ErrFileServiceNotFound, + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + resolver, err := api_storage.NewFileServiceResolver(services) + require.NoError(t, err) + + // Act + fileService, err := resolver.Resolve(testCase.resolveName) + + // Assert + if testCase.expected.errIs != nil { + require.ErrorIs(t, err, testCase.expected.errIs) + require.Nil(t, fileService) + return + } + + require.NoError(t, err) + require.Same(t, testCase.expected.service, fileService) + }) + } +} + +func TestFileServiceResolver_CopiesServices(t *testing.T) { + t.Parallel() + + // Arrange + var ( + controller = gomock.NewController(t) + workService = mocks.NewMockFileService(controller) + services = api_storage.FileServiceMap{ + storage.FileServiceWork: workService, + } + ) + + resolver, err := api_storage.NewFileServiceResolver(services) + require.NoError(t, err) + + delete(services, storage.FileServiceWork) + + // Act + fileService, err := resolver.Resolve(storage.FileServiceWork) + + // Assert + require.NoError(t, err) + require.Same(t, workService, fileService) +} + +func TestNewDefaultFileServices(t *testing.T) { + t.Parallel() + + // Arrange + var ( + workDir = t.TempDir() + collectorsBasePath = t.TempDir() + configuration = newTestStorageConfiguration(workDir, collectorsBasePath) + ) + + require.NoError(t, os.MkdirAll(configuration.TempDirectory(), 0o750)) + require.NoError(t, os.MkdirAll(configuration.RetainedFilesDirectory(), 0o750)) + + // Act + fileServices, err := api_storage.NewDefaultFileServices(configuration) + + // Assert + require.NoError(t, err) + require.Contains(t, fileServices, storage.FileServiceWork) + require.Contains(t, fileServices, storage.FileServiceIngest) + require.Contains(t, fileServices, storage.FileServiceRetained) + require.Contains(t, fileServices, storage.FileServiceCollectors) + require.Len(t, fileServices, 4) + + for _, fileService := range fileServices { + storageFileService, ok := fileService.(*storage.StorageFileService) + require.True(t, ok) + + localStore, ok := storageFileService.Storage.(*storage.LocalStore) + require.True(t, ok) + require.NoError(t, localStore.Close()) + } +} + +func TestNewDefaultFileServices_ReturnsError(t *testing.T) { + t.Parallel() + + type testData struct { + name string + setup func(t *testing.T) config.Configuration + } + + tests := []testData{ + { + name: "missing work directory", + setup: func(t *testing.T) config.Configuration { + return newTestStorageConfiguration(filepath.Join(t.TempDir(), "missing"), t.TempDir()) + }, + }, + { + name: "missing temp directory", + setup: func(t *testing.T) config.Configuration { + workDir := t.TempDir() + + return newTestStorageConfiguration(workDir, t.TempDir()) + }, + }, + { + name: "missing retained directory", + setup: func(t *testing.T) config.Configuration { + workDir := t.TempDir() + configuration := newTestStorageConfiguration(workDir, t.TempDir()) + require.NoError(t, os.MkdirAll(configuration.TempDirectory(), 0o750)) + + return configuration + }, + }, + { + name: "missing collectors directory", + setup: func(t *testing.T) config.Configuration { + workDir := t.TempDir() + configuration := newTestStorageConfiguration(workDir, filepath.Join(t.TempDir(), "missing")) + require.NoError(t, os.MkdirAll(configuration.TempDirectory(), 0o750)) + require.NoError(t, os.MkdirAll(configuration.RetainedFilesDirectory(), 0o750)) + + return configuration + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + configuration := testCase.setup(t) + + // Act + fileServices, err := api_storage.NewDefaultFileServices(configuration) + + // Assert + require.Error(t, err) + require.Nil(t, fileServices) + }) + } +} diff --git a/cmd/api/src/services/storage/mocks/fileserviceresolver.go b/cmd/api/src/services/storage/mocks/fileserviceresolver.go new file mode 100644 index 000000000000..626dafa98deb --- /dev/null +++ b/cmd/api/src/services/storage/mocks/fileserviceresolver.go @@ -0,0 +1,72 @@ +// Copyright 2026 Specter Ops, Inc. +// +// Licensed under the Apache License, Version 2.0 +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 + +// Code generated by MockGen. DO NOT EDIT. +// Source: github.com/specterops/bloodhound/cmd/api/src/services/storage (interfaces: FileServiceResolver) +// +// Generated by this command: +// +// mockgen -copyright_file=../../../../../LICENSE.header -destination=./mocks/fileserviceresolver.go -package=mocks . FileServiceResolver +// + +// Package mocks is a generated GoMock package. +package mocks + +import ( + reflect "reflect" + + storage "github.com/specterops/bloodhound/packages/go/storage" + gomock "go.uber.org/mock/gomock" +) + +// MockFileServiceResolver is a mock of FileServiceResolver interface. +type MockFileServiceResolver struct { + ctrl *gomock.Controller + recorder *MockFileServiceResolverMockRecorder + isgomock struct{} +} + +// MockFileServiceResolverMockRecorder is the mock recorder for MockFileServiceResolver. +type MockFileServiceResolverMockRecorder struct { + mock *MockFileServiceResolver +} + +// NewMockFileServiceResolver creates a new mock instance. +func NewMockFileServiceResolver(ctrl *gomock.Controller) *MockFileServiceResolver { + mock := &MockFileServiceResolver{ctrl: ctrl} + mock.recorder = &MockFileServiceResolverMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockFileServiceResolver) EXPECT() *MockFileServiceResolverMockRecorder { + return m.recorder +} + +// Resolve mocks base method. +func (m *MockFileServiceResolver) Resolve(name storage.FileServiceName) (storage.FileService, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Resolve", name) + ret0, _ := ret[0].(storage.FileService) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Resolve indicates an expected call of Resolve. +func (mr *MockFileServiceResolverMockRecorder) Resolve(name any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Resolve", reflect.TypeOf((*MockFileServiceResolver)(nil).Resolve), name) +} diff --git a/cmd/api/src/services/upload/upload.go b/cmd/api/src/services/upload/upload.go index c6efd57a1119..45e9aa606e92 100644 --- a/cmd/api/src/services/upload/upload.go +++ b/cmd/api/src/services/upload/upload.go @@ -17,12 +17,13 @@ package upload import ( + "context" "errors" "fmt" "io" "log/slog" "net/http" - "os" + "time" "github.com/specterops/bloodhound/cmd/api/src/model" "github.com/specterops/bloodhound/cmd/api/src/model/ingest" @@ -30,11 +31,12 @@ import ( "github.com/specterops/bloodhound/packages/go/bhlog/attr" "github.com/specterops/bloodhound/packages/go/headers" "github.com/specterops/bloodhound/packages/go/mediatypes" + "github.com/specterops/bloodhound/packages/go/storage" ) var ErrInvalidJSON = errors.New("file is not valid json") -func SaveIngestFile(location string, request *http.Request, validator IngestValidator, jobID int64) (IngestTaskParams, error) { +func SaveIngestFile(ctx context.Context, fileService storage.FileService, request *http.Request, validator IngestValidator, jobID int64) (IngestTaskParams, error) { fileData := request.Body var ( @@ -53,7 +55,7 @@ func SaveIngestFile(location string, request *http.Request, validator IngestVali return IngestTaskParams{}, fmt.Errorf("invalid content type for ingest file") } - if tempFileName, err := WriteAndValidateFileWithPrefix(fileData, location, ingestFileTempPrefix(jobID), validationFn); err != nil { + if tempFileName, err := WriteAndValidateFile(ctx, fileService, fileData, fmt.Sprintf("file_upload_job%d_", jobID), validationFn); err != nil { return IngestTaskParams{}, err } else { return IngestTaskParams{ @@ -61,58 +63,94 @@ func SaveIngestFile(location string, request *http.Request, validator IngestVali FileType: fileType, }, nil } - -} - -func ingestFileTempPrefix(jobID int64) string { - return fmt.Sprintf("file_upload_job%d_", jobID) -} - -func WriteAndValidateFile(fileData io.Reader, location string, jobID int64, validationFunc FileValidator) (string, error) { - return WriteAndValidateFileWithPrefix(fileData, location, ingestFileTempPrefix(jobID), validationFunc) } -func WriteAndValidateFileWithPrefix(fileData io.Reader, location string, tempFilePrefix string, validationFunc FileValidator) (string, error) { - var ( - tempFile *os.File - tempFileName string - err error - validationErr error - ) - - // Write a temp file. If it passes validation, keep it around and return the filename. Otherwise destroy it. - if tempFile, err = os.CreateTemp(location, tempFilePrefix); err != nil { - return "", fmt.Errorf("error creating ingest file: %w", err) +func CleanupTempFile(ctx context.Context, fileService storage.FileService, tempFileName string) { + if tempFileName == "" { + return } - // Save this for later - tempFileName = tempFile.Name() - - // Run validation on the file to see if we even want to keep it - _, validationErr = validationFunc(fileData, tempFile) + cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second) + defer cancel() - // We close the file next, not last. We can't defer this if we might want to delete it. - // Note: fileData does not need to be closed because the HTTP server manages it's lifecyle - if closeErr := tempFile.Close(); closeErr != nil { - slog.Error( - "Error closing temp file", - slog.String("file", tempFileName), - attr.Error(closeErr), + if err := fileService.DeleteFile(cleanupCtx, tempFileName); err != nil { + slog.ErrorContext( + cleanupCtx, + "Failed to delete temp file", + slog.String("temp_file_name", tempFileName), + attr.Error(err), ) } +} + +func WriteAndValidateFile(ctx context.Context, fileService storage.FileService, fileData io.Reader, prefix string, validationFunc FileValidator) (string, error) { + if validationFunc == nil { + return "", fmt.Errorf("validation function is required") + } + + // Create a pipe: pr (read end) and pw (write end). + // Data written to pw can be read from pr. + pr, pw := io.Pipe() + + // validationErrCh carries the result of the validation goroutine. + // Using a buffered channel (size 1) ensures the goroutine never blocks on send, + // and gives the main goroutine a synchronization point to wait for the result. + validationErrCh := make(chan error, 1) + + // Start validation in a separate goroutine. + // validationFunc reads from pr, writes to io.Discard. + // This validates the stream without storing the data twice. + go func() { + _, err := validationFunc(pr, io.Discard) + pr.Close() + validationErrCh <- err + }() + + // TeeReader: as we read fileData, a copy is written to pw and flows to the goroutine via pr. + teeReader := io.TeeReader(fileData, pw) + + // Write to storage while validation happens concurrently. + tempFileName, writeErr := fileService.WriteTempFile(ctx, prefix, teeReader, storage.WriteOptions{}) + + // Closing pw signals EOF to the validation goroutine so it can finish. + pw.Close() + + select { + case validationErr := <-validationErrCh: + // Context cancelation wins over validation errors when both are ready + if err := ctx.Err(); err != nil { + CleanupTempFile(ctx, fileService, tempFileName) + return "", err + } - // If the validation was not successful, after we close the file, we remove it and return the error - if validationErr != nil { - if removeErr := os.Remove(tempFileName); removeErr != nil { - slog.Error( - "Error deleting temp file", - slog.String("file", tempFileName), - attr.Error(removeErr), + if validationErr != nil { + slog.ErrorContext( + ctx, + "Validation failed", + slog.String("temp_file_name", + tempFileName), + attr.Error(validationErr), ) + CleanupTempFile(ctx, fileService, tempFileName) + return "", validationErr } - return "", validationErr + case <-ctx.Done(): + CleanupTempFile(ctx, fileService, tempFileName) + return "", ctx.Err() + } + + // Check if writing failed; the temp file should be cleaned up. + if writeErr != nil { + slog.ErrorContext( + ctx, + "Write failed", + slog.String("temp_file_name", tempFileName), + attr.Error(writeErr), + ) + CleanupTempFile(ctx, fileService, tempFileName) + return "", writeErr } - // Assuming no other errors, return the name of the closed temp file + slog.InfoContext(ctx, "File written and validated", slog.String("temp_file_name", tempFileName)) return tempFileName, nil } diff --git a/cmd/api/src/services/upload/upload_test.go b/cmd/api/src/services/upload/upload_test.go index 6393eeead557..8581cd411eab 100644 --- a/cmd/api/src/services/upload/upload_test.go +++ b/cmd/api/src/services/upload/upload_test.go @@ -18,17 +18,20 @@ package upload import ( "bytes" + "context" "errors" "fmt" "io" "os" - "path/filepath" "strings" "testing" "github.com/specterops/bloodhound/cmd/api/src/model/ingest" + "github.com/specterops/bloodhound/packages/go/storage" + storagemocks "github.com/specterops/bloodhound/packages/go/storage/mocks" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" ) func buildValidator(t *testing.T, expectedContent string, validationErr error) FileValidator { @@ -39,9 +42,6 @@ func buildValidator(t *testing.T, expectedContent string, validationErr error) F require.NoError(t, err) require.Equal(t, expectedContent, string(content)) - _, err = dst.Write(content) - require.NoError(t, err) - return ingest.OriginalMetadata{}, validationErr } } @@ -156,39 +156,110 @@ func TestWriteAndValidateJSON_NormalizationError(t *testing.T) { assert.ErrorIs(t, err, ErrInvalidJSON) } -func TestWriteAndValidateFileWithPrefix(t *testing.T) { +func TestUpload_WriteAndValidateFile(t *testing.T) { t.Parallel() var ( - errValidation = errors.New("validation failed") - tempFilePrefix = "custom_prefix_" + errValidation = errors.New("validation failed") + errWrite = errors.New("write failed") ) - t.Run("writes validated file using caller supplied prefix", func(t *testing.T) { - t.Parallel() + type expected struct { + errIs error + fileName string + } - tempDirectory := t.TempDir() - fileName, err := WriteAndValidateFileWithPrefix(strings.NewReader("content"), tempDirectory, tempFilePrefix, buildValidator(t, "content", nil)) - require.NoError(t, err) - require.Contains(t, filepath.Base(fileName), tempFilePrefix) + type testData struct { + name string + tempFileName string + writeErr error + validationErr error + expected expected + expectDelete bool + } - fileContent, err := os.ReadFile(fileName) - require.NoError(t, err) - require.Equal(t, "content", string(fileContent)) - }) + tests := []testData{ + { + name: "writes and validates file", + tempFileName: "prefix/tmp-file", + expected: expected{ + fileName: "prefix/tmp-file", + }, + }, + { + name: "validation error deletes temp file", + tempFileName: "prefix/tmp-file", + validationErr: errValidation, + expected: expected{ + errIs: errValidation, + }, + expectDelete: true, + }, + { + name: "write error deletes temp file", + tempFileName: "prefix/tmp-file", + writeErr: errWrite, + expected: expected{ + errIs: errWrite, + }, + expectDelete: true, + }, + { + name: "validation error takes precedence over write error", + tempFileName: "prefix/tmp-file", + writeErr: errWrite, + validationErr: errValidation, + expected: expected{ + errIs: errValidation, + }, + expectDelete: true, + }, + { + name: "write error without temp path does not delete empty path", + tempFileName: "", + writeErr: errWrite, + expected: expected{ + errIs: errWrite, + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + ctx := context.Background() + mockFileService := storagemocks.NewMockFileService(gomock.NewController(t)) + validator := buildValidator(t, "content", testCase.validationErr) + + mockFileService.EXPECT(). + WriteTempFile(ctx, "prefix", gomock.Any(), storage.WriteOptions{}). + DoAndReturn(func(_ context.Context, _ string, reader io.Reader, _ storage.WriteOptions) (string, error) { + content, err := io.ReadAll(reader) + require.NoError(t, err) + require.Equal(t, "content", string(content)) + + return testCase.tempFileName, testCase.writeErr + }) + if testCase.expectDelete { + mockFileService.EXPECT().DeleteFile(gomock.Any(), testCase.tempFileName).Return(nil) + } - t.Run("validation error deletes temp file", func(t *testing.T) { - t.Parallel() + // Act + actualFileName, err := WriteAndValidateFile(ctx, mockFileService, strings.NewReader("content"), "prefix", validator) - tempDirectory := t.TempDir() - fileName, err := WriteAndValidateFileWithPrefix(strings.NewReader("content"), tempDirectory, tempFilePrefix, buildValidator(t, "content", errValidation)) - require.ErrorIs(t, err, errValidation) - require.Empty(t, fileName) + // Assert + if testCase.expected.errIs != nil { + require.ErrorIs(t, err, testCase.expected.errIs) + require.Empty(t, actualFileName) + return + } - files, err := filepath.Glob(filepath.Join(tempDirectory, tempFilePrefix+"*")) - require.NoError(t, err) - require.Empty(t, files) - }) + require.NoError(t, err) + require.Equal(t, testCase.expected.fileName, actualFileName) + }) + } } // ErrorReader is a mock reader that always returns an error diff --git a/cmd/api/src/test/lab/fixtures/api.go b/cmd/api/src/test/lab/fixtures/api.go index 4348e097e10d..95145c8d154e 100644 --- a/cmd/api/src/test/lab/fixtures/api.go +++ b/cmd/api/src/test/lab/fixtures/api.go @@ -60,14 +60,15 @@ func NewCustomApiFixture(cfgFixture *lab.Fixture[config.Configuration]) *lab.Fix defer wg.Done() initializer := bootstrap.Initializer[*database.BloodhoundDB, *graph.DatabaseSwitch]{ - Configuration: cfg, - DBConnector: services.ConnectDatabases, - Entrypoint: func(ctx context.Context, cfg config.Configuration, databaseConnections bootstrap.DatabaseConnections[*database.BloodhoundDB, *graph.DatabaseSwitch]) ([]daemons.Daemon, error) { + Configuration: cfg, + DBConnector: services.ConnectDatabases, + RuntimeDependencies: services.CreateRuntimeDependencies, + Entrypoint: func(ctx context.Context, cfg config.Configuration, databaseConnections bootstrap.DatabaseConnections[*database.BloodhoundDB, *graph.DatabaseSwitch], dependencies bootstrap.RuntimeDependencies) ([]daemons.Daemon, error) { if err := databaseConnections.RDMS.Wipe(ctx); err != nil { return nil, err } - return services.Entrypoint(ctx, cfg, databaseConnections) + return services.Entrypoint(ctx, cfg, databaseConnections, dependencies) }, } diff --git a/cmd/api/src/utils/util.go b/cmd/api/src/utils/util.go index 47f15e620259..333c1c1ad604 100644 --- a/cmd/api/src/utils/util.go +++ b/cmd/api/src/utils/util.go @@ -49,6 +49,19 @@ const ( ClientTypeOpenHound ClientType = 3 ) +func (c ClientType) String() string { + switch c { + case ClientTypeSharpHound: + return "SharpHound" + case ClientTypeAzureHound: + return "AzureHound" + case ClientTypeOpenHound: + return "OpenHound" + default: + return fmt.Sprintf("UnknownClientType(%d)", int(c)) + } +} + type ClientVersion struct { ClientType ClientType Major int @@ -57,6 +70,7 @@ type ClientVersion struct { Extra int Prerelease string BuildMetadata string + Original string } // IsValidClientVersion checks the version from a user agent to ensure it's a valid UserAgent and that @@ -104,7 +118,7 @@ func ParseClientVersion(userAgent string) (ClientVersion, error) { } // clientVersionFromSemver maps a parsed semver.Version into a ClientVersion struct. -func clientVersionFromSemver(clientType ClientType, parsedVersion *semver.Version, extra int) ClientVersion { +func clientVersionFromSemver(clientType ClientType, parsedVersion *semver.Version, extra int, original string) ClientVersion { return ClientVersion{ ClientType: clientType, Major: int(parsedVersion.Major()), @@ -113,6 +127,7 @@ func clientVersionFromSemver(clientType ClientType, parsedVersion *semver.Versio Extra: extra, Prerelease: parsedVersion.Prerelease(), BuildMetadata: parsedVersion.Metadata(), + Original: original, } } @@ -152,7 +167,7 @@ func parseCollectorVersion(clientType ClientType, rawVersion string) (ClientVers return version, ErrInvalidCollectorVersion } - return clientVersionFromSemver(clientType, parsedVersion, 0), nil + return clientVersionFromSemver(clientType, parsedVersion, 0, rawVersion), nil } // parseSharpHoundVersion parses a raw version string in SharpHound's X.Y.Z.W[-rcN] format. @@ -178,7 +193,7 @@ func parseSharpHoundVersion(rawVersion string) (ClientVersion, error) { return version, ErrInvalidSharpHoundVersion } - return clientVersionFromSemver(ClientTypeSharpHound, parsedVersion, extra), nil + return clientVersionFromSemver(ClientTypeSharpHound, parsedVersion, extra, rawVersion), nil } // splitSharpHoundVersion splits a SharpHound version string into a semver-compatible version diff --git a/go.mod b/go.mod index f3ec35911c07..6ce544feee3c 100644 --- a/go.mod +++ b/go.mod @@ -21,8 +21,9 @@ require ( cuelang.org/go v0.16.0 github.com/Masterminds/semver/v3 v3.4.0 github.com/RoaringBitmap/roaring/v2 v2.16.0 - github.com/aws/aws-sdk-go-v2/config v1.32.13 + github.com/aws/aws-sdk-go-v2/config v1.32.20 github.com/aws/aws-sdk-go-v2/feature/rds/auth v1.6.21 + github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.22 github.com/bloodhoundad/azurehound/v2 v2.12.1 github.com/cespare/xxhash/v2 v2.3.0 github.com/channelmeter/iso8601duration v0.0.0-20150204201828-8da3af7a2a61 @@ -65,6 +66,9 @@ require ( ) require ( + github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.11 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.18 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.25 // indirect github.com/mfridman/interpolate v0.0.2 // indirect github.com/sethvargo/go-retry v0.3.0 // indirect ) @@ -102,19 +106,20 @@ require ( github.com/antlr4-go/antlr/v4 v4.13.1 // indirect github.com/ashanbrown/forbidigo/v2 v2.3.0 // indirect github.com/ashanbrown/makezero/v2 v2.1.0 // indirect - github.com/aws/aws-sdk-go-v2 v1.41.5 // indirect - github.com/aws/aws-sdk-go-v2/credentials v1.19.13 // indirect - github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.21 // indirect - github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.21 // indirect - github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.21 // indirect - github.com/aws/aws-sdk-go-v2/internal/ini v1.8.6 // indirect - github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.7 // indirect - github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.21 // indirect - github.com/aws/aws-sdk-go-v2/service/signin v1.0.9 // indirect - github.com/aws/aws-sdk-go-v2/service/sso v1.30.14 // indirect - github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.18 // indirect - github.com/aws/aws-sdk-go-v2/service/sts v1.41.10 // indirect - github.com/aws/smithy-go v1.24.2 // indirect + github.com/aws/aws-sdk-go-v2 v1.41.9 + github.com/aws/aws-sdk-go-v2/credentials v1.19.19 + github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.25 // indirect + github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.25 // indirect + github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.25 // indirect + github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.26 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.10 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.25 // indirect + github.com/aws/aws-sdk-go-v2/service/s3 v1.102.2 + github.com/aws/aws-sdk-go-v2/service/signin v1.1.1 // indirect + github.com/aws/aws-sdk-go-v2/service/sso v1.30.19 // indirect + github.com/aws/aws-sdk-go-v2/service/ssooidc v1.36.2 // indirect + github.com/aws/aws-sdk-go-v2/service/sts v1.42.3 // indirect + github.com/aws/smithy-go v1.26.0 github.com/axiomhq/hyperloglog v0.2.6 // indirect github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect github.com/beevik/etree v1.6.0 // indirect diff --git a/go.sum b/go.sum index 89574c64a7ee..87c9a5a6e31d 100644 --- a/go.sum +++ b/go.sum @@ -558,72 +558,82 @@ github.com/aws/aws-sdk-go-v2 v0.18.0/go.mod h1:JWVYvqSMppoMJC0x5wdwiImzgXTI9FuZw github.com/aws/aws-sdk-go-v2 v1.17.5/go.mod h1:uzbQtefpm44goOPmdKyAlXSNcwlRgF3ePWVW6EtJvvw= github.com/aws/aws-sdk-go-v2 v1.25.2/go.mod h1:Evoc5AsmtveRt1komDwIsjHFyrP5tDuF1D1U+6z6pNo= github.com/aws/aws-sdk-go-v2 v1.36.3/go.mod h1:LLXuLpgzEbD766Z5ECcRmi8AzSwfZItDtmABVkRLGzg= -github.com/aws/aws-sdk-go-v2 v1.41.5 h1:dj5kopbwUsVUVFgO4Fi5BIT3t4WyqIDjGKCangnV/yY= -github.com/aws/aws-sdk-go-v2 v1.41.5/go.mod h1:mwsPRE8ceUUpiTgF7QmQIJ7lgsKUPQOUl3o72QBrE1o= +github.com/aws/aws-sdk-go-v2 v1.41.9 h1:/rYeyO2+HrMztAmxAq9++XJtFMqSIpSsNA0yDGALYq4= +github.com/aws/aws-sdk-go-v2 v1.41.9/go.mod h1:+HsoOEX80qAVUitj1A2DhCNTjmb3edVyuDypb6LNEeo= +github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.11 h1:h5+3VT69KUBK24grGuuA5saDJTj2IIjLb9au668Fo5I= +github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.11/go.mod h1:dnakxebH6UwFvcvujL0LVggYQ8nEvBGjU4G/V79Nv94= github.com/aws/aws-sdk-go-v2/config v1.18.14/go.mod h1:0pI6JQBHKwd0JnwAZS3VCapLKMO++UL2BOkWwyyzTnA= github.com/aws/aws-sdk-go-v2/config v1.27.4/go.mod h1:zq2FFXK3A416kiukwpsd+rD4ny6JC7QSkp4QdN1Mp2g= github.com/aws/aws-sdk-go-v2/config v1.29.12/go.mod h1:xse1YTjmORlb/6fhkWi8qJh3cvZi4JoVNhc+NbJt4kI= -github.com/aws/aws-sdk-go-v2/config v1.32.13 h1:5KgbxMaS2coSWRrx9TX/QtWbqzgQkOdEa3sZPhBhCSg= -github.com/aws/aws-sdk-go-v2/config v1.32.13/go.mod h1:8zz7wedqtCbw5e9Mi2doEwDyEgHcEE9YOJp6a8jdSMY= +github.com/aws/aws-sdk-go-v2/config v1.32.20 h1:8VMDnWc/kEzxsI/1ngGM9mG81a8IGmIHD8KLcYGwagc= +github.com/aws/aws-sdk-go-v2/config v1.32.20/go.mod h1:PuwEpciweIXGULWeOeSTXtSbH4CW9mWdWrhdCKQI1sM= github.com/aws/aws-sdk-go-v2/credentials v1.13.14/go.mod h1:85ckagDuzdIOnZRwws1eLKnymJs3ZM1QwVC1XcuNGOY= github.com/aws/aws-sdk-go-v2/credentials v1.17.4/go.mod h1:+30tpwrkOgvkJL1rUZuRLoxcJwtI/OkeBLYnHxJtVe0= github.com/aws/aws-sdk-go-v2/credentials v1.17.65/go.mod h1:4zyjAuGOdikpNYiSGpsGz8hLGmUzlY8pc8r9QQ/RXYQ= -github.com/aws/aws-sdk-go-v2/credentials v1.19.13 h1:mA59E3fokBvyEGHKFdnpNNrvaR351cqiHgRg+JzOSRI= -github.com/aws/aws-sdk-go-v2/credentials v1.19.13/go.mod h1:yoTXOQKea18nrM69wGF9jBdG4WocSZA1h38A+t/MAsk= +github.com/aws/aws-sdk-go-v2/credentials v1.19.19 h1:yuFzSV1U0aRNYCQGVaTY2zW2M/L93pYHnXnrJUphYhU= +github.com/aws/aws-sdk-go-v2/credentials v1.19.19/go.mod h1:7y63L1kGzeoDlJaQ3Z578KrnmfBut96JjvJUzGwR+YE= github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.12.23/go.mod h1:mOtmAg65GT1HIL/HT/PynwPbS+UG0BgCZ6vhkPqnxWo= github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.15.2/go.mod h1:iRlGzMix0SExQEviAyptRWRGdYNo3+ufW/lCzvKVTUc= github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.30/go.mod h1:Jpne2tDnYiFascUEs2AWHJL9Yp7A5ZVy3TNyxaAjD6M= -github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.21 h1:NUS3K4BTDArQqNu2ih7yeDLaS3bmHD0YndtA6UP884g= -github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.21/go.mod h1:YWNWJQNjKigKY1RHVJCuupeWDrrHjRqHm0N9rdrWzYI= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.25 h1:0w6dCiO8iez+YKwRhRBlL1CH/E3GTfdkuzrwj1by8vo= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.25/go.mod h1:9FDWUothyr5RCRAHc45XOiVCzUR8n/IhCYX+uVqw6vk= github.com/aws/aws-sdk-go-v2/feature/rds/auth v1.6.21 h1:HFn8sVT87KWnGs2Q2gO/brPZc2bR0RXD++cYKRmABzk= github.com/aws/aws-sdk-go-v2/feature/rds/auth v1.6.21/go.mod h1:BGZ/K6gLGJt8K36j6gcsD7WVxmWt0MGBYtr57iLweio= +github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.22 h1:eBnBL5sAT9aIg014yU6DG4YT/ov8r9jnOPW/t7rr8Ig= +github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.22/go.mod h1:PjXTcjNr3cv5pfRG+PakZ9Yypx2vLRjCAcm1kEfHSjk= github.com/aws/aws-sdk-go-v2/internal/configsources v1.1.29/go.mod h1:Dip3sIGv485+xerzVv24emnjX5Sg88utCL8fwGmCeWg= github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.2/go.mod h1:wRQv0nN6v9wDXuWThpovGQjqF1HFdcgWjporw14lS8k= github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.34/go.mod h1:p4VfIceZokChbA9FzMbRGz5OV+lekcVtHlPKEO0gSZY= -github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.21 h1:Rgg6wvjjtX8bNHcvi9OnXWwcE0a2vGpbwmtICOsvcf4= -github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.21/go.mod h1:A/kJFst/nm//cyqonihbdpQZwiUhhzpqTsdbhDdRF9c= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.25 h1:Uii3frf9ztec/ABM2/FSH9/z7PLzxfpG8h4RpkUFflQ= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.25/go.mod h1:G6kntsA2GorAxDPbap6xgB2F+amSLUF8GJTi7PUoX44= github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.4.23/go.mod h1:mr6c4cHC+S/MMkrjtSlG4QA36kOznDep+0fga5L/fGQ= github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.2/go.mod h1:tyF5sKccmDz0Bv4NrstEr+/9YkSPJHrcO7UsUKf7pWM= github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.34/go.mod h1:dFZsC0BLo346mvKQLWmoJxT+Sjp+qcVR1tRVHQGOH9Q= -github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.21 h1:PEgGVtPoB6NTpPrBgqSE5hE/o47Ij9qk/SEZFbUOe9A= -github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.21/go.mod h1:p+hz+PRAYlY3zcpJhPwXlLC4C+kqn70WIHwnzAfs6ps= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.25 h1:r1+/l6m+WaUJF9HISEsNOLHSNj5EXYQxK8VX6Cz9NlA= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.25/go.mod h1:cKf+D+NMDK1LndD7BowHbBZPgR9V0/5HubH0PFWvA+c= github.com/aws/aws-sdk-go-v2/internal/ini v1.3.30/go.mod h1:vsbq62AOBwQ1LJ/GWKFxX8beUEYeRp/Agitrxee2/qM= github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0/go.mod h1:8tu/lYfQfFe6IGnaOdrpVgEL2IrrDOf6/m9RQum4NkY= github.com/aws/aws-sdk-go-v2/internal/ini v1.8.3/go.mod h1:H5O/EsxDWyU+LP/V8i5sm8cxoZgc2fdNR9bxlOFrQTo= -github.com/aws/aws-sdk-go-v2/internal/ini v1.8.6 h1:qYQ4pzQ2Oz6WpQ8T3HvGHnZydA72MnLuFK9tJwmrbHw= -github.com/aws/aws-sdk-go-v2/internal/ini v1.8.6/go.mod h1:O3h0IK87yXci+kg6flUKzJnWeziQUKciKrLjcatSNcY= +github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.26 h1:A1PmWU2zfkIm9EyFlJncFXL4W4phML+h8KjltUsCvNQ= +github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.26/go.mod h1:dY4MRzXEizrD4hqtpKvWVGPX7QleSGGVY+EBolo1RmM= github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.1/go.mod h1:JKpmtYhhPs7D97NL/ltqz7yCkERFW5dOlHyVl66ZYF8= github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.3/go.mod h1:0yKJC/kb8sAnmlYa6Zs3QVYqaC8ug2AbnNChv5Ox3uA= -github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.7 h1:5EniKhLZe4xzL7a+fU3C2tfUN4nWIqlLesfrjkuPFTY= -github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.7/go.mod h1:x0nZssQ3qZSnIcePWLvcoFisRXJzcTVvYpAAdYX8+GI= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.10 h1:d5/908OJ4bXg8lyjeMPvXetEKqoDoLi5Owy1zNue3yg= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.10/go.mod h1:a57l7Hwh+FWI+we50g5NPJHYUKeJKfXbc4w8SyXu8Ig= +github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.18 h1:W/EyPFl9A5rXrtoilfwHYEvzHER+K4SpBPtMXi24Mos= +github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.18/go.mod h1:UG50K+pvd/uy6xExbobg0rjqFBFZe6I3l75EPDZw4tg= github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.9.23/go.mod h1:9uPh+Hrz2Vn6oMnQYiUi/zbh3ovbnQk19YKINkQny44= github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.2/go.mod h1:Ru7vg1iQ7cR4i7SZ/JTLYN9kaXtbL69UdgG0OQWQxW0= github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.15/go.mod h1:SwFBy2vjtA0vZbjjaFtfN045boopadnoVPhu4Fv66vY= -github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.21 h1:c31//R3xgIJMSC8S6hEVq+38DcvUlgFY0FM6mSI5oto= -github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.21/go.mod h1:r6+pf23ouCB718FUxaqzZdbpYFyDtehyZcmP5KL9FkA= -github.com/aws/aws-sdk-go-v2/service/signin v1.0.9 h1:QKZH0S178gCmFEgst8hN0mCX1KxLgHBKKY/CLqwP8lg= -github.com/aws/aws-sdk-go-v2/service/signin v1.0.9/go.mod h1:7yuQJoT+OoH8aqIxw9vwF+8KpvLZ8AWmvmUWHsGQZvI= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.25 h1:dD3dhHNglpd98gs72my22Ndqi1hqQGllFFg1F+twfxg= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.25/go.mod h1:0yAbjPfd64gG7mj85RW+fMEYdfBgCRZw8g/oWcL1pjc= +github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.25 h1:2pQEbwf+/6EDbiit/GcBE2K4IUpMZymaA0kOz3xK978= +github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.25/go.mod h1:KvT6NCcQ0EZ+ZkVRrlBMt04Po3ok23YELEp7WimhLhM= +github.com/aws/aws-sdk-go-v2/service/s3 v1.102.2 h1:ie4ElCmUKS26pzrZcIk/lmt4yWjAqLLcawstyQCh298= +github.com/aws/aws-sdk-go-v2/service/s3 v1.102.2/go.mod h1:zjsomFeX5duj+4PlMB+o4JoWTIx+G0XMyzjYrUbQkN0= +github.com/aws/aws-sdk-go-v2/service/signin v1.1.1 h1:1VwbP3qMNfxUDEXWki4rCE5iA+44VA1lokTz9HasGzw= +github.com/aws/aws-sdk-go-v2/service/signin v1.1.1/go.mod h1:vUtyoSj0OPji3kjIVSc/GlKuWEiL33f/WFxl6dmpy/A= github.com/aws/aws-sdk-go-v2/service/sso v1.12.3/go.mod h1:jtLIhd+V+lft6ktxpItycqHqiVXrPIRjWIsFIlzMriw= github.com/aws/aws-sdk-go-v2/service/sso v1.20.1/go.mod h1:RsYqzYr2F2oPDdpy+PdhephuZxTfjHQe7SOBcZGoAU8= github.com/aws/aws-sdk-go-v2/service/sso v1.25.2/go.mod h1:qs4a9T5EMLl/Cajiw2TcbNt2UNo/Hqlyp+GiuG4CFDI= -github.com/aws/aws-sdk-go-v2/service/sso v1.30.14 h1:GcLE9ba5ehAQma6wlopUesYg/hbcOhFNWTjELkiWkh4= -github.com/aws/aws-sdk-go-v2/service/sso v1.30.14/go.mod h1:WSvS1NLr7JaPunCXqpJnWk1Bjo7IxzZXrZi1QQCkuqM= +github.com/aws/aws-sdk-go-v2/service/sso v1.30.19 h1:N6pIsdFOW1Kd9S4KyFKXdGRBojPPxkP32+uHFWLv4Hc= +github.com/aws/aws-sdk-go-v2/service/sso v1.30.19/go.mod h1:3gt5WJArFooNmyLONS+h/R4J+o86II8du38IgCwj9dE= github.com/aws/aws-sdk-go-v2/service/ssooidc v1.14.3/go.mod h1:zVwRrfdSmbRZWkUkWjOItY7SOalnFnq/Yg2LVPqDjwc= github.com/aws/aws-sdk-go-v2/service/ssooidc v1.23.1/go.mod h1:YjAPFn4kGFqKC54VsHs5fn5B6d+PCY2tziEa3U/GB5Y= github.com/aws/aws-sdk-go-v2/service/ssooidc v1.30.0/go.mod h1:MlYRNmYu/fGPoxBQVvBYr9nyr948aY/WLUvwBMBJubs= -github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.18 h1:mP49nTpfKtpXLt5SLn8Uv8z6W+03jYVoOSAl/c02nog= -github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.18/go.mod h1:YO8TrYtFdl5w/4vmjL8zaBSsiNp3w0L1FfKVKenZT7w= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.36.2 h1:hc+lBYiiTr8Zk4MTzIsQ92MeDWCIDvWGmzKUWOaBcOg= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.36.2/go.mod h1:hU6fqB3OJA6/ePheD47LQnxvjYk6br6PtQxs+Q9ojvk= github.com/aws/aws-sdk-go-v2/service/sts v1.18.4/go.mod h1:1mKZHLLpDMHTNSYPJ7qrcnCQdHCWsNQaT0xRvq2u80s= github.com/aws/aws-sdk-go-v2/service/sts v1.28.1/go.mod h1:uQ7YYKZt3adCRrdCBREm1CD3efFLOUNH77MrUCvx5oA= github.com/aws/aws-sdk-go-v2/service/sts v1.33.17/go.mod h1:cQnB8CUnxbMU82JvlqjKR2HBOm3fe9pWorWBza6MBJ4= -github.com/aws/aws-sdk-go-v2/service/sts v1.41.10 h1:p8ogvvLugcR/zLBXTXrTkj0RYBUdErbMnAFFp12Lm/U= -github.com/aws/aws-sdk-go-v2/service/sts v1.41.10/go.mod h1:60dv0eZJfeVXfbT1tFJinbHrDfSJ2GZl4Q//OSSNAVw= +github.com/aws/aws-sdk-go-v2/service/sts v1.42.3 h1:ErklX/7uhSbkAAeyQD/Y1OoQ9hO3SJXQNEgksORW3Js= +github.com/aws/aws-sdk-go-v2/service/sts v1.42.3/go.mod h1:ULe4HCzfKPiR6R3HEurE3b1upEkuk8AkMrOKtaOxKO8= github.com/aws/smithy-go v1.13.5/go.mod h1:Tg+OJXh4MB2R/uN61Ko2f6hTZwB/ZYGOtib8J3gBHzA= github.com/aws/smithy-go v1.20.1/go.mod h1:krry+ya/rV9RDcV/Q16kpu6ypI4K2czasz0NC3qS14E= github.com/aws/smithy-go v1.22.2/go.mod h1:irrKGvNn1InZwb2d7fkIRNucdfwR8R+Ts3wxYa/cJHg= github.com/aws/smithy-go v1.22.3/go.mod h1:t1ufH5HMublsJYulve2RKmHDC15xu1f26kHCp/HgceI= -github.com/aws/smithy-go v1.24.2 h1:FzA3bu/nt/vDvmnkg+R8Xl46gmzEDam6mZ1hzmwXFng= -github.com/aws/smithy-go v1.24.2/go.mod h1:YE2RhdIuDbA5E5bTdciG9KrW3+TiEONeUWCqxX9i1Fc= +github.com/aws/smithy-go v1.26.0 h1:9ouqbi+NyKP7fV3Te7UElCwdAb6Y8uk7LGwPE5tVe/s= +github.com/aws/smithy-go v1.26.0/go.mod h1:YE2RhdIuDbA5E5bTdciG9KrW3+TiEONeUWCqxX9i1Fc= github.com/axiomhq/hyperloglog v0.2.6 h1:sRhvvF3RIXWQgAXaTphLp4yJiX4S0IN3MWTaAgZoRJw= github.com/axiomhq/hyperloglog v0.2.6/go.mod h1:YjX/dQqCR/7QYX0g8mu8UZAjpIenz1FKM71UEsjFoTo= github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k= diff --git a/packages/go/storage/localstore.go b/packages/go/storage/localstore.go new file mode 100644 index 000000000000..31b597c0e99a --- /dev/null +++ b/packages/go/storage/localstore.go @@ -0,0 +1,444 @@ +// Copyright 2026 Specter Ops, Inc. +// +// Licensed under the Apache License, Version 2.0 +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 + +package storage + +import ( + "context" + "errors" + "fmt" + "io" + "io/fs" + "mime" + "os" + "path" + "strings" + "sync" + + "github.com/specterops/bloodhound/packages/go/mediatypes" +) + +var ErrIsDirectory = errors.New("is a directory") + +// ctxReader wraps an io.Reader so that context cancellation is observed between +// reads. io.Copy calls Read in a loop and each call checks ctx.Err() before delegating. +type ctxReader struct { + ctx context.Context + reader io.Reader +} + +func (s *ctxReader) Read(buffer []byte) (int, error) { + if err := s.ctx.Err(); err != nil { + return 0, err + } + return s.reader.Read(buffer) +} + +// LocalStore is a Storage implementation backed by a sandboxed directory on the +// local filesystem. All operations are confined to the root supplied to +// NewLocalStore. Callers use forward-slash-separated logical paths (e.g. +// "archives/clients/abc/file") and the kernel enforces that resolution stays +// within the root, including through symlinks. Passing a path with a leading slash +// will cause an error. +type LocalStore struct { + root *os.Root + closeOnce sync.Once + closeErr error +} + +func NewLocalStore(rootPath string) (*LocalStore, error) { + rootDirectory, err := os.OpenRoot(rootPath) + if err != nil { + return nil, err + } + return &LocalStore{ + root: rootDirectory, + }, nil +} + +func (s *LocalStore) Close() error { + s.closeOnce.Do(func() { s.closeErr = s.root.Close() }) + return s.closeErr +} + +func detectContentType(name string) string { + if ext := path.Ext(name); ext != "" { + if contentType := mime.TypeByExtension(ext); contentType != "" { + return contentType + } + } + return mediatypes.ApplicationOctetStream.String() +} + +// syncDir is used to ensure that once the entry in a directory +// is changed, to sync the change with the parent directory. While +// this error does not mean the bytes of the file were not saved, +// it may be helpful for troubleshooting if the otherwise successfully +// saved file does not appear. +func syncDir(root *os.Root, dir string) error { + dirFile, err := root.Open(dir) + if err != nil { + return err + } + defer dirFile.Close() + + return dirFile.Sync() +} + +// writeAtomic streams src into a temp file under dir(name), then publishes it at name. +// If failIfExists is true, publish uses link+unlink and returns an error satisfying +// errors.Is(err, fs.ErrExist) on collision. Otherwise, publish uses rename and silently +// replaces any existing file at name. The temp file is removed on every failure path. +// However, failed or crash-lost temp cleanup may leak .tmp-* which can be picked up by +// a sweeper. +func (s *LocalStore) writeAtomic(ctx context.Context, name string, src io.Reader, failIfExists bool) error { + var ( + dir = path.Dir(name) + tmpName string + tmp *os.File + id string + closed bool + err error + ) + + if err = ctx.Err(); err != nil { + return err + } + + if dir != "." { + if err = s.root.MkdirAll(dir, 0o750); err != nil { + return err + } + } + + if id, err = randomID(); err != nil { + return err + } + tmpName = path.Join(dir, ".tmp-"+id) + + if tmp, err = s.root.OpenFile(tmpName, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o640); err != nil { + return err + } + + // Remove the temp file on any failure path. On the success path we rename it + // before this runs, so Remove returns ErrNotExist and we ignore it. + defer func() { + if err != nil { + if !closed { + _ = tmp.Close() + } + _ = s.root.Remove(tmpName) + } + }() + + if _, err = io.Copy(tmp, &ctxReader{ctx: ctx, reader: src}); err != nil { + return err + } + if err = tmp.Sync(); err != nil { // flush data blocks + return err + } + if err = tmp.Close(); err != nil { + return err + } + closed = true + + if failIfExists { + // Link fails with fs.ErrExist if name already exists. Atomic and race-free + // against concurrent writers on the same filesystem. + if err = s.root.Link(tmpName, name); err != nil { + return err + } + if err = syncDir(s.root, dir); err != nil { + return err + } + + // Publish is durable, a failed Remove leaks a .tmp-... that a sweeper can reclaim + _ = s.root.Remove(tmpName) + return nil + } + if err = s.root.Rename(tmpName, name); err != nil { + return err + } + if err = syncDir(s.root, dir); err != nil { + return err + } + return nil +} + +// Open returns an fs.File, which must be closed. +func (s *LocalStore) Open(name string) (fs.File, error) { + return s.root.Open(name) // *os.File Satisfies fs.File +} + +func (s *LocalStore) Stat(ctx context.Context, name string) (FileInfo, error) { + if err := ctx.Err(); err != nil { + return FileInfo{}, err + } + stat, err := s.root.Stat(name) + if err != nil { + return FileInfo{}, err + } + if stat.IsDir() { + return FileInfo{}, fmt.Errorf("stat %q: %w", name, ErrIsDirectory) + } + + return FileInfo{ + Path: name, + Size: stat.Size(), + LastModified: stat.ModTime(), + IsDir: stat.IsDir(), + ContentType: detectContentType(name), + }, nil +} + +func (s *LocalStore) Get(ctx context.Context, name string) (io.ReadCloser, FileInfo, error) { + if err := ctx.Err(); err != nil { + return nil, FileInfo{}, err + } + file, err := s.root.Open(name) + if err != nil { + return nil, FileInfo{}, err + } + stat, err := file.Stat() + if err != nil { + _ = file.Close() + return nil, FileInfo{}, err + } + if stat.IsDir() { + _ = file.Close() + return nil, FileInfo{}, fmt.Errorf("get: %q, %w", name, ErrIsDirectory) + } + if err := ctx.Err(); err != nil { + _ = file.Close() + return nil, FileInfo{}, err + } + + return file, FileInfo{ + Path: name, + Size: stat.Size(), + LastModified: stat.ModTime(), + IsDir: stat.IsDir(), + ContentType: detectContentType(name), + }, nil +} + +func (s *LocalStore) Put(ctx context.Context, name string, reader io.Reader, options WriteOptions) error { + return s.writeAtomic(ctx, name, reader, options.FailIfExists) +} + +func (s *LocalStore) Exists(ctx context.Context, name string) (bool, error) { + if err := ctx.Err(); err != nil { + return false, err + } + stat, err := s.root.Stat(name) + if errors.Is(err, os.ErrNotExist) { + return false, nil + } + if err != nil { + return false, err + } + if stat.IsDir() { + return false, nil + } + return true, nil +} + +func (s *LocalStore) Delete(ctx context.Context, name string) error { + if err := ctx.Err(); err != nil { + return err + } + + stat, err := s.root.Stat(name) + if errors.Is(err, os.ErrNotExist) { + return nil + } + if err != nil { + return err + } + if stat.IsDir() { + return fmt.Errorf("delete %q: %w", name, ErrIsDirectory) + } + if err := s.root.Remove(name); err != nil { + return err + } + return syncDir(s.root, path.Dir(name)) +} + +func (s *LocalStore) List(ctx context.Context, name string, options ListOptions) ([]FileInfo, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + listName := name + if strings.TrimSpace(listName) == "" || listName == "/" { + listName = "." + } + fsys := s.root.FS() + if options.Recursive { + out := []FileInfo{} + err := fs.WalkDir(fsys, listName, func(entryPath string, entry fs.DirEntry, err error) error { + if cerr := ctx.Err(); cerr != nil { + return cerr + } + if err != nil { + if errors.Is(err, fs.ErrNotExist) && entryPath == listName { + return fs.SkipAll + } + return err + } + if entry.IsDir() { + return nil + } + info, err := entry.Info() + if err != nil { + return err + } + out = append(out, FileInfo{ + Path: entryPath, + Size: info.Size(), + LastModified: info.ModTime(), + IsDir: entry.IsDir(), + ContentType: detectContentType(entryPath), + }) + if options.Limit > 0 && len(out) >= options.Limit { + return fs.SkipAll + } + return nil + }) + if err != nil { + return nil, err + } + return out, nil + } + entries, err := fs.ReadDir(fsys, listName) + if errors.Is(err, fs.ErrNotExist) { + return []FileInfo{}, nil + } + if err != nil { + return nil, err + } + out := make([]FileInfo, 0, len(entries)) + for _, entry := range entries { + if err := ctx.Err(); err != nil { + return nil, err + } + if entry.IsDir() { + continue + } + info, err := entry.Info() + if err != nil { + return nil, err + } + entryPath := path.Join(listName, entry.Name()) + if listName == "." { + entryPath = entry.Name() + } + out = append(out, FileInfo{ + Path: entryPath, + Size: info.Size(), + LastModified: info.ModTime(), + IsDir: entry.IsDir(), + ContentType: detectContentType(entryPath), + }) + if options.Limit > 0 && len(out) >= options.Limit { + break + } + } + return out, nil +} + +// Copy duplicates srcName to destName atomically: on failure the destination is unchanged; +// on success, callers observe either the old content or new, never a partial write. +func (s *LocalStore) Copy(ctx context.Context, srcName, dstName string, options WriteOptions) error { + if err := ctx.Err(); err != nil { + return err + } + src, err := s.root.Open(srcName) + if err != nil { + return err + } + defer src.Close() + info, err := src.Stat() + if err != nil { + return err + } + if info.IsDir() { + return fmt.Errorf("copy %q: %w", srcName, ErrIsDirectory) + } + return s.writeAtomic(ctx, dstName, src, options.FailIfExists) +} + +func (s *LocalStore) moveSyncDir(srcDir, dstDir string) error { + // Starting with destDir because we are ensuring the final location is durable. + if err := syncDir(s.root, dstDir); err != nil { + return err + } + + // Only have to sync once if the move was within the same directory. + if srcDir != dstDir { + if err := syncDir(s.root, srcDir); err != nil { + return err + } + } + + return nil +} + +// Move is able to move a file from srcName to dstName using a two-step Link+Remove. A crash between link +// and unlink leaves both names present. +func (s *LocalStore) Move(ctx context.Context, srcName, dstName string, options WriteOptions) error { + if err := ctx.Err(); err != nil { + return err + } + info, err := s.root.Stat(srcName) + if err != nil { + return err + } + if info.IsDir() { + return fmt.Errorf("move %q: %w", srcName, ErrIsDirectory) + } + + srcDir := path.Dir(srcName) + dstDir := path.Dir(dstName) + if dstDir != "." { + if err := s.root.MkdirAll(dstDir, 0o750); err != nil { + return err + } + } + if options.FailIfExists { + if err := s.root.Link(srcName, dstName); err != nil { + return err + } + + // Ensure the destination is durable before removing the source. + if err := syncDir(s.root, dstDir); err != nil { + return err + } + + if err := s.root.Remove(srcName); err != nil { + return err + } + + // Ensure the source directory is durable after the remove. + if err := syncDir(s.root, srcDir); err != nil { + return err + } + return nil + } + if err := s.root.Rename(srcName, dstName); err != nil { + return err + } + return s.moveSyncDir(srcDir, dstDir) +} diff --git a/packages/go/storage/localstore_test.go b/packages/go/storage/localstore_test.go new file mode 100644 index 000000000000..1326bd169a56 --- /dev/null +++ b/packages/go/storage/localstore_test.go @@ -0,0 +1,1712 @@ +// Copyright 2026 Specter Ops, Inc. +// +// Licensed under the Apache License, Version 2.0 +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 + +package storage_test + +import ( + "context" + "errors" + "io" + "io/fs" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/specterops/bloodhound/packages/go/storage" + "github.com/stretchr/testify/require" +) + +// errorReader is a type to simulate a failed reader +type errorReader struct { + err error +} + +func (s errorReader) Read([]byte) (int, error) { + return 0, s.err +} + +// partialErrorReader creates the error case that some bytes were written and then failed. +type partialErrorReader struct { + err error + done bool +} + +func (s *partialErrorReader) Read(p []byte) (int, error) { + if s.done { + return 0, s.err + } + + s.done = true + return copy(p, "partial"), nil +} + +// requireNoTempFiles is used to walk down directories and ensure no temp files exist. +func requireNoTempFiles(t *testing.T, rootPath string) { + t.Helper() + + err := filepath.WalkDir(rootPath, func(filePath string, entry fs.DirEntry, walkErr error) error { + require.NoError(t, walkErr) + + if entry.IsDir() { + return nil + } + + require.False(t, strings.HasPrefix(entry.Name(), ".tmp-"), "unexpected temp file: %s", filePath) + return nil + }) + require.NoError(t, err) +} + +func newTestLocalStore(t *testing.T) (string, *storage.LocalStore) { + t.Helper() + + rootPath := t.TempDir() + localStore, err := storage.NewLocalStore(rootPath) + require.NoError(t, err) + + t.Cleanup(func() { + require.NoError(t, localStore.Close()) + }) + + return rootPath, localStore +} + +func writeTestFile(t *testing.T, rootPath string, name string, data string) { + t.Helper() + + filePath := filepath.Join(rootPath, filepath.FromSlash(name)) + require.NoError(t, os.MkdirAll(filepath.Dir(filePath), 0o750)) + require.NoError(t, os.WriteFile(filePath, []byte(data), 0o640)) +} + +func readTestFile(t *testing.T, rootPath string, name string) string { + t.Helper() + + data, err := os.ReadFile(filepath.Join(rootPath, filepath.FromSlash(name))) + require.NoError(t, err) + return string(data) +} + +func requireReadFile(t *testing.T, filePath string) []byte { + t.Helper() + + data, err := os.ReadFile(filePath) + require.NoError(t, err) + return data +} + +func fileInfoPaths(fileInfos []storage.FileInfo) []string { + paths := make([]string, 0, len(fileInfos)) + for _, fileInfo := range fileInfos { + paths = append(paths, fileInfo.Path) + } + return paths +} + +func TestNewLocalStore(t *testing.T) { + t.Parallel() + + type testData struct { + name string + setupRoot func(t *testing.T) string + wantErr bool + } + + tests := []testData{ + { + name: "opens existing directory", + setupRoot: func(t *testing.T) string { + return t.TempDir() + }, + }, + { + name: "missing root returns error", + setupRoot: func(t *testing.T) string { + return filepath.Join(t.TempDir(), "missing") + }, + wantErr: true, + }, + { + name: "file root returns error", + setupRoot: func(t *testing.T) string { + filePath := filepath.Join(t.TempDir(), "file") + require.NoError(t, os.WriteFile(filePath, []byte("data"), 0o640)) + return filePath + }, + wantErr: true, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + rootPath := testCase.setupRoot(t) + + // Act + localStore, err := storage.NewLocalStore(rootPath) + + // Assert + if testCase.wantErr { + require.Error(t, err) + require.Nil(t, localStore) + return + } + + require.NoError(t, err) + require.NotNil(t, localStore) + require.NoError(t, localStore.Close()) + }) + } +} + +func TestLocalStore_Close(t *testing.T) { + t.Parallel() + + // Arrange + rootPath := t.TempDir() + localStore, err := storage.NewLocalStore(rootPath) + require.NoError(t, err) + + // Act / Assert + require.NoError(t, localStore.Close()) + require.NoError(t, localStore.Close()) +} + +func TestLocalStore_Put(t *testing.T) { + t.Parallel() + + var errRead = errors.New("read failed") + + type expected struct { + errIs error + errContains string + fileContent map[string]string + missing []string + } + + type testData struct { + name string + setup func(t *testing.T, rootPath string) + buildContext func(t *testing.T) context.Context + buildReader func(t *testing.T) io.Reader + fileName string + options storage.WriteOptions + expected expected + verify func(t *testing.T, rootPath string) + } + + tests := []testData{ + { + name: "writes new file", + fileName: "file.json", + buildReader: func(t *testing.T) io.Reader { + return strings.NewReader(`{"ok":true}`) + }, + expected: expected{ + fileContent: map[string]string{"file.json": `{"ok":true}`}, + }, + verify: func(t *testing.T, rootPath string) { + requireNoTempFiles(t, rootPath) + }, + }, + { + name: "writes new nested file", + fileName: "nested/file.json", + buildReader: func(t *testing.T) io.Reader { + return strings.NewReader(`{"ok":true}`) + }, + expected: expected{ + fileContent: map[string]string{"nested/file.json": `{"ok":true}`}, + }, + verify: func(t *testing.T, rootPath string) { + requireNoTempFiles(t, rootPath) + }, + }, + { + name: "empty file name returns error", + fileName: "", + buildReader: func(t *testing.T) io.Reader { + return strings.NewReader("data") + }, + expected: expected{ + errContains: "empty path", + }, + verify: func(t *testing.T, rootPath string) { + requireNoTempFiles(t, rootPath) + }, + }, + { + name: "overwrites existing file by default", + fileName: "nested/file.json", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "nested/file.json", "old") + }, + buildReader: func(t *testing.T) io.Reader { + return strings.NewReader("new") + }, + expected: expected{ + fileContent: map[string]string{"nested/file.json": "new"}, + }, + verify: func(t *testing.T, rootPath string) { + requireNoTempFiles(t, rootPath) + }, + }, + { + name: "fail if exists preserves existing file", + fileName: "nested/file.json", + options: storage.WriteOptions{FailIfExists: true}, + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "nested/file.json", "old") + }, + buildReader: func(t *testing.T) io.Reader { + return strings.NewReader("new") + }, + expected: expected{ + errIs: fs.ErrExist, + fileContent: map[string]string{"nested/file.json": "old"}, + }, + verify: func(t *testing.T, rootPath string) { + requireNoTempFiles(t, rootPath) + }, + }, + { + name: "canceled context does not save file", + fileName: "file.json", + buildReader: func(t *testing.T) io.Reader { + return strings.NewReader("canceled") + }, + buildContext: func(t *testing.T) context.Context { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + return ctx + }, + expected: expected{ + errIs: context.Canceled, + missing: []string{"file.json"}, + }, + verify: func(t *testing.T, rootPath string) { + requireNoTempFiles(t, rootPath) + }, + }, + { + name: "failed reader returns error and removes temp file", + fileName: "file.json", + buildReader: func(t *testing.T) io.Reader { + return errorReader{err: errRead} + }, + expected: expected{ + errIs: errRead, + missing: []string{"file.json"}, + }, + verify: func(t *testing.T, rootPath string) { + requireNoTempFiles(t, rootPath) + }, + }, + { + name: "partial reader error does not publish partial file", + fileName: "file.json", + buildReader: func(t *testing.T) io.Reader { + return &partialErrorReader{err: errRead} + }, + expected: expected{ + errIs: errRead, + missing: []string{"file.json"}, + }, + verify: func(t *testing.T, rootPath string) { + requireNoTempFiles(t, rootPath) + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + rootPath, localStore := newTestLocalStore(t) + + ctx := context.Background() + if testCase.buildContext != nil { + ctx = testCase.buildContext(t) + } + + var reader io.Reader = strings.NewReader("") + if testCase.buildReader != nil { + reader = testCase.buildReader(t) + } + + if testCase.setup != nil { + testCase.setup(t, rootPath) + } + + // Act + err := localStore.Put(ctx, testCase.fileName, reader, testCase.options) + + // Assert + switch { + case testCase.expected.errIs != nil: + require.ErrorIs(t, err, testCase.expected.errIs) + case testCase.expected.errContains != "": + require.Error(t, err) + require.Contains(t, err.Error(), testCase.expected.errContains) + default: + require.NoError(t, err) + } + + for fileName, expectedContent := range testCase.expected.fileContent { + actualContent, err := os.ReadFile(filepath.Join(rootPath, filepath.FromSlash(fileName))) + require.NoError(t, err) + require.Equal(t, expectedContent, string(actualContent)) + } + + for _, missingFile := range testCase.expected.missing { + _, err := os.Stat(filepath.Join(rootPath, filepath.FromSlash(missingFile))) + require.ErrorIs(t, err, os.ErrNotExist) + } + + if testCase.verify != nil { + testCase.verify(t, rootPath) + } + }) + } +} + +func TestLocalStore_Get(t *testing.T) { + t.Parallel() + + type expected struct { + errIs error + errContains string + content string + } + + type testData struct { + name string + setup func(t *testing.T, rootPath string) + buildContext func(t *testing.T) context.Context + fileName string + expected expected + verify func(t *testing.T, fileInfo storage.FileInfo) + } + + tests := []testData{ + { + name: "gets file content and metadata", + fileName: "file.json", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "file.json", `{"ok":true}`) + }, + expected: expected{ + content: `{"ok":true}`, + }, + verify: func(t *testing.T, fileInfo storage.FileInfo) { + require.Equal(t, "file.json", fileInfo.Path) + require.Equal(t, int64(len(`{"ok":true}`)), fileInfo.Size) + require.Equal(t, "application/json", fileInfo.ContentType) + require.False(t, fileInfo.IsDir) + require.False(t, fileInfo.LastModified.IsZero()) + }, + }, + { + name: "gets nested file content and metadata", + fileName: "nested/file.json", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "nested/file.json", `{"ok":true}`) + }, + expected: expected{ + content: `{"ok":true}`, + }, + verify: func(t *testing.T, fileInfo storage.FileInfo) { + require.Equal(t, "nested/file.json", fileInfo.Path) + require.Equal(t, int64(len(`{"ok":true}`)), fileInfo.Size) + require.Equal(t, "application/json", fileInfo.ContentType) + require.False(t, fileInfo.IsDir) + require.False(t, fileInfo.LastModified.IsZero()) + }, + }, + { + name: "missing file returns error", + fileName: "file.json", + expected: expected{ + errIs: os.ErrNotExist, + }, + }, + { + name: "get directory returns error", + fileName: "nested", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "nested/file.json", `{"ok":true}`) + }, + expected: expected{ + errIs: storage.ErrIsDirectory, + }, + }, + { + name: "canceled context returns canceled", + fileName: "nested/file.json", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "nested/file.json", `{"ok":true}`) + }, + buildContext: func(t *testing.T) context.Context { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + return ctx + }, + expected: expected{ + errIs: context.Canceled, + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + rootPath, localStore := newTestLocalStore(t) + + ctx := context.Background() + if testCase.buildContext != nil { + ctx = testCase.buildContext(t) + } + + if testCase.setup != nil { + testCase.setup(t, rootPath) + } + + // Act + readCloser, fileInfo, err := localStore.Get(ctx, testCase.fileName) + + // Assert + switch { + case testCase.expected.errIs != nil: + require.ErrorIs(t, err, testCase.expected.errIs) + require.Nil(t, readCloser) + return + case testCase.expected.errContains != "": + require.Error(t, err) + require.Contains(t, err.Error(), testCase.expected.errContains) + require.Nil(t, readCloser) + return + default: + require.NoError(t, err) + require.NotNil(t, readCloser) + } + + content, err := io.ReadAll(readCloser) + require.NoError(t, err) + require.NoError(t, readCloser.Close()) + + require.Equal(t, testCase.expected.content, string(content)) + if testCase.verify != nil { + testCase.verify(t, fileInfo) + } + }) + } +} + +func TestLocalStore_Stat(t *testing.T) { + t.Parallel() + + type expected struct { + errIs error + errContains string + fileInfo storage.FileInfo + } + + type testData struct { + name string + setup func(t *testing.T, rootPath string) + buildContext func(t *testing.T) context.Context + fileName string + expected expected + } + + tests := []testData{ + { + name: "gets file metadata", + fileName: "file.json", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "file.json", `{"ok":true}`) + }, + expected: expected{ + fileInfo: storage.FileInfo{ + Path: "file.json", + Size: int64(len(`{"ok":true}`)), + ContentType: "application/json", + IsDir: false, + }, + }, + }, + { + name: "gets nested file metadata", + fileName: "nested/file.json", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "nested/file.json", `{"ok":true}`) + }, + expected: expected{ + fileInfo: storage.FileInfo{ + Path: "nested/file.json", + Size: int64(len(`{"ok":true}`)), + ContentType: "application/json", + IsDir: false, + }, + }, + }, + { + name: "missing file returns error", + fileName: "file.json", + expected: expected{ + errIs: os.ErrNotExist, + }, + }, + { + name: "get directory returns error", + fileName: "nested", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "nested/file.json", `{"ok":true}`) + }, + expected: expected{ + errIs: storage.ErrIsDirectory, + }, + }, + { + name: "canceled context returns canceled", + fileName: "nested/file.json", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "nested/file.json", `{"ok":true}`) + }, + buildContext: func(t *testing.T) context.Context { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + return ctx + }, + expected: expected{ + errIs: context.Canceled, + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + rootPath, localStore := newTestLocalStore(t) + + ctx := context.Background() + if testCase.buildContext != nil { + ctx = testCase.buildContext(t) + } + + if testCase.setup != nil { + testCase.setup(t, rootPath) + } + + // Act + fileInfo, err := localStore.Stat(ctx, testCase.fileName) + + // Assert + switch { + case testCase.expected.errIs != nil: + require.ErrorIs(t, err, testCase.expected.errIs) + return + case testCase.expected.errContains != "": + require.Error(t, err) + require.Contains(t, err.Error(), testCase.expected.errContains) + return + default: + require.NoError(t, err) + } + + require.Equal(t, testCase.expected.fileInfo.Path, fileInfo.Path) + require.Equal(t, testCase.expected.fileInfo.Size, fileInfo.Size) + require.Equal(t, testCase.expected.fileInfo.ContentType, fileInfo.ContentType) + require.False(t, fileInfo.IsDir) + require.False(t, fileInfo.LastModified.IsZero()) + }) + } +} + +func TestLocalStore_Exists(t *testing.T) { + t.Parallel() + + type expected struct { + errIs error + errContains string + exists bool + } + + type testData struct { + name string + setup func(t *testing.T, rootPath string) + buildContext func(t *testing.T) context.Context + fileName string + expected expected + } + + tests := []testData{ + { + name: "exists file returns exists", + fileName: "file.json", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "file.json", `{"ok":true}`) + }, + expected: expected{ + exists: true, + }, + }, + { + name: "nested file returns exists", + fileName: "nested/file.json", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "nested/file.json", `{"ok":true}`) + }, + expected: expected{ + exists: true, + }, + }, + { + name: "missing file returns false", + fileName: "file.json", + expected: expected{ + exists: false, + }, + }, + { + name: "exists on directory returns false", + fileName: "nested", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "nested/file.json", `{"ok":true}`) + }, + expected: expected{ + exists: false, + }, + }, + { + name: "canceled context returns canceled", + fileName: "nested/file.json", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "nested/file.json", `{"ok":true}`) + }, + buildContext: func(t *testing.T) context.Context { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + return ctx + }, + expected: expected{ + errIs: context.Canceled, + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + rootPath, localStore := newTestLocalStore(t) + + ctx := context.Background() + if testCase.buildContext != nil { + ctx = testCase.buildContext(t) + } + + if testCase.setup != nil { + testCase.setup(t, rootPath) + } + + // Act + exists, err := localStore.Exists(ctx, testCase.fileName) + + // Assert + switch { + case testCase.expected.errIs != nil: + require.ErrorIs(t, err, testCase.expected.errIs) + return + case testCase.expected.errContains != "": + require.Error(t, err) + require.Contains(t, err.Error(), testCase.expected.errContains) + return + default: + require.NoError(t, err) + } + + require.Equal(t, testCase.expected.exists, exists) + }) + } +} + +func TestLocalStore_Delete(t *testing.T) { + t.Parallel() + + type expected struct { + errIs error + errContains string + checkRemoved bool + shouldRemove bool + } + + type testData struct { + name string + setup func(t *testing.T, rootPath string) + buildContext func(t *testing.T) context.Context + fileName string + expected expected + verify func(t *testing.T, rootPath string) + } + + tests := []testData{ + { + name: "delete file removes file", + fileName: "file.json", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "file.json", `{"ok":true}`) + }, + expected: expected{ + checkRemoved: true, + shouldRemove: true, + }, + verify: func(t *testing.T, rootPath string) { + requireNoTempFiles(t, rootPath) + }, + }, + { + name: "delete nested file removes file", + fileName: "nested/file.json", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "nested/file.json", `{"ok":true}`) + }, + expected: expected{ + checkRemoved: true, + shouldRemove: true, + }, + verify: func(t *testing.T, rootPath string) { + requireNoTempFiles(t, rootPath) + }, + }, + { + name: "delete nested file does not removes directory with items in it", + fileName: "nested/file.json", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "nested/file.json", `{"ok":true}`) + writeTestFile(t, rootPath, "nested/file2.json", `{"ok":true}`) + }, + expected: expected{ + checkRemoved: true, + shouldRemove: true, + }, + verify: func(t *testing.T, rootPath string) { + requireNoTempFiles(t, rootPath) + stat, err := os.Stat(filepath.Join(rootPath, "nested")) + require.Nil(t, err) + require.True(t, stat.IsDir()) + }, + }, + { + name: "missing file does not return error", + fileName: "file.json", + }, + { + name: "delete directory returns error on directory with items", + fileName: "nested", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "nested/file.json", `{"ok":true}`) + }, + expected: expected{ + errIs: storage.ErrIsDirectory, + }, + }, + { + name: "delete directory returns error on empty directory", + fileName: "nested", + setup: func(t *testing.T, rootPath string) { + require.NoError(t, os.MkdirAll(filepath.Join(rootPath, filepath.FromSlash("nested")), 0o750)) + }, + expected: expected{ + errIs: storage.ErrIsDirectory, + }, + }, + { + name: "canceled context returns canceled", + fileName: "nested/file.json", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "nested/file.json", `{"ok":true}`) + }, + buildContext: func(t *testing.T) context.Context { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + return ctx + }, + expected: expected{ + checkRemoved: true, + shouldRemove: false, + errIs: context.Canceled, + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + rootPath, localStore := newTestLocalStore(t) + + ctx := context.Background() + if testCase.buildContext != nil { + ctx = testCase.buildContext(t) + } + + if testCase.setup != nil { + testCase.setup(t, rootPath) + } + + // Act + err := localStore.Delete(ctx, testCase.fileName) + + // Assert + switch { + case testCase.expected.errIs != nil: + require.ErrorIs(t, err, testCase.expected.errIs) + case testCase.expected.errContains != "": + require.Error(t, err) + require.Contains(t, err.Error(), testCase.expected.errContains) + default: + require.NoError(t, err) + } + + if testCase.expected.checkRemoved { + _, err := os.Stat(filepath.Join(rootPath, filepath.FromSlash(testCase.fileName))) + if testCase.expected.shouldRemove { + require.NotNil(t, err) + require.ErrorIs(t, err, os.ErrNotExist) + } else { + require.Nil(t, err) + } + } + + if testCase.verify != nil { + testCase.verify(t, rootPath) + } + }) + } +} + +func TestLocalStore_List(t *testing.T) { + t.Parallel() + + type expected struct { + errIs error + errContains string + paths []string + } + + type testData struct { + name string + setup func(t *testing.T, rootPath string) + buildContext func(t *testing.T) context.Context + listName string + options storage.ListOptions + expected expected + verify func(t *testing.T, fileInfos []storage.FileInfo) + } + + setupFiles := func(t *testing.T, rootPath string) { + t.Helper() + + writeTestFile(t, rootPath, "file.json", `{"ok":true}`) + writeTestFile(t, rootPath, "nested/file.json", `{"nested":true}`) + writeTestFile(t, rootPath, "nested/file.txt", "nested text") + writeTestFile(t, rootPath, "nested/deeper/file.json", `{"deeper":true}`) + } + + tests := []testData{ + { + name: "empty path lists root files only", + listName: "", + setup: setupFiles, + expected: expected{ + paths: []string{"file.json"}, + }, + verify: func(t *testing.T, fileInfos []storage.FileInfo) { + require.Len(t, fileInfos, 1) + require.Equal(t, "application/json", fileInfos[0].ContentType) + require.Equal(t, int64(len(`{"ok":true}`)), fileInfos[0].Size) + require.False(t, fileInfos[0].IsDir) + require.False(t, fileInfos[0].LastModified.IsZero()) + }, + }, + { + name: "slash lists root files only", + listName: "/", + setup: setupFiles, + expected: expected{ + paths: []string{"file.json"}, + }, + }, + { + name: "lists subdirectory files only", + listName: "nested", + setup: setupFiles, + expected: expected{ + paths: []string{ + "nested/file.json", + "nested/file.txt", + }, + }, + }, + { + name: "recursively lists root files", + listName: "", + options: storage.ListOptions{Recursive: true}, + setup: setupFiles, + expected: expected{ + paths: []string{ + "file.json", + "nested/deeper/file.json", + "nested/file.json", + "nested/file.txt", + }, + }, + }, + { + name: "recursively lists subdirectory files", + listName: "nested", + options: storage.ListOptions{Recursive: true}, + setup: setupFiles, + expected: expected{ + paths: []string{ + "nested/deeper/file.json", + "nested/file.json", + "nested/file.txt", + }, + }, + }, + { + name: "limit stops list after requested number of files", + listName: "", + options: storage.ListOptions{Recursive: true, Limit: 2}, + setup: setupFiles, + verify: func(t *testing.T, fileInfos []storage.FileInfo) { + require.Len(t, fileInfos, 2) + require.Subset(t, []string{ + "file.json", + "nested/deeper/file.json", + "nested/file.json", + "nested/file.txt", + }, fileInfoPaths(fileInfos)) + }, + }, + { + name: "limit stops non-recursive list after requested number of files", + listName: "nested", + options: storage.ListOptions{Limit: 1}, + setup: setupFiles, + verify: func(t *testing.T, fileInfos []storage.FileInfo) { + require.Len(t, fileInfos, 1) + require.Subset(t, []string{ + "nested/file.json", + "nested/file.txt", + }, fileInfoPaths(fileInfos)) + }, + }, + { + name: "missing path returns empty list", + listName: "missing", + setup: setupFiles, + expected: expected{ + paths: []string{}, + }, + }, + { + name: "missing recursive path returns empty list", + listName: "missing", + options: storage.ListOptions{Recursive: true}, + setup: setupFiles, + expected: expected{ + paths: []string{}, + }, + }, + { + name: "canceled context returns canceled", + listName: "", + setup: setupFiles, + buildContext: func(t *testing.T) context.Context { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + return ctx + }, + expected: expected{ + errIs: context.Canceled, + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + rootPath, localStore := newTestLocalStore(t) + + ctx := context.Background() + if testCase.buildContext != nil { + ctx = testCase.buildContext(t) + } + + if testCase.setup != nil { + testCase.setup(t, rootPath) + } + + // Act + fileInfos, err := localStore.List(ctx, testCase.listName, testCase.options) + + // Assert + switch { + case testCase.expected.errIs != nil: + require.ErrorIs(t, err, testCase.expected.errIs) + require.Nil(t, fileInfos) + return + case testCase.expected.errContains != "": + require.Error(t, err) + require.Contains(t, err.Error(), testCase.expected.errContains) + require.Nil(t, fileInfos) + return + default: + require.NoError(t, err) + } + + if testCase.expected.paths != nil { + require.ElementsMatch(t, testCase.expected.paths, fileInfoPaths(fileInfos)) + } + + for _, fileInfo := range fileInfos { + require.False(t, fileInfo.IsDir) + require.False(t, fileInfo.LastModified.IsZero()) + } + + if testCase.verify != nil { + testCase.verify(t, fileInfos) + } + }) + } +} + +func TestLocalStore_Copy(t *testing.T) { + t.Parallel() + + type expected struct { + errIs error + errContains string + fileContent map[string]string + missing []string + } + + type testData struct { + name string + setup func(t *testing.T, rootPath string) + buildContext func(t *testing.T) context.Context + sourceName string + destinationName string + options storage.WriteOptions + expected expected + verify func(t *testing.T, rootPath string) + } + + tests := []testData{ + { + name: "copies file to new destination", + sourceName: "source.json", + destinationName: "destination.json", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "source.json", `{"ok":true}`) + }, + expected: expected{ + fileContent: map[string]string{ + "source.json": `{"ok":true}`, + "destination.json": `{"ok":true}`, + }, + }, + verify: func(t *testing.T, rootPath string) { + requireNoTempFiles(t, rootPath) + }, + }, + { + name: "copies file to nested destination", + sourceName: "source.json", + destinationName: "nested/destination.json", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "source.json", `{"ok":true}`) + }, + expected: expected{ + fileContent: map[string]string{ + "source.json": `{"ok":true}`, + "nested/destination.json": `{"ok":true}`, + }, + }, + verify: func(t *testing.T, rootPath string) { + requireNoTempFiles(t, rootPath) + }, + }, + { + name: "copy overwrites destination by default", + sourceName: "source.json", + destinationName: "destination.json", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "source.json", "new") + writeTestFile(t, rootPath, "destination.json", "old") + }, + expected: expected{ + fileContent: map[string]string{ + "source.json": "new", + "destination.json": "new", + }, + }, + verify: func(t *testing.T, rootPath string) { + requireNoTempFiles(t, rootPath) + }, + }, + { + name: "fail if exists preserves existing destination", + sourceName: "source.json", + destinationName: "destination.json", + options: storage.WriteOptions{FailIfExists: true}, + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "source.json", "new") + writeTestFile(t, rootPath, "destination.json", "old") + }, + expected: expected{ + errIs: fs.ErrExist, + fileContent: map[string]string{ + "source.json": "new", + "destination.json": "old", + }, + }, + verify: func(t *testing.T, rootPath string) { + requireNoTempFiles(t, rootPath) + }, + }, + { + name: "missing source returns not exist", + sourceName: "missing.json", + destinationName: "destination.json", + expected: expected{ + errIs: os.ErrNotExist, + missing: []string{"destination.json"}, + }, + verify: func(t *testing.T, rootPath string) { + requireNoTempFiles(t, rootPath) + }, + }, + { + name: "directory source returns directory error", + sourceName: "nested", + destinationName: "destination.json", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "nested/file.json", `{"ok":true}`) + }, + expected: expected{ + errIs: storage.ErrIsDirectory, + missing: []string{"destination.json"}, + }, + verify: func(t *testing.T, rootPath string) { + requireNoTempFiles(t, rootPath) + }, + }, + { + name: "canceled context preserves source and does not copy", + sourceName: "source.json", + destinationName: "destination.json", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "source.json", `{"ok":true}`) + }, + buildContext: func(t *testing.T) context.Context { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + return ctx + }, + expected: expected{ + errIs: context.Canceled, + fileContent: map[string]string{ + "source.json": `{"ok":true}`, + }, + missing: []string{"destination.json"}, + }, + verify: func(t *testing.T, rootPath string) { + requireNoTempFiles(t, rootPath) + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + rootPath, localStore := newTestLocalStore(t) + + ctx := context.Background() + if testCase.buildContext != nil { + ctx = testCase.buildContext(t) + } + + if testCase.setup != nil { + testCase.setup(t, rootPath) + } + + // Act + err := localStore.Copy(ctx, testCase.sourceName, testCase.destinationName, testCase.options) + + // Assert + switch { + case testCase.expected.errIs != nil: + require.ErrorIs(t, err, testCase.expected.errIs) + case testCase.expected.errContains != "": + require.Error(t, err) + require.Contains(t, err.Error(), testCase.expected.errContains) + default: + require.NoError(t, err) + } + + for fileName, expectedContent := range testCase.expected.fileContent { + require.Equal(t, expectedContent, readTestFile(t, rootPath, fileName)) + } + + for _, missingFile := range testCase.expected.missing { + _, err := os.Stat(filepath.Join(rootPath, filepath.FromSlash(missingFile))) + require.ErrorIs(t, err, os.ErrNotExist) + } + + if testCase.verify != nil { + testCase.verify(t, rootPath) + } + }) + } +} + +func TestLocalStore_Move(t *testing.T) { + t.Parallel() + + type expected struct { + errIs error + errContains string + fileContent map[string]string + missing []string + } + + type testData struct { + name string + setup func(t *testing.T, rootPath string) + buildContext func(t *testing.T) context.Context + sourceName string + destinationName string + options storage.WriteOptions + expected expected + verify func(t *testing.T, rootPath string) + } + + tests := []testData{ + { + name: "moves file to new destination", + sourceName: "source.json", + destinationName: "destination.json", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "source.json", `{"ok":true}`) + }, + expected: expected{ + fileContent: map[string]string{ + "destination.json": `{"ok":true}`, + }, + missing: []string{"source.json"}, + }, + verify: func(t *testing.T, rootPath string) { + requireNoTempFiles(t, rootPath) + }, + }, + { + name: "moves file to nested destination", + sourceName: "source.json", + destinationName: "nested/destination.json", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "source.json", `{"ok":true}`) + }, + expected: expected{ + fileContent: map[string]string{ + "nested/destination.json": `{"ok":true}`, + }, + missing: []string{"source.json"}, + }, + verify: func(t *testing.T, rootPath string) { + requireNoTempFiles(t, rootPath) + }, + }, + { + name: "move overwrites destination by default", + sourceName: "source.json", + destinationName: "destination.json", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "source.json", "new") + writeTestFile(t, rootPath, "destination.json", "old") + }, + expected: expected{ + fileContent: map[string]string{ + "destination.json": "new", + }, + missing: []string{"source.json"}, + }, + verify: func(t *testing.T, rootPath string) { + requireNoTempFiles(t, rootPath) + }, + }, + { + name: "fail if exists preserves source and destination", + sourceName: "source.json", + destinationName: "destination.json", + options: storage.WriteOptions{FailIfExists: true}, + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "source.json", "new") + writeTestFile(t, rootPath, "destination.json", "old") + }, + expected: expected{ + errIs: fs.ErrExist, + fileContent: map[string]string{ + "source.json": "new", + "destination.json": "old", + }, + }, + verify: func(t *testing.T, rootPath string) { + requireNoTempFiles(t, rootPath) + }, + }, + { + name: "missing source returns not exist", + sourceName: "missing.json", + destinationName: "destination.json", + expected: expected{ + errIs: os.ErrNotExist, + missing: []string{"destination.json"}, + }, + verify: func(t *testing.T, rootPath string) { + requireNoTempFiles(t, rootPath) + }, + }, + { + name: "directory source returns directory error", + sourceName: "nested", + destinationName: "destination", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "nested/file.json", `{"ok":true}`) + }, + expected: expected{ + errIs: storage.ErrIsDirectory, + fileContent: map[string]string{ + "nested/file.json": `{"ok":true}`, + }, + missing: []string{"destination"}, + }, + verify: func(t *testing.T, rootPath string) { + requireNoTempFiles(t, rootPath) + }, + }, + { + name: "canceled context preserves source and does not move", + sourceName: "source.json", + destinationName: "destination.json", + setup: func(t *testing.T, rootPath string) { + writeTestFile(t, rootPath, "source.json", `{"ok":true}`) + }, + buildContext: func(t *testing.T) context.Context { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + return ctx + }, + expected: expected{ + errIs: context.Canceled, + fileContent: map[string]string{ + "source.json": `{"ok":true}`, + }, + missing: []string{"destination.json"}, + }, + verify: func(t *testing.T, rootPath string) { + requireNoTempFiles(t, rootPath) + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + rootPath, localStore := newTestLocalStore(t) + + ctx := context.Background() + if testCase.buildContext != nil { + ctx = testCase.buildContext(t) + } + + if testCase.setup != nil { + testCase.setup(t, rootPath) + } + + // Act + err := localStore.Move(ctx, testCase.sourceName, testCase.destinationName, testCase.options) + + // Assert + switch { + case testCase.expected.errIs != nil: + require.ErrorIs(t, err, testCase.expected.errIs) + case testCase.expected.errContains != "": + require.Error(t, err) + require.Contains(t, err.Error(), testCase.expected.errContains) + default: + require.NoError(t, err) + } + + for fileName, expectedContent := range testCase.expected.fileContent { + require.Equal(t, expectedContent, readTestFile(t, rootPath, fileName)) + } + + for _, missingFile := range testCase.expected.missing { + _, err := os.Stat(filepath.Join(rootPath, filepath.FromSlash(missingFile))) + require.ErrorIs(t, err, os.ErrNotExist) + } + + if testCase.verify != nil { + testCase.verify(t, rootPath) + } + }) + } +} + +func TestLocalStore_PathSafety(t *testing.T) { + t.Parallel() + + type testData struct { + name string + setup func(t *testing.T, rootPath string, outsidePath string) string + operation func(ctx context.Context, localStore *storage.LocalStore, unsafeName string) error + } + + tests := []testData{ + { + name: "put rejects parent traversal destination", + operation: func(ctx context.Context, localStore *storage.LocalStore, unsafeName string) error { + return localStore.Put(ctx, unsafeName, strings.NewReader("changed"), storage.WriteOptions{}) + }, + }, + { + name: "put rejects absolute destination", + setup: func(t *testing.T, rootPath string, outsidePath string) string { + return outsidePath + }, + operation: func(ctx context.Context, localStore *storage.LocalStore, unsafeName string) error { + return localStore.Put(ctx, unsafeName, strings.NewReader("changed"), storage.WriteOptions{}) + }, + }, + { + name: "get rejects absolute source", + setup: func(t *testing.T, rootPath string, outsidePath string) string { + return outsidePath + }, + operation: func(ctx context.Context, localStore *storage.LocalStore, unsafeName string) error { + readCloser, _, err := localStore.Get(ctx, unsafeName) + if readCloser != nil { + require.NoError(t, readCloser.Close()) + } + return err + }, + }, + { + name: "stat rejects absolute source", + setup: func(t *testing.T, rootPath string, outsidePath string) string { + return outsidePath + }, + operation: func(ctx context.Context, localStore *storage.LocalStore, unsafeName string) error { + _, err := localStore.Stat(ctx, unsafeName) + return err + }, + }, + { + name: "delete rejects absolute target", + setup: func(t *testing.T, rootPath string, outsidePath string) string { + return outsidePath + }, + operation: func(ctx context.Context, localStore *storage.LocalStore, unsafeName string) error { + return localStore.Delete(ctx, unsafeName) + }, + }, + { + name: "get rejects parent traversal source", + operation: func(ctx context.Context, localStore *storage.LocalStore, unsafeName string) error { + readCloser, _, err := localStore.Get(ctx, unsafeName) + if readCloser != nil { + require.NoError(t, readCloser.Close()) + } + return err + }, + }, + { + name: "stat rejects parent traversal source", + operation: func(ctx context.Context, localStore *storage.LocalStore, unsafeName string) error { + _, err := localStore.Stat(ctx, unsafeName) + return err + }, + }, + { + name: "delete rejects parent traversal target", + operation: func(ctx context.Context, localStore *storage.LocalStore, unsafeName string) error { + return localStore.Delete(ctx, unsafeName) + }, + }, + { + name: "list rejects parent traversal path", + operation: func(ctx context.Context, localStore *storage.LocalStore, unsafeName string) error { + _, err := localStore.List(ctx, unsafeName, storage.ListOptions{}) + return err + }, + }, + { + name: "copy rejects parent traversal source", + operation: func(ctx context.Context, localStore *storage.LocalStore, unsafeName string) error { + return localStore.Copy(ctx, unsafeName, "destination.json", storage.WriteOptions{}) + }, + }, + { + name: "copy rejects parent traversal destination", + operation: func(ctx context.Context, localStore *storage.LocalStore, unsafeName string) error { + return localStore.Copy(ctx, "source.json", unsafeName, storage.WriteOptions{}) + }, + }, + { + name: "copy rejects absolute source", + setup: func(t *testing.T, rootPath string, outsidePath string) string { + return outsidePath + }, + operation: func(ctx context.Context, localStore *storage.LocalStore, unsafeName string) error { + return localStore.Copy(ctx, unsafeName, "destination.json", storage.WriteOptions{}) + }, + }, + { + name: "copy rejects absolute destination", + setup: func(t *testing.T, rootPath string, outsidePath string) string { + return outsidePath + }, + operation: func(ctx context.Context, localStore *storage.LocalStore, unsafeName string) error { + return localStore.Copy(ctx, "source.json", unsafeName, storage.WriteOptions{}) + }, + }, + { + name: "move rejects parent traversal source", + operation: func(ctx context.Context, localStore *storage.LocalStore, unsafeName string) error { + return localStore.Move(ctx, unsafeName, "destination.json", storage.WriteOptions{}) + }, + }, + { + name: "move rejects parent traversal destination", + operation: func(ctx context.Context, localStore *storage.LocalStore, unsafeName string) error { + return localStore.Move(ctx, "source.json", unsafeName, storage.WriteOptions{}) + }, + }, + { + name: "move rejects absolute source", + setup: func(t *testing.T, rootPath string, outsidePath string) string { + return outsidePath + }, + operation: func(ctx context.Context, localStore *storage.LocalStore, unsafeName string) error { + return localStore.Move(ctx, unsafeName, "destination.json", storage.WriteOptions{}) + }, + }, + { + name: "move rejects absolute destination", + setup: func(t *testing.T, rootPath string, outsidePath string) string { + return outsidePath + }, + operation: func(ctx context.Context, localStore *storage.LocalStore, unsafeName string) error { + return localStore.Move(ctx, "source.json", unsafeName, storage.WriteOptions{}) + }, + }, + { + name: "put rejects symlink escape path", + setup: func(t *testing.T, rootPath string, outsidePath string) string { + linkPath := filepath.Join(rootPath, "outside") + if err := os.Symlink(filepath.Dir(outsidePath), linkPath); err != nil { + t.Skipf("unable to create symlink: %v", err) + } + return "outside/outside.txt" + }, + operation: func(ctx context.Context, localStore *storage.LocalStore, unsafeName string) error { + return localStore.Put(ctx, unsafeName, strings.NewReader("changed"), storage.WriteOptions{}) + }, + }, + { + name: "get rejects symlink escape path", + setup: func(t *testing.T, rootPath string, outsidePath string) string { + linkPath := filepath.Join(rootPath, "outside") + if err := os.Symlink(filepath.Dir(outsidePath), linkPath); err != nil { + t.Skipf("unable to create symlink: %v", err) + } + return "outside/outside.txt" + }, + operation: func(ctx context.Context, localStore *storage.LocalStore, unsafeName string) error { + readCloser, _, err := localStore.Get(ctx, unsafeName) + if readCloser != nil { + require.NoError(t, readCloser.Close()) + } + return err + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + var ( + ctx = context.Background() + rootPath, localStore = newTestLocalStore(t) + outsideDir = t.TempDir() + outsidePath = filepath.Join(outsideDir, "outside.txt") + ) + + require.NoError(t, os.WriteFile(outsidePath, []byte("outside"), 0o640)) + writeTestFile(t, rootPath, "source.json", "source") + + unsafeName := "" + if testCase.setup != nil { + unsafeName = testCase.setup(t, rootPath, outsidePath) + } else { + relativePath, err := filepath.Rel(rootPath, outsidePath) + require.NoError(t, err) + unsafeName = filepath.ToSlash(relativePath) + } + + // Act + err := testCase.operation(ctx, localStore, unsafeName) + + // Assert + require.Error(t, err) + require.Equal(t, "outside", string(requireReadFile(t, outsidePath))) + require.Equal(t, "source", readTestFile(t, rootPath, "source.json")) + requireNoTempFiles(t, rootPath) + }) + } +} diff --git a/packages/go/storage/mocks/storage.go b/packages/go/storage/mocks/storage.go new file mode 100644 index 000000000000..1179e89cf7e3 --- /dev/null +++ b/packages/go/storage/mocks/storage.go @@ -0,0 +1,332 @@ +// Copyright 2026 Specter Ops, Inc. +// +// Licensed under the Apache License, Version 2.0 +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 + +// Code generated by MockGen. DO NOT EDIT. +// Source: github.com/specterops/bloodhound/packages/go/storage (interfaces: Storage,FileService) +// +// Generated by this command: +// +// mockgen -copyright_file=../../../LICENSE.header -destination=./mocks/storage.go -package=mocks . Storage,FileService +// + +// Package mocks is a generated GoMock package. +package mocks + +import ( + context "context" + io "io" + reflect "reflect" + + storage "github.com/specterops/bloodhound/packages/go/storage" + gomock "go.uber.org/mock/gomock" +) + +// MockStorage is a mock of Storage interface. +type MockStorage struct { + ctrl *gomock.Controller + recorder *MockStorageMockRecorder + isgomock struct{} +} + +// MockStorageMockRecorder is the mock recorder for MockStorage. +type MockStorageMockRecorder struct { + mock *MockStorage +} + +// NewMockStorage creates a new mock instance. +func NewMockStorage(ctrl *gomock.Controller) *MockStorage { + mock := &MockStorage{ctrl: ctrl} + mock.recorder = &MockStorageMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockStorage) EXPECT() *MockStorageMockRecorder { + return m.recorder +} + +// Copy mocks base method. +func (m *MockStorage) Copy(ctx context.Context, srcName, dstName string, options storage.WriteOptions) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Copy", ctx, srcName, dstName, options) + ret0, _ := ret[0].(error) + return ret0 +} + +// Copy indicates an expected call of Copy. +func (mr *MockStorageMockRecorder) Copy(ctx, srcName, dstName, options any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Copy", reflect.TypeOf((*MockStorage)(nil).Copy), ctx, srcName, dstName, options) +} + +// Delete mocks base method. +func (m *MockStorage) Delete(ctx context.Context, name string) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Delete", ctx, name) + ret0, _ := ret[0].(error) + return ret0 +} + +// Delete indicates an expected call of Delete. +func (mr *MockStorageMockRecorder) Delete(ctx, name any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Delete", reflect.TypeOf((*MockStorage)(nil).Delete), ctx, name) +} + +// Exists mocks base method. +func (m *MockStorage) Exists(ctx context.Context, name string) (bool, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Exists", ctx, name) + ret0, _ := ret[0].(bool) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Exists indicates an expected call of Exists. +func (mr *MockStorageMockRecorder) Exists(ctx, name any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Exists", reflect.TypeOf((*MockStorage)(nil).Exists), ctx, name) +} + +// Get mocks base method. +func (m *MockStorage) Get(ctx context.Context, name string) (io.ReadCloser, storage.FileInfo, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Get", ctx, name) + ret0, _ := ret[0].(io.ReadCloser) + ret1, _ := ret[1].(storage.FileInfo) + ret2, _ := ret[2].(error) + return ret0, ret1, ret2 +} + +// Get indicates an expected call of Get. +func (mr *MockStorageMockRecorder) Get(ctx, name any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get", reflect.TypeOf((*MockStorage)(nil).Get), ctx, name) +} + +// List mocks base method. +func (m *MockStorage) List(ctx context.Context, name string, options storage.ListOptions) ([]storage.FileInfo, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "List", ctx, name, options) + ret0, _ := ret[0].([]storage.FileInfo) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// List indicates an expected call of List. +func (mr *MockStorageMockRecorder) List(ctx, name, options any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "List", reflect.TypeOf((*MockStorage)(nil).List), ctx, name, options) +} + +// Move mocks base method. +func (m *MockStorage) Move(ctx context.Context, srcName, dstName string, options storage.WriteOptions) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Move", ctx, srcName, dstName, options) + ret0, _ := ret[0].(error) + return ret0 +} + +// Move indicates an expected call of Move. +func (mr *MockStorageMockRecorder) Move(ctx, srcName, dstName, options any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Move", reflect.TypeOf((*MockStorage)(nil).Move), ctx, srcName, dstName, options) +} + +// Put mocks base method. +func (m *MockStorage) Put(ctx context.Context, name string, reader io.Reader, options storage.WriteOptions) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Put", ctx, name, reader, options) + ret0, _ := ret[0].(error) + return ret0 +} + +// Put indicates an expected call of Put. +func (mr *MockStorageMockRecorder) Put(ctx, name, reader, options any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Put", reflect.TypeOf((*MockStorage)(nil).Put), ctx, name, reader, options) +} + +// Stat mocks base method. +func (m *MockStorage) Stat(ctx context.Context, name string) (storage.FileInfo, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Stat", ctx, name) + ret0, _ := ret[0].(storage.FileInfo) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Stat indicates an expected call of Stat. +func (mr *MockStorageMockRecorder) Stat(ctx, name any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Stat", reflect.TypeOf((*MockStorage)(nil).Stat), ctx, name) +} + +// MockFileService is a mock of FileService interface. +type MockFileService struct { + ctrl *gomock.Controller + recorder *MockFileServiceMockRecorder + isgomock struct{} +} + +// MockFileServiceMockRecorder is the mock recorder for MockFileService. +type MockFileServiceMockRecorder struct { + mock *MockFileService +} + +// NewMockFileService creates a new mock instance. +func NewMockFileService(ctrl *gomock.Controller) *MockFileService { + mock := &MockFileService{ctrl: ctrl} + mock.recorder = &MockFileServiceMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockFileService) EXPECT() *MockFileServiceMockRecorder { + return m.recorder +} + +// DeleteFile mocks base method. +func (m *MockFileService) DeleteFile(ctx context.Context, name string) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DeleteFile", ctx, name) + ret0, _ := ret[0].(error) + return ret0 +} + +// DeleteFile indicates an expected call of DeleteFile. +func (mr *MockFileServiceMockRecorder) DeleteFile(ctx, name any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteFile", reflect.TypeOf((*MockFileService)(nil).DeleteFile), ctx, name) +} + +// GetFile mocks base method. +func (m *MockFileService) GetFile(ctx context.Context, name string) (io.ReadCloser, storage.FileInfo, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetFile", ctx, name) + ret0, _ := ret[0].(io.ReadCloser) + ret1, _ := ret[1].(storage.FileInfo) + ret2, _ := ret[2].(error) + return ret0, ret1, ret2 +} + +// GetFile indicates an expected call of GetFile. +func (mr *MockFileServiceMockRecorder) GetFile(ctx, name any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetFile", reflect.TypeOf((*MockFileService)(nil).GetFile), ctx, name) +} + +// GetFileInfo mocks base method. +func (m *MockFileService) GetFileInfo(ctx context.Context, name string) (storage.FileInfo, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetFileInfo", ctx, name) + ret0, _ := ret[0].(storage.FileInfo) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetFileInfo indicates an expected call of GetFileInfo. +func (mr *MockFileServiceMockRecorder) GetFileInfo(ctx, name any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetFileInfo", reflect.TypeOf((*MockFileService)(nil).GetFileInfo), ctx, name) +} + +// ListFiles mocks base method. +func (m *MockFileService) ListFiles(ctx context.Context, name string, opts storage.ListOptions) ([]storage.FileInfo, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ListFiles", ctx, name, opts) + ret0, _ := ret[0].([]storage.FileInfo) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ListFiles indicates an expected call of ListFiles. +func (mr *MockFileServiceMockRecorder) ListFiles(ctx, name, opts any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListFiles", reflect.TypeOf((*MockFileService)(nil).ListFiles), ctx, name, opts) +} + +// MoveFile mocks base method. +func (m *MockFileService) MoveFile(ctx context.Context, srcName, dstName string, opts storage.WriteOptions) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "MoveFile", ctx, srcName, dstName, opts) + ret0, _ := ret[0].(error) + return ret0 +} + +// MoveFile indicates an expected call of MoveFile. +func (mr *MockFileServiceMockRecorder) MoveFile(ctx, srcName, dstName, opts any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MoveFile", reflect.TypeOf((*MockFileService)(nil).MoveFile), ctx, srcName, dstName, opts) +} + +// ReadFile mocks base method. +func (m *MockFileService) ReadFile(ctx context.Context, name string) ([]byte, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ReadFile", ctx, name) + ret0, _ := ret[0].([]byte) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ReadFile indicates an expected call of ReadFile. +func (mr *MockFileServiceMockRecorder) ReadFile(ctx, name any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReadFile", reflect.TypeOf((*MockFileService)(nil).ReadFile), ctx, name) +} + +// WriteFile mocks base method. +func (m *MockFileService) WriteFile(ctx context.Context, name string, data []byte, opts storage.WriteOptions) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "WriteFile", ctx, name, data, opts) + ret0, _ := ret[0].(error) + return ret0 +} + +// WriteFile indicates an expected call of WriteFile. +func (mr *MockFileServiceMockRecorder) WriteFile(ctx, name, data, opts any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WriteFile", reflect.TypeOf((*MockFileService)(nil).WriteFile), ctx, name, data, opts) +} + +// WriteFileFromReader mocks base method. +func (m *MockFileService) WriteFileFromReader(ctx context.Context, name string, reader io.Reader, opts storage.WriteOptions) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "WriteFileFromReader", ctx, name, reader, opts) + ret0, _ := ret[0].(error) + return ret0 +} + +// WriteFileFromReader indicates an expected call of WriteFileFromReader. +func (mr *MockFileServiceMockRecorder) WriteFileFromReader(ctx, name, reader, opts any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WriteFileFromReader", reflect.TypeOf((*MockFileService)(nil).WriteFileFromReader), ctx, name, reader, opts) +} + +// WriteTempFile mocks base method. +func (m *MockFileService) WriteTempFile(ctx context.Context, prefix string, reader io.Reader, opts storage.WriteOptions) (string, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "WriteTempFile", ctx, prefix, reader, opts) + ret0, _ := ret[0].(string) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// WriteTempFile indicates an expected call of WriteTempFile. +func (mr *MockFileServiceMockRecorder) WriteTempFile(ctx, prefix, reader, opts any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WriteTempFile", reflect.TypeOf((*MockFileService)(nil).WriteTempFile), ctx, prefix, reader, opts) +} diff --git a/packages/go/storage/s3.go b/packages/go/storage/s3.go new file mode 100644 index 000000000000..2f744182f755 --- /dev/null +++ b/packages/go/storage/s3.go @@ -0,0 +1,519 @@ +// Copyright 2026 Specter Ops, Inc. +// +// Licensed under the Apache License, Version 2.0 +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 + +package storage + +import ( + "context" + "errors" + "fmt" + "io" + "io/fs" + "mime" + "net/http" + "net/url" + "os" + "path" + "sort" + "strings" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/feature/s3/manager" + "github.com/aws/aws-sdk-go-v2/service/s3" + s3types "github.com/aws/aws-sdk-go-v2/service/s3/types" + "github.com/aws/smithy-go" + smithyhttp "github.com/aws/smithy-go/transport/http" + "github.com/specterops/bloodhound/packages/go/mediatypes" +) + +const ( + defaultUploadPartSize int64 = 8 * 1024 * 1024 + defaultUploadConcurrency = 4 + + defaultMultipartCopyCutoff int64 = 128 * 1024 * 1024 + defaultMultipartCopyPartSize int64 = 64 * 1024 * 1024 +) + +type Store struct { + bucket string + prefix string + client *s3.Client + // uploader has been depreciated, but its replacement, transfermanager, does not currently support copy + // rather than having two different dependencies, for now we are just using the sdk until the replacement + // can support all the features we need. + uploader *manager.Uploader + partSize int64 + copyCutoff int64 +} + +func NewS3Store(bucket, prefix string, client *s3.Client) *Store { + // Used to help break up large uploads for Put and Copy + uploader := manager.NewUploader(client, func(s *manager.Uploader) { + s.PartSize = defaultUploadPartSize + s.Concurrency = defaultUploadConcurrency + }) + + return &Store{ + bucket: bucket, + prefix: strings.Trim(prefix, "/"), + client: client, + uploader: uploader, + partSize: defaultMultipartCopyPartSize, + copyCutoff: defaultMultipartCopyCutoff, + } +} + +func normalizePath(name string) (string, error) { + name = strings.TrimSpace(name) + name = strings.ReplaceAll(name, "\\", "/") + name = strings.TrimLeft(name, "/") + + if name == "" { + return "", errors.New("invalid empty path") + } + + // recommendation of avoiding period-only path segments + for _, segment := range strings.Split(name, "/") { + if segment == "." || segment == ".." { + return "", errors.New("path contains relative segment") + } + } + + cleaned := path.Clean(name) + if cleaned == "." || cleaned == "" || strings.HasPrefix(cleaned, "/") { + return "", errors.New("invalid empty path") + } + + return cleaned, nil +} + +func s3DetectContentType(name string) string { + ext := path.Ext(name) + if ext == "" { + return mediatypes.ApplicationOctetStream.String() + } + + contentType := mime.TypeByExtension(ext) + if contentType == "" { + return mediatypes.ApplicationOctetStream.String() + } + + return contentType +} + +func mapExistsError(err error) error { + if err == nil { + return nil + } + + var httpErr *smithyhttp.ResponseError + if errors.As(err, &httpErr) && httpErr.HTTPStatusCode() == http.StatusPreconditionFailed { + return fmt.Errorf("%w: %w", fs.ErrExist, err) + } + return err +} + +func mapNotFoundError(err error) error { + if err == nil { + return nil + } + + var apiErr smithy.APIError + if errors.As(err, &apiErr) { + switch apiErr.ErrorCode() { + case "NoSuchKey", "NotFound": + return fmt.Errorf("%w: %w", os.ErrNotExist, err) + } + } + + var httpErr *smithyhttp.ResponseError + if errors.As(err, &httpErr) && httpErr.HTTPStatusCode() == http.StatusNotFound { + return fmt.Errorf("%w: %w", os.ErrNotExist, err) + } + + return err +} + +func (s *Store) key(name string) (string, error) { + normalizedPath, err := normalizePath(name) + if err != nil { + return "", err + } + + if s.prefix == "" { + return normalizedPath, nil + } + return s.prefix + "/" + normalizedPath, nil +} + +func (s *Store) Put(ctx context.Context, name string, reader io.Reader, options WriteOptions) error { + key, err := s.key(name) + if err != nil { + return err + } + + input := &s3.PutObjectInput{ + Bucket: aws.String(s.bucket), + Key: aws.String(key), + Body: reader, + } + + if options.FailIfExists { + input.IfNoneMatch = aws.String("*") + } + + if options.ContentType != "" { + input.ContentType = aws.String(options.ContentType) + } + + if len(options.Metadata) > 0 { + input.Metadata = options.Metadata + } + + _, err = s.uploader.Upload(ctx, input) + return mapExistsError(err) +} + +func (s *Store) Get(ctx context.Context, name string) (io.ReadCloser, FileInfo, error) { + key, err := s.key(name) + if err != nil { + return nil, FileInfo{}, err + } + + normalizedPath, err := normalizePath(name) + if err != nil { + return nil, FileInfo{}, err + } + out, err := s.client.GetObject(ctx, &s3.GetObjectInput{ + Bucket: aws.String(s.bucket), + Key: aws.String(key), + }) + if err != nil { + return nil, FileInfo{}, mapNotFoundError(err) + } + + info := FileInfo{ + Path: normalizedPath, + Size: aws.ToInt64(out.ContentLength), + ContentType: aws.ToString(out.ContentType), + ETag: aws.ToString(out.ETag), + LastModified: aws.ToTime(out.LastModified), + } + // Set in case S3 does not provide a content type + if info.ContentType == "" { + info.ContentType = s3DetectContentType(normalizedPath) + } + + return out.Body, info, nil +} + +func (s *Store) Stat(ctx context.Context, name string) (FileInfo, error) { + key, err := s.key(name) + if err != nil { + return FileInfo{}, err + } + out, err := s.client.HeadObject(ctx, &s3.HeadObjectInput{ + Bucket: aws.String(s.bucket), + Key: aws.String(key), + }) + if err != nil { + return FileInfo{}, mapNotFoundError(err) + } + + normalizedPath, err := normalizePath(name) + if err != nil { + return FileInfo{}, err + } + + info := FileInfo{ + Path: normalizedPath, + Size: aws.ToInt64(out.ContentLength), + ContentType: aws.ToString(out.ContentType), + ETag: aws.ToString(out.ETag), + LastModified: aws.ToTime(out.LastModified), + } + + // Set in case S3 does not provide a content type + if info.ContentType == "" { + info.ContentType = s3DetectContentType(normalizedPath) + } + + return info, nil +} + +func (s *Store) Exists(ctx context.Context, name string) (bool, error) { + _, err := s.Stat(ctx, name) + if err == nil { + return true, nil + } + if errors.Is(err, os.ErrNotExist) { + return false, nil + } + return false, err +} + +func (s *Store) Delete(ctx context.Context, name string) error { + key, err := s.key(name) + if err != nil { + return err + } + _, err = s.client.DeleteObject(ctx, &s3.DeleteObjectInput{ + Bucket: aws.String(s.bucket), + Key: aws.String(key), + }) + var noSuchKey *s3types.NoSuchKey + if errors.As(err, &noSuchKey) { + return nil + } + return err +} + +func isRootPath(name string) bool { + name = strings.TrimSpace(name) + return name == "" || name == "/" +} + +func stripPrefix(prefix, key string) string { + if prefix == "" { + return key + } + return strings.TrimPrefix(key, prefix+"/") +} + +func (s *Store) logicalPathFromKey(key string) string { + return stripPrefix(s.prefix, key) +} + +func (s *Store) listPrefix(name string) (string, error) { + key, err := s.key(name) + if err != nil { + return "", err + } + + if key == "" { + return "", nil + } + + return strings.TrimSuffix(key, "/") + "/", nil +} + +func (s *Store) List(ctx context.Context, name string, options ListOptions) ([]FileInfo, error) { + var ( + results []FileInfo + normalizedName string + keyPrefix string + err error + ) + + // This is done so items at the root path can be listed + if isRootPath(name) { + if s.prefix != "" { + keyPrefix = s.prefix + "/" + } else { + keyPrefix = "" + } + normalizedName = "" + } else { + normalizedName, err = normalizePath(name) + if err != nil { + return nil, err + } + + keyPrefix, err = s.listPrefix(normalizedName) + if err != nil { + return nil, err + } + } + + input := &s3.ListObjectsV2Input{ + Bucket: aws.String(s.bucket), + Prefix: aws.String(keyPrefix), + } + + if !options.Recursive { + input.Delimiter = aws.String("/") + } + + paginator := s3.NewListObjectsV2Paginator(s.client, input) + + for paginator.HasMorePages() { + page, err := paginator.NextPage(ctx) + if err != nil { + return nil, err + } + + for _, obj := range page.Contents { + if err = ctx.Err(); err != nil { + return []FileInfo{}, err + } + if options.Limit > 0 && len(results) >= options.Limit { + return results, nil + } + logicalPath := s.logicalPathFromKey(aws.ToString(obj.Key)) + if logicalPath == normalizedName { + continue + } + results = append(results, FileInfo{ + Path: logicalPath, + ContentType: s3DetectContentType(logicalPath), + Size: aws.ToInt64(obj.Size), + ETag: aws.ToString(obj.ETag), + LastModified: aws.ToTime(obj.LastModified), + }) + } + } + + return results, nil +} + +func (s *Store) copySmall(ctx context.Context, srcKey, dstKey string, options WriteOptions) error { + input := &s3.CopyObjectInput{ + Bucket: aws.String(s.bucket), + CopySource: aws.String(url.PathEscape(s.bucket + "/" + srcKey)), + Key: aws.String(dstKey), + } + + if options.FailIfExists { + input.IfNoneMatch = aws.String("*") + } + _, err := s.client.CopyObject(ctx, input) + return mapExistsError(err) +} + +func (s *Store) copyMultipart(ctx context.Context, srcKey, dstKey string, sourceSize int64, options WriteOptions) error { + createInput := &s3.CreateMultipartUploadInput{ + Bucket: aws.String(s.bucket), + Key: aws.String(dstKey), + } + + // Cannot safely emulate FailIfExists, must do an explicit pre-check, which may lead to a race condition + if options.FailIfExists { + exists, err := s.Exists(ctx, s.logicalPathFromKey(dstKey)) + if err != nil { + return err + } + if exists { + return fmt.Errorf("destination already exists: %w", fs.ErrExist) + } + } + + createOutput, err := s.client.CreateMultipartUpload(ctx, createInput) + if err != nil { + return err + } + + completedParts := make([]s3types.CompletedPart, 0) + uploadID := aws.ToString(createOutput.UploadId) + + defer func() { + if err != nil { + _, _ = s.client.AbortMultipartUpload(ctx, &s3.AbortMultipartUploadInput{ + Bucket: aws.String(s.bucket), + Key: aws.String(dstKey), + UploadId: aws.String(uploadID), + }) + } + }() + + var ( + partNumber int32 = 1 + start int64 = 0 + ) + + for start < sourceSize { + end := start + s.partSize - 1 + if end >= sourceSize { + end = sourceSize - 1 + } + + rangeHeader := fmt.Sprintf("bytes=%d-%d", start, end) + + partOutput, uploadErr := s.client.UploadPartCopy(ctx, &s3.UploadPartCopyInput{ + Bucket: aws.String(s.bucket), + Key: aws.String(dstKey), + UploadId: aws.String(uploadID), + PartNumber: aws.Int32(partNumber), + CopySource: aws.String(url.PathEscape(s.bucket + "/" + srcKey)), + CopySourceRange: aws.String(rangeHeader), + }) + if uploadErr != nil { + err = uploadErr + return err + } + + completedParts = append(completedParts, s3types.CompletedPart{ + ETag: partOutput.CopyPartResult.ETag, + PartNumber: aws.Int32(partNumber), + }) + + partNumber++ + start = end + 1 + } + + sort.Slice(completedParts, func(left, right int) bool { + return aws.ToInt32(completedParts[left].PartNumber) < aws.ToInt32(completedParts[right].PartNumber) + }) + + _, err = s.client.CompleteMultipartUpload(ctx, &s3.CompleteMultipartUploadInput{ + Bucket: aws.String(s.bucket), + Key: aws.String(dstKey), + UploadId: aws.String(uploadID), + MultipartUpload: &s3types.CompletedMultipartUpload{ + Parts: completedParts, + }, + }) + return err +} + +// Copy currently honors FailIfExists for the destination object. +// Metadata and ContentType are preserved from the source object. +// If a large copy is done (> copyCutoff) we use multipart copy, which cannot +// guarentee FailIfExists due to a possible race condition. This function does +// best effort by checking to see if the file exists before the copy starts. +func (s *Store) Copy(ctx context.Context, srcName, dstName string, options WriteOptions) error { + srcKey, err := s.key(srcName) + if err != nil { + return err + } + + dstKey, err := s.key(dstName) + if err != nil { + return err + } + + head, err := s.client.HeadObject(ctx, &s3.HeadObjectInput{ + Bucket: aws.String(s.bucket), + Key: aws.String(srcKey), + }) + if err != nil { + return err + } + + sourceSize := aws.ToInt64(head.ContentLength) + if sourceSize < s.copyCutoff { + return s.copySmall(ctx, srcKey, dstKey, options) + } + + return s.copyMultipart(ctx, srcKey, dstKey, sourceSize, options) + +} + +func (s *Store) Move(ctx context.Context, srcPath, dstPath string, options WriteOptions) error { + if err := s.Copy(ctx, srcPath, dstPath, options); err != nil { + return err + } + return s.Delete(ctx, srcPath) +} diff --git a/packages/go/storage/s3_test.go b/packages/go/storage/s3_test.go new file mode 100644 index 000000000000..9a71e9ff9aa5 --- /dev/null +++ b/packages/go/storage/s3_test.go @@ -0,0 +1,754 @@ +// Copyright 2026 Specter Ops, Inc. +// +// Licensed under the Apache License, Version 2.0 +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 + +package storage + +import ( + "context" + "errors" + "fmt" + "io" + "io/fs" + "net/http" + "os" + "strings" + "testing" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/credentials" + "github.com/aws/aws-sdk-go-v2/service/s3" + smithyhttp "github.com/aws/smithy-go/transport/http" + "github.com/stretchr/testify/require" +) + +type testHTTPResponse struct { + statusCode int + headers map[string]string + body string +} + +type testHTTPClient struct { + requests []*http.Request + responses []testHTTPResponse +} + +func (s *testHTTPClient) Do(request *http.Request) (*http.Response, error) { + if len(s.responses) == 0 { + return nil, fmt.Errorf("unexpected request: %s %s", request.Method, request.URL.String()) + } + + response := s.responses[0] + s.responses = s.responses[1:] + s.requests = append(s.requests, request) + + headers := http.Header{} + for key, value := range response.headers { + headers.Set(key, value) + } + + return &http.Response{ + StatusCode: response.statusCode, + Header: headers, + Body: io.NopCloser(strings.NewReader(response.body)), + ContentLength: int64(len(response.body)), + Request: request, + }, nil +} + +func newTestStore(httpClient *testHTTPClient) *Store { + cfg := aws.Config{ + Region: "us-east-1", + Credentials: credentials.NewStaticCredentialsProvider( + "test-access-key", + "test-secret-key", + "", + ), + HTTPClient: httpClient, + Retryer: func() aws.Retryer { + return aws.NopRetryer{} + }, + } + + client := s3.NewFromConfig(cfg, func(options *s3.Options) { + options.BaseEndpoint = aws.String("https://s3.test") + options.UsePathStyle = true + }) + + return NewS3Store("test-bucket", "prefix", client) +} + +func TestNormalizePath(t *testing.T) { + t.Parallel() + + type expected struct { + path string + error bool + } + + type testData struct { + name string + input string + expected expected + } + + testCases := []testData{ + { + name: "trims whitespace", + input: " file.json ", + expected: expected{ + path: "file.json", + }, + }, + { + name: "removes leading slash", + input: "/dir/file.json", + expected: expected{ + path: "dir/file.json", + }, + }, + { + name: "normalizes windows separators", + input: `dir\file.json`, + expected: expected{ + path: "dir/file.json", + }, + }, + { + name: "rejects empty path", + input: " ", + expected: expected{ + error: true, + }, + }, + { + name: "rejects escaping path", + input: "../file.json", + expected: expected{ + error: true, + }, + }, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + actualPath, err := normalizePath(testCase.input) + if testCase.expected.error { + require.Error(t, err) + require.Empty(t, actualPath) + return + } + + require.NoError(t, err) + require.Equal(t, testCase.expected.path, actualPath) + }) + } +} + +func TestDetectContentType(t *testing.T) { + t.Parallel() + + type testData struct { + name string + input string + expected string + } + + testCases := []testData{ + { + name: "json extension", + input: "file.json", + expected: "application/json", + }, + { + name: "unknown extension", + input: "file.unknownext", + expected: "application/octet-stream", + }, + { + name: "no extension", + input: "file", + expected: "application/octet-stream", + }, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + require.Equal(t, testCase.expected, detectContentType(testCase.input)) + }) + } +} + +func TestMapExistsError(t *testing.T) { + t.Parallel() + + errWrapped := errors.New("wrapped") + + type expected struct { + errIs error + err error + } + + type testData struct { + name string + input error + expected expected + } + + testCases := []testData{ + { + name: "nil error returns nil", + expected: expected{ + err: nil, + }, + }, + { + name: "precondition failure maps to exists", + input: &smithyhttp.ResponseError{ + Response: &smithyhttp.Response{Response: &http.Response{StatusCode: http.StatusPreconditionFailed}}, + Err: errWrapped, + }, + expected: expected{ + errIs: fs.ErrExist, + }, + }, + { + name: "other status returns original error", + input: &smithyhttp.ResponseError{ + Response: &smithyhttp.Response{Response: &http.Response{StatusCode: http.StatusInternalServerError}}, + Err: errWrapped, + }, + expected: expected{ + errIs: errWrapped, + }, + }, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + err := mapExistsError(testCase.input) + if testCase.expected.errIs != nil { + require.ErrorIs(t, err, testCase.expected.errIs) + return + } + + require.Equal(t, testCase.expected.err, err) + }) + } +} + +func TestStore_PathHelpers(t *testing.T) { + t.Parallel() + + type expected struct { + key string + listPrefix string + logicalPath string + rootPath bool + strippedPath string + keyErr bool + listPrefixErr bool + } + + type testData struct { + name string + storePrefix string + input string + keyInput string + stripPrefix string + stripKey string + logicalKey string + expected expected + } + + testCases := []testData{ + { + name: "prefix is applied to keys", + storePrefix: "prefix", + input: "dir/file.json", + keyInput: "dir/file.json", + stripPrefix: "prefix", + stripKey: "prefix/dir/file.json", + logicalKey: "prefix/dir/file.json", + expected: expected{ + key: "prefix/dir/file.json", + listPrefix: "prefix/dir/file.json/", + logicalPath: "dir/file.json", + strippedPath: "dir/file.json", + }, + }, + { + name: "empty prefix uses normalized key", + storePrefix: "", + input: "/dir/file.json", + keyInput: "/dir/file.json", + stripKey: "dir/file.json", + logicalKey: "dir/file.json", + expected: expected{ + key: "dir/file.json", + listPrefix: "dir/file.json/", + logicalPath: "dir/file.json", + strippedPath: "dir/file.json", + }, + }, + { + name: "root path is detected", + storePrefix: "prefix", + input: "/", + keyInput: "/", + expected: expected{ + rootPath: true, + keyErr: true, + listPrefixErr: true, + }, + }, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + store := &Store{prefix: testCase.storePrefix} + + actualKey, keyErr := store.key(testCase.keyInput) + if testCase.expected.keyErr { + require.Error(t, keyErr) + } else { + require.NoError(t, keyErr) + require.Equal(t, testCase.expected.key, actualKey) + } + + actualListPrefix, listPrefixErr := store.listPrefix(testCase.input) + if testCase.expected.listPrefixErr { + require.Error(t, listPrefixErr) + } else { + require.NoError(t, listPrefixErr) + require.Equal(t, testCase.expected.listPrefix, actualListPrefix) + } + + require.Equal(t, testCase.expected.logicalPath, store.logicalPathFromKey(testCase.logicalKey)) + require.Equal(t, testCase.expected.rootPath, isRootPath(testCase.input)) + require.Equal(t, testCase.expected.strippedPath, stripPrefix(testCase.stripPrefix, testCase.stripKey)) + }) + } +} + +func TestStore_Stat(t *testing.T) { + t.Parallel() + + lastModified := time.Date(2026, time.May, 20, 12, 30, 0, 0, time.UTC) + + type expected struct { + errIs error + fileInfo FileInfo + request string + } + + type testData struct { + name string + fileName string + responses []testHTTPResponse + expected expected + } + + testCases := []testData{ + { + name: "returns object metadata", + fileName: "dir/file.json", + responses: []testHTTPResponse{ + { + statusCode: http.StatusOK, + headers: map[string]string{ + "Content-Length": "12", + "Content-Type": "application/custom-json", + "ETag": `"etag-value"`, + "Last-Modified": lastModified.Format(http.TimeFormat), + }, + }, + }, + expected: expected{ + fileInfo: FileInfo{ + Path: "dir/file.json", + Size: 12, + ContentType: "application/custom-json", + ETag: `"etag-value"`, + LastModified: lastModified, + }, + request: "/test-bucket/prefix/dir/file.json", + }, + }, + { + name: "falls back to extension content type", + fileName: "dir/file.json", + responses: []testHTTPResponse{ + { + statusCode: http.StatusOK, + headers: map[string]string{ + "Content-Length": "12", + "ETag": `"etag-value"`, + "Last-Modified": lastModified.Format(http.TimeFormat), + }, + }, + }, + expected: expected{ + fileInfo: FileInfo{ + Path: "dir/file.json", + Size: 12, + ContentType: "application/json", + ETag: `"etag-value"`, + LastModified: lastModified, + }, + request: "/test-bucket/prefix/dir/file.json", + }, + }, + { + name: "not found error maps to os not exist", + fileName: "missing.json", + responses: []testHTTPResponse{ + { + statusCode: http.StatusNotFound, + }, + }, + expected: expected{ + errIs: os.ErrNotExist, + request: "/test-bucket/prefix/missing.json", + }, + }, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + httpClient := &testHTTPClient{responses: testCase.responses} + store := newTestStore(httpClient) + + actualFileInfo, err := store.Stat(context.Background(), testCase.fileName) + if testCase.expected.errIs != nil { + require.ErrorIs(t, err, testCase.expected.errIs) + require.Empty(t, actualFileInfo) + } else { + require.NoError(t, err) + require.Equal(t, testCase.expected.fileInfo, actualFileInfo) + } + + require.Len(t, httpClient.requests, 1) + require.Equal(t, http.MethodHead, httpClient.requests[0].Method) + require.Equal(t, testCase.expected.request, httpClient.requests[0].URL.Path) + }) + } +} + +func TestStore_Exists(t *testing.T) { + t.Parallel() + + type expected struct { + exists bool + responseErrorAs bool + } + + type testData struct { + name string + fileName string + responses []testHTTPResponse + expected expected + } + + testCases := []testData{ + { + name: "head success returns true", + fileName: "file.json", + responses: []testHTTPResponse{ + { + statusCode: http.StatusOK, + headers: map[string]string{ + "Content-Length": "12", + }, + }, + }, + expected: expected{ + exists: true, + }, + }, + { + name: "head not found status returns false", + fileName: "missing.json", + responses: []testHTTPResponse{ + { + statusCode: http.StatusNotFound, + }, + }, + expected: expected{ + exists: false, + }, + }, + { + name: "head no such key error returns false", + fileName: "missing.json", + responses: []testHTTPResponse{ + { + statusCode: http.StatusNotFound, + body: `NoSuchKeymissing`, + }, + }, + expected: expected{ + exists: false, + }, + }, + { + name: "service error is returned", + fileName: "file.json", + responses: []testHTTPResponse{ + { + statusCode: http.StatusInternalServerError, + }, + }, + expected: expected{ + responseErrorAs: true, + }, + }, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + httpClient := &testHTTPClient{responses: testCase.responses} + store := newTestStore(httpClient) + + actualExists, err := store.Exists(context.Background(), testCase.fileName) + if testCase.expected.responseErrorAs { + var responseError *smithyhttp.ResponseError + require.ErrorAs(t, err, &responseError) + require.False(t, actualExists) + return + } + + require.NoError(t, err) + require.Equal(t, testCase.expected.exists, actualExists) + }) + } +} + +func TestStore_Get(t *testing.T) { + t.Parallel() + + lastModified := time.Date(2026, time.May, 20, 12, 30, 0, 0, time.UTC) + + type expected struct { + errIs error + fileInfo FileInfo + content string + } + + type testData struct { + name string + fileName string + responses []testHTTPResponse + expected expected + } + + testCases := []testData{ + { + name: "returns object content and metadata", + fileName: "dir/file.json", + responses: []testHTTPResponse{ + { + statusCode: http.StatusOK, + headers: map[string]string{ + "Content-Length": "7", + "Content-Type": "application/json", + "ETag": `"etag-value"`, + "Last-Modified": lastModified.Format(http.TimeFormat), + }, + body: `{"x":1}`, + }, + }, + expected: expected{ + fileInfo: FileInfo{ + Path: "dir/file.json", + Size: 7, + ContentType: "application/json", + ETag: `"etag-value"`, + LastModified: lastModified, + }, + content: `{"x":1}`, + }, + }, + { + name: "not found maps to os not exist", + fileName: "missing.json", + responses: []testHTTPResponse{ + { + statusCode: http.StatusNotFound, + body: `NoSuchKeymissing`, + }, + }, + expected: expected{ + errIs: os.ErrNotExist, + }, + }, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + httpClient := &testHTTPClient{responses: testCase.responses} + store := newTestStore(httpClient) + + readCloser, actualFileInfo, err := store.Get(context.Background(), testCase.fileName) + if testCase.expected.errIs != nil { + require.ErrorIs(t, err, testCase.expected.errIs) + require.Nil(t, readCloser) + require.Empty(t, actualFileInfo) + return + } + defer readCloser.Close() + + actualContent, err := io.ReadAll(readCloser) + require.NoError(t, err) + require.Equal(t, testCase.expected.fileInfo, actualFileInfo) + require.Equal(t, testCase.expected.content, string(actualContent)) + }) + } +} + +func TestStore_List(t *testing.T) { + t.Parallel() + + lastModified := time.Date(2026, time.May, 20, 12, 30, 0, 0, time.UTC) + listBody := fmt.Sprintf(` + + test-bucket + prefix/dir/ + 2 + + prefix/dir/file-a.json + %s + "etag-a" + 10 + STANDARD + + + prefix/dir/file-b.zip + %s + "etag-b" + 20 + STANDARD + +`, lastModified.Format(time.RFC3339), lastModified.Format(time.RFC3339)) + + type expected struct { + files []FileInfo + query map[string]string + } + + type testData struct { + name string + path string + options ListOptions + expected expected + } + + testCases := []testData{ + { + name: "non recursive list sets delimiter", + path: "dir", + expected: expected{ + files: []FileInfo{ + { + Path: "dir/file-a.json", + ContentType: "application/json", + Size: 10, + ETag: `"etag-a"`, + LastModified: lastModified, + }, + { + Path: "dir/file-b.zip", + ContentType: "application/zip", + Size: 20, + ETag: `"etag-b"`, + LastModified: lastModified, + }, + }, + query: map[string]string{ + "delimiter": "/", + "prefix": "prefix/dir/", + }, + }, + }, + { + name: "recursive list omits delimiter", + path: "dir", + options: ListOptions{ + Recursive: true, + }, + expected: expected{ + files: []FileInfo{ + { + Path: "dir/file-a.json", + ContentType: "application/json", + Size: 10, + ETag: `"etag-a"`, + LastModified: lastModified, + }, + { + Path: "dir/file-b.zip", + ContentType: "application/zip", + Size: 20, + ETag: `"etag-b"`, + LastModified: lastModified, + }, + }, + query: map[string]string{ + "prefix": "prefix/dir/", + }, + }, + }, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + httpClient := &testHTTPClient{responses: []testHTTPResponse{ + { + statusCode: http.StatusOK, + body: listBody, + }, + }} + store := newTestStore(httpClient) + + actualFiles, err := store.List(context.Background(), testCase.path, testCase.options) + require.NoError(t, err) + require.Equal(t, testCase.expected.files, actualFiles) + require.Len(t, httpClient.requests, 1) + require.Equal(t, http.MethodGet, httpClient.requests[0].Method) + + query := httpClient.requests[0].URL.Query() + require.Equal(t, "2", query.Get("list-type")) + require.Equal(t, testCase.expected.query["prefix"], query.Get("prefix")) + require.Equal(t, testCase.expected.query["delimiter"], query.Get("delimiter")) + }) + } +} diff --git a/packages/go/storage/storage.go b/packages/go/storage/storage.go new file mode 100644 index 000000000000..4937193d693f --- /dev/null +++ b/packages/go/storage/storage.go @@ -0,0 +1,230 @@ +// Copyright 2025 Specter Ops, Inc. +// +// Licensed under the Apache License, Version 2.0 +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 + +package storage + +import ( + "bytes" + "context" + "crypto/rand" + "encoding/hex" + "errors" + "io" + "time" +) + +//go:generate go run go.uber.org/mock/mockgen -copyright_file=../../../LICENSE.header -destination=./mocks/storage.go -package=mocks . Storage,FileService + +type FileServiceName string + +const ( + FileServiceIngest FileServiceName = "ingest" + FileServiceRetained FileServiceName = "retained" + FileServiceCollectors FileServiceName = "collectors" + FileServiceJobLogs FileServiceName = "job_logs" + FileServiceWork FileServiceName = "work" + FileServiceCollectorArtifacts FileServiceName = "collector_artifacts" +) + +var ErrFileServiceNotFound = errors.New("file service not found") + +type FileInfo struct { + Path string + Size int64 + ContentType string + ETag string + LastModified time.Time + IsDir bool +} + +type WriteOptions struct { + ContentType string + Metadata map[string]string + // FailIfExists causes write to return an error wrapping fs.ErrExist when the + // destination already exists, instead of silently replacing it. + FailIfExists bool +} + +type ListOptions struct { + Recursive bool + Limit int +} + +// Storage serves as a storage abstraction that can be used to store and manage files +// in a variety of storage backends. +type Storage interface { + // Put writes a file at the given path. + Put(ctx context.Context, name string, reader io.Reader, options WriteOptions) error + + // Get opens a file for reading. + Get(ctx context.Context, name string) (io.ReadCloser, FileInfo, error) + + // Stat returns metadata for a given path. + Stat(ctx context.Context, name string) (FileInfo, error) + + // Delete removes a file. + Delete(ctx context.Context, name string) error + + // Exists checks whether a file exists. + Exists(ctx context.Context, name string) (bool, error) + + // List returns a list of files in a given path. + List(ctx context.Context, name string, options ListOptions) ([]FileInfo, error) + + // Copy duplicates an object. + Copy(ctx context.Context, srcName, dstName string, options WriteOptions) error + + // Move moves an object. Is done by a copy and a delete. + Move(ctx context.Context, srcName, dstName string, options WriteOptions) error +} + +// FileService serves as an abstraction to handle files with different storage backends. This serves as +// a list of general functions that each file service must implement. +type FileService interface { + // GetFile returns a io.ReadCloser and FileInfo for the named file that is requested. + GetFile(ctx context.Context, name string) (io.ReadCloser, FileInfo, error) + + // GetFileInfo returns just the FileInfo of the file requested. + GetFileInfo(ctx context.Context, name string) (FileInfo, error) + + // ReadFile returns the byte information of the file that is requested. + ReadFile(ctx context.Context, name string) ([]byte, error) + + // WriteFile takes the name and byte information as well as WriteOptions to write to the + // storage backend. + WriteFile(ctx context.Context, name string, data []byte, opts WriteOptions) error + + // WriteFileFromReader takes the name, io.Reader, and WriteOptions to write to the + // storage backend. + WriteFileFromReader(ctx context.Context, name string, reader io.Reader, opts WriteOptions) error + + // DeleteFile deletes a file at a specific name from the storage backend. If the file + // is not found, no error is returned. + DeleteFile(ctx context.Context, name string) error + + // WriteTempFile handles the creation of a temp file when given an io.Reader. A prefix + // can also be used to define how the temp file is created. WriteOptions can also be + // specified. + WriteTempFile(ctx context.Context, prefix string, reader io.Reader, opts WriteOptions) (string, error) + + // MoveFile takes a srcName and dstName to move a file from one location to another on + // the storage backend. WriteOptions can also be specified in the case of collisions. + MoveFile(ctx context.Context, srcName, dstName string, opts WriteOptions) error + + // ListFiles lists the files at a given location in the storage backend. This can be done + // recursively, or with a limit on the specified directory. + ListFiles(ctx context.Context, name string, opts ListOptions) ([]FileInfo, error) +} + +type StorageFileService struct { + Storage Storage +} + +func NewFileService(storage Storage) *StorageFileService { + return &StorageFileService{Storage: storage} +} + +func randomID() (string, error) { + var b [16]byte + if _, err := rand.Read(b[:]); err != nil { + return "", err + } + return hex.EncodeToString(b[:]), nil +} + +func (s *StorageFileService) GetFile(ctx context.Context, name string) (io.ReadCloser, FileInfo, error) { + return s.Storage.Get(ctx, name) +} + +func (s *StorageFileService) GetFileInfo(ctx context.Context, name string) (FileInfo, error) { + return s.Storage.Stat(ctx, name) +} + +func (s *StorageFileService) ReadFile(ctx context.Context, name string) ([]byte, error) { + rc, _, err := s.Storage.Get(ctx, name) + if err != nil { + return nil, err + } + data, readErr := io.ReadAll(rc) + if closeErr := rc.Close(); closeErr != nil { + if readErr != nil { + return nil, errors.Join(readErr, closeErr) + } + return nil, closeErr + } + return data, readErr +} + +func (s *StorageFileService) WriteFile(ctx context.Context, name string, data []byte, opts WriteOptions) error { + return s.Storage.Put(ctx, name, bytes.NewReader(data), opts) +} + +func (s *StorageFileService) WriteFileFromReader(ctx context.Context, name string, reader io.Reader, opts WriteOptions) error { + return s.Storage.Put(ctx, name, reader, opts) +} + +func (s *StorageFileService) DeleteFile(ctx context.Context, name string) error { + return s.Storage.Delete(ctx, name) +} + +func (s *StorageFileService) WriteTempFile(ctx context.Context, prefix string, reader io.Reader, opts WriteOptions) (string, error) { + id, err := randomID() + if err != nil { + return "", err + } + + tempPath := prefix + "tmp-" + id + if err := s.Storage.Put(ctx, tempPath, reader, opts); err != nil { + return "", err + } + + return tempPath, nil +} + +func (s *StorageFileService) MoveFile(ctx context.Context, srcName, dstName string, options WriteOptions) error { + return s.Storage.Move(ctx, srcName, dstName, options) +} + +func (s *StorageFileService) ListFiles(ctx context.Context, name string, options ListOptions) ([]FileInfo, error) { + return s.Storage.List(ctx, name, options) +} + +func MoveFileBetweenServices( + ctx context.Context, + sourceService FileService, + destinationService FileService, + sourceName string, + destinationName string, + opts WriteOptions, +) error { + sourceFile, _, err := sourceService.GetFile(ctx, sourceName) + if err != nil { + return err + } + + if err := destinationService.WriteFileFromReader(ctx, destinationName, sourceFile, opts); err != nil { + if closeErr := sourceFile.Close(); closeErr != nil { + return errors.Join(err, closeErr) + } + return err + } + + if err := sourceFile.Close(); err != nil { + return err + } + + return sourceService.DeleteFile(ctx, sourceName) +} diff --git a/packages/go/storage/storage_test.go b/packages/go/storage/storage_test.go new file mode 100644 index 000000000000..26c9ce649b03 --- /dev/null +++ b/packages/go/storage/storage_test.go @@ -0,0 +1,845 @@ +// Copyright 2026 Specter Ops, Inc. +// +// Licensed under the Apache License, Version 2.0 +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 + +package storage_test + +import ( + "context" + "errors" + "io" + "strings" + "testing" + + "github.com/specterops/bloodhound/packages/go/storage" + "github.com/specterops/bloodhound/packages/go/storage/mocks" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" +) + +type trackingReadCloser struct { + reader io.Reader + closeErr error + closed bool + closeCall int +} + +func (s *trackingReadCloser) Read(p []byte) (int, error) { + return s.reader.Read(p) +} + +func (s *trackingReadCloser) Close() error { + s.closed = true + s.closeCall++ + return s.closeErr +} + +type readErrorSource struct { + err error +} + +func (s readErrorSource) Read([]byte) (int, error) { + return 0, s.err +} + +func TestNewFileService(t *testing.T) { + t.Parallel() + + // Arrange + mockStorage := mocks.NewMockStorage(gomock.NewController(t)) + + // Act + fileService := storage.NewFileService(mockStorage) + + // Assert + require.Same(t, mockStorage, fileService.Storage) +} + +func TestStorageFileService_GetFile(t *testing.T) { + t.Parallel() + + var ( + errGet = errors.New("get failed") + fileInfo = storage.FileInfo{Path: "file.json", Size: 4, ContentType: "application/json"} + readCloser = &trackingReadCloser{reader: strings.NewReader("data")} + ) + + type expected struct { + errIs error + fileInfo storage.FileInfo + readCloser io.ReadCloser + } + + type testData struct { + name string + getErr error + expected expected + } + + tests := []testData{ + { + name: "gets file from storage", + expected: expected{ + fileInfo: fileInfo, + readCloser: readCloser, + }, + }, + { + name: "get error returns error", + getErr: errGet, + expected: expected{ + errIs: errGet, + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + var ( + ctx = context.Background() + mockStorage = mocks.NewMockStorage(gomock.NewController(t)) + fileService = storage.NewFileService(mockStorage) + ) + + mockStorage.EXPECT(). + Get(ctx, "file.json"). + Return(testCase.expected.readCloser, testCase.expected.fileInfo, testCase.getErr) + + // Act + actualReadCloser, actualFileInfo, err := fileService.GetFile(ctx, "file.json") + + // Assert + if testCase.expected.errIs != nil { + require.ErrorIs(t, err, testCase.expected.errIs) + require.Nil(t, actualReadCloser) + require.Empty(t, actualFileInfo) + return + } + + require.NoError(t, err) + require.Same(t, testCase.expected.readCloser, actualReadCloser) + require.Equal(t, testCase.expected.fileInfo, actualFileInfo) + }) + } +} + +func TestStorageFileService_ReadFile(t *testing.T) { + t.Parallel() + + var ( + errGet = errors.New("get failed") + errRead = errors.New("read failed") + ) + + type expected struct { + errIs error + content string + closed bool + } + + type testData struct { + name string + setupMock func(t *testing.T, ctx context.Context, mockStorage *mocks.MockStorage) *trackingReadCloser + expected expected + } + + tests := []testData{ + { + name: "reads file content and closes reader", + setupMock: func(t *testing.T, ctx context.Context, mockStorage *mocks.MockStorage) *trackingReadCloser { + readCloser := &trackingReadCloser{reader: strings.NewReader("content")} + mockStorage.EXPECT(). + Get(ctx, "file.json"). + Return(readCloser, storage.FileInfo{Path: "file.json"}, nil) + + return readCloser + }, + expected: expected{ + content: "content", + closed: true, + }, + }, + { + name: "get error returns error", + setupMock: func(t *testing.T, ctx context.Context, mockStorage *mocks.MockStorage) *trackingReadCloser { + mockStorage.EXPECT(). + Get(ctx, "file.json"). + Return(nil, storage.FileInfo{}, errGet) + + return nil + }, + expected: expected{ + errIs: errGet, + }, + }, + { + name: "reader error returns error and closes reader", + setupMock: func(t *testing.T, ctx context.Context, mockStorage *mocks.MockStorage) *trackingReadCloser { + readCloser := &trackingReadCloser{reader: readErrorSource{err: errRead}} + mockStorage.EXPECT(). + Get(ctx, "file.json"). + Return(readCloser, storage.FileInfo{Path: "file.json"}, nil) + + return readCloser + }, + expected: expected{ + errIs: errRead, + closed: true, + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + var ( + ctx = context.Background() + mockStorage = mocks.NewMockStorage(gomock.NewController(t)) + fileService = storage.NewFileService(mockStorage) + readCloser = testCase.setupMock(t, ctx, mockStorage) + ) + + // Act + content, err := fileService.ReadFile(ctx, "file.json") + + // Assert + if testCase.expected.errIs != nil { + require.ErrorIs(t, err, testCase.expected.errIs) + } else { + require.NoError(t, err) + } + + require.Equal(t, testCase.expected.content, string(content)) + if readCloser != nil { + require.Equal(t, testCase.expected.closed, readCloser.closed) + } + }) + } +} + +func TestStorageFileService_WriteFile(t *testing.T) { + t.Parallel() + + var errPut = errors.New("put failed") + + type expected struct { + errIs error + } + + type testData struct { + name string + putErr error + expected expected + } + + tests := []testData{ + { + name: "writes file content", + }, + { + name: "put error returns error", + putErr: errPut, + expected: expected{ + errIs: errPut, + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + var ( + ctx = context.Background() + options = storage.WriteOptions{ContentType: "application/json"} + mockStorage = mocks.NewMockStorage(gomock.NewController(t)) + fileService = storage.NewFileService(mockStorage) + ) + + mockStorage.EXPECT(). + Put(ctx, "file.json", gomock.Any(), options). + DoAndReturn(func(_ context.Context, _ string, reader io.Reader, _ storage.WriteOptions) error { + content, err := io.ReadAll(reader) + require.NoError(t, err) + require.Equal(t, `{"ok":true}`, string(content)) + + return testCase.putErr + }) + + // Act + err := fileService.WriteFile(ctx, "file.json", []byte(`{"ok":true}`), options) + + // Assert + if testCase.expected.errIs != nil { + require.ErrorIs(t, err, testCase.expected.errIs) + } else { + require.NoError(t, err) + } + }) + } +} + +func TestStorageFileService_WriteFileFromReader(t *testing.T) { + t.Parallel() + + var errPut = errors.New("put failed") + + type expected struct { + errIs error + } + + type testData struct { + name string + putErr error + expected expected + } + + tests := []testData{ + { + name: "writes reader", + }, + { + name: "put error returns error", + putErr: errPut, + expected: expected{ + errIs: errPut, + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + var ( + ctx = context.Background() + options = storage.WriteOptions{FailIfExists: true} + reader = strings.NewReader("content") + mockStorage = mocks.NewMockStorage(gomock.NewController(t)) + fileService = storage.NewFileService(mockStorage) + ) + + mockStorage.EXPECT(). + Put(ctx, "file.json", reader, options). + Return(testCase.putErr) + + // Act + err := fileService.WriteFileFromReader(ctx, "file.json", reader, options) + + // Assert + if testCase.expected.errIs != nil { + require.ErrorIs(t, err, testCase.expected.errIs) + } else { + require.NoError(t, err) + } + }) + } +} + +func TestStorageFileService_DeleteFile(t *testing.T) { + t.Parallel() + + var errDelete = errors.New("delete failed") + + type expected struct { + errIs error + } + + type testData struct { + name string + deleteErr error + expected expected + } + + tests := []testData{ + { + name: "deletes file", + }, + { + name: "delete error returns error", + deleteErr: errDelete, + expected: expected{ + errIs: errDelete, + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + var ( + ctx = context.Background() + mockStorage = mocks.NewMockStorage(gomock.NewController(t)) + fileService = storage.NewFileService(mockStorage) + ) + + mockStorage.EXPECT(). + Delete(ctx, "file.json"). + Return(testCase.deleteErr) + + // Act + err := fileService.DeleteFile(ctx, "file.json") + + // Assert + if testCase.expected.errIs != nil { + require.ErrorIs(t, err, testCase.expected.errIs) + } else { + require.NoError(t, err) + } + }) + } +} + +func TestStorageFileService_WriteTempFile(t *testing.T) { + t.Parallel() + + var errPut = errors.New("put failed") + + type expected struct { + errIs error + } + + type testData struct { + name string + putErr error + expected expected + } + + tests := []testData{ + { + name: "writes temp file and returns temp path", + }, + { + name: "put error returns error", + putErr: errPut, + expected: expected{ + errIs: errPut, + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + var ( + ctx = context.Background() + options = storage.WriteOptions{ContentType: "text/plain"} + mockStorage = mocks.NewMockStorage(gomock.NewController(t)) + fileService = storage.NewFileService(mockStorage) + writtenPath string + ) + + mockStorage.EXPECT(). + Put(ctx, gomock.Any(), gomock.Any(), options). + DoAndReturn(func(_ context.Context, name string, reader io.Reader, _ storage.WriteOptions) error { + content, err := io.ReadAll(reader) + require.NoError(t, err) + require.Equal(t, "content", string(content)) + require.Regexp(t, `^prefix_tmp-[0-9a-f]{32}$`, name) + + writtenPath = name + return testCase.putErr + }) + + // Act + tempPath, err := fileService.WriteTempFile(ctx, "prefix_", strings.NewReader("content"), options) + + // Assert + if testCase.expected.errIs != nil { + require.ErrorIs(t, err, testCase.expected.errIs) + require.Empty(t, tempPath) + } else { + require.NoError(t, err) + require.Equal(t, writtenPath, tempPath) + } + }) + } +} + +func TestStorageFileService_MoveFile(t *testing.T) { + t.Parallel() + + var errMove = errors.New("move failed") + + type expected struct { + errIs error + } + + type testData struct { + name string + moveErr error + expected expected + } + + tests := []testData{ + { + name: "moves file", + }, + { + name: "move error returns error", + moveErr: errMove, + expected: expected{ + errIs: errMove, + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + var ( + ctx = context.Background() + options = storage.WriteOptions{FailIfExists: true} + mockStorage = mocks.NewMockStorage(gomock.NewController(t)) + fileService = storage.NewFileService(mockStorage) + ) + + mockStorage.EXPECT(). + Move(ctx, "source.json", "destination.json", options). + Return(testCase.moveErr) + + // Act + err := fileService.MoveFile(ctx, "source.json", "destination.json", options) + + // Assert + if testCase.expected.errIs != nil { + require.ErrorIs(t, err, testCase.expected.errIs) + } else { + require.NoError(t, err) + } + }) + } +} + +func TestStorageFileService_ListFiles(t *testing.T) { + t.Parallel() + + var ( + errList = errors.New("list failed") + fileInfos = []storage.FileInfo{{Path: "file.json"}} + ) + + type expected struct { + errIs error + fileInfos []storage.FileInfo + } + + type testData struct { + name string + listErr error + expected expected + } + + tests := []testData{ + { + name: "lists files", + expected: expected{ + fileInfos: fileInfos, + }, + }, + { + name: "list error returns error", + listErr: errList, + expected: expected{ + errIs: errList, + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + var ( + ctx = context.Background() + options = storage.ListOptions{Recursive: true, Limit: 10} + mockStorage = mocks.NewMockStorage(gomock.NewController(t)) + fileService = storage.NewFileService(mockStorage) + ) + + mockStorage.EXPECT(). + List(ctx, "prefix", options). + Return(testCase.expected.fileInfos, testCase.listErr) + + // Act + actualFileInfos, err := fileService.ListFiles(ctx, "prefix", options) + + // Assert + if testCase.expected.errIs != nil { + require.ErrorIs(t, err, testCase.expected.errIs) + require.Nil(t, actualFileInfos) + return + } + + require.NoError(t, err) + require.Equal(t, testCase.expected.fileInfos, actualFileInfos) + }) + } +} + +func TestMoveFileBetweenServices(t *testing.T) { + t.Parallel() + + var ( + errGet = errors.New("get failed") + errWrite = errors.New("write failed") + errDelete = errors.New("delete failed") + errClose = errors.New("close failed") + ) + + type expected struct { + errIs error + additionalErrIs error + closed bool + } + + type testData struct { + name string + setupMock func( + t *testing.T, + ctx context.Context, + sourceService *mocks.MockFileService, + destinationService *mocks.MockFileService, + options storage.WriteOptions, + ) *trackingReadCloser + expected expected + } + + tests := []testData{ + { + name: "moves file between services", + setupMock: func( + t *testing.T, + ctx context.Context, + sourceService *mocks.MockFileService, + destinationService *mocks.MockFileService, + options storage.WriteOptions, + ) *trackingReadCloser { + readCloser := &trackingReadCloser{reader: strings.NewReader("content")} + + gomock.InOrder( + sourceService.EXPECT(). + GetFile(ctx, "source.json"). + Return(readCloser, storage.FileInfo{Path: "source.json"}, nil), + destinationService.EXPECT(). + WriteFileFromReader(ctx, "destination.json", readCloser, options). + DoAndReturn(func(_ context.Context, _ string, reader io.Reader, _ storage.WriteOptions) error { + content, err := io.ReadAll(reader) + require.NoError(t, err) + require.Equal(t, "content", string(content)) + + return nil + }), + sourceService.EXPECT(). + DeleteFile(ctx, "source.json"). + DoAndReturn(func(_ context.Context, _ string) error { + require.True(t, readCloser.closed) + require.Equal(t, 1, readCloser.closeCall) + return nil + }), + ) + + return readCloser + }, + expected: expected{ + closed: true, + }, + }, + { + name: "get error returns error", + setupMock: func( + t *testing.T, + ctx context.Context, + sourceService *mocks.MockFileService, + destinationService *mocks.MockFileService, + options storage.WriteOptions, + ) *trackingReadCloser { + sourceService.EXPECT(). + GetFile(ctx, "source.json"). + Return(nil, storage.FileInfo{}, errGet) + + return nil + }, + expected: expected{ + errIs: errGet, + }, + }, + { + name: "write error returns error and does not delete source", + setupMock: func( + t *testing.T, + ctx context.Context, + sourceService *mocks.MockFileService, + destinationService *mocks.MockFileService, + options storage.WriteOptions, + ) *trackingReadCloser { + readCloser := &trackingReadCloser{reader: strings.NewReader("content")} + + gomock.InOrder( + sourceService.EXPECT(). + GetFile(ctx, "source.json"). + Return(readCloser, storage.FileInfo{Path: "source.json"}, nil), + destinationService.EXPECT(). + WriteFileFromReader(ctx, "destination.json", readCloser, options). + Return(errWrite), + ) + + return readCloser + }, + expected: expected{ + errIs: errWrite, + closed: true, + }, + }, + { + name: "write and close errors returns both errors and does not delete source", + setupMock: func( + t *testing.T, + ctx context.Context, + sourceService *mocks.MockFileService, + destinationService *mocks.MockFileService, + options storage.WriteOptions, + ) *trackingReadCloser { + readCloser := &trackingReadCloser{ + reader: strings.NewReader("content"), + closeErr: errClose, + } + + gomock.InOrder( + sourceService.EXPECT(). + GetFile(ctx, "source.json"). + Return(readCloser, storage.FileInfo{Path: "source.json"}, nil), + destinationService.EXPECT(). + WriteFileFromReader(ctx, "destination.json", readCloser, options). + Return(errWrite), + ) + + return readCloser + }, + expected: expected{ + errIs: errWrite, + additionalErrIs: errClose, + closed: true, + }, + }, + { + name: "delete error returns error", + setupMock: func( + t *testing.T, + ctx context.Context, + sourceService *mocks.MockFileService, + destinationService *mocks.MockFileService, + options storage.WriteOptions, + ) *trackingReadCloser { + readCloser := &trackingReadCloser{reader: strings.NewReader("content")} + + gomock.InOrder( + sourceService.EXPECT(). + GetFile(ctx, "source.json"). + Return(readCloser, storage.FileInfo{Path: "source.json"}, nil), + destinationService.EXPECT(). + WriteFileFromReader(ctx, "destination.json", readCloser, options). + Return(nil), + sourceService.EXPECT(). + DeleteFile(ctx, "source.json"). + Return(errDelete), + ) + + return readCloser + }, + expected: expected{ + errIs: errDelete, + closed: true, + }, + }, + { + name: "close error returns error and does not delete source", + setupMock: func( + t *testing.T, + ctx context.Context, + sourceService *mocks.MockFileService, + destinationService *mocks.MockFileService, + options storage.WriteOptions, + ) *trackingReadCloser { + readCloser := &trackingReadCloser{ + reader: strings.NewReader("content"), + closeErr: errClose, + } + + gomock.InOrder( + sourceService.EXPECT(). + GetFile(ctx, "source.json"). + Return(readCloser, storage.FileInfo{Path: "source.json"}, nil), + destinationService.EXPECT(). + WriteFileFromReader(ctx, "destination.json", readCloser, options). + Return(nil), + ) + + return readCloser + }, + expected: expected{ + errIs: errClose, + closed: true, + }, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + // Arrange + var ( + ctx = context.Background() + options = storage.WriteOptions{FailIfExists: true} + controller = gomock.NewController(t) + sourceService = mocks.NewMockFileService(controller) + destinationService = mocks.NewMockFileService(controller) + readCloser = testCase.setupMock(t, ctx, sourceService, destinationService, options) + ) + + // Act + err := storage.MoveFileBetweenServices(ctx, sourceService, destinationService, "source.json", "destination.json", options) + + // Assert + if testCase.expected.errIs != nil { + require.ErrorIs(t, err, testCase.expected.errIs) + } else { + require.NoError(t, err) + } + if testCase.expected.additionalErrIs != nil { + require.ErrorIs(t, err, testCase.expected.additionalErrIs) + } + + if readCloser != nil { + require.Equal(t, testCase.expected.closed, readCloser.closed) + } + }) + } +}