diff --git a/cmd/api/src/api/registration/registration.go b/cmd/api/src/api/registration/registration.go index 62bc07922b89..93cb3ecf6667 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,6 +90,6 @@ func RegisterFossRoutes( routerInst.PathPrefix("/ui", 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..e22bcd4320d5 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/cmd/api/src/services/storage" "github.com/specterops/bloodhound/packages/go/bhlog/attr" "github.com/specterops/bloodhound/packages/go/headers" ) 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..eb44586a5a19 --- /dev/null +++ b/cmd/api/src/api/tools/ingest_test.go @@ -0,0 +1,554 @@ +// 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/cmd/api/src/services/storage" + storagemocks "github.com/specterops/bloodhound/cmd/api/src/services/storage/mocks" + "github.com/specterops/bloodhound/packages/go/headers" + "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.True(t, ingestControl.retainedFileLock.TryLock()) + ingestControl.retainedFileLock.Unlock() + }) + + 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.True(t, ingestControl.retainedFileLock.TryLock()) + ingestControl.retainedFileLock.Unlock() + }) +} + +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..427d38efdedc 100644 --- a/cmd/api/src/api/v2/collectors.go +++ b/cmd/api/src/api/v2/collectors.go @@ -26,6 +26,7 @@ import ( "github.com/gorilla/mux" "github.com/specterops/bloodhound/cmd/api/src/api" "github.com/specterops/bloodhound/cmd/api/src/config" + "github.com/specterops/bloodhound/cmd/api/src/services/storage" "github.com/specterops/bloodhound/packages/go/bhlog/attr" ) @@ -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..3059626f219d 100644 --- a/cmd/api/src/api/v2/collectors_test.go +++ b/cmd/api/src/api/v2/collectors_test.go @@ -28,7 +28,8 @@ 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" + "github.com/specterops/bloodhound/cmd/api/src/services/storage" + storagemocks "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" @@ -40,7 +41,8 @@ import ( func TestResources_DownloadCollectorByVersion(t *testing.T) { type mock struct { - mockFS *fsmocks.MockService + mockFS *storagemocks.MockFileService + mockFileServiceResolver *storagemocks.MockFileServiceResolver } type expected struct { responseBody string @@ -99,7 +101,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 +140,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 +159,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 +174,8 @@ func TestResources_DownloadCollectorByVersion(t *testing.T) { ctrl := gomock.NewController(t) mock := &mock{ - mockFS: fsmocks.NewMockService(ctrl), + mockFS: storagemocks.NewMockFileService(ctrl), + mockFileServiceResolver: storagemocks.NewMockFileServiceResolver(ctrl), } testCase.setupMocks(t, mock) @@ -162,8 +187,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 +212,8 @@ func TestResources_DownloadCollectorByVersion(t *testing.T) { func TestResources_DownloadCollectorChecksumByVersion(t *testing.T) { type mock struct { - mockFS *fsmocks.MockService + mockFS *storagemocks.MockFileService + mockFileServiceResolver *storagemocks.MockFileServiceResolver } type expected struct { responseBody string @@ -246,7 +272,27 @@ 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(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{ + Path: "/api/v2/collectors/azurehound/latest/checksum", + }, + 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, @@ -265,7 +311,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 +330,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 +352,14 @@ func TestResources_DownloadCollectorChecksumByVersion(t *testing.T) { ctrl := gomock.NewController(t) mock := &mock{ - mockFS: fsmocks.NewMockService(ctrl), + mockFS: storagemocks.NewMockFileService(ctrl), + mockFileServiceResolver: storagemocks.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..fd0372dd221e 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" @@ -38,6 +37,7 @@ import ( "github.com/specterops/bloodhound/packages/go/headers" "github.com/specterops/bloodhound/cmd/api/src/services/job" + "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" ) @@ -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..ec9bdf849773 100644 --- a/cmd/api/src/api/v2/fileingest_test.go +++ b/cmd/api/src/api/v2/fileingest_test.go @@ -43,6 +43,8 @@ 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" + "github.com/specterops/bloodhound/cmd/api/src/services/storage" + storagemocks "github.com/specterops/bloodhound/cmd/api/src/services/storage/mocks" "github.com/specterops/bloodhound/packages/go/headers" "github.com/specterops/bloodhound/cmd/api/src/utils/test" @@ -479,9 +481,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 *storagemocks.MockFileServiceResolver } type expected struct { responseBody string @@ -489,10 +501,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 +586,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 +643,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 +669,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 +702,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 +733,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 +766,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 +817,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 +838,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: storagemocks.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..50b5298ecbb0 100644 --- a/cmd/api/src/bootstrap/initializer.go +++ b/cmd/api/src/bootstrap/initializer.go @@ -33,13 +33,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, connections DatabaseConnections[DBType, GraphType]) (RuntimeDependencies, error) +type InitializerLogic[DBType database.Database, GraphType graph.Database] func(ctx context.Context, cfg config.Configuration, databaseConnections DatabaseConnections[DBType, GraphType], deps 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 { @@ -65,29 +67,34 @@ func (s Initializer[DBType, GraphType]) Launch(parentCtx context.Context, handle defer databaseConnections.RDMS.Close(ctx) defer databaseConnections.Graph.Close(ctx) - // Daemons that start prior to blocking db migration - if s.PreMigrationDaemons != nil { - if daemonInstances, err := s.PreMigrationDaemons(ctx, s.Configuration, databaseConnections); err != nil { + if dependencies, err := s.RuntimeDependencies(ctx, s.Configuration, databaseConnections); err != nil { + return fmt.Errorf("failed to initialize runtime dependencies: %w", err) + } else { + // Daemons that start prior to blocking db migration + if s.PreMigrationDaemons != nil { + if daemonInstances, err := s.PreMigrationDaemons(ctx, s.Configuration, databaseConnections, dependencies); err != nil { + return fmt.Errorf("failed to start services: %w", err) + } else { + daemonManager.Start(ctx, daemonInstances...) + } + } + + // Daemons that start after blocking db migration + if daemonInstances, err := s.Entrypoint(ctx, s.Configuration, databaseConnections, dependencies); err != nil { return fmt.Errorf("failed to start services: %w", err) } else { daemonManager.Start(ctx, daemonInstances...) } - } - // Daemons that start after blocking db migration - if daemonInstances, err := s.Entrypoint(ctx, s.Configuration, databaseConnections); err != nil { - return fmt.Errorf("failed to start services: %w", err) - } else { - daemonManager.Start(ctx, daemonInstances...) - } + // Log successful start and wait for a signal to exit + slog.InfoContext(ctx, "Server started successfully") + <-ctx.Done() - // Log successful start and wait for a signal to exit - slog.InfoContext(ctx, "Server started successfully") - <-ctx.Done() + slog.InfoContext(ctx, "Shutting down") + // TODO: Refactor this pattern in favor of context handling + daemonManager.Stop() - slog.InfoContext(ctx, "Shutting down") - // TODO: Refactor this pattern in favor of context handling - daemonManager.Stop() + } return nil } 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 8be699c1ec07..980100605e0d 100644 --- a/cmd/api/src/bootstrap/util.go +++ b/cmd/api/src/bootstrap/util.go @@ -25,6 +25,7 @@ import ( "github.com/jackc/pgx/v5/pgxpool" "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" @@ -32,6 +33,14 @@ import ( "github.com/specterops/dawgs/util/size" ) +// RuntimeDependencies holds values that must be created prior to the entrypoint. 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) { @@ -57,6 +66,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 } 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 731dc4a63b2b..1abc3af27e93 100644 --- a/cmd/api/src/config/config.go +++ b/cmd/api/src/config/config.go @@ -146,6 +146,10 @@ type Configuration struct { EnableAuditLogStdout bool `json:"enable_audit_log_stdout"` } +func (s Configuration) ScratchDirectory() string { + return filepath.Join(s.WorkDir, "ingest_scratch") +} + func (s Configuration) TempDirectory() string { return filepath.Join(s.WorkDir, "tmp") } diff --git a/cmd/api/src/daemons/api/toolapi/api.go b/cmd/api/src/daemons/api/toolapi/api.go index decfb07a80dd..3a8dc549084b 100644 --- a/cmd/api/src/daemons/api/toolapi/api.go +++ b/cmd/api/src/daemons/api/toolapi/api.go @@ -33,6 +33,7 @@ import ( "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/storage" "github.com/specterops/bloodhound/packages/go/bhlog/attr" "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..51df798e8b92 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" @@ -27,9 +29,15 @@ import ( "strings" "sync" + "github.com/specterops/bloodhound/cmd/api/src/services/storage" "github.com/specterops/bloodhound/packages/go/bhlog/attr" ) +const ( + orphanMinimumAge = 24 * time.Hour + scratchMinimumAge = 24 * time.Hour +) + // FileOperations is an interface for describing filesystem actions. This implementation exists due to deficiencies in // the fs package where the fs.FS interface does not support destructive operations like RemoveAll. type FileOperations interface { @@ -55,112 +63,244 @@ 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 + } - if err := s.fileOps.RemoveAll(fullPath); err != nil { - slog.ErrorContext( - ctx, - "Failed removing orphaned file", - slog.String("full_path", fullPath), - attr.Error(err), - ) - } + cutoff := time.Now().Add(-minimumAge) - numDeleted += 1 + for _, entry := range entries { + if ctx.Err() != nil { + return } - if numDeleted > 0 { - slog.InfoContext( + if entry.IsDir() { + continue + } + + 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 (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 + } + if entry.IsDir() { + continue + } + + fullPath := filepath.Join(s.tempDirectoryRootPath, entry.Name()) + if _, ok := expectedLocalFiles[filepath.Clean(fullPath)]; ok { + 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..91661eb6831c --- /dev/null +++ b/cmd/api/src/daemons/datapipe/cleanup_internal_test.go @@ -0,0 +1,302 @@ +// 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/cmd/api/src/services/storage" + storagemocks "github.com/specterops/bloodhound/cmd/api/src/services/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, "old.json")).Return(nil) + + sweeper.clearLegacyLocalIngestFiles(context.Background(), []string{"expected.json", "../old.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..d2a9294edb04 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/cmd/api/src/services/storage" + storagemocks "github.com/specterops/bloodhound/cmd/api/src/services/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 38440046b754..73483ba39520 100644 --- a/cmd/api/src/daemons/datapipe/datapipe_integration_test.go +++ b/cmd/api/src/daemons/datapipe/datapipe_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/cache" @@ -112,15 +113,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..c4ae3a7ee3f2 100644 --- a/cmd/api/src/daemons/datapipe/delete_integration_test.go +++ b/cmd/api/src/daemons/datapipe/delete_integration_test.go @@ -28,6 +28,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/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..1782f43c9542 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" + "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" @@ -48,21 +49,23 @@ type BHCEPipeline struct { cfg config.Configuration orphanedFileSweeper *OrphanFileSweeper ingestSchema upload.IngestSchema + fileServiceResolver storage.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 storage.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 +158,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 +167,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/services/entrypoint.go b/cmd/api/src/services/entrypoint.go index 5c217551427b..fe15ec715315 100644 --- a/cmd/api/src/services/entrypoint.go +++ b/cmd/api/src/services/entrypoint.go @@ -41,6 +41,7 @@ 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" + "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" @@ -74,14 +75,37 @@ 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 deps = bootstrap.RuntimeDependencies{} + if fileServices, err := storage.NewDefaultFileServices(cfg); err != nil { + return deps, fmt.Errorf("failed to initialize file services: %w", err) + } else if fileServiceResolver, err := storage.NewFileServiceResolver(fileServices); err != nil { + return deps, 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 deps, fmt.Errorf("failed to resolve FileServiceRetained which is needed for the PreMigrationDaemons: %w", err) + } else { + deps.FileServiceResolver = fileServiceResolver + return deps, 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], deps bootstrap.RuntimeDependencies) ([]daemons.Daemon, error) { + if retainedFileService, err := deps.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], deps bootstrap.RuntimeDependencies) ([]daemons.Daemon, error) { dogtagsService := dogtags.NewDefaultService() @@ -137,7 +161,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, deps.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 +171,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, deps.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..d85e4b3d1a67 --- /dev/null +++ b/cmd/api/src/services/entrypoint_test.go @@ -0,0 +1,118 @@ +// 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" + "github.com/specterops/bloodhound/cmd/api/src/services/storage" + storagemocks "github.com/specterops/bloodhound/cmd/api/src/services/storage/mocks" + "github.com/specterops/dawgs/graph" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" +) + +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 := storage.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 := storage.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) + }) +} + +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"), + } +} 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 f016fc4de89a..58400fb47fc0 100644 --- a/cmd/api/src/services/graphify/graphify_integration_test.go +++ b/cmd/api/src/services/graphify/graphify_integration_test.go @@ -33,6 +33,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" @@ -107,9 +108,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 new file mode 100644 index 000000000000..4377355d82c6 --- /dev/null +++ b/cmd/api/src/services/graphify/ingest_storage.go @@ -0,0 +1,203 @@ +// 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 graphify + +import ( + "archive/zip" + "context" + "fmt" + "io" + "log/slog" + "os" + + "github.com/specterops/bloodhound/cmd/api/src/model" + "github.com/specterops/bloodhound/cmd/api/src/services/storage" + "github.com/specterops/bloodhound/packages/go/bhlog/attr" + "github.com/specterops/bloodhound/packages/go/bomenc" + "github.com/specterops/dawgs/util" +) + +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 + ) + + 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) + } + + 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) + } + + 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 OpenScratchReadSeeker(ctx context.Context, scratchDirectory string, fileService storage.FileService, storedFileName string) (*os.File, string, error) { + var ( + scratchPath string + scratchFile *os.File + err error + ) + + if scratchPath, err = SpoolToScratch(ctx, scratchDirectory, fileService, storedFileName); err != nil { + return nil, "", err + } + + if scratchFile, err = os.Open(scratchPath); err != nil { + _ = os.Remove(scratchPath) + return nil, "", fmt.Errorf("error opening ingest scratch file %q: %w", scratchPath, 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() + + if normalizedFile, err = bomenc.NormalizeToUTF8(sourceFile); err != 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 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 + } + + // Zip Path: + scratchPath, err := SpoolToScratch(ctx, scratchDirectory, fileService, storedFileName) + if err != nil { + return []IngestFileData{ + { + Name: providedFileName, + Path: storedFileName, + Errors: []string{fmt.Sprintf("Error spooling archive to scratch: %v", err)}, + }, + }, err + } + defer os.Remove(scratchPath) + + 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 + } + + processedFileData := IngestFileData{ + Name: archiveFile.Name, + ParentFile: providedFileName, + } + + if extractedPath, err := WriteArchiveFileToStorage(ctx, fileService, archiveFile, prefix); err != nil { + processedFileData.Errors = []string{err.Error()} + + // TODO MC: should a single extract break the full process? Any file data with error can be skipped directly + // 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 deleting archive", + slog.String("path", storedFileName), + attr.Error(err), + ) + } + + 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 new file mode 100644 index 000000000000..f0995ff25f28 --- /dev/null +++ b/cmd/api/src/services/graphify/ingest_storage_test.go @@ -0,0 +1,633 @@ +// 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 graphify_test + +import ( + "archive/zip" + "bytes" + "context" + "errors" + "io" + "os" + "strings" + "testing" + + "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/storage" + storagemocks "github.com/specterops/bloodhound/cmd/api/src/services/storage/mocks" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" +) + +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 string + isDir bool +} + +func buildIngestStorageZip(t *testing.T, entries ...ingestStorageZipEntry) []byte { + t.Helper() + + var archive bytes.Buffer + archiveWriter := zip.NewWriter(&archive) + + for _, entry := range entries { + if entry.isDir { + _, err := archiveWriter.Create(entry.name) + require.NoError(t, err) + continue + } + + entryWriter, err := archiveWriter.Create(entry.name) + require.NoError(t, err) + + _, err = entryWriter.Write([]byte(entry.content)) + require.NoError(t, err) + } + + 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 TestSpoolToScratch(t *testing.T) { + t.Parallel() + + var ( + errGet = errors.New("get failed") + errRead = errors.New("read failed") + ) + + 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: "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, + }, + }, + { + 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 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 ( + errGet = errors.New("get failed") + errWrite = errors.New("write failed") + ) + + type expected struct { + errIs error + 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: "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: "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 continues", + 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"}, + ) + + 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)) + + return "prefix/tmp-two", nil + }), + ) + mockFileService.EXPECT(). + DeleteFile(ctx, "stored.json"). + Return(nil) + }, + expected: expected{ + 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) + }, + }, + } + + 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)) + ) + + 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) + } else { + require.NoError(t, err) + } + require.Equal(t, testCase.expected.fileData, fileData) + + if testCase.verify != nil { + testCase.verify(t, scratchDirectory) + } + }) + } +} + +func requireReadFile(t *testing.T, filePath string) []byte { + t.Helper() + + data, err := os.ReadFile(filePath) + require.NoError(t, err) + return data +} diff --git a/cmd/api/src/services/graphify/service.go b/cmd/api/src/services/graphify/service.go index 6138cd65915a..ff5d9e860d06 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" ) @@ -39,23 +40,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 60d72dbbffd9..5f83bf2be0be 100644 --- a/cmd/api/src/services/graphify/tasks.go +++ b/cmd/api/src/services/graphify/tasks.go @@ -17,10 +17,9 @@ package graphify import ( - "archive/zip" "context" "errors" - "io" + "fmt" "io/fs" "log/slog" "os" @@ -29,9 +28,9 @@ 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/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/dawgs/graph" ) @@ -57,115 +56,12 @@ type IngestFileData struct { UserDataErrs []string } -// extractIngestFiles will take a path and extract zips if necessary, returning the paths for files to process -// along with any errors and the number of failed files (in the case of a zip archive) -func (s *GraphifyService) extractIngestFiles(path string, providedFileName string, fileType model.FileType) ([]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: path, - Errors: []string{}, - }, - }, nil - } else if archive, err := zip.OpenReader(path); err != nil { - return []IngestFileData{}, err - } else { - var ( - errs = errorlist.NewBuilder() - fileData = make([]IngestFileData, 0) - ) - - defer func() { - if err := archive.Close(); err != nil { - slog.ErrorContext( - s.ctx, - "Error closing archive", - slog.String("path", path), - attr.Error(err), - ) - } - if err := os.Remove(path); err != nil { - slog.ErrorContext( - s.ctx, - "Error deleting archive", - slog.String("path", path), - attr.Error(err), - ) - } - }() - - for _, f := range archive.File { - // skip directories - if f.FileInfo().IsDir() { - continue - } - - fileName, err := s.extractToTempFile(f) - if err != nil { - fileData = append(fileData, IngestFileData{ - Name: f.Name, - ParentFile: providedFileName, - Errors: []string{err.Error()}, - }) - - errs.Add(err) - } else { - fileData = append(fileData, IngestFileData{ - Name: f.Name, - ParentFile: providedFileName, - Path: fileName, - }) - } - } - - return fileData, errs.Build() - } -} - -func (s *GraphifyService) extractToTempFile(f *zip.File) (string, error) { - // Given a single artifact in an archive, extract it out to a temporary file - tempFile, err := os.CreateTemp(s.cfg.TempDirectory(), "bh") - if err != nil { - return "", err - } - - success := false - defer func() { - // Always close the tempFile, but... - tempFile.Close() - if !success { - // ... only delete if it wasn't successful. Otherwise we leave it around to be processed - os.Remove(tempFile.Name()) - } - }() - - srcFile, err := f.Open() - if err != nil { - return "", err - } - defer srcFile.Close() - - // this creates a normalized file to feed to the copy - if normFile, err := bomenc.NormalizeToUTF8(srcFile); err != nil { - return "", err - // and this is what actually copies it to disk - } else if _, err := io.Copy(tempFile, normFile); err != nil { - return "", err - } else { - // let the deferred method above know we shouldn't delete it and return the filename - success = true - return tempFile.Name(), 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 := s.extractIngestFiles(task.StoredFileName, task.OriginalFileName, task.FileType); 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() @@ -179,7 +75,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 @@ -224,10 +124,10 @@ func (s *GraphifyService) NewIngestContext(ctx context.Context, ingestTime time. return NewIngestContext(ctx, opts...) } -func processSingleFile(ctx context.Context, fileData IngestFileData, ingestContext *IngestContext, readOpts ReadOptions) error { +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, err := os.Open(fileData.Path) + file, scratchPath, err := OpenScratchReadSeeker(ctx, tempDirectory, fileService, fileData.Path) if err != nil { slog.ErrorContext( ctx, @@ -239,14 +139,29 @@ func processSingleFile(ctx context.Context, fileData IngestFileData, ingestConte } defer func() { - file.Close() + 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(fileData.Path); err != nil && !errors.Is(err, fs.ErrNotExist) { - slog.ErrorContext( + 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("filepath", fileData.Path), + slog.String("storage_path", fileData.Path), attr.Error(err), ) } @@ -256,7 +171,7 @@ func processSingleFile(ctx context.Context, fileData IngestFileData, ingestConte slog.ErrorContext( ctx, "Error reading ingest file", - slog.String("filepath", fileData.Path), + slog.String("storage_path", fileData.Path), attr.Error(err), ) return err @@ -281,6 +196,12 @@ func (s *GraphifyService) ProcessTasks(updateJob UpdateJobFunc) { return } + ingestFileService, err := s.fileServiceResolver.Resolve(storage.FileServiceIngest) + if err != nil { + slog.ErrorContext(s.ctx, "Error resolve ingest file service", attr.Error(err)) + return + } + start := time.Now() slog.InfoContext(s.ctx, "Ingest run starting", @@ -306,7 +227,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/localstore.go b/cmd/api/src/services/storage/localstore.go new file mode 100644 index 000000000000..4d7588666104 --- /dev/null +++ b/cmd/api/src/services/storage/localstore.go @@ -0,0 +1,379 @@ +// 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" +) + +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 + r io.Reader +} + +func (s *ctxReader) Read(p []byte) (int, error) { + if err := s.ctx.Err(); err != nil { + return 0, err + } + return s.r.Read(p) +} + +// 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(root string) (*LocalStore, error) { + r, err := os.OpenRoot(root) + if err != nil { + return nil, err + } + return &LocalStore{ + root: r, + }, 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 ct := mime.TypeByExtension(ext); ct != "" { + return ct + } + } + return "application/octet-stream" +} + +// 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. +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, r: 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 + } + // Publish is durable, a failed Remove leaks a .tmp-... that a sweeper can reclaim + _ = s.root.Remove(tmpName) + return nil + } + err = s.root.Rename(tmpName, name) + return err +} + +// 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) + } + return s.root.Remove(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(p string, d fs.DirEntry, err error) error { + if cerr := ctx.Err(); cerr != nil { + return cerr + } + if err != nil { + if errors.Is(err, fs.ErrNotExist) && p == listName { + return fs.SkipAll + } + return err + } + if d.IsDir() { + return nil + } + info, err := d.Info() + if err != nil { + return err + } + out = append(out, FileInfo{ + Path: p, + Size: info.Size(), + LastModified: info.ModTime(), + IsDir: d.IsDir(), + ContentType: detectContentType(p), + }) + 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) +} + +// 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) + } + + dir := path.Dir(dstName) + if dir != "." { + if err := s.root.MkdirAll(dir, 0o750); err != nil { + return err + } + } + if options.FailIfExists { + if err := s.root.Link(srcName, dstName); err != nil { + return err + } + return s.root.Remove(srcName) + } + return s.root.Rename(srcName, dstName) +} diff --git a/cmd/api/src/services/storage/localstore_test.go b/cmd/api/src/services/storage/localstore_test.go new file mode 100644 index 000000000000..a71f802d991e --- /dev/null +++ b/cmd/api/src/services/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/cmd/api/src/services/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/cmd/api/src/services/storage/mocks/fs.go b/cmd/api/src/services/storage/mocks/fs.go new file mode 100644 index 000000000000..176cd183e81a --- /dev/null +++ b/cmd/api/src/services/storage/mocks/fs.go @@ -0,0 +1,356 @@ +// 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: Storage,FileService,FileServiceResolver) +// +// Generated by this command: +// +// mockgen -copyright_file=../../../../../LICENSE.header -destination=./mocks/fs.go -package=mocks . Storage,FileService,FileServiceResolver +// + +// Package mocks is a generated GoMock package. +package mocks + +import ( + context "context" + io "io" + reflect "reflect" + + storage "github.com/specterops/bloodhound/cmd/api/src/services/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) +} + +// 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) +} + +// 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/storage/storage.go b/cmd/api/src/services/storage/storage.go new file mode 100644 index 000000000000..1060888c51ba --- /dev/null +++ b/cmd/api/src/services/storage/storage.go @@ -0,0 +1,314 @@ +// 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" + "fmt" + "io" + "time" + + "github.com/specterops/bloodhound/cmd/api/src/config" +) + +//go:generate go run go.uber.org/mock/mockgen -copyright_file=../../../../../LICENSE.header -destination=./mocks/fs.go -package=mocks . Storage,FileService,FileServiceResolver + +type FileServiceName string + +const ( + FileServiceIngest FileServiceName = "ingest" + FileServiceRetained FileServiceName = "retained" + FileServiceCollectors FileServiceName = "collectors" + FileServiceJobLogs FileServiceName = "job_logs" + FileServiceWork FileServiceName = "work" +) + +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 +} + +// 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 +} + +// Serves as an abstraction to hanlde files with different storage backends. This functions +// are general functions that each file service must implement. +type FileService interface { + // GetFile returns a io.ReadCloser and FileInfo for the named filed that is requested. + GetFile(ctx context.Context, name string) (io.ReadCloser, 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) ReadFile(ctx context.Context, name string) ([]byte, error) { + rc, _, err := s.Storage.Get(ctx, name) + if err != nil { + return nil, err + } + defer rc.Close() + + return io.ReadAll(rc) +} + +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) +} + +// FileServiceResolver is an interface that is used to resolve the actual FileService needed for +// a specific use case. This is ultimately map backed. +type FileServiceResolver interface { + // Resolve returns a FileService interface if a FileService is found with the given name. + // Otherwise, an error is returned. + Resolve(name FileServiceName) (FileService, error) +} + +type fileServiceResolver struct { + services map[FileServiceName]FileService +} + +func NewFileServiceResolver(services map[FileServiceName]FileService) (FileServiceResolver, error) { + var ( + serviceName FileServiceName + fileService FileService + copiedServices = make(map[FileServiceName]FileService, 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 FileServiceName) (FileService, error) { + var ( + fileService FileService + found bool + ) + + if name == "" { + return nil, fmt.Errorf("%w: empty name", ErrFileServiceNotFound) + } + + fileService, found = s.services[name] + if !found { + return nil, fmt.Errorf("%w: %s", ErrFileServiceNotFound, name) + } + + return fileService, nil +} + +// 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) (map[FileServiceName]FileService, error) { + var ( + fileServices = make(map[FileServiceName]FileService) + ) + workStore, err := NewLocalStore(cfg.WorkDir) + if err != nil { + return nil, err + } + + ingestStore, err := NewLocalStore(cfg.TempDirectory()) + if err != nil { + return nil, err + } + + retainStore, err := NewLocalStore(cfg.RetainedFilesDirectory()) + if err != nil { + return nil, err + } + + collectorsStore, err := NewLocalStore(cfg.CollectorsBasePath) + if err != nil { + return nil, err + } + + workService := NewFileService(workStore) + fileServices[FileServiceWork] = workService + + ingestService := NewFileService(ingestStore) + fileServices[FileServiceIngest] = ingestService + + retainService := NewFileService(retainStore) + fileServices[FileServiceRetained] = retainService + + collectorsService := NewFileService(collectorsStore) + fileServices[FileServiceCollectors] = collectorsService + + return fileServices, nil +} diff --git a/cmd/api/src/services/storage/storage_test.go b/cmd/api/src/services/storage/storage_test.go new file mode 100644 index 000000000000..f4e2269e57c4 --- /dev/null +++ b/cmd/api/src/services/storage/storage_test.go @@ -0,0 +1,1115 @@ +// 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" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/specterops/bloodhound/cmd/api/src/config" + "github.com/specterops/bloodhound/cmd/api/src/services/storage" + "github.com/specterops/bloodhound/cmd/api/src/services/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 newTestStorageConfiguration(workDir string, collectorsBasePath string) config.Configuration { + return config.Configuration{ + WorkDir: workDir, + CollectorsBasePath: collectorsBasePath, + } +} + +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) + } + }) + } +} + +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) map[storage.FileServiceName]storage.FileService + expected expected + } + + tests := []testData{ + { + name: "creates resolver with services", + buildServices: func(t *testing.T, controller *gomock.Controller) map[storage.FileServiceName]storage.FileService { + return map[storage.FileServiceName]storage.FileService{ + storage.FileServiceWork: mocks.NewMockFileService(controller), + } + }, + }, + { + name: "empty service name returns error", + buildServices: func(t *testing.T, controller *gomock.Controller) map[storage.FileServiceName]storage.FileService { + return map[storage.FileServiceName]storage.FileService{ + "": mocks.NewMockFileService(controller), + } + }, + expected: expected{ + errContains: "file service name is required", + }, + }, + { + name: "nil service returns error", + buildServices: func(t *testing.T, controller *gomock.Controller) map[storage.FileServiceName]storage.FileService { + return map[storage.FileServiceName]storage.FileService{ + 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 := 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 := map[storage.FileServiceName]storage.FileService{ + 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 := 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 = map[storage.FileServiceName]storage.FileService{ + storage.FileServiceWork: workService, + } + ) + + resolver, err := 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 := 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 := storage.NewDefaultFileServices(configuration) + + // Assert + require.Error(t, err) + require.Nil(t, fileServices) + }) + } +} diff --git a/cmd/api/src/services/upload/streamdecoder.go b/cmd/api/src/services/upload/streamdecoder.go index 113fb07c367f..b5b700e8a6de 100644 --- a/cmd/api/src/services/upload/streamdecoder.go +++ b/cmd/api/src/services/upload/streamdecoder.go @@ -295,13 +295,13 @@ func (s ValidationReport) BuildAPIError() []string { func (s ValidationReport) Error() string { var sb strings.Builder if len(s.CriticalErrors) > 0 { - sb.WriteString(fmt.Sprintf("(%d) critical error(s): [%s]", len(s.CriticalErrors), formatAggregateErrors(s.CriticalErrors))) + fmt.Fprintf(&sb, "(%d) critical error(s): [%s]", len(s.CriticalErrors), formatAggregateErrors(s.CriticalErrors)) if len(s.ValidationErrors) > 0 { sb.WriteString(", ") } } if len(s.ValidationErrors) > 0 { - sb.WriteString(fmt.Sprintf("(%d) validation error(s): [%s]", len(s.ValidationErrors), formatAggregateErrors(s.ValidationErrors))) + fmt.Fprintf(&sb, "(%d) validation error(s): [%s]", len(s.ValidationErrors), formatAggregateErrors(s.ValidationErrors)) } return sb.String() } @@ -310,7 +310,7 @@ func formatSchemaValidationError(arrayName string, index int, err error) string var sb strings.Builder if ve, ok := err.(*jsonschema.ValidationError); ok { numberOfViolations := len(ve.Causes) - sb.WriteString(fmt.Sprintf("%s[%d] schema validation failed with %d error(s): ", arrayName, index, numberOfViolations)) + fmt.Fprintf(&sb, "%s[%d] schema validation failed with %d error(s): ", arrayName, index, numberOfViolations) sb.WriteString("[") @@ -328,17 +328,13 @@ func formatSchemaValidationError(arrayName string, index int, err error) string switch { // Case: property value is an object (not allowed) case isPropertyError && isTypeError(cause, "object"): - sb.WriteString(fmt.Sprintf( - "Invalid property '%s': objects are not allowed in the property bag. Use only strings, numbers, booleans, nulls, or arrays of these types.", - propertyName, - )) + fmt.Fprintf(&sb, "Invalid property '%s': objects are not allowed in the property bag. Use only strings, numbers, booleans, nulls, or arrays of these types.", + propertyName) // Case: array contains a nested object (also not allowed) case isPropertyError && isNotError(cause): - sb.WriteString(fmt.Sprintf( - "Invalid property '%s': array contains an object. Arrays must contain only primitive values (string, number, boolean, or null).", - propertyName, - )) + fmt.Fprintf(&sb, "Invalid property '%s': array contains an object. Arrays must contain only primitive values (string, number, boolean, or null).", + propertyName) default: sb.WriteString(cause.Error()) diff --git a/cmd/api/src/services/upload/upload.go b/cmd/api/src/services/upload/upload.go index 484b8fda3bb4..b072ad0069ba 100644 --- a/cmd/api/src/services/upload/upload.go +++ b/cmd/api/src/services/upload/upload.go @@ -17,15 +17,16 @@ package upload import ( + "context" "errors" "fmt" "io" "log/slog" "net/http" - "os" "github.com/specterops/bloodhound/cmd/api/src/model" "github.com/specterops/bloodhound/cmd/api/src/model/ingest" + "github.com/specterops/bloodhound/cmd/api/src/services/storage" "github.com/specterops/bloodhound/cmd/api/src/utils" "github.com/specterops/bloodhound/packages/go/bhlog/attr" "github.com/specterops/bloodhound/packages/go/headers" @@ -34,7 +35,7 @@ import ( 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 +54,7 @@ func SaveIngestFile(location string, request *http.Request, validator IngestVali return IngestTaskParams{}, fmt.Errorf("invalid content type for ingest file") } - if tempFileName, err := WriteAndValidateFile(fileData, location, 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,44 +62,68 @@ func SaveIngestFile(location string, request *http.Request, validator IngestVali FileType: fileType, }, nil } - } -func WriteAndValidateFile(fileData io.Reader, location string, jobID int64, validationFunc FileValidator) (string, error) { - // Write a temp file. If it passes validation, keep it around and return the filename. Otherwise destroy it. - tempFile, err := os.CreateTemp(location, fmt.Sprintf("file_upload_job%d_", jobID)) - if err != nil { - return "", fmt.Errorf("error creating ingest file: %w", err) - } +func WriteAndValidateFile(ctx context.Context, fileService storage.FileService, fileData io.Reader, prefix string, validationFunc FileValidator) (string, error) { + // Create a pipe: pr (read end) and pw (write end). + // Data written to pw can be read from pr. + pr, pw := io.Pipe() - // Save this for later - tempFileName := tempFile.Name() + // 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) - // Run validation on the file to see if we even want to keep it - _, validationErr := validationFunc(fileData, tempFile) + // 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 + }() - // 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), - ) - } + // 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() - // If the validation was not successful, after we close the file, we remove it and return the error + // Wait for the validation goroutine to finish and collect its result. + validationErr := <-validationErrCh + + // Check if validation failed first — the temp file should be cleaned up. if validationErr != nil { - if removeErr := os.Remove(tempFileName); removeErr != nil { - slog.Error( - "Error deleting temp file", - slog.String("file", tempFileName), - attr.Error(removeErr), - ) + slog.ErrorContext( + ctx, + "Validation failed", + slog.String("temp_file_name", + tempFileName), + attr.Error(validationErr), + ) + if tempFileName != "" { + fileService.DeleteFile(ctx, tempFileName) } return "", validationErr } - // Assuming no other errors, return the name of the closed temp file + // 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), + ) + if tempFileName != "" { + fileService.DeleteFile(ctx, tempFileName) + } + return "", writeErr + } + + 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 8ce9c4b6030f..d372337fcb22 100644 --- a/cmd/api/src/services/upload/upload_test.go +++ b/cmd/api/src/services/upload/upload_test.go @@ -18,6 +18,7 @@ package upload import ( "bytes" + "context" "errors" "fmt" "io" @@ -26,9 +27,25 @@ import ( "testing" "github.com/specterops/bloodhound/cmd/api/src/model/ingest" + "github.com/specterops/bloodhound/cmd/api/src/services/storage" + storagemocks "github.com/specterops/bloodhound/cmd/api/src/services/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 { + t.Helper() + + return func(src io.Reader, dst io.Writer) (ingest.OriginalMetadata, error) { + content, err := io.ReadAll(src) + require.NoError(t, err) + require.Equal(t, expectedContent, string(content)) + + return ingest.OriginalMetadata{}, validationErr + } +} + func TestWriteAndValidateZip(t *testing.T) { t.Run("valid zip file is ok", func(t *testing.T) { writer := bytes.Buffer{} @@ -139,6 +156,112 @@ func TestWriteAndValidateJSON_NormalizationError(t *testing.T) { assert.ErrorIs(t, err, ErrInvalidJSON) } +func TestUpload_WriteAndValidateFile(t *testing.T) { + t.Parallel() + + var ( + errValidation = errors.New("validation failed") + errWrite = errors.New("write failed") + ) + + type expected struct { + errIs error + fileName string + } + + type testData struct { + name string + tempFileName string + writeErr error + validationErr error + expected expected + expectDelete bool + } + + 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(ctx, testCase.tempFileName).Return(nil) + } + + // Act + actualFileName, err := WriteAndValidateFile(ctx, mockFileService, strings.NewReader("content"), "prefix", validator) + + // Assert + if testCase.expected.errIs != nil { + require.ErrorIs(t, err, testCase.expected.errIs) + require.Empty(t, actualFileName) + return + } + + require.NoError(t, err) + require.Equal(t, testCase.expected.fileName, actualFileName) + }) + } +} + // ErrorReader is a mock reader that always returns an error type ErrorReader struct { err error diff --git a/cmd/api/src/test/lab/fixtures/api.go b/cmd/api/src/test/lab/fixtures/api.go index 4348e097e10d..df01d5a0ea89 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], deps 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, deps) }, } diff --git a/go.mod b/go.mod index 0d7bad7d3e11..2674cd85d8c2 100644 --- a/go.mod +++ b/go.mod @@ -100,21 +100,21 @@ 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/config v1.32.13 // 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 v1.41.7 // indirect + github.com/aws/aws-sdk-go-v2/config v1.32.17 // indirect + github.com/aws/aws-sdk-go-v2/credentials v1.19.16 // indirect + github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.23 // indirect github.com/aws/aws-sdk-go-v2/feature/rds/auth v1.6.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/internal/configsources v1.4.23 // indirect + github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.23 // indirect + github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.24 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.9 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.23 // indirect + github.com/aws/aws-sdk-go-v2/service/signin v1.0.11 // indirect + github.com/aws/aws-sdk-go-v2/service/sso v1.30.17 // indirect + github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.21 // indirect + github.com/aws/aws-sdk-go-v2/service/sts v1.42.1 // indirect + github.com/aws/smithy-go v1.25.1 // indirect 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 6c9becae9bec..51962d3ef47b 100644 --- a/go.sum +++ b/go.sum @@ -558,72 +558,72 @@ 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.7 h1:DWpAJt66FmnnaRIOT/8ASTucrvuDPZASqhhLey6tLY8= +github.com/aws/aws-sdk-go-v2 v1.41.7/go.mod h1:4LAfZOPHNVNQEckOACQx60Y8pSRjIkNZQz1w92xpMJc= 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.17 h1:FpL4/758/diKwqbytU0prpuiu60fgXKUWCpDJtApclU= +github.com/aws/aws-sdk-go-v2/config v1.32.17/go.mod h1:OXqUMzgXytfoF9JaKkhrOYsyh72t9G+MJH8mMRaexOE= 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.16 h1:r3RJBuU7X9ibt8RHbMjWE6y60QbKBiII6wSrXnapxSU= +github.com/aws/aws-sdk-go-v2/credentials v1.19.16/go.mod h1:6cx7zqDENJDbBIIWX6P8s0h6hqHC8Avbjh9Dseo27ug= 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.23 h1:UuSfcORqNSz/ey3VPRS8TcVH2Ikf0/sC+Hdj400QI6U= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.23/go.mod h1:+G/OSGiOFnSOkYloKj/9M35s74LgVAdJBSD5lsFfqKg= 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/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.23 h1:GpT/TrnBYuE5gan2cZbTtvP+JlHsutdmlV2YfEyNde0= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.23/go.mod h1:xYWD6BS9ywC5bS3sz9Xh04whO/hzK2plt2Zkyrp4JuA= 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.23 h1:bpd8vxhlQi2r1hiueOw02f/duEPTMK59Q4QMAoTTtTo= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.23/go.mod h1:15DfR2nw+CRHIk0tqNyifu3G1YdAOy68RftkhMDDwYk= 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.24 h1:OQqn11BtaYv1WLUowvcA30MpzIu8Ti4pcLPIIyoKZrA= +github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.24/go.mod h1:X5ZJyfwVrWA96GzPmUCWFQaEARPR7gCrpq2E92PJwAE= 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.9 h1:FLudkZLt5ci0ozzgkVo8BJGwvqNaZbTWb3UcucAateA= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.9/go.mod h1:w7wZ/s9qK7c8g4al+UyoF1Sp/Z45UwMGcqIzLWVQHWk= 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.23 h1:pbrxO/kuIwgEsOPLkaHu0O+m4fNgLU8B3vxQ+72jTPw= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.23/go.mod h1:/CMNUqoj46HpS3MNRDEDIwcgEnrtZlKRaHNaHxIFpNA= +github.com/aws/aws-sdk-go-v2/service/signin v1.0.11 h1:TdJ+HdzOBhU8+iVAOGUTU63VXopcumCOF1paFulHWZc= +github.com/aws/aws-sdk-go-v2/service/signin v1.0.11/go.mod h1:R82ZRExE/nheo0N+T8zHPcLRTcH8MGsnR3BiVGX0TwI= 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.17 h1:7byT8HUWrgoRp6sXjxtZwgOKfhss5fW6SkLBtqzgRoE= +github.com/aws/aws-sdk-go-v2/service/sso v1.30.17/go.mod h1:xNWknVi4Ezm1vg1QsB/5EWpAJURq22uqd38U8qKvOJc= 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.35.21 h1:+1Kl1zx6bWi4X7cKi3VYh29h8BvsCoHQEQ6ST9X8w7w= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.21/go.mod h1:4vIRDq+CJB2xFAXZ+YgGUTiEft7oAQlhIs71xcSeuVg= 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.1 h1:F/M5Y9I3nwr2IEpshZgh1GeHpOItExNM9L1euNuh/fk= +github.com/aws/aws-sdk-go-v2/service/sts v1.42.1/go.mod h1:mTNxImtovCOEEuD65mKW7DCsL+2gjEH+RPEAexAzAio= 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.25.1 h1:J8ERsGSU7d+aCmdQur5Txg6bVoYelvQJgtZehD12GkI= +github.com/aws/smithy-go v1.25.1/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/schemagen/generator/sql.go b/packages/go/schemagen/generator/sql.go index a2e944e408fc..c225ad47bb20 100644 --- a/packages/go/schemagen/generator/sql.go +++ b/packages/go/schemagen/generator/sql.go @@ -178,7 +178,7 @@ func GenerateExtensionSQLAzure(dir string, azSchema model.Azure) error { func GenerateExtensionSQL(name string, displayName string, version string, namespace, dir string, fileName string, nodeKinds []model.StringEnum, relationshipKinds []model.StringEnum, pathfindingRelationshipKinds []model.StringEnum) error { var sb strings.Builder - sb.WriteString(fmt.Sprintf("-- Code generated by Cuelang code gen. DO NOT EDIT!\n-- Cuelang source: %s/", SchemaSourceName)) + fmt.Fprintf(&sb, "-- Code generated by Cuelang code gen. DO NOT EDIT!\n-- Cuelang source: %s/", SchemaSourceName) sb.WriteString(` CREATE OR REPLACE FUNCTION genscript_upsert_kind(node_kind_name TEXT) RETURNS SMALLINT AS $$ @@ -300,20 +300,20 @@ $$ LANGUAGE plpgsql; sb.WriteString("\nDO $$\nDECLARE\n\textension_id INT;\n\tenvironment_id INT;\nBEGIN\n\tLOCK schema_extensions, schema_node_kinds, schema_relationship_kinds, kind;\n\n") - sb.WriteString(fmt.Sprintf("\tIF NOT EXISTS (SELECT id FROM schema_extensions WHERE name = '%s') THEN\n", name)) - sb.WriteString(fmt.Sprintf("\t\tINSERT INTO schema_extensions (name, display_name, version, is_builtin, namespace, created_at, updated_at) VALUES ('%s', '%s', '%s', true, '%s', NOW(), NOW()) RETURNING id INTO extension_id;\n", name, displayName, version, namespace)) + fmt.Fprintf(&sb, "\tIF NOT EXISTS (SELECT id FROM schema_extensions WHERE name = '%s') THEN\n", name) + fmt.Fprintf(&sb, "\t\tINSERT INTO schema_extensions (name, display_name, version, is_builtin, namespace, created_at, updated_at) VALUES ('%s', '%s', '%s', true, '%s', NOW(), NOW()) RETURNING id INTO extension_id;\n", name, displayName, version, namespace) sb.WriteString("\tELSE\n") - sb.WriteString(fmt.Sprintf("\t\tUPDATE schema_extensions SET display_name = '%s', version = '%s', namespace = '%s', updated_at = NOW() WHERE name = '%s' RETURNING id INTO extension_id;\n", displayName, version, namespace, name)) + fmt.Fprintf(&sb, "\t\tUPDATE schema_extensions SET display_name = '%s', version = '%s', namespace = '%s', updated_at = NOW() WHERE name = '%s' RETURNING id INTO extension_id;\n", displayName, version, namespace, name) sb.WriteString("\tEND IF;\n\n") sb.WriteString("\t-- Insert Node Kinds\n") for _, kind := range nodeKinds { - sb.WriteString(fmt.Sprintf("\tPERFORM genscript_upsert_kind('%s');\n", kind.GetRepresentation())) + fmt.Fprintf(&sb, "\tPERFORM genscript_upsert_kind('%s');\n", kind.GetRepresentation()) } sb.WriteString("\n\t-- Insert Relationship Kinds\n") for _, kind := range relationshipKinds { - sb.WriteString(fmt.Sprintf("\tPERFORM genscript_upsert_kind('%s');\n", kind.GetRepresentation())) + fmt.Fprintf(&sb, "\tPERFORM genscript_upsert_kind('%s');\n", kind.GetRepresentation()) } sb.WriteString("\n") @@ -321,7 +321,7 @@ $$ LANGUAGE plpgsql; for _, kind := range nodeKinds { iconInfo, found := nodeIcons[kind.GetRepresentation()] - sb.WriteString(fmt.Sprintf("\tPERFORM genscript_upsert_schema_node_kind(extension_id, '%s', '%s', '', %t, '%s', '%s');\n", kind.GetRepresentation(), kind.GetRepresentation(), found, iconInfo.Icon, iconInfo.Color)) + fmt.Fprintf(&sb, "\tPERFORM genscript_upsert_schema_node_kind(extension_id, '%s', '%s', '', %t, '%s', '%s');\n", kind.GetRepresentation(), kind.GetRepresentation(), found, iconInfo.Icon, iconInfo.Color) } traversableMap := make(map[string]struct{}) @@ -335,7 +335,7 @@ $$ LANGUAGE plpgsql; for _, kind := range relationshipKinds { _, traversable := traversableMap[kind.GetRepresentation()] - sb.WriteString(fmt.Sprintf("\tPERFORM genscript_upsert_schema_relationship_kind(extension_id, '%s', '', %t);\n", kind.GetRepresentation(), traversable)) + fmt.Fprintf(&sb, "\tPERFORM genscript_upsert_schema_relationship_kind(extension_id, '%s', '', %t);\n", kind.GetRepresentation(), traversable) } sb.WriteString("\n")