From 9d767cc75196cda99585358c83a55b9b1306e07c Mon Sep 17 00:00:00 2001 From: Michael Cuomo Date: Fri, 3 Apr 2026 16:16:14 -0400 Subject: [PATCH 01/26] wip: start abstraction --- cmd/api/src/api/v2/collectors.go | 4 +- cmd/api/src/api/v2/collectors_test.go | 22 ++-- cmd/api/src/api/v2/fileingest.go | 2 +- cmd/api/src/api/v2/fileingest_test.go | 52 ++++++-- cmd/api/src/api/v2/model.go | 6 +- cmd/api/src/services/storage/mocks/fs.go | 132 +++++++++++++++++++ cmd/api/src/services/storage/osstore.go | 155 +++++++++++++++++++++++ cmd/api/src/services/storage/storage.go | 148 ++++++++++++++++++++++ cmd/api/src/services/upload/upload.go | 74 ++++++----- 9 files changed, 538 insertions(+), 57 deletions(-) create mode 100644 cmd/api/src/services/storage/mocks/fs.go create mode 100644 cmd/api/src/services/storage/osstore.go create mode 100644 cmd/api/src/services/storage/storage.go diff --git a/cmd/api/src/api/v2/collectors.go b/cmd/api/src/api/v2/collectors.go index 7c1567a379e1..940a4f0b432e 100644 --- a/cmd/api/src/api/v2/collectors.go +++ b/cmd/api/src/api/v2/collectors.go @@ -89,7 +89,7 @@ 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 data, err := s.FileService.ReadFile(request.Context(), filepath.Join(s.Config.CollectorsDirectory(), 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 +122,7 @@ 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 data, err := s.FileService.ReadFile(request.Context(), filepath.Join(s.Config.CollectorsDirectory(), 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..ec642e9f79bc 100644 --- a/cmd/api/src/api/v2/collectors_test.go +++ b/cmd/api/src/api/v2/collectors_test.go @@ -28,7 +28,7 @@ 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" + 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 +40,7 @@ import ( func TestResources_DownloadCollectorByVersion(t *testing.T) { type mock struct { - mockFS *fsmocks.MockService + mockFS *storagemocks.MockFileService } type expected struct { responseBody string @@ -99,7 +99,7 @@ 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.mockFS.EXPECT().ReadFile(gomock.Any(), "azurehound/azurehound-latest.zip").Return([]byte{}, errors.New("error")) }, expected: expected{ responseCode: http.StatusInternalServerError, @@ -118,7 +118,7 @@ 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.mockFS.EXPECT().ReadFile(gomock.Any(), "azurehound/azurehound-v1.0.0.zip").Return([]byte{}, nil) }, expected: expected{ responseCode: http.StatusOK, @@ -136,7 +136,7 @@ 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.mockFS.EXPECT().ReadFile(gomock.Any(), "azurehound/azurehound-latest.zip").Return([]byte{}, nil) }, expected: expected{ responseCode: http.StatusOK, @@ -150,7 +150,7 @@ func TestResources_DownloadCollectorByVersion(t *testing.T) { ctrl := gomock.NewController(t) mock := &mock{ - mockFS: fsmocks.NewMockService(ctrl), + mockFS: storagemocks.NewMockFileService(ctrl), } testCase.setupMocks(t, mock) @@ -187,7 +187,7 @@ func TestResources_DownloadCollectorByVersion(t *testing.T) { func TestResources_DownloadCollectorChecksumByVersion(t *testing.T) { type mock struct { - mockFS *fsmocks.MockService + mockFS *storagemocks.MockFileService } type expected struct { responseBody string @@ -246,7 +246,7 @@ func TestResources_DownloadCollectorChecksumByVersion(t *testing.T) { } }, setupMocks: func(t *testing.T, mock *mock) { - mock.mockFS.EXPECT().ReadFile("azurehound/azurehound-latest.zip.sha256").Return([]byte{}, errors.New("error")) + mock.mockFS.EXPECT().ReadFile(gomock.Any(), "azurehound/azurehound-latest.zip.sha256").Return([]byte{}, errors.New("error")) }, expected: expected{ responseCode: http.StatusInternalServerError, @@ -265,7 +265,7 @@ 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.mockFS.EXPECT().ReadFile(gomock.Any(), "azurehound/azurehound-v1.0.0.zip.sha256").Return([]byte{}, nil) }, expected: expected{ responseCode: http.StatusOK, @@ -283,7 +283,7 @@ func TestResources_DownloadCollectorChecksumByVersion(t *testing.T) { } }, setupMocks: func(t *testing.T, mock *mock) { - mock.mockFS.EXPECT().ReadFile("azurehound/azurehound-latest.zip.sha256").Return([]byte{}, nil) + mock.mockFS.EXPECT().ReadFile(gomock.Any(), "azurehound/azurehound-latest.zip.sha256").Return([]byte{}, nil) }, expected: expected{ responseCode: http.StatusOK, @@ -304,7 +304,7 @@ func TestResources_DownloadCollectorChecksumByVersion(t *testing.T) { ctrl := gomock.NewController(t) mock := &mock{ - mockFS: fsmocks.NewMockService(ctrl), + mockFS: storagemocks.NewMockFileService(ctrl), } testCase.setupMocks(t, mock) diff --git a/cmd/api/src/api/v2/fileingest.go b/cmd/api/src/api/v2/fileingest.go index 4df9be834f4b..83df8c8eaaec 100644 --- a/cmd/api/src/api/v2/fileingest.go +++ b/cmd/api/src/api/v2/fileingest.go @@ -146,7 +146,7 @@ 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 ingestTaskParams, err := upload.SaveIngestFile(request.Context(), s.FileService, 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 ( diff --git a/cmd/api/src/api/v2/fileingest_test.go b/cmd/api/src/api/v2/fileingest_test.go index 0f675ecda18c..723d13c293b2 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" @@ -481,7 +483,8 @@ func TestResources_ListAcceptedFileUploadTypes(t *testing.T) { func TestResources_ProcessIngestTask(t *testing.T) { type mock struct { - mockDatabase *dbmocks.MockDatabase + mockDatabase *dbmocks.MockDatabase + mockFileService *storagemocks.MockFileService } type expected struct { responseBody string @@ -489,10 +492,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{ @@ -574,6 +578,9 @@ func TestResources_ProcessIngestTask(t *testing.T) { t.Helper() mock.mockDatabase.EXPECT().GetIngestJob(gomock.Any(), int64(1)).Return(model.IngestJob{Status: model.JobStatusRunning}, nil) }, + // Use a real OsStore so that actual file I/O gives the concurrent validation + // goroutine time to set validationErr before the main goroutine reads it. + fileServiceOvrd: storage.NewLocalFileService(storage.NewOsStore()), 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"}`, @@ -603,6 +610,8 @@ func TestResources_ProcessIngestTask(t *testing.T) { t.Helper() mock.mockDatabase.EXPECT().GetIngestJob(gomock.Any(), int64(1)).Return(model.IngestJob{Status: model.JobStatusRunning}, nil) }, + // Use a real OsStore for the same reason as the ErrInvalidJSON case above. + fileServiceOvrd: storage.NewLocalFileService(storage.NewOsStore()), 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"}`, @@ -626,6 +635,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.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{}, errors.New("error")) }, expected: expected{ @@ -651,6 +664,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.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 +694,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.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 +726,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.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 +776,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.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 +796,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), } request := testCase.buildRequest() testCase.setupMocks(t, mocks) + fileService := storage.FileService(mocks.mockFileService) + if testCase.fileServiceOvrd != nil { + fileService = testCase.fileServiceOvrd + } + resources := v2.Resources{ - DB: mocks.mockDatabase, - Config: config.Configuration{}, + DB: mocks.mockDatabase, + Config: config.Configuration{}, + FileService: fileService, } 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..e48cb8c1ce83 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 + FileService storage.FileService OpenGraphSchemaService OpenGraphSchemaService DogTags dogtags.Service } @@ -145,7 +145,7 @@ func NewResources( Authorizer: authorizer, Authenticator: authenticator, IngestSchema: ingestSchema, - FileService: &fs.Client{}, + FileService: storage.NewLocalFileService(storage.NewOsStore()), DogTags: dogtagsService, OpenGraphSchemaService: openGraphSchemaService, } 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..451b68f6a7fb --- /dev/null +++ b/cmd/api/src/services/storage/mocks/fs.go @@ -0,0 +1,132 @@ +// 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: FileService) +// +// Generated by this command: +// +// mockgen -copyright_file=../../../../../LICENSE.header -destination=./mocks/fs.go -package=mocks . FileService +// + +// 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" +) + +// 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) +} + +// 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) +} + +// TempPath mocks base method. +func (m *MockFileService) TempPath(prefix, pattern string) (string, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "TempPath", prefix, pattern) + ret0, _ := ret[0].(string) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// TempPath indicates an expected call of TempPath. +func (mr *MockFileServiceMockRecorder) TempPath(prefix, pattern any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "TempPath", reflect.TypeOf((*MockFileService)(nil).TempPath), prefix, pattern) +} + +// 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) +} + +// WriteTempFile mocks base method. +func (m *MockFileService) WriteTempFile(ctx context.Context, prefix string, reader io.Reader, opts storage.WriteOptions) (string, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "WriteTempFile", ctx, prefix, reader, opts) + ret0, _ := ret[0].(string) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// WriteTempFile indicates an expected call of WriteTempFile. +func (mr *MockFileServiceMockRecorder) WriteTempFile(ctx, prefix, reader, opts any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WriteTempFile", reflect.TypeOf((*MockFileService)(nil).WriteTempFile), ctx, prefix, reader, opts) +} diff --git a/cmd/api/src/services/storage/osstore.go b/cmd/api/src/services/storage/osstore.go new file mode 100644 index 000000000000..0d1b8229e0c6 --- /dev/null +++ b/cmd/api/src/services/storage/osstore.go @@ -0,0 +1,155 @@ +// 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" + "io" + "os" + "path/filepath" +) + +// OsStore is a Storage implementation that operates directly on the OS filesystem +// using absolute paths. It is intended for use in CE deployments where all paths +// are already fully-qualified (e.g. temp directory paths from config). +type OsStore struct{} + +func NewOsStore() *OsStore { + return &OsStore{} +} + +func (s *OsStore) Put(ctx context.Context, path string, reader io.Reader, options WriteOptions) error { + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + return err + } + + tmp, err := os.CreateTemp(filepath.Dir(path), ".tmp-*") + if err != nil { + return err + } + tmpName := tmp.Name() + + cleanup := func() { + tmp.Close() + os.Remove(tmpName) + } + + if _, err := io.Copy(tmp, reader); err != nil { + cleanup() + return err + } + + if err := tmp.Close(); err != nil { + _ = os.Remove(tmpName) + return err + } + + return os.Rename(tmpName, path) +} + +func (s *OsStore) Get(ctx context.Context, path string) (io.ReadCloser, FileInfo, error) { + file, err := os.Open(path) + if err != nil { + return nil, FileInfo{}, err + } + + stat, err := file.Stat() + if err != nil { + _ = file.Close() + return nil, FileInfo{}, err + } + + return file, FileInfo{ + Path: path, + Size: stat.Size(), + LastModified: stat.ModTime(), + IsDir: stat.IsDir(), + }, nil +} + +func (s *OsStore) Stat(ctx context.Context, path string) (FileInfo, error) { + stat, err := os.Stat(path) + if err != nil { + return FileInfo{}, err + } + + return FileInfo{ + Path: path, + Size: stat.Size(), + LastModified: stat.ModTime(), + IsDir: stat.IsDir(), + }, nil +} + +func (s *OsStore) Exists(ctx context.Context, path string) (bool, error) { + _, err := os.Stat(path) + if errors.Is(err, os.ErrNotExist) { + return false, nil + } + if err != nil { + return false, err + } + return true, nil +} + +func (s *OsStore) Delete(ctx context.Context, path string) error { + err := os.Remove(path) + if errors.Is(err, os.ErrNotExist) { + return nil + } + return err +} + +func (s *OsStore) List(ctx context.Context, path string, options ListOptions) ([]FileInfo, error) { + entries, err := os.ReadDir(path) + if errors.Is(err, os.ErrNotExist) { + return []FileInfo{}, nil + } + if err != nil { + return nil, err + } + + var out []FileInfo + for _, entry := range entries { + info, err := entry.Info() + if err != nil { + return nil, err + } + out = append(out, FileInfo{ + Path: filepath.Join(path, entry.Name()), + Size: info.Size(), + LastModified: info.ModTime(), + IsDir: entry.IsDir(), + }) + } + return out, nil +} + +func (s *OsStore) Copy(ctx context.Context, srcPath, dstPath string) error { + src, err := os.Open(srcPath) + if err != nil { + return err + } + defer src.Close() + + return s.Put(ctx, dstPath, src, WriteOptions{}) +} + +func (s *OsStore) Move(ctx context.Context, srcPath, dstPath string) error { + return os.Rename(srcPath, dstPath) +} diff --git a/cmd/api/src/services/storage/storage.go b/cmd/api/src/services/storage/storage.go new file mode 100644 index 000000000000..282721dd9a01 --- /dev/null +++ b/cmd/api/src/services/storage/storage.go @@ -0,0 +1,148 @@ +// 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" + "fmt" + "io" + "path" + "time" +) + +//go:generate go run go.uber.org/mock/mockgen -copyright_file=../../../../../LICENSE.header -destination=./mocks/fs.go -package=mocks . FileService + +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 +} + +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, path string, reader io.Reader, options WriteOptions) error + + // Get opens a file for reading. + Get(ctx context.Context, path string) (io.ReadCloser, FileInfo, error) + + // Stat returns metadata for a given path. + Stat(ctx context.Context, path string) (FileInfo, error) + + // Delete removes a file. + Delete(ctx context.Context, path string) error + + // Exists checks whether a file exists. + Exists(ctx context.Context, path string) (bool, error) + + // List returns a list of files in a given path. + List(ctx context.Context, path string, options ListOptions) ([]FileInfo, error) + + // Copy duplicates an object. + Copy(ctx context.Context, srcPath, dstPath string) error + + // Move moves an object. + // Is done by a copy and a delete. + Move(ctx context.Context, srcPath, dstPath string) error +} + +type FileService interface { + ReadFile(ctx context.Context, name string) ([]byte, error) + WriteFile(ctx context.Context, name string, data []byte, opts WriteOptions) error + DeleteFile(ctx context.Context, name string) error + TempPath(prefix, pattern string) (string, error) + WriteTempFile(ctx context.Context, prefix string, reader io.Reader, opts WriteOptions) (string, error) +} + +type LocalFileService struct { + Storage Storage +} + +func NewLocalFileService(storage Storage) *LocalFileService { + return &LocalFileService{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 *LocalFileService) 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 *LocalFileService) WriteFile(ctx context.Context, name string, data []byte, opts WriteOptions) error { + return s.Storage.Put(ctx, name, bytes.NewReader(data), opts) +} + +func (s *LocalFileService) DeleteFile(ctx context.Context, name string) error { + return s.Storage.Delete(ctx, name) +} + +func (s *LocalFileService) TempPath(prefix, pattern string) (string, error) { + id, err := randomID() + if err != nil { + return "", err + } + + name := pattern + if name == "" { + name = "tmp-*" + } + + return path.Join(prefix, fmt.Sprintf("%s-%s", name, id)), nil +} + +func (s *LocalFileService) WriteTempFile(ctx context.Context, prefix string, reader io.Reader, opts WriteOptions) (string, error) { + id, err := randomID() + if err != nil { + return "", err + } + + tempPath := path.Join(prefix, "tmp-"+id) + if err := s.Storage.Put(ctx, tempPath, reader, opts); err != nil { + return "", err + } + + return tempPath, nil +} diff --git a/cmd/api/src/services/upload/upload.go b/cmd/api/src/services/upload/upload.go index 484b8fda3bb4..d7713bd97b4b 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,53 @@ 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{}) - // If the validation was not successful, after we close the file, we remove it and return the error + // Closing pw signals EOF to the validation goroutine so it can finish. + pw.Close() + + // 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("tempFileName", tempFileName), attr.Error(validationErr)) + 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("tempFileName", tempFileName), attr.Error(writeErr)) + fileService.DeleteFile(ctx, tempFileName) + return "", writeErr + } + + slog.InfoContext(ctx, "File written and validated", slog.String("tempFileName", tempFileName)) return tempFileName, nil } From 02f8ca80a1c7c85988918a649b69972717793872 Mon Sep 17 00:00:00 2001 From: Michael Cuomo Date: Wed, 22 Apr 2026 19:41:00 -0400 Subject: [PATCH 02/26] wip: starting to refine localstore --- cmd/api/src/services/storage/localstore.go | 432 +++++++++++++++++++++ cmd/api/src/services/storage/osstore.go | 189 ++++++--- cmd/api/src/services/storage/storage.go | 16 - 3 files changed, 568 insertions(+), 69 deletions(-) create mode 100644 cmd/api/src/services/storage/localstore.go diff --git a/cmd/api/src/services/storage/localstore.go b/cmd/api/src/services/storage/localstore.go new file mode 100644 index 000000000000..d9e34c89f8c2 --- /dev/null +++ b/cmd/api/src/services/storage/localstore.go @@ -0,0 +1,432 @@ +// 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" + "io" + "mime" + "os" + "path/filepath" + "strings" +) + +type LocalStore struct { + root string +} + +func NewLocalStore(root string) *LocalStore { + return &LocalStore{ + root: root, + } +} + +func (s *LocalStore) fullPath(path string) (string, error) { + path = filepath.ToSlash(strings.TrimSpace(path)) + path = strings.TrimPrefix(path, "/") + + cleaned := filepath.Clean(path) + + if cleaned == "." { + return "", errors.New("invalid empty path") + } + + if cleaned == ".." || strings.HasPrefix(cleaned, ".."+string(filepath.Separator)) { + return "", errors.New("path escapes storage root") + } + + full := filepath.Join(s.root, cleaned) + return full, nil +} + +func normalizeLogicalPath(path string) string { + path = filepath.ToSlash(strings.TrimSpace(path)) + path = strings.TrimPrefix(path, "/") + return filepath.Clean(path) +} + +func detectContentType(path string) string { + ext := filepath.Ext(path) + if ext == "" { + return "application/octet-stream" + } + + ct := mime.TypeByExtension(ext) + if ct == "" { + return "application/octet-stream" + } + + return ct +} + +func (s *LocalStore) Put(ctx context.Context, path string, reader io.Reader, options WriteOptions) error { + full, err := s.fullPath(path) + if err != nil { + return err + } + + dir := filepath.Dir(full) + if err := os.MkdirAll(dir, 0o755); err != nil { + return err + } + + tmp, err := os.CreateTemp(dir, ".tmp-*") + if err != nil { + return err + } + tmpName := tmp.Name() + + cleanup := func() { + tmp.Close() + os.Remove(tmpName) + } + + buf := make([]byte, 32*1024) + for { + select { + case <-ctx.Done(): + cleanup() + return ctx.Err() + default: + } + + n, readErr := reader.Read(buf) + if n > 0 { + if _, err := tmp.Write(buf[:n]); err != nil { + cleanup() + return err + } + } + + if readErr == io.EOF { + break + } + if readErr != nil { + cleanup() + return readErr + } + } + + if err := tmp.Close(); err != nil { + _ = os.Remove(tmpName) + return err + } + + if err := os.Rename(tmpName, full); err != nil { + _ = os.Remove(tmpName) + return err + } + + return nil +} + + +func (s *LocalStore) Get(ctx context.Context, path string) (io.ReadCloser, FileInfo, error) { + full, err := s.fullPath(path) + if err != nil { + return nil, FileInfo{}, err + } + + file, err := os.Open(full) + if err != nil { + return nil, FileInfo{}, err + } + + stat, err := file.Stat() + if err != nil { + _ = file.Close() + return nil, FileInfo{}, err + } + + info := FileInfo{ + Path: normalizeLogicalPath(path), + Size: stat.Size(), + LastModified: stat.ModTime(), + IsDir: stat.IsDir(), + ContentType: detectContentType(path), + } + + select { + case <-ctx.Done(): + _ = file.Close() + return nil, FileInfo{}, ctx.Err() + default: + } + + return file, info, nil +} + +func (s *LocalStore) Stat(ctx context.Context, path string) (FileInfo, error) { + full, err := s.fullPath(path) + if err != nil { + return FileInfo{}, err + } + + select { + case <-ctx.Done(): + return FileInfo{}, ctx.Err() + default: + } + + stat, err := os.Stat(full) + if err != nil { + return FileInfo{}, err + } + + return FileInfo{ + Path: normalizeLogicalPath(path), + Size: stat.Size(), + LastModified: stat.ModTime(), + IsDir: stat.IsDir(), + ContentType: detectContentType(path), + }, nil +} + +func (s *LocalStore) Exists(ctx context.Context, path string) (bool, error) { + _, err := s.Stat(ctx, path) + if err == nil { + return true, nil + } + if errors.Is(err, os.ErrNotExist) { + return false, nil + } + return false, err +} + +func (s *LocalStore) Delete(ctx context.Context, path string) error { + full, err := s.fullPath(path) + if err != nil { + return err + } + + select { + case <-ctx.Done(): + return ctx.Err() + default: + } + + err = os.Remove(full) + if errors.Is(err, os.ErrNotExist) { + return nil + } + return err +} + +func (s *LocalStore) Move(ctx context.Context, srcPath, dstPath string) error { + srcFull, err := s.fullPath(srcPath) + if err != nil { + return err + } + + dstFull, err := s.fullPath(dstPath) + if err != nil { + return err + } + + if err := os.MkdirAll(filepath.Dir(dstFull), 0o755); err != nil { + return err + } + + select { + case <-ctx.Done(): + return ctx.Err() + default: + } + + return os.Rename(srcFull, dstFull) +} + +func (s *LocalStore) listRecursive(ctx context.Context, fullPrefix string) ([]FileInfo, error) { + var out []FileInfo + + err := filepath.WalkDir(fullPrefix, func(path string, d os.DirEntry, err error) error { + if err != nil { + return err + } + + select { + case <-ctx.Done(): + return ctx.Err() + default: + } + + if d.IsDir() { + return nil + } + + info, err := d.Info() + if err != nil { + return err + } + + relPath, err := filepath.Rel(s.root, path) + if err != nil { + return err + } + + out = append(out, FileInfo{ + Path: normalizeLogicalPath(relPath), + Size: info.Size(), + ContentType: detectContentType(path), + LastModified: info.ModTime(), + }) + + return nil + }) + + if errors.Is(err, os.ErrNotExist) { + return []FileInfo{}, nil + } + if err != nil { + return nil, err + } + + return out, nil +} + +func (s *LocalStore) listShallow(ctx context.Context, fullPrefix string) ([]FileInfo, error) { + entries, err := os.ReadDir(fullPrefix) + if errors.Is(err, os.ErrNotExist) { + return []FileInfo{}, nil + } + if err != nil { + return nil, err + } + + out := make([]FileInfo, 0, len(entries)) + + for _, entry := range entries { + select { + case <-ctx.Done(): + return nil, ctx.Err() + default: + } + + if entry.IsDir() { + continue + } + + info, err := entry.Info() + if err != nil { + return nil, err + } + + full := filepath.Join(fullPrefix, entry.Name()) + rel, err := filepath.Rel(s.root, full) + if err != nil { + return nil, err + } + + out = append(out, FileInfo{ + Path: normalizeLogicalPath(rel), + Size: info.Size(), + ContentType: detectContentType(full), + LastModified: info.ModTime(), + }) + } + + return out, nil +} + +func (s *LocalStore) List(ctx context.Context, path string, options ListOptions) ([]FileInfo, error) { + full, err := s.fullPath(path) + if err != nil { + return nil, err + } + + if options.Recursive { + return s.listRecursive(ctx, full) + } + + return s.listShallow(ctx, full) +} + +func (s *LocalStore) Copy(ctx context.Context, srcPath, dstPath string) error { + srcFull, err := s.fullPath(srcPath) + if err != nil { + return err + } + + dstFull, err := s.fullPath(dstPath) + if err != nil { + return err + } + + srcFile, err := os.Open(srcFull) + if err != nil { + return err + } + defer srcFile.Close() + + srcInfo, err := srcFile.Stat() + if err != nil { + return err + } + if srcInfo.IsDir() { + return errors.New("cannot copy a directory") + } + + if err := os.MkdirAll(filepath.Dir(dstFull), 0o755); err != nil { + return err + } + + tmp, err := os.CreateTemp(filepath.Dir(dstFull), ".tmp-*") + if err != nil { + return err + } + tmpName := tmp.Name() + + cleanup := func() { + _ = tmp.Close() + _ = os.Remove(tmpName) + } + + buf := make([]byte, 32*1024) + for { + select { + case <-ctx.Done(): + cleanup() + return ctx.Err() + default: + } + + n, readErr := srcFile.Read(buf) + if n > 0 { + if _, err := tmp.Write(buf[:n]); err != nil { + cleanup() + return err + } + } + + if readErr == io.EOF { + break + } + if readErr != nil { + cleanup() + return readErr + } + } + + if err := tmp.Close(); err != nil { + _ = os.Remove(tmpName) + return err + } + + return os.Rename(tmpName, dstFull) +} + diff --git a/cmd/api/src/services/storage/osstore.go b/cmd/api/src/services/storage/osstore.go index 0d1b8229e0c6..7f7660f24556 100644 --- a/cmd/api/src/services/storage/osstore.go +++ b/cmd/api/src/services/storage/osstore.go @@ -20,54 +20,43 @@ import ( "context" "errors" "io" + "io/fs" "os" - "path/filepath" + "path" ) -// OsStore is a Storage implementation that operates directly on the OS filesystem -// using absolute paths. It is intended for use in CE deployments where all paths -// are already fully-qualified (e.g. temp directory paths from config). -type OsStore struct{} - -func NewOsStore() *OsStore { - return &OsStore{} +// OsStore is a Storage implementation backed by a sandboxed directory on the +// local filesystem. All operations are confined to the root supplied to +// NewOsStore; 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. +type OsStore struct { + root *os.Root } -func (s *OsStore) Put(ctx context.Context, path string, reader io.Reader, options WriteOptions) error { - if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { - return err - } - - tmp, err := os.CreateTemp(filepath.Dir(path), ".tmp-*") +func NewOsStore(root string) (*OsStore, error) { + r, err := os.OpenRoot(root) if err != nil { - return err - } - tmpName := tmp.Name() - - cleanup := func() { - tmp.Close() - os.Remove(tmpName) - } - - if _, err := io.Copy(tmp, reader); err != nil { - cleanup() - return err + return nil, err } + return &OsStore{ + root: r, + }, nil +} - if err := tmp.Close(); err != nil { - _ = os.Remove(tmpName) - return err - } +func (s *OsStore) Close() error { + return s.root.Close() +} - return os.Rename(tmpName, path) +func (s *OsStore) Open(name string) (fs.File, error) { + return s.root.Open(name) // *os.File Satisfies fs.File } -func (s *OsStore) Get(ctx context.Context, path string) (io.ReadCloser, FileInfo, error) { - file, err := os.Open(path) +func (s *OsStore) Get(ctx context.Context, name string) (io.ReadCloser, FileInfo, error) { + file, err := s.root.Open(name) if err != nil { return nil, FileInfo{}, err } - stat, err := file.Stat() if err != nil { _ = file.Close() @@ -75,29 +64,80 @@ func (s *OsStore) Get(ctx context.Context, path string) (io.ReadCloser, FileInfo } return file, FileInfo{ - Path: path, + Path: name, Size: stat.Size(), LastModified: stat.ModTime(), IsDir: stat.IsDir(), }, nil } -func (s *OsStore) Stat(ctx context.Context, path string) (FileInfo, error) { - stat, err := os.Stat(path) +func (s *OsStore) Put(ctx context.Context, name string, reader io.Reader, options WriteOptions) error { + var ( + dir = path.Dir(name) + tmpName string + tmp *os.File + closed bool + err error + ) + if dir != "." { + if err := s.root.MkdirAll(dir, 0o750); err != nil { + return err + } + } + + id, err := randomID() + if err != nil { + return err + } + tmpName = path.Join(dir, ".tmp-"+id) + + tmp, err = s.root.OpenFile(tmpName, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o640) + if 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, reader); 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 + + err = s.root.Rename(tmpName, name) + return err +} + +func (s *OsStore) Stat(ctx context.Context, name string) (FileInfo, error) { + stat, err := s.root.Stat(name) if err != nil { return FileInfo{}, err } return FileInfo{ - Path: path, + Path: name, Size: stat.Size(), LastModified: stat.ModTime(), IsDir: stat.IsDir(), }, nil } -func (s *OsStore) Exists(ctx context.Context, path string) (bool, error) { - _, err := os.Stat(path) +func (s *OsStore) Exists(ctx context.Context, name string) (bool, error) { + _, err := s.root.Stat(name) if errors.Is(err, os.ErrNotExist) { return false, nil } @@ -107,49 +147,92 @@ func (s *OsStore) Exists(ctx context.Context, path string) (bool, error) { return true, nil } -func (s *OsStore) Delete(ctx context.Context, path string) error { - err := os.Remove(path) +func (s *OsStore) Delete(ctx context.Context, name string) error { + err := s.root.Remove(name) if errors.Is(err, os.ErrNotExist) { return nil } return err } -func (s *OsStore) List(ctx context.Context, path string, options ListOptions) ([]FileInfo, error) { - entries, err := os.ReadDir(path) - if errors.Is(err, os.ErrNotExist) { +func (s *OsStore) List(ctx context.Context, name string, options ListOptions) ([]FileInfo, error) { + fsys := s.root.FS() + if options.Recursive { + out := []FileInfo{} + err := fs.WalkDir(fsys, name, func(p string, d fs.DirEntry, err error) error { + if err != nil { + if errors.Is(err, fs.ErrNotExist) && p == name { + return fs.SkipAll + } + return err + } + if p == name && d.IsDir() { + return nil // skip the prefix directory itself + } + info, err := d.Info() + if err != nil { + return err + } + out = append(out, FileInfo{ + Path: p, + Size: info.Size(), + LastModified: info.ModTime(), + IsDir: d.IsDir(), + }) + 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, name) + if errors.Is(err, fs.ErrNotExist) { return []FileInfo{}, nil } if err != nil { return nil, err } - - var out []FileInfo + out := make([]FileInfo, 0, len(entries)) for _, entry := range entries { info, err := entry.Info() if err != nil { return nil, err } out = append(out, FileInfo{ - Path: filepath.Join(path, entry.Name()), + Path: path.Join(name, entry.Name()), Size: info.Size(), LastModified: info.ModTime(), IsDir: entry.IsDir(), }) + if options.Limit > 0 && len(out) >= options.Limit { + break + } } return out, nil } -func (s *OsStore) Copy(ctx context.Context, srcPath, dstPath string) error { - src, err := os.Open(srcPath) +// Copy duplicates srcName to destName atomically: on failure the destination is unchanged; +// on success, callers observe either the old content or hew, never a partial write. +func (s *OsStore) Copy(ctx context.Context, srcName, dstName string) error { + src, err := s.root.Open(srcName) if err != nil { return err } defer src.Close() - return s.Put(ctx, dstPath, src, WriteOptions{}) + return s.Put(ctx, dstName, src, WriteOptions{}) } -func (s *OsStore) Move(ctx context.Context, srcPath, dstPath string) error { - return os.Rename(srcPath, dstPath) +func (s *OsStore) Move(ctx context.Context, srcName, dstName string) error { + dir := path.Dir(dstName) + if dir != "." { + if err := s.root.MkdirAll(dir, 0o750); err != nil { + return err + } + } + return s.root.Rename(srcName, dstName) } diff --git a/cmd/api/src/services/storage/storage.go b/cmd/api/src/services/storage/storage.go index 282721dd9a01..a408d9f2e25d 100644 --- a/cmd/api/src/services/storage/storage.go +++ b/cmd/api/src/services/storage/storage.go @@ -21,7 +21,6 @@ import ( "context" "crypto/rand" "encoding/hex" - "fmt" "io" "path" "time" @@ -81,7 +80,6 @@ type FileService interface { ReadFile(ctx context.Context, name string) ([]byte, error) WriteFile(ctx context.Context, name string, data []byte, opts WriteOptions) error DeleteFile(ctx context.Context, name string) error - TempPath(prefix, pattern string) (string, error) WriteTempFile(ctx context.Context, prefix string, reader io.Reader, opts WriteOptions) (string, error) } @@ -119,20 +117,6 @@ func (s *LocalFileService) DeleteFile(ctx context.Context, name string) error { return s.Storage.Delete(ctx, name) } -func (s *LocalFileService) TempPath(prefix, pattern string) (string, error) { - id, err := randomID() - if err != nil { - return "", err - } - - name := pattern - if name == "" { - name = "tmp-*" - } - - return path.Join(prefix, fmt.Sprintf("%s-%s", name, id)), nil -} - func (s *LocalFileService) WriteTempFile(ctx context.Context, prefix string, reader io.Reader, opts WriteOptions) (string, error) { id, err := randomID() if err != nil { From cfe85186d3684545618ab2098ad0fe795dc634b1 Mon Sep 17 00:00:00 2001 From: Michael Cuomo Date: Thu, 23 Apr 2026 19:51:57 -0400 Subject: [PATCH 03/26] fix: refine the logic for localstore --- cmd/api/src/api/v2/fileingest_test.go | 4 +- cmd/api/src/api/v2/model.go | 2 +- cmd/api/src/services/storage/localstore.go | 505 +++++++++------------ cmd/api/src/services/storage/osstore.go | 238 ---------- cmd/api/src/services/storage/storage.go | 3 + 5 files changed, 220 insertions(+), 532 deletions(-) delete mode 100644 cmd/api/src/services/storage/osstore.go diff --git a/cmd/api/src/api/v2/fileingest_test.go b/cmd/api/src/api/v2/fileingest_test.go index 723d13c293b2..afeaaa769aab 100644 --- a/cmd/api/src/api/v2/fileingest_test.go +++ b/cmd/api/src/api/v2/fileingest_test.go @@ -580,7 +580,7 @@ func TestResources_ProcessIngestTask(t *testing.T) { }, // Use a real OsStore so that actual file I/O gives the concurrent validation // goroutine time to set validationErr before the main goroutine reads it. - fileServiceOvrd: storage.NewLocalFileService(storage.NewOsStore()), + fileServiceOvrd: storage.NewLocalFileService(storage.NewLocalStore()), 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"}`, @@ -611,7 +611,7 @@ func TestResources_ProcessIngestTask(t *testing.T) { mock.mockDatabase.EXPECT().GetIngestJob(gomock.Any(), int64(1)).Return(model.IngestJob{Status: model.JobStatusRunning}, nil) }, // Use a real OsStore for the same reason as the ErrInvalidJSON case above. - fileServiceOvrd: storage.NewLocalFileService(storage.NewOsStore()), + fileServiceOvrd: storage.NewLocalFileService(storage.NewLocalStore()), 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"}`, diff --git a/cmd/api/src/api/v2/model.go b/cmd/api/src/api/v2/model.go index e48cb8c1ce83..03934efe8413 100644 --- a/cmd/api/src/api/v2/model.go +++ b/cmd/api/src/api/v2/model.go @@ -145,7 +145,7 @@ func NewResources( Authorizer: authorizer, Authenticator: authenticator, IngestSchema: ingestSchema, - FileService: storage.NewLocalFileService(storage.NewOsStore()), + FileService: storage.NewLocalFileService(storage.NewLocalStore()), DogTags: dogtagsService, OpenGraphSchemaService: openGraphSchemaService, } diff --git a/cmd/api/src/services/storage/localstore.go b/cmd/api/src/services/storage/localstore.go index d9e34c89f8c2..87f36d3e86e9 100644 --- a/cmd/api/src/services/storage/localstore.go +++ b/cmd/api/src/services/storage/localstore.go @@ -19,414 +19,337 @@ package storage import ( "context" "errors" + "fmt" "io" + "io/fs" "mime" "os" - "path/filepath" - "strings" + "path" + "sync" ) -type LocalStore struct { - root string +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 NewLocalStore(root string) *LocalStore { - return &LocalStore{ - root: root, +func (s *ctxReader) Read(p []byte) (int, error) { + if err := s.ctx.Err(); err != nil { + return 0, err } + return s.r.Read(p) } -func (s *LocalStore) fullPath(path string) (string, error) { - path = filepath.ToSlash(strings.TrimSpace(path)) - path = strings.TrimPrefix(path, "/") - - cleaned := filepath.Clean(path) - - if cleaned == "." { - return "", errors.New("invalid empty path") - } +// 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 +} - if cleaned == ".." || strings.HasPrefix(cleaned, ".."+string(filepath.Separator)) { - return "", errors.New("path escapes storage root") +func NewLocalStore(root string) (*LocalStore, error) { + r, err := os.OpenRoot(root) + if err != nil { + return nil, err } - - full := filepath.Join(s.root, cleaned) - return full, nil + return &LocalStore{ + root: r, + }, nil } -func normalizeLogicalPath(path string) string { - path = filepath.ToSlash(strings.TrimSpace(path)) - path = strings.TrimPrefix(path, "/") - return filepath.Clean(path) +func (s *LocalStore) Close() error { + s.closeOnce.Do(func() { s.closeErr = s.root.Close() }) + return s.closeErr } -func detectContentType(path string) string { - ext := filepath.Ext(path) - if ext == "" { - return "application/octet-stream" - } - - ct := mime.TypeByExtension(ext) - if ct == "" { - return "application/octet-stream" +func detectContentType(name string) string { + if ext := path.Ext(name); ext != "" { + if ct := mime.TypeByExtension(ext); ct != "" { + return ct + } } - - return ct + return "application/octet-stream" } -func (s *LocalStore) Put(ctx context.Context, path string, reader io.Reader, options WriteOptions) error { - full, err := s.fullPath(path) - if err != nil { +// 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 } - dir := filepath.Dir(full) - if err := os.MkdirAll(dir, 0o755); err != nil { - return err + if dir != "." { + if err = s.root.MkdirAll(dir, 0o750); err != nil { + return err + } } - tmp, err := os.CreateTemp(dir, ".tmp-*") - if err != nil { + if id, err = randomID(); err != nil { return err } - tmpName := tmp.Name() + tmpName = path.Join(dir, ".tmp-"+id) - cleanup := func() { - tmp.Close() - os.Remove(tmpName) + if tmp, err = s.root.OpenFile(tmpName, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o640); err != nil { + return err } - buf := make([]byte, 32*1024) - for { - select { - case <-ctx.Done(): - cleanup() - return ctx.Err() - default: - } - - n, readErr := reader.Read(buf) - if n > 0 { - if _, err := tmp.Write(buf[:n]); err != nil { - cleanup() - 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 readErr == io.EOF { - break - } - if readErr != nil { - cleanup() - return readErr - } + if _, err = io.Copy(tmp, &ctxReader{ctx: ctx, r: src}); err != nil { + return err } - - if err := tmp.Close(); err != nil { - _ = os.Remove(tmpName) + if err = tmp.Sync(); err != nil { // flush data blocks return err } - - if err := os.Rename(tmpName, full); err != nil { - _ = os.Remove(tmpName) + if err = tmp.Close(); err != nil { return err } + closed = true - return nil + 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) Get(ctx context.Context, path string) (io.ReadCloser, FileInfo, error) { - full, err := s.fullPath(path) - if err != nil { - return nil, FileInfo{}, err +func (s *LocalStore) Stat(ctx context.Context, name string) (FileInfo, error) { + if err := ctx.Err(); err != nil { + return FileInfo{}, err } - - file, err := os.Open(full) - if err != nil { - return nil, FileInfo{}, err + stat, err := s.root.Stat(name) + if stat.IsDir() { + return FileInfo{}, fmt.Errorf("stat %q: %w", name, ErrIsDirectory) } - - stat, err := file.Stat() if err != nil { - _ = file.Close() - return nil, FileInfo{}, err + return FileInfo{}, err } - info := FileInfo{ - Path: normalizeLogicalPath(path), + return FileInfo{ + Path: name, Size: stat.Size(), LastModified: stat.ModTime(), IsDir: stat.IsDir(), - ContentType: detectContentType(path), - } - - select { - case <-ctx.Done(): - _ = file.Close() - return nil, FileInfo{}, ctx.Err() - default: - } - - return file, info, nil + ContentType: detectContentType(name), + }, nil } -func (s *LocalStore) Stat(ctx context.Context, path string) (FileInfo, error) { - full, err := s.fullPath(path) - if err != nil { - return FileInfo{}, err +func (s *LocalStore) Get(ctx context.Context, name string) (io.ReadCloser, FileInfo, error) { + if err := ctx.Err(); err != nil { + return nil, FileInfo{}, err } - - select { - case <-ctx.Done(): - return FileInfo{}, ctx.Err() - default: + file, err := s.root.Open(name) + if err != nil { + return nil, FileInfo{}, err } - - stat, err := os.Stat(full) + stat, err := file.Stat() if err != nil { - return FileInfo{}, err + _ = file.Close() + return nil, FileInfo{}, err + } + if stat.IsDir() { + return nil, FileInfo{}, fmt.Errorf("get: %q, %w", name, ErrIsDirectory) + } + if err := ctx.Err(); err != nil { + _ = file.Close() + return nil, FileInfo{}, err } - return FileInfo{ - Path: normalizeLogicalPath(path), + return file, FileInfo{ + Path: name, Size: stat.Size(), LastModified: stat.ModTime(), IsDir: stat.IsDir(), - ContentType: detectContentType(path), + ContentType: detectContentType(name), }, nil } -func (s *LocalStore) Exists(ctx context.Context, path string) (bool, error) { - _, err := s.Stat(ctx, path) - if err == nil { - return true, 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 } - return false, err -} - -func (s *LocalStore) Delete(ctx context.Context, path string) error { - full, err := s.fullPath(path) if err != nil { - return err + return false, err } - - select { - case <-ctx.Done(): - return ctx.Err() - default: + if stat.IsDir() { + return false, nil } + return true, nil +} - err = os.Remove(full) +func (s *LocalStore) Delete(ctx context.Context, name string) error { + if err := ctx.Err(); err != nil { + return err + } + err := s.root.Remove(name) if errors.Is(err, os.ErrNotExist) { return nil } return err } -func (s *LocalStore) Move(ctx context.Context, srcPath, dstPath string) error { - srcFull, err := s.fullPath(srcPath) - if err != nil { - return err - } - - dstFull, err := s.fullPath(dstPath) - if err != nil { - return err - } - - if err := os.MkdirAll(filepath.Dir(dstFull), 0o755); err != nil { - return err - } - - select { - case <-ctx.Done(): - return ctx.Err() - default: +func (s *LocalStore) List(ctx context.Context, name string, options ListOptions) ([]FileInfo, error) { + if err := ctx.Err(); err != nil { + return nil, err } - - return os.Rename(srcFull, dstFull) -} - -func (s *LocalStore) listRecursive(ctx context.Context, fullPrefix string) ([]FileInfo, error) { - var out []FileInfo - - err := filepath.WalkDir(fullPrefix, func(path string, d os.DirEntry, err error) error { - if err != nil { - return err - } - - select { - case <-ctx.Done(): - return ctx.Err() - default: - } - - if d.IsDir() { + fsys := s.root.FS() + if options.Recursive { + out := []FileInfo{} + err := fs.WalkDir(fsys, name, 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 == name { + 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 - } - - info, err := d.Info() - if err != nil { - return err - } - - relPath, err := filepath.Rel(s.root, path) + }) if err != nil { - return err + return nil, err } - - out = append(out, FileInfo{ - Path: normalizeLogicalPath(relPath), - Size: info.Size(), - ContentType: detectContentType(path), - LastModified: info.ModTime(), - }) - - return nil - }) - - if errors.Is(err, os.ErrNotExist) { - return []FileInfo{}, nil + return out, nil } - if err != nil { - return nil, err - } - - return out, nil -} - -func (s *LocalStore) listShallow(ctx context.Context, fullPrefix string) ([]FileInfo, error) { - entries, err := os.ReadDir(fullPrefix) - if errors.Is(err, os.ErrNotExist) { + entries, err := fs.ReadDir(fsys, name) + 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 { - select { - case <-ctx.Done(): - return nil, ctx.Err() - default: + if err := ctx.Err(); err != nil { + return nil, err } - if entry.IsDir() { continue } - info, err := entry.Info() if err != nil { return nil, err } - - full := filepath.Join(fullPrefix, entry.Name()) - rel, err := filepath.Rel(s.root, full) - if err != nil { - return nil, err - } - + entryPath := path.Join(name, entry.Name()) out = append(out, FileInfo{ - Path: normalizeLogicalPath(rel), + Path: entryPath, Size: info.Size(), - ContentType: detectContentType(full), LastModified: info.ModTime(), + IsDir: entry.IsDir(), + ContentType: detectContentType(entryPath), }) + if options.Limit > 0 && len(out) >= options.Limit { + break + } } - return out, nil } -func (s *LocalStore) List(ctx context.Context, path string, options ListOptions) ([]FileInfo, error) { - full, err := s.fullPath(path) - if err != nil { - return nil, err - } - - if options.Recursive { - return s.listRecursive(ctx, full) - } - - return s.listShallow(ctx, full) -} - -func (s *LocalStore) Copy(ctx context.Context, srcPath, dstPath string) error { - srcFull, err := s.fullPath(srcPath) - if err != nil { - return err - } - - dstFull, err := s.fullPath(dstPath) - if err != 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 } - - srcFile, err := os.Open(srcFull) + src, err := s.root.Open(srcName) if err != nil { return err } - defer srcFile.Close() - - srcInfo, err := srcFile.Stat() + defer src.Close() + info, err := src.Stat() if err != nil { return err } - if srcInfo.IsDir() { - return errors.New("cannot copy a directory") - } - - if err := os.MkdirAll(filepath.Dir(dstFull), 0o755); err != nil { - return err + if info.IsDir() { + return fmt.Errorf("copy %q: %w", srcName, ErrIsDirectory) } + return s.writeAtomic(ctx, dstName, src, options.FailIfExists) +} - tmp, err := os.CreateTemp(filepath.Dir(dstFull), ".tmp-*") - if err != nil { +// Move is able to move a file from srcName to dstName using a two-step Link+Remove. A crash between link +// and unlink leaves both names present. +func (s *LocalStore) Move(ctx context.Context, srcName, dstName string, options WriteOptions) error { + if err := ctx.Err(); err != nil { return err } - tmpName := tmp.Name() - - cleanup := func() { - _ = tmp.Close() - _ = os.Remove(tmpName) - } - - buf := make([]byte, 32*1024) - for { - select { - case <-ctx.Done(): - cleanup() - return ctx.Err() - default: - } - - n, readErr := srcFile.Read(buf) - if n > 0 { - if _, err := tmp.Write(buf[:n]); err != nil { - cleanup() - return err - } - } - - if readErr == io.EOF { - break - } - if readErr != nil { - cleanup() - return readErr + dir := path.Dir(dstName) + if dir != "." { + if err := s.root.MkdirAll(dir, 0o750); err != nil { + return err } } - - if err := tmp.Close(); err != nil { - _ = os.Remove(tmpName) - return err + if options.FailIfExists { + if err := s.root.Link(srcName, dstName); err != nil { + return err + } + return s.root.Remove(srcName) } - - return os.Rename(tmpName, dstFull) + return s.root.Rename(srcName, dstName) } - diff --git a/cmd/api/src/services/storage/osstore.go b/cmd/api/src/services/storage/osstore.go deleted file mode 100644 index 7f7660f24556..000000000000 --- a/cmd/api/src/services/storage/osstore.go +++ /dev/null @@ -1,238 +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 - -package storage - -import ( - "context" - "errors" - "io" - "io/fs" - "os" - "path" -) - -// OsStore is a Storage implementation backed by a sandboxed directory on the -// local filesystem. All operations are confined to the root supplied to -// NewOsStore; 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. -type OsStore struct { - root *os.Root -} - -func NewOsStore(root string) (*OsStore, error) { - r, err := os.OpenRoot(root) - if err != nil { - return nil, err - } - return &OsStore{ - root: r, - }, nil -} - -func (s *OsStore) Close() error { - return s.root.Close() -} - -func (s *OsStore) Open(name string) (fs.File, error) { - return s.root.Open(name) // *os.File Satisfies fs.File -} - -func (s *OsStore) Get(ctx context.Context, name string) (io.ReadCloser, FileInfo, error) { - 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 - } - - return file, FileInfo{ - Path: name, - Size: stat.Size(), - LastModified: stat.ModTime(), - IsDir: stat.IsDir(), - }, nil -} - -func (s *OsStore) Put(ctx context.Context, name string, reader io.Reader, options WriteOptions) error { - var ( - dir = path.Dir(name) - tmpName string - tmp *os.File - closed bool - err error - ) - if dir != "." { - if err := s.root.MkdirAll(dir, 0o750); err != nil { - return err - } - } - - id, err := randomID() - if err != nil { - return err - } - tmpName = path.Join(dir, ".tmp-"+id) - - tmp, err = s.root.OpenFile(tmpName, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o640) - if 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, reader); 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 - - err = s.root.Rename(tmpName, name) - return err -} - -func (s *OsStore) Stat(ctx context.Context, name string) (FileInfo, error) { - stat, err := s.root.Stat(name) - if err != nil { - return FileInfo{}, err - } - - return FileInfo{ - Path: name, - Size: stat.Size(), - LastModified: stat.ModTime(), - IsDir: stat.IsDir(), - }, nil -} - -func (s *OsStore) Exists(ctx context.Context, name string) (bool, error) { - _, err := s.root.Stat(name) - if errors.Is(err, os.ErrNotExist) { - return false, nil - } - if err != nil { - return false, err - } - return true, nil -} - -func (s *OsStore) Delete(ctx context.Context, name string) error { - err := s.root.Remove(name) - if errors.Is(err, os.ErrNotExist) { - return nil - } - return err -} - -func (s *OsStore) List(ctx context.Context, name string, options ListOptions) ([]FileInfo, error) { - fsys := s.root.FS() - if options.Recursive { - out := []FileInfo{} - err := fs.WalkDir(fsys, name, func(p string, d fs.DirEntry, err error) error { - if err != nil { - if errors.Is(err, fs.ErrNotExist) && p == name { - return fs.SkipAll - } - return err - } - if p == name && d.IsDir() { - return nil // skip the prefix directory itself - } - info, err := d.Info() - if err != nil { - return err - } - out = append(out, FileInfo{ - Path: p, - Size: info.Size(), - LastModified: info.ModTime(), - IsDir: d.IsDir(), - }) - 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, name) - 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 { - info, err := entry.Info() - if err != nil { - return nil, err - } - out = append(out, FileInfo{ - Path: path.Join(name, entry.Name()), - Size: info.Size(), - LastModified: info.ModTime(), - IsDir: entry.IsDir(), - }) - 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 hew, never a partial write. -func (s *OsStore) Copy(ctx context.Context, srcName, dstName string) error { - src, err := s.root.Open(srcName) - if err != nil { - return err - } - defer src.Close() - - return s.Put(ctx, dstName, src, WriteOptions{}) -} - -func (s *OsStore) Move(ctx context.Context, srcName, dstName string) error { - dir := path.Dir(dstName) - if dir != "." { - if err := s.root.MkdirAll(dir, 0o750); err != nil { - return err - } - } - return s.root.Rename(srcName, dstName) -} diff --git a/cmd/api/src/services/storage/storage.go b/cmd/api/src/services/storage/storage.go index a408d9f2e25d..0ab89e70d890 100644 --- a/cmd/api/src/services/storage/storage.go +++ b/cmd/api/src/services/storage/storage.go @@ -40,6 +40,9 @@ type FileInfo struct { 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 { From 1119be2bf28ad240f15c78090c0e4518cb7b7964 Mon Sep 17 00:00:00 2001 From: Michael Cuomo Date: Thu, 30 Apr 2026 22:23:31 -0400 Subject: [PATCH 04/26] wip: cleanup storage interface --- cmd/api/src/api/v2/model.go | 8 +++++++- cmd/api/src/services/storage/localstore.go | 6 +++--- cmd/api/src/services/storage/storage.go | 19 +++++++++---------- 3 files changed, 19 insertions(+), 14 deletions(-) diff --git a/cmd/api/src/api/v2/model.go b/cmd/api/src/api/v2/model.go index 03934efe8413..4e0a18b17420 100644 --- a/cmd/api/src/api/v2/model.go +++ b/cmd/api/src/api/v2/model.go @@ -17,6 +17,8 @@ package v2 import ( + "log/slog" + "github.com/gorilla/schema" "github.com/specterops/bloodhound/cmd/api/src/api" "github.com/specterops/bloodhound/cmd/api/src/auth" @@ -133,6 +135,10 @@ func NewResources( dogtagsService dogtags.Service, openGraphSchemaService OpenGraphSchemaService, ) Resources { + localStore, err := storage.NewLocalStore("/") + if err != nil { + slog.Error("Error creating local store") + } return Resources{ Decoder: schema.NewDecoder(), DB: rdms, @@ -145,7 +151,7 @@ func NewResources( Authorizer: authorizer, Authenticator: authenticator, IngestSchema: ingestSchema, - FileService: storage.NewLocalFileService(storage.NewLocalStore()), + FileService: storage.NewLocalFileService(localStore), DogTags: dogtagsService, OpenGraphSchemaService: openGraphSchemaService, } diff --git a/cmd/api/src/services/storage/localstore.go b/cmd/api/src/services/storage/localstore.go index 87f36d3e86e9..0359a5916d1a 100644 --- a/cmd/api/src/services/storage/localstore.go +++ b/cmd/api/src/services/storage/localstore.go @@ -159,12 +159,12 @@ func (s *LocalStore) Stat(ctx context.Context, name string) (FileInfo, error) { return FileInfo{}, err } stat, err := s.root.Stat(name) - if stat.IsDir() { - return FileInfo{}, fmt.Errorf("stat %q: %w", name, ErrIsDirectory) - } if err != nil { return FileInfo{}, err } + if stat.IsDir() { + return FileInfo{}, fmt.Errorf("stat %q: %w", name, ErrIsDirectory) + } return FileInfo{ Path: name, diff --git a/cmd/api/src/services/storage/storage.go b/cmd/api/src/services/storage/storage.go index 0ab89e70d890..98a38c3b5576 100644 --- a/cmd/api/src/services/storage/storage.go +++ b/cmd/api/src/services/storage/storage.go @@ -54,29 +54,28 @@ type ListOptions struct { // in a variety of storage backends. type Storage interface { // Put writes a file at the given path. - Put(ctx context.Context, path string, reader io.Reader, options WriteOptions) error + Put(ctx context.Context, name string, reader io.Reader, options WriteOptions) error // Get opens a file for reading. - Get(ctx context.Context, path string) (io.ReadCloser, FileInfo, error) + Get(ctx context.Context, name string) (io.ReadCloser, FileInfo, error) // Stat returns metadata for a given path. - Stat(ctx context.Context, path string) (FileInfo, error) + Stat(ctx context.Context, name string) (FileInfo, error) // Delete removes a file. - Delete(ctx context.Context, path string) error + Delete(ctx context.Context, name string) error // Exists checks whether a file exists. - Exists(ctx context.Context, path string) (bool, error) + Exists(ctx context.Context, name string) (bool, error) // List returns a list of files in a given path. - List(ctx context.Context, path string, options ListOptions) ([]FileInfo, error) + List(ctx context.Context, name string, options ListOptions) ([]FileInfo, error) // Copy duplicates an object. - Copy(ctx context.Context, srcPath, dstPath string) error + 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, srcPath, dstPath string) error + // Move moves an object. Is done by a copy and a delete. + Move(ctx context.Context, srcName, dstName string, options WriteOptions) error } type FileService interface { From acfc81f704a06289a65768401827a786aa5930a6 Mon Sep 17 00:00:00 2001 From: Michael Cuomo Date: Fri, 1 May 2026 17:06:49 -0400 Subject: [PATCH 05/26] fix: close directory on error for Get --- cmd/api/src/services/storage/localstore.go | 1 + 1 file changed, 1 insertion(+) diff --git a/cmd/api/src/services/storage/localstore.go b/cmd/api/src/services/storage/localstore.go index 0359a5916d1a..62c831fb30e5 100644 --- a/cmd/api/src/services/storage/localstore.go +++ b/cmd/api/src/services/storage/localstore.go @@ -189,6 +189,7 @@ func (s *LocalStore) Get(ctx context.Context, name string) (io.ReadCloser, FileI 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 { From 37268b1976c4fa61ffe7205e760df2f1fc9ab2dc Mon Sep 17 00:00:00 2001 From: Michael Cuomo Date: Wed, 6 May 2026 15:12:11 -0400 Subject: [PATCH 06/26] feat: pipe file service resolver around resources --- cmd/api/src/api/registration/registration.go | 4 +- cmd/api/src/api/v2/collectors.go | 11 ++- cmd/api/src/api/v2/collectors_test.go | 8 +- cmd/api/src/api/v2/fileingest.go | 5 +- cmd/api/src/api/v2/fileingest_test.go | 10 +- cmd/api/src/api/v2/model.go | 15 ++- cmd/api/src/services/entrypoint.go | 7 +- cmd/api/src/services/storage/storage.go | 99 +++++++++++++++++++- cmd/api/src/services/upload/upload.go | 2 +- 9 files changed, 137 insertions(+), 24 deletions(-) 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/v2/collectors.go b/cmd/api/src/api/v2/collectors.go index 940a4f0b432e..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(request.Context(), 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(request.Context(), 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 ec642e9f79bc..642cb7bbde35 100644 --- a/cmd/api/src/api/v2/collectors_test.go +++ b/cmd/api/src/api/v2/collectors_test.go @@ -162,8 +162,8 @@ func TestResources_DownloadCollectorByVersion(t *testing.T) { } resources := v2.Resources{ - CollectorManifests: collectorManifests, - FileService: mock.mockFS, + CollectorManifests: collectorManifests, + FileServiceResolver: mock.mockFS, } response := httptest.NewRecorder() @@ -309,8 +309,8 @@ func TestResources_DownloadCollectorChecksumByVersion(t *testing.T) { testCase.setupMocks(t, mock) resources := v2.Resources{ - CollectorManifests: collectorManifests, - FileService: mock.mockFS, + CollectorManifests: collectorManifests, + FileServiceResolver: mock.mockFS, } response := httptest.NewRecorder() diff --git a/cmd/api/src/api/v2/fileingest.go b/cmd/api/src/api/v2/fileingest.go index 83df8c8eaaec..1d5c1b354994 100644 --- a/cmd/api/src/api/v2/fileingest.go +++ b/cmd/api/src/api/v2/fileingest.go @@ -38,6 +38,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 +147,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(request.Context(), s.FileService, 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 ( diff --git a/cmd/api/src/api/v2/fileingest_test.go b/cmd/api/src/api/v2/fileingest_test.go index afeaaa769aab..3c126ea9f98d 100644 --- a/cmd/api/src/api/v2/fileingest_test.go +++ b/cmd/api/src/api/v2/fileingest_test.go @@ -580,7 +580,7 @@ func TestResources_ProcessIngestTask(t *testing.T) { }, // Use a real OsStore so that actual file I/O gives the concurrent validation // goroutine time to set validationErr before the main goroutine reads it. - fileServiceOvrd: storage.NewLocalFileService(storage.NewLocalStore()), + fileServiceOvrd: storage.NewFileService(storage.NewLocalStore()), 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"}`, @@ -611,7 +611,7 @@ func TestResources_ProcessIngestTask(t *testing.T) { mock.mockDatabase.EXPECT().GetIngestJob(gomock.Any(), int64(1)).Return(model.IngestJob{Status: model.JobStatusRunning}, nil) }, // Use a real OsStore for the same reason as the ErrInvalidJSON case above. - fileServiceOvrd: storage.NewLocalFileService(storage.NewLocalStore()), + fileServiceOvrd: storage.NewFileService(storage.NewLocalStore()), 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"}`, @@ -809,9 +809,9 @@ func TestResources_ProcessIngestTask(t *testing.T) { } resources := v2.Resources{ - DB: mocks.mockDatabase, - Config: config.Configuration{}, - FileService: fileService, + DB: mocks.mockDatabase, + Config: config.Configuration{}, + FileServiceResolver: fileService, } 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 4e0a18b17420..93b9a2aff18b 100644 --- a/cmd/api/src/api/v2/model.go +++ b/cmd/api/src/api/v2/model.go @@ -17,8 +17,6 @@ package v2 import ( - "log/slog" - "github.com/gorilla/schema" "github.com/specterops/bloodhound/cmd/api/src/api" "github.com/specterops/bloodhound/cmd/api/src/auth" @@ -117,7 +115,7 @@ type Resources struct { Authorizer auth.Authorizer Authenticator api.Authenticator IngestSchema upload.IngestSchema - FileService storage.FileService + FileServiceResolver storage.FileServiceResolver OpenGraphSchemaService OpenGraphSchemaService DogTags dogtags.Service } @@ -132,13 +130,14 @@ func NewResources( authorizer auth.Authorizer, authenticator api.Authenticator, ingestSchema upload.IngestSchema, + fileServiceResolver storage.FileServiceResolver, dogtagsService dogtags.Service, openGraphSchemaService OpenGraphSchemaService, ) Resources { - localStore, err := storage.NewLocalStore("/") - if err != nil { - slog.Error("Error creating local store") - } + // localStore, err := storage.NewLocalStore("/tmp/") + // if err != nil { + // slog.Error("Error creating local store") + // } return Resources{ Decoder: schema.NewDecoder(), DB: rdms, @@ -151,7 +150,7 @@ func NewResources( Authorizer: authorizer, Authenticator: authenticator, IngestSchema: ingestSchema, - FileService: storage.NewLocalFileService(localStore), + FileServiceResolver: fileServiceResolver, DogTags: dogtagsService, OpenGraphSchemaService: openGraphSchemaService, } diff --git a/cmd/api/src/services/entrypoint.go b/cmd/api/src/services/entrypoint.go index 5c217551427b..9436a060dc10 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" @@ -132,6 +133,10 @@ func Entrypoint(ctx context.Context, cfg config.Configuration, connections boots return nil, fmt.Errorf("failed to save collector manifests: %w", err) } else if ingestSchema, err := upload.LoadIngestSchema(); err != nil { return nil, fmt.Errorf("failed to load OpenGraph schema: %w", err) + } else if fileServices, err := storage.NewDefaultFileServices(cfg); err != nil { + return nil, fmt.Errorf("failed to intialize file services: %w", err) + } else if fileServiceResolver, err := storage.NewFileServiceResolver(fileServices); err != nil { + return nil, fmt.Errorf("failed to initialize file service resolver: %w", err) } else { startDelay := 0 * time.Second @@ -147,7 +152,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, fileServiceResolver, dogtagsService, openGraphSchemaService) // Set neo4j batch and flush sizes neo4jParameters := appcfg.GetNeo4jParameters(ctx, connections.RDMS) diff --git a/cmd/api/src/services/storage/storage.go b/cmd/api/src/services/storage/storage.go index 98a38c3b5576..f1517694b487 100644 --- a/cmd/api/src/services/storage/storage.go +++ b/cmd/api/src/services/storage/storage.go @@ -21,13 +21,28 @@ import ( "context" "crypto/rand" "encoding/hex" + "errors" + "fmt" "io" "path" "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 . FileService +type FileServiceName string + +const ( + FileServiceIngest FileServiceName = "ingest" + 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 @@ -89,7 +104,7 @@ type LocalFileService struct { Storage Storage } -func NewLocalFileService(storage Storage) *LocalFileService { +func NewFileService(storage Storage) *LocalFileService { return &LocalFileService{Storage: storage} } @@ -132,3 +147,85 @@ func (s *LocalFileService) WriteTempFile(ctx context.Context, prefix string, rea return tempPath, nil } + +// 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(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 +} + +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 + } + + collectorsStore, err := NewLocalStore(cfg.CollectorsBasePath) + if err != nil { + return nil, err + } + + workService := NewFileService(workStore) + fileServices[FileServiceWork] = workService + + ingestService := NewFileService(ingestStore) + fileServices[FileServiceIngest] = ingestService + + collectersService := NewFileService(collectorsStore) + fileServices[FileServiceCollectors] = collectersService + + return fileServices, nil +} diff --git a/cmd/api/src/services/upload/upload.go b/cmd/api/src/services/upload/upload.go index d7713bd97b4b..c4626ac1aeab 100644 --- a/cmd/api/src/services/upload/upload.go +++ b/cmd/api/src/services/upload/upload.go @@ -54,7 +54,7 @@ func SaveIngestFile(ctx context.Context, fileService storage.FileService, reques return IngestTaskParams{}, fmt.Errorf("invalid content type for ingest file") } - if tempFileName, err := WriteAndValidateFile(ctx, fileService, fileData, fmt.Sprintf("file_upload_job%d_", jobID), validationFn); err != nil { + if tempFileName, err := WriteAndValidateFile(ctx, fileService, fileData, fmt.Sprintf("tmp/file_upload_job%d_", jobID), validationFn); err != nil { return IngestTaskParams{}, err } else { return IngestTaskParams{ From 540f470b5b0dba2bff0ffd72f4ce89490e644326 Mon Sep 17 00:00:00 2001 From: Michael Cuomo Date: Thu, 7 May 2026 09:57:28 -0400 Subject: [PATCH 07/26] feat: add get and move to FileService --- cmd/api/src/services/storage/storage.go | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/cmd/api/src/services/storage/storage.go b/cmd/api/src/services/storage/storage.go index f1517694b487..26a3dced8404 100644 --- a/cmd/api/src/services/storage/storage.go +++ b/cmd/api/src/services/storage/storage.go @@ -94,10 +94,12 @@ type Storage interface { } type FileService interface { + GetFile(ctx context.Context, name string) (io.ReadCloser, FileInfo, error) ReadFile(ctx context.Context, name string) ([]byte, error) WriteFile(ctx context.Context, name string, data []byte, opts WriteOptions) error DeleteFile(ctx context.Context, name string) error WriteTempFile(ctx context.Context, prefix string, reader io.Reader, opts WriteOptions) (string, error) + MoveFile(ctx context.Context, srcName, dstName string, opts WriteOptions) error } type LocalFileService struct { @@ -116,6 +118,10 @@ func randomID() (string, error) { return hex.EncodeToString(b[:]), nil } +func (s *LocalFileService) GetFile(ctx context.Context, name string) (io.ReadCloser, FileInfo, error) { + return s.Storage.Get(ctx, name) +} + func (s *LocalFileService) ReadFile(ctx context.Context, name string) ([]byte, error) { rc, _, err := s.Storage.Get(ctx, name) if err != nil { @@ -148,6 +154,10 @@ func (s *LocalFileService) WriteTempFile(ctx context.Context, prefix string, rea return tempPath, nil } +func (s *LocalFileService) MoveFile(ctx context.Context, srcName, dstName string, options WriteOptions) error { + return s.Storage.Move(ctx, srcName, dstName, options) +} + // 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 { From 8b4cc0b177c7c0cc314e1017fa0d7df2811b3ed3 Mon Sep 17 00:00:00 2001 From: Michael Cuomo Date: Thu, 7 May 2026 13:52:29 -0400 Subject: [PATCH 08/26] fix: centralize ingest logic, add fileservice to datapipe --- cmd/api/src/api/tools/ingest.go | 2 + cmd/api/src/api/v2/fileingest.go | 3 +- cmd/api/src/daemons/datapipe/cleanup.go | 2 + cmd/api/src/daemons/datapipe/pipeline.go | 7 +- cmd/api/src/services/entrypoint.go | 2 +- .../src/services/graphify/ingest_storage.go | 188 ++++++++++++++++++ cmd/api/src/services/graphify/service.go | 33 +-- cmd/api/src/services/graphify/tasks.go | 120 ++++------- cmd/api/src/services/upload/upload.go | 2 +- 9 files changed, 258 insertions(+), 101 deletions(-) create mode 100644 cmd/api/src/services/graphify/ingest_storage.go diff --git a/cmd/api/src/api/tools/ingest.go b/cmd/api/src/api/tools/ingest.go index 9cb7700db532..62cf792ff8d1 100644 --- a/cmd/api/src/api/tools/ingest.go +++ b/cmd/api/src/api/tools/ingest.go @@ -87,6 +87,7 @@ func (s *IngestControl) FetchRetainedIngestFiles(response http.ResponseWriter, r retainedFilesDirectory := s.cfg.RetainedFilesDirectory() + // TODO MC: should there be a retained fileService? 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. @@ -228,6 +229,7 @@ func (s *IngestControl) DisableIngestFileRetention(response http.ResponseWriter, retainedFilesDirectory := s.cfg.RetainedFilesDirectory() + // TODO MC: should there be a retained fileService if retainedFilesDirEntries, err := os.ReadDir(retainedFilesDirectory); err != nil { slog.WarnContext( request.Context(), diff --git a/cmd/api/src/api/v2/fileingest.go b/cmd/api/src/api/v2/fileingest.go index 1d5c1b354994..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" @@ -172,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/daemons/datapipe/cleanup.go b/cmd/api/src/daemons/datapipe/cleanup.go index 08eb8feeb91c..b99e60f95c34 100644 --- a/cmd/api/src/daemons/datapipe/cleanup.go +++ b/cmd/api/src/daemons/datapipe/cleanup.go @@ -30,6 +30,8 @@ import ( "github.com/specterops/bloodhound/packages/go/bhlog/attr" ) +// TODO MC: update cleanup to take into account the new file service + // 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 { diff --git a/cmd/api/src/daemons/datapipe/pipeline.go b/cmd/api/src/daemons/datapipe/pipeline.go index c54af65283d6..760ed4b80bad 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,12 +49,13 @@ 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, @@ -61,8 +63,9 @@ func NewPipeline(ctx context.Context, cfg config.Configuration, db database.Data cfg: cfg, orphanedFileSweeper: NewOrphanFileSweeper(NewOSFileOperations(), cfg.TempDirectory()), 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, } } diff --git a/cmd/api/src/services/entrypoint.go b/cmd/api/src/services/entrypoint.go index 9436a060dc10..2857fc48d7ad 100644 --- a/cmd/api/src/services/entrypoint.go +++ b/cmd/api/src/services/entrypoint.go @@ -142,7 +142,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, 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) 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..48d982b07e49 --- /dev/null +++ b/cmd/api/src/services/graphify/ingest_storage.go @@ -0,0 +1,188 @@ +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, tempDirectory 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(tempDirectory, "ingest-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, tempDirectory string, fileService storage.FileService, storedFileName string) (*os.File, string, error) { + var ( + scratchPath string + scratchFile *os.File + err error + ) + + if scratchPath, err = SpoolToScratch(ctx, tempDirectory, 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{}, // TODO MC: include write options? + ) + if err != nil { + return "", fmt.Errorf("write archive file %q to stroage: %w", archiveFile.Name, err) + } + + return extractedPath, nil +} + +func ExtractIngestFiles(ctx context.Context, tempDirectory 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, tempDirectory, 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, fmt.Sprintf("tmp/file_upload_job%d_", jobID)); err != nil { + if extractedPath, err := WriteArchiveFileToStorage(ctx, fileService, archiveFile, prefix); err != nil { + processedFileData.Errors = []string{err.Error()} + + errs.Add(fmt.Errorf( + "error extracting file %s in archive %s: %w", + archiveFile.Name, + storedFileName, + err, + )) + } else { + processedFileData.Path = extractedPath + } + + fileData = append(fileData, processedFileData) + } + + if err := fileService.DeleteFile(ctx, storedFileName); err != nil { + slog.ErrorContext( + ctx, + "Error deleting archive", + slog.String("path", storedFileName), + attr.Error(err), + ) + } + + return fileData, errs.Combined() +} 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..bd819ca27de0 100644 --- a/cmd/api/src/services/graphify/tasks.go +++ b/cmd/api/src/services/graphify/tasks.go @@ -20,6 +20,7 @@ import ( "archive/zip" "context" "errors" + "fmt" "io" "io/fs" "log/slog" @@ -29,6 +30,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/packages/go/bhlog/attr" "github.com/specterops/bloodhound/packages/go/bhlog/measure" "github.com/specterops/bloodhound/packages/go/bomenc" @@ -57,73 +59,6 @@ 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") @@ -162,10 +97,10 @@ func (s *GraphifyService) extractToTempFile(f *zip.File) (string, error) { // 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.TempDirectory(), fileService, task.StoredFileName, task.OriginalFileName, task.FileType, fmt.Sprintf("tmp/file_upload_job%d_", ic.JobId)); err != nil { + return fileData, err } else { errs := errorlist.NewBuilder() @@ -179,7 +114,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.TempDirectory(), data, ic, readOpts); err != nil { var ( graphifyError errorlist.Error resolutionErr endpoint.ResolutionError @@ -224,10 +163,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 +178,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 +210,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 +235,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 +266,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/upload/upload.go b/cmd/api/src/services/upload/upload.go index c4626ac1aeab..d7713bd97b4b 100644 --- a/cmd/api/src/services/upload/upload.go +++ b/cmd/api/src/services/upload/upload.go @@ -54,7 +54,7 @@ func SaveIngestFile(ctx context.Context, fileService storage.FileService, reques return IngestTaskParams{}, fmt.Errorf("invalid content type for ingest file") } - if tempFileName, err := WriteAndValidateFile(ctx, fileService, fileData, fmt.Sprintf("tmp/file_upload_job%d_", 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{ From 23f408aa31c49c007964345702fa238f80d58cca Mon Sep 17 00:00:00 2001 From: Michael Cuomo Date: Thu, 7 May 2026 17:24:59 -0400 Subject: [PATCH 09/26] fix: alter cleanup logic to handle orphaned files --- cmd/api/src/bootstrap/util.go | 4 + cmd/api/src/config/config.go | 4 + cmd/api/src/daemons/datapipe/cleanup.go | 249 +++++++++++------- cmd/api/src/daemons/datapipe/pipeline.go | 6 +- .../src/services/graphify/ingest_storage.go | 26 +- cmd/api/src/services/graphify/tasks.go | 4 +- cmd/api/src/services/storage/storage.go | 5 + 7 files changed, 185 insertions(+), 113 deletions(-) diff --git a/cmd/api/src/bootstrap/util.go b/cmd/api/src/bootstrap/util.go index 8be699c1ec07..e239aa90d1a1 100644 --- a/cmd/api/src/bootstrap/util.go +++ b/cmd/api/src/bootstrap/util.go @@ -57,6 +57,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/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/datapipe/cleanup.go b/cmd/api/src/daemons/datapipe/cleanup.go index b99e60f95c34..efae977d9feb 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,11 +29,17 @@ import ( "strings" "sync" + "github.com/specterops/bloodhound/cmd/api/src/services/storage" "github.com/specterops/bloodhound/packages/go/bhlog/attr" ) // TODO MC: update cleanup to take into account the new file service +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 { @@ -57,112 +65,161 @@ 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 } -func NewOrphanFileSweeper(fileOps FileOperations, tempDirectoryRootPath string) *OrphanFileSweeper { +func NewOrphanFileSweeper(fileOps FileOperations, tempDirectoryRootPath string, scratchDirectoryRootPath string) *OrphanFileSweeper { return &OrphanFileSweeper{ - lock: &sync.Mutex{}, - fileOps: fileOps, - tempDirectoryRootPath: strings.TrimSuffix(tempDirectoryRootPath, string(filepath.Separator)), + lock: &sync.Mutex{}, + fileOps: fileOps, + tempDirectoryRootPath: strings.TrimSuffix(tempDirectoryRootPath, string(filepath.Separator)), + scratchDirectoryRootPath: strings.TrimSuffix(scratchDirectoryRootPath, string(filepath.Separator)), } } -// 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() { +func ClearLocalIngestScratch(ctx context.Context, scratchDirectory string, minimumAge time.Duration) { + entries, err := os.ReadDir(scratchDirectory) + if err != nil { + slog.ErrorContext(ctx, "Error reading scratch directory", slog.String("scratch_directory", scratchDirectory), attr.Error(err)) return } - // Release the lock once finished - defer s.lock.Unlock() + cutoff := time.Now().Add(-minimumAge) + + for _, entry := range entries { + if ctx.Err() != nil { + return + } + + 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())) + } +} + +// addExpectedStoragePath handles saved file paths that were saved in absolute prior to this release +func (s *OrphanFileSweeper) addExpectedStoragePath(expectedFiles map[string]struct{}, expectedFileName string) { + if expectedFileName == "" { + return + } - 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 { - slog.ErrorContext( - ctx, - "Failed reading work directory", - slog.String("path", s.tempDirectoryRootPath), - 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()) - - if err := s.fileOps.RemoveAll(fullPath); err != nil { - slog.ErrorContext( - ctx, - "Failed removing orphaned file", - slog.String("full_path", fullPath), - attr.Error(err), - ) - } - - numDeleted += 1 - } - - if numDeleted > 0 { - slog.InfoContext( - ctx, - "Finished removing orphaned ingest files", - slog.Int("num_deleted", numDeleted), - ) + 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 } + + expectedFiles[path.Clean(filepath.ToSlash(expectedFileName))] = 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 _, 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, "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 { + if filepath.IsAbs(expectedFileName) { + expectedLocalFiles[filepath.Clean(expectedFileName)] = struct{}{} + } + } + + 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/pipeline.go b/cmd/api/src/daemons/datapipe/pipeline.go index 760ed4b80bad..7186f84e5d74 100644 --- a/cmd/api/src/daemons/datapipe/pipeline.go +++ b/cmd/api/src/daemons/datapipe/pipeline.go @@ -61,7 +61,7 @@ func NewPipeline(ctx context.Context, cfg config.Configuration, db database.Data graphdb: graphDB, cache: cache, cfg: cfg, - orphanedFileSweeper: NewOrphanFileSweeper(NewOSFileOperations(), cfg.TempDirectory()), + orphanedFileSweeper: NewOrphanFileSweeper(NewOSFileOperations(), cfg.TempDirectory(), cfg.ScratchDirectory()), ingestSchema: ingestSchema, fileServiceResolver: fileServiceResolver, jobService: job.NewJobService(ctx, db), @@ -158,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)) @@ -165,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/services/graphify/ingest_storage.go b/cmd/api/src/services/graphify/ingest_storage.go index 48d982b07e49..b3a66949e82e 100644 --- a/cmd/api/src/services/graphify/ingest_storage.go +++ b/cmd/api/src/services/graphify/ingest_storage.go @@ -15,7 +15,7 @@ import ( "github.com/specterops/dawgs/util" ) -func SpoolToScratch(ctx context.Context, tempDirectory string, fileService storage.FileService, storedFileName string) (string, error) { +func SpoolToScratch(ctx context.Context, scratchDirectory string, fileService storage.FileService, storedFileName string) (string, error) { var ( sourceFile io.ReadCloser scratchFile *os.File @@ -30,7 +30,7 @@ func SpoolToScratch(ctx context.Context, tempDirectory string, fileService stora } defer sourceFile.Close() - scratchFile, err = os.CreateTemp(tempDirectory, "ingest-archive-*") + scratchFile, err = os.CreateTemp(scratchDirectory, "archive-*") if err != nil { return "", fmt.Errorf("create ingest scratch file: %w", err) } @@ -60,14 +60,14 @@ func SpoolToScratch(ctx context.Context, tempDirectory string, fileService stora return scratchPath, nil } -func OpenScratchReadSeeker(ctx context.Context, tempDirectory string, fileService storage.FileService, storedFileName string) (*os.File, string, error) { +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, tempDirectory, fileService, storedFileName); err != nil { + if scratchPath, err = SpoolToScratch(ctx, scratchDirectory, fileService, storedFileName); err != nil { return nil, "", err } @@ -108,7 +108,7 @@ func WriteArchiveFileToStorage(ctx context.Context, fileService storage.FileServ return extractedPath, nil } -func ExtractIngestFiles(ctx context.Context, tempDirectory string, fileService storage.FileService, storedFileName, providedFileName string, fileType model.FileType, prefix string) ([]IngestFileData, error) { +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{ @@ -120,7 +120,7 @@ func ExtractIngestFiles(ctx context.Context, tempDirectory string, fileService s } // Zip Path: - scratchPath, err := SpoolToScratch(ctx, tempDirectory, fileService, storedFileName) + scratchPath, err := SpoolToScratch(ctx, scratchDirectory, fileService, storedFileName) if err != nil { return []IngestFileData{ { @@ -158,16 +158,16 @@ func ExtractIngestFiles(ctx context.Context, tempDirectory string, fileService s ParentFile: providedFileName, } - //if extractedPath, err := WriteArchiveFileToStorage(ctx, fileService, archiveFile, fmt.Sprintf("tmp/file_upload_job%d_", jobID)); err != nil { if extractedPath, err := WriteArchiveFileToStorage(ctx, fileService, archiveFile, prefix); err != nil { processedFileData.Errors = []string{err.Error()} - errs.Add(fmt.Errorf( - "error extracting file %s in archive %s: %w", - archiveFile.Name, - storedFileName, - err, - )) + // 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 } diff --git a/cmd/api/src/services/graphify/tasks.go b/cmd/api/src/services/graphify/tasks.go index bd819ca27de0..35c3a78168b2 100644 --- a/cmd/api/src/services/graphify/tasks.go +++ b/cmd/api/src/services/graphify/tasks.go @@ -99,7 +99,7 @@ func (s *GraphifyService) extractToTempFile(f *zip.File) (string, error) { // archive, the number of files that failed to ingest as JSON, and an error func (s *GraphifyService) ProcessIngestFile(ic *IngestContext, fileService storage.FileService, task model.IngestTask) ([]IngestFileData, error) { // Try to pre-process the file. If any of them fail, stop processing and return the error - if fileData, err := ExtractIngestFiles(ic.Ctx, s.cfg.TempDirectory(), fileService, task.StoredFileName, task.OriginalFileName, task.FileType, fmt.Sprintf("tmp/file_upload_job%d_", ic.JobId)); err != nil { + if fileData, err := ExtractIngestFiles(ic.Ctx, s.cfg.ScratchDirectory(), fileService, task.StoredFileName, task.OriginalFileName, task.FileType, fmt.Sprintf("tmp/file_upload_job%d_", ic.JobId)); err != nil { return fileData, err } else { errs := errorlist.NewBuilder() @@ -118,7 +118,7 @@ func (s *GraphifyService) ProcessIngestFile(ic *IngestContext, fileService stora continue } - if err := processSingleFile(ic.Ctx, fileService, s.cfg.TempDirectory(), data, ic, readOpts); err != nil { + if err := processSingleFile(ic.Ctx, fileService, s.cfg.ScratchDirectory(), data, ic, readOpts); err != nil { var ( graphifyError errorlist.Error resolutionErr endpoint.ResolutionError diff --git a/cmd/api/src/services/storage/storage.go b/cmd/api/src/services/storage/storage.go index 26a3dced8404..8ae2b99ab1c2 100644 --- a/cmd/api/src/services/storage/storage.go +++ b/cmd/api/src/services/storage/storage.go @@ -100,6 +100,7 @@ type FileService interface { DeleteFile(ctx context.Context, name string) error WriteTempFile(ctx context.Context, prefix string, reader io.Reader, opts WriteOptions) (string, error) MoveFile(ctx context.Context, srcName, dstName string, opts WriteOptions) error + ListFiles(ctx context.Context, name string, opts ListOptions) ([]FileInfo, error) } type LocalFileService struct { @@ -158,6 +159,10 @@ func (s *LocalFileService) MoveFile(ctx context.Context, srcName, dstName string return s.Storage.Move(ctx, srcName, dstName, options) } +func (s *LocalFileService) ListFiles(ctx context.Context, name string, options ListOptions) ([]FileInfo, error) { + return s.Storage.List(ctx, name, options) +} + // 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 { From aebea94f0a8ee7dae47e40a4d02c3d4d1002154f Mon Sep 17 00:00:00 2001 From: Michael Cuomo Date: Thu, 7 May 2026 17:50:11 -0400 Subject: [PATCH 10/26] feat: add retained file service --- cmd/api/src/services/storage/storage.go | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/cmd/api/src/services/storage/storage.go b/cmd/api/src/services/storage/storage.go index 8ae2b99ab1c2..e253493be278 100644 --- a/cmd/api/src/services/storage/storage.go +++ b/cmd/api/src/services/storage/storage.go @@ -36,6 +36,7 @@ type FileServiceName string const ( FileServiceIngest FileServiceName = "ingest" + FileServiceRetained FileServiceName = "retained" FileServiceCollectors FileServiceName = "collectors" FileServiceJobLogs FileServiceName = "job_logs" FileServiceWork FileServiceName = "work" @@ -228,6 +229,11 @@ func NewDefaultFileServices(cfg config.Configuration) (map[FileServiceName]FileS 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 @@ -239,6 +245,9 @@ func NewDefaultFileServices(cfg config.Configuration) (map[FileServiceName]FileS ingestService := NewFileService(ingestStore) fileServices[FileServiceIngest] = ingestService + retainService := NewFileService(retainStore) + fileServices[FileServiceRetained] = retainService + collectersService := NewFileService(collectorsStore) fileServices[FileServiceCollectors] = collectersService From d61a5941be8a484d72e29268faef5bf369c1684a Mon Sep 17 00:00:00 2001 From: Michael Cuomo Date: Fri, 8 May 2026 11:32:37 -0400 Subject: [PATCH 11/26] fix: remove temp prefix from ingest --- cmd/api/src/services/graphify/tasks.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cmd/api/src/services/graphify/tasks.go b/cmd/api/src/services/graphify/tasks.go index 35c3a78168b2..a84fd5f6d628 100644 --- a/cmd/api/src/services/graphify/tasks.go +++ b/cmd/api/src/services/graphify/tasks.go @@ -99,7 +99,7 @@ func (s *GraphifyService) extractToTempFile(f *zip.File) (string, error) { // archive, the number of files that failed to ingest as JSON, and an error func (s *GraphifyService) ProcessIngestFile(ic *IngestContext, fileService storage.FileService, task model.IngestTask) ([]IngestFileData, error) { // Try to pre-process the file. If any of them fail, stop processing and return the error - if fileData, err := ExtractIngestFiles(ic.Ctx, s.cfg.ScratchDirectory(), fileService, task.StoredFileName, task.OriginalFileName, task.FileType, fmt.Sprintf("tmp/file_upload_job%d_", ic.JobId)); err != nil { + 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() From a12ec84b713f31dd6eb144e2d02b8bd6dcc45896 Mon Sep 17 00:00:00 2001 From: Michael Cuomo Date: Mon, 11 May 2026 13:11:17 -0400 Subject: [PATCH 12/26] feat: add write from reader, as well as move between services --- cmd/api/src/services/storage/storage.go | 27 +++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/cmd/api/src/services/storage/storage.go b/cmd/api/src/services/storage/storage.go index e253493be278..121af9426078 100644 --- a/cmd/api/src/services/storage/storage.go +++ b/cmd/api/src/services/storage/storage.go @@ -98,6 +98,7 @@ type FileService interface { GetFile(ctx context.Context, name string) (io.ReadCloser, FileInfo, error) ReadFile(ctx context.Context, name string) ([]byte, error) WriteFile(ctx context.Context, name string, data []byte, opts WriteOptions) error + WriteFileFromReader(ctx context.Context, name string, reader io.Reader, opts WriteOptions) error DeleteFile(ctx context.Context, name string) error WriteTempFile(ctx context.Context, prefix string, reader io.Reader, opts WriteOptions) (string, error) MoveFile(ctx context.Context, srcName, dstName string, opts WriteOptions) error @@ -138,6 +139,10 @@ func (s *LocalFileService) WriteFile(ctx context.Context, name string, data []by return s.Storage.Put(ctx, name, bytes.NewReader(data), opts) } +func (s *LocalFileService) WriteFileFromReader(ctx context.Context, name string, reader io.Reader, opts WriteOptions) error { + return s.Storage.Put(ctx, name, reader, opts) +} + func (s *LocalFileService) DeleteFile(ctx context.Context, name string) error { return s.Storage.Delete(ctx, name) } @@ -164,6 +169,28 @@ func (s *LocalFileService) ListFiles(ctx context.Context, name string, options L return s.Storage.List(ctx, name, options) } +// TODO MC: is this functinality necessary? +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 + } + defer sourceFile.Close() + + if err := destinationService.WriteFileFromReader(ctx, destinationName, sourceFile, opts); 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 { From acc6bc8a204f1374b8bd87c5e2414f8817de075e Mon Sep 17 00:00:00 2001 From: Michael Cuomo Date: Thu, 14 May 2026 16:46:54 -0400 Subject: [PATCH 13/26] fix: localstore.go fix for recursive List if "" is passed --- cmd/api/src/services/storage/localstore.go | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/cmd/api/src/services/storage/localstore.go b/cmd/api/src/services/storage/localstore.go index 62c831fb30e5..2e4281cf0d9d 100644 --- a/cmd/api/src/services/storage/localstore.go +++ b/cmd/api/src/services/storage/localstore.go @@ -25,6 +25,7 @@ import ( "mime" "os" "path" + "strings" "sync" ) @@ -242,15 +243,19 @@ func (s *LocalStore) List(ctx context.Context, name string, options ListOptions) 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, name, func(p string, d fs.DirEntry, err error) error { + 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 == name { + if errors.Is(err, fs.ErrNotExist) && p == listName { return fs.SkipAll } return err @@ -279,7 +284,7 @@ func (s *LocalStore) List(ctx context.Context, name string, options ListOptions) } return out, nil } - entries, err := fs.ReadDir(fsys, name) + entries, err := fs.ReadDir(fsys, listName) if errors.Is(err, fs.ErrNotExist) { return []FileInfo{}, nil } @@ -298,7 +303,10 @@ func (s *LocalStore) List(ctx context.Context, name string, options ListOptions) if err != nil { return nil, err } - entryPath := path.Join(name, entry.Name()) + entryPath := path.Join(listName, entry.Name()) + if listName == "." { + entryPath = entry.Name() + } out = append(out, FileInfo{ Path: entryPath, Size: info.Size(), From 4d32ea0ebb895ac255ac26502f3035405bda9a25 Mon Sep 17 00:00:00 2001 From: Michael Cuomo Date: Thu, 14 May 2026 17:06:08 -0400 Subject: [PATCH 14/26] feat: adding RuntimeDependencies and moving FileService to earlier startup, Adjust Cleanup to use FileServiceRetained --- cmd/api/src/api/tools/ingest.go | 207 +++++++++-------------- cmd/api/src/bootstrap/initializer.go | 43 +++-- cmd/api/src/bootstrap/util.go | 10 ++ cmd/api/src/cmd/bhapi/main.go | 1 + cmd/api/src/daemons/api/toolapi/api.go | 5 +- cmd/api/src/daemons/datapipe/cleanup.go | 77 ++++++++- cmd/api/src/daemons/datapipe/pipeline.go | 2 +- cmd/api/src/services/entrypoint.go | 40 +++-- 8 files changed, 222 insertions(+), 163 deletions(-) diff --git a/cmd/api/src/api/tools/ingest.go b/cmd/api/src/api/tools/ingest.go index 62cf792ff8d1..9a7872505d8a 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,102 +118,33 @@ func (s *IngestControl) FetchRetainedIngestFiles(response http.ResponseWriter, r return } - retainedFilesDirectory := s.cfg.RetainedFilesDirectory() - - // TODO MC: should there be a retained fileService? - 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) - ) - - for _, retainedFilesDirEntry := range retainedFilesDirEntries { - retainedFilePath := filepath.Join(retainedFilesDirectory, retainedFilesDirEntry.Name()) - - 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 - } + 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) - 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 - } - } - } + gzipWriter := gzip.NewWriter(response) + defer gzipWriter.Close() + + tarWriter := tar.NewWriter(gzipWriter) + defer tarWriter.Close() + + 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 + } } } @@ -216,6 +180,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 } @@ -224,32 +189,24 @@ 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() - - // TODO MC: should there be a retained fileService - 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), - ) - } + files, err := s.retainedFileService.ListFiles(ctx, "", storage.ListOptions{Recursive: true}) + if err != nil { + slog.WarnContext(ctx, "Failed listing retained files", 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/bootstrap/initializer.go b/cmd/api/src/bootstrap/initializer.go index abb0ba1070e6..d6c0a2a4eb92 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, conections 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/util.go b/cmd/api/src/bootstrap/util.go index e239aa90d1a1..e4e9787a129c 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,15 @@ 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 hve access to the FileServiceRetained, the +// FileServiceResolver is created prior ot 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) { 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/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 efae977d9feb..e2f5bec4e45a 100644 --- a/cmd/api/src/daemons/datapipe/cleanup.go +++ b/cmd/api/src/daemons/datapipe/cleanup.go @@ -33,8 +33,6 @@ import ( "github.com/specterops/bloodhound/packages/go/bhlog/attr" ) -// TODO MC: update cleanup to take into account the new file service - const ( orphanMinimumAge = 24 * time.Hour scratchMinimumAge = 24 * time.Hour @@ -69,14 +67,37 @@ type OrphanFileSweeper struct { fileOps FileOperations tempDirectoryRootPath string scratchDirectoryRootPath string + excludedStoragePrefixes []string +} + +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 } -func NewOrphanFileSweeper(fileOps FileOperations, tempDirectoryRootPath string, scratchDirectoryRootPath string) *OrphanFileSweeper { +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), } } @@ -111,6 +132,48 @@ func ClearLocalIngestScratch(ctx context.Context, scratchDirectory string, minim } } +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 bn legacy absolute filesystem paths or file-service logical paths. Logical paths are +// normalized, resolved under tempDirectoryRootPath, and ignored if they are empty or would excape 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) { if expectedFileName == "" { @@ -156,6 +219,10 @@ func (s *OrphanFileSweeper) clearStoredIngestFiles(ctx context.Context, ingestFi } logicalPath := path.Clean(file.Path) + if s.isExcludedStoragePath(logicalPath) { + continue + } + if _, ok := expectedFiles[logicalPath]; ok { continue } @@ -174,9 +241,7 @@ func (s *OrphanFileSweeper) clearLegacyLocalIngestFiles(ctx context.Context, exp expectedLocalFiles := map[string]struct{}{} for _, expectedFileName := range expectedFileNames { - if filepath.IsAbs(expectedFileName) { - expectedLocalFiles[filepath.Clean(expectedFileName)] = struct{}{} - } + s.addExpectedLocalPath(expectedLocalFiles, expectedFileName) } entries, err := s.fileOps.ReadDir(s.tempDirectoryRootPath) diff --git a/cmd/api/src/daemons/datapipe/pipeline.go b/cmd/api/src/daemons/datapipe/pipeline.go index 7186f84e5d74..cc81261489e9 100644 --- a/cmd/api/src/daemons/datapipe/pipeline.go +++ b/cmd/api/src/daemons/datapipe/pipeline.go @@ -61,7 +61,7 @@ func NewPipeline(ctx context.Context, cfg config.Configuration, db database.Data graphdb: graphDB, cache: cache, cfg: cfg, - orphanedFileSweeper: NewOrphanFileSweeper(NewOSFileOperations(), cfg.TempDirectory(), cfg.ScratchDirectory()), + orphanedFileSweeper: NewOrphanFileSweeper(NewOSFileOperations(), cfg.TempDirectory(), cfg.ScratchDirectory(), "retained"), ingestSchema: ingestSchema, fileServiceResolver: fileServiceResolver, jobService: job.NewJobService(ctx, db), diff --git a/cmd/api/src/services/entrypoint.go b/cmd/api/src/services/entrypoint.go index 2857fc48d7ad..fab5457cb4e5 100644 --- a/cmd/api/src/services/entrypoint.go +++ b/cmd/api/src/services/entrypoint.go @@ -75,14 +75,36 @@ 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 intialize 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() @@ -133,16 +155,12 @@ func Entrypoint(ctx context.Context, cfg config.Configuration, connections boots return nil, fmt.Errorf("failed to save collector manifests: %w", err) } else if ingestSchema, err := upload.LoadIngestSchema(); err != nil { return nil, fmt.Errorf("failed to load OpenGraph schema: %w", err) - } else if fileServices, err := storage.NewDefaultFileServices(cfg); err != nil { - return nil, fmt.Errorf("failed to intialize file services: %w", err) - } else if fileServiceResolver, err := storage.NewFileServiceResolver(fileServices); err != nil { - return nil, fmt.Errorf("failed to initialize file service resolver: %w", err) } else { startDelay := 0 * time.Second var ( cl = changelog.NewChangelog(connections.Graph, connections.RDMS, changelog.DefaultOptions()) - pipeline = datapipe.NewPipeline(ctx, cfg, connections.RDMS, connections.Graph, graphQueryCache, ingestSchema, fileServiceResolver, 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) @@ -152,7 +170,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, fileServiceResolver, 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) From 82b7c2ad6c6ad571c6006a7ac66362571d0e33bd Mon Sep 17 00:00:00 2001 From: Michael Cuomo Date: Thu, 14 May 2026 18:49:40 -0400 Subject: [PATCH 15/26] chore: clean up typos --- cmd/api/src/bootstrap/util.go | 4 ++-- cmd/api/src/daemons/datapipe/cleanup.go | 4 ++-- cmd/api/src/services/entrypoint.go | 4 ++-- cmd/api/src/services/storage/storage.go | 6 +++--- 4 files changed, 9 insertions(+), 9 deletions(-) diff --git a/cmd/api/src/bootstrap/util.go b/cmd/api/src/bootstrap/util.go index e4e9787a129c..f736d48b2625 100644 --- a/cmd/api/src/bootstrap/util.go +++ b/cmd/api/src/bootstrap/util.go @@ -35,8 +35,8 @@ import ( // 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 hve access to the FileServiceRetained, the -// FileServiceResolver is created prior ot the PreMigrationDaemons and the Entrypoint. This could +// 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 diff --git a/cmd/api/src/daemons/datapipe/cleanup.go b/cmd/api/src/daemons/datapipe/cleanup.go index e2f5bec4e45a..115525672b1d 100644 --- a/cmd/api/src/daemons/datapipe/cleanup.go +++ b/cmd/api/src/daemons/datapipe/cleanup.go @@ -145,8 +145,8 @@ func (s *OrphanFileSweeper) isExcludedStoragePath(logicalPath string) bool { } // addExpectedLocalPath records an expected file path for the legacy local temp-directory cleanup pass. -// Expected names may bn legacy absolute filesystem paths or file-service logical paths. Logical paths are -// normalized, resolved under tempDirectoryRootPath, and ignored if they are empty or would excape the temp root. +// 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 == "" { diff --git a/cmd/api/src/services/entrypoint.go b/cmd/api/src/services/entrypoint.go index fab5457cb4e5..0cbdef1a114e 100644 --- a/cmd/api/src/services/entrypoint.go +++ b/cmd/api/src/services/entrypoint.go @@ -80,7 +80,7 @@ func ConnectDatabases(ctx context.Context, cfg config.Configuration) (bootstrap. 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 intialize file services: %w", err) + 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 @@ -96,7 +96,7 @@ func CreateRuntimeDependencies(ctx context.Context, cfg config.Configuration, co // 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], 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) + return nil, fmt.Errorf("error resolving FileServiceRetained: %w", err) } else { return []daemons.Daemon{ toolapi.NewDaemon(ctx, connections, cfg, schema.DefaultGraphSchema(), retainedFileService), diff --git a/cmd/api/src/services/storage/storage.go b/cmd/api/src/services/storage/storage.go index 121af9426078..881593cfb770 100644 --- a/cmd/api/src/services/storage/storage.go +++ b/cmd/api/src/services/storage/storage.go @@ -169,7 +169,7 @@ func (s *LocalFileService) ListFiles(ctx context.Context, name string, options L return s.Storage.List(ctx, name, options) } -// TODO MC: is this functinality necessary? +// TODO MC: is this functionality necessary? func MoveFileBetweenServices( ctx context.Context, sourceService FileService, @@ -275,8 +275,8 @@ func NewDefaultFileServices(cfg config.Configuration) (map[FileServiceName]FileS retainService := NewFileService(retainStore) fileServices[FileServiceRetained] = retainService - collectersService := NewFileService(collectorsStore) - fileServices[FileServiceCollectors] = collectersService + collectorsService := NewFileService(collectorsStore) + fileServices[FileServiceCollectors] = collectorsService return fileServices, nil } From d7ab5c2f573d9ef10c5e032fb7367933d60bd9a3 Mon Sep 17 00:00:00 2001 From: Michael Cuomo Date: Fri, 15 May 2026 16:55:00 -0400 Subject: [PATCH 16/26] test: bring existing tests inline with new FileService --- cmd/api/src/api/v2/collectors_test.go | 23 +- cmd/api/src/api/v2/fileingest_test.go | 48 +++- cmd/api/src/daemons/datapipe/cleanup_test.go | 117 ++++++-- .../datapipe/datapipe_integration_test.go | 10 +- .../datapipe/delete_integration_test.go | 29 +- .../datapipe/pipeline_integration_test.go | 27 +- cmd/api/src/services/fs/fs.go | 38 --- cmd/api/src/services/fs/mocks/fs.go | 87 ------ .../graphify/graphify_integration_test.go | 8 +- .../src/services/graphify/ingest_storage.go | 15 + .../graphify/tasks_integration_test.go | 41 ++- cmd/api/src/services/storage/mocks/fs.go | 256 ++++++++++++++++-- cmd/api/src/services/storage/storage.go | 3 +- cmd/api/src/services/upload/upload.go | 6 +- cmd/api/src/test/lab/fixtures/api.go | 4 +- 15 files changed, 494 insertions(+), 218 deletions(-) delete mode 100644 cmd/api/src/services/fs/fs.go delete mode 100644 cmd/api/src/services/fs/mocks/fs.go diff --git a/cmd/api/src/api/v2/collectors_test.go b/cmd/api/src/api/v2/collectors_test.go index 642cb7bbde35..a3ca271b11b6 100644 --- a/cmd/api/src/api/v2/collectors_test.go +++ b/cmd/api/src/api/v2/collectors_test.go @@ -28,6 +28,7 @@ import ( "github.com/gorilla/mux" v2 "github.com/specterops/bloodhound/cmd/api/src/api/v2" "github.com/specterops/bloodhound/cmd/api/src/config" + "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" @@ -40,7 +41,8 @@ import ( func TestResources_DownloadCollectorByVersion(t *testing.T) { type mock struct { - mockFS *storagemocks.MockFileService + mockFS *storagemocks.MockFileService + mockFileServiceResolver *storagemocks.MockFileServiceResolver } type expected struct { responseBody string @@ -99,6 +101,7 @@ func TestResources_DownloadCollectorByVersion(t *testing.T) { } }, setupMocks: func(t *testing.T, mock *mock) { + mock.mockFileServiceResolver.EXPECT().Resolve(storage.FileServiceCollectors).Return(mock.mockFS, nil) mock.mockFS.EXPECT().ReadFile(gomock.Any(), "azurehound/azurehound-latest.zip").Return([]byte{}, errors.New("error")) }, expected: expected{ @@ -118,6 +121,7 @@ func TestResources_DownloadCollectorByVersion(t *testing.T) { } }, setupMocks: func(t *testing.T, mock *mock) { + mock.mockFileServiceResolver.EXPECT().Resolve(storage.FileServiceCollectors).Return(mock.mockFS, nil) mock.mockFS.EXPECT().ReadFile(gomock.Any(), "azurehound/azurehound-v1.0.0.zip").Return([]byte{}, nil) }, expected: expected{ @@ -136,6 +140,7 @@ func TestResources_DownloadCollectorByVersion(t *testing.T) { } }, setupMocks: func(t *testing.T, mock *mock) { + mock.mockFileServiceResolver.EXPECT().Resolve(storage.FileServiceCollectors).Return(mock.mockFS, nil) mock.mockFS.EXPECT().ReadFile(gomock.Any(), "azurehound/azurehound-latest.zip").Return([]byte{}, nil) }, expected: expected{ @@ -150,7 +155,8 @@ func TestResources_DownloadCollectorByVersion(t *testing.T) { ctrl := gomock.NewController(t) mock := &mock{ - mockFS: storagemocks.NewMockFileService(ctrl), + mockFS: storagemocks.NewMockFileService(ctrl), + mockFileServiceResolver: storagemocks.NewMockFileServiceResolver(ctrl), } testCase.setupMocks(t, mock) @@ -163,7 +169,7 @@ func TestResources_DownloadCollectorByVersion(t *testing.T) { resources := v2.Resources{ CollectorManifests: collectorManifests, - FileServiceResolver: mock.mockFS, + FileServiceResolver: mock.mockFileServiceResolver, } response := httptest.NewRecorder() @@ -187,7 +193,8 @@ func TestResources_DownloadCollectorByVersion(t *testing.T) { func TestResources_DownloadCollectorChecksumByVersion(t *testing.T) { type mock struct { - mockFS *storagemocks.MockFileService + mockFS *storagemocks.MockFileService + mockFileServiceResolver *storagemocks.MockFileServiceResolver } type expected struct { responseBody string @@ -246,6 +253,7 @@ func TestResources_DownloadCollectorChecksumByVersion(t *testing.T) { } }, setupMocks: func(t *testing.T, mock *mock) { + mock.mockFileServiceResolver.EXPECT().Resolve(storage.FileServiceCollectors).Return(mock.mockFS, nil) mock.mockFS.EXPECT().ReadFile(gomock.Any(), "azurehound/azurehound-latest.zip.sha256").Return([]byte{}, errors.New("error")) }, expected: expected{ @@ -265,6 +273,7 @@ func TestResources_DownloadCollectorChecksumByVersion(t *testing.T) { } }, setupMocks: func(t *testing.T, mock *mock) { + mock.mockFileServiceResolver.EXPECT().Resolve(storage.FileServiceCollectors).Return(mock.mockFS, nil) mock.mockFS.EXPECT().ReadFile(gomock.Any(), "azurehound/azurehound-v1.0.0.zip.sha256").Return([]byte{}, nil) }, expected: expected{ @@ -283,6 +292,7 @@ func TestResources_DownloadCollectorChecksumByVersion(t *testing.T) { } }, setupMocks: func(t *testing.T, mock *mock) { + mock.mockFileServiceResolver.EXPECT().Resolve(storage.FileServiceCollectors).Return(mock.mockFS, nil) mock.mockFS.EXPECT().ReadFile(gomock.Any(), "azurehound/azurehound-latest.zip.sha256").Return([]byte{}, nil) }, expected: expected{ @@ -304,13 +314,14 @@ func TestResources_DownloadCollectorChecksumByVersion(t *testing.T) { ctrl := gomock.NewController(t) mock := &mock{ - mockFS: storagemocks.NewMockFileService(ctrl), + mockFS: storagemocks.NewMockFileService(ctrl), + mockFileServiceResolver: storagemocks.NewMockFileServiceResolver(ctrl), } testCase.setupMocks(t, mock) resources := v2.Resources{ CollectorManifests: collectorManifests, - FileServiceResolver: mock.mockFS, + FileServiceResolver: mock.mockFileServiceResolver, } response := httptest.NewRecorder() diff --git a/cmd/api/src/api/v2/fileingest_test.go b/cmd/api/src/api/v2/fileingest_test.go index 3c126ea9f98d..d773bafe4893 100644 --- a/cmd/api/src/api/v2/fileingest_test.go +++ b/cmd/api/src/api/v2/fileingest_test.go @@ -481,10 +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 - mockFileService *storagemocks.MockFileService + mockDatabase *dbmocks.MockDatabase + mockFileService *storagemocks.MockFileService + trueFileService storage.FileService + mockFileServiceResolver *storagemocks.MockFileServiceResolver } type expected struct { responseBody string @@ -577,10 +586,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.trueFileService, nil) }, - // Use a real OsStore so that actual file I/O gives the concurrent validation + // 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: storage.NewFileService(storage.NewLocalStore()), + 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"}`, @@ -609,9 +619,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 OsStore for the same reason as the ErrInvalidJSON case above. - fileServiceOvrd: storage.NewFileService(storage.NewLocalStore()), + // 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"}`, @@ -634,11 +645,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 "/tmp/test", err + 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{ @@ -664,6 +678,7 @@ 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 @@ -694,6 +709,7 @@ 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 @@ -726,6 +742,7 @@ 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 @@ -776,6 +793,7 @@ 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 @@ -796,22 +814,22 @@ func TestResources_ProcessIngestTask(t *testing.T) { ctrl := gomock.NewController(t) mocks := &mock{ - mockDatabase: dbmocks.NewMockDatabase(ctrl), - mockFileService: storagemocks.NewMockFileService(ctrl), + mockDatabase: dbmocks.NewMockDatabase(ctrl), + mockFileService: storagemocks.NewMockFileService(ctrl), + mockFileServiceResolver: storagemocks.NewMockFileServiceResolver(ctrl), } - request := testCase.buildRequest() - testCase.setupMocks(t, mocks) - - fileService := storage.FileService(mocks.mockFileService) if testCase.fileServiceOvrd != nil { - fileService = testCase.fileServiceOvrd + mocks.trueFileService = testCase.fileServiceOvrd } + request := testCase.buildRequest() + testCase.setupMocks(t, mocks) + resources := v2.Resources{ DB: mocks.mockDatabase, Config: config.Configuration{}, - FileServiceResolver: fileService, + FileServiceResolver: mocks.mockFileServiceResolver, } err := os.Mkdir(resources.Config.TempDirectory(), 0755) 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..661e3ec9d44b 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" @@ -45,12 +46,15 @@ func TestDeleteAllData(t *testing.T) { path.Join(testSuite.WorkDir, "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") @@ -93,12 +97,15 @@ func TestDeleteAllData_Alternative(t *testing.T) { path.Join(testSuite.WorkDir, "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") @@ -138,12 +145,15 @@ func TestDeleteSourcelessData(t *testing.T) { path.Join(testSuite.WorkDir, "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") @@ -182,12 +192,15 @@ func TestDeleteSourceKindsData(t *testing.T) { path.Join(testSuite.WorkDir, "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_integration_test.go b/cmd/api/src/daemons/datapipe/pipeline_integration_test.go index cd80757a98be..7c78a4b3b415 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" @@ -47,12 +48,15 @@ func TestDeleteData_Sourceless(t *testing.T) { path.Join(testSuite.WorkDir, "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) @@ -96,12 +100,15 @@ func TestDeleteData_SourceKinds(t *testing.T) { path.Join(testSuite.WorkDir, "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) @@ -145,12 +152,15 @@ func TestDeleteData_All(t *testing.T) { path.Join(testSuite.WorkDir, "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,11 +203,14 @@ 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{ + fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestCtx, ingestFileService, model.IngestTask{ StoredFileName: path.Join(testSuite.WorkDir, "oneGoodOneInvalidRel.json"), FileType: model.FileTypeJson, }) 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 index b3a66949e82e..56ea5d467631 100644 --- a/cmd/api/src/services/graphify/ingest_storage.go +++ b/cmd/api/src/services/graphify/ingest_storage.go @@ -1,3 +1,18 @@ +// 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 ( diff --git a/cmd/api/src/services/graphify/tasks_integration_test.go b/cmd/api/src/services/graphify/tasks_integration_test.go index 04d37a098c66..fc8d4c49d5f4 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" ) @@ -49,12 +50,15 @@ func TestVersion5IngestJSON(t *testing.T) { path.Join(testSuite.WorkDir, "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 @@ -85,12 +89,15 @@ func TestVersion5IngestZIP(t *testing.T) { path.Join(testSuite.WorkDir, "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 @@ -133,12 +140,15 @@ func TestVersion6ADCSJSON(t *testing.T) { path.Join(testSuite.WorkDir, "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 @@ -169,12 +179,15 @@ func TestVersion6ADCSZIP(t *testing.T) { path.Join(testSuite.WorkDir, "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 @@ -217,12 +230,15 @@ func TestVersion6AllJSON(t *testing.T) { path.Join(testSuite.WorkDir, "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 @@ -253,12 +269,15 @@ func TestVersion6AllZIP(t *testing.T) { path.Join(testSuite.WorkDir, "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 @@ -296,12 +315,15 @@ func TestVersion6IngestJSON(t *testing.T) { path.Join(testSuite.WorkDir, "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 @@ -332,12 +354,15 @@ func TestVersion6IngestZIP(t *testing.T) { path.Join(testSuite.WorkDir, "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/mocks/fs.go b/cmd/api/src/services/storage/mocks/fs.go index 451b68f6a7fb..176cd183e81a 100644 --- a/cmd/api/src/services/storage/mocks/fs.go +++ b/cmd/api/src/services/storage/mocks/fs.go @@ -15,11 +15,11 @@ // SPDX-License-Identifier: Apache-2.0 // Code generated by MockGen. DO NOT EDIT. -// Source: github.com/specterops/bloodhound/cmd/api/src/services/storage (interfaces: FileService) +// 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 . FileService +// mockgen -copyright_file=../../../../../LICENSE.header -destination=./mocks/fs.go -package=mocks . Storage,FileService,FileServiceResolver // // Package mocks is a generated GoMock package. @@ -34,6 +34,147 @@ import ( 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 @@ -72,34 +213,64 @@ func (mr *MockFileServiceMockRecorder) DeleteFile(ctx, name any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteFile", reflect.TypeOf((*MockFileService)(nil).DeleteFile), ctx, name) } -// ReadFile mocks base method. -func (m *MockFileService) ReadFile(ctx context.Context, name string) ([]byte, error) { +// 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, "ReadFile", ctx, name) - ret0, _ := ret[0].([]byte) + 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 } -// ReadFile indicates an expected call of ReadFile. -func (mr *MockFileServiceMockRecorder) ReadFile(ctx, name any) *gomock.Call { +// 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, "ReadFile", reflect.TypeOf((*MockFileService)(nil).ReadFile), ctx, name) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListFiles", reflect.TypeOf((*MockFileService)(nil).ListFiles), ctx, name, opts) } -// TempPath mocks base method. -func (m *MockFileService) TempPath(prefix, pattern string) (string, error) { +// 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, "TempPath", prefix, pattern) - ret0, _ := ret[0].(string) + 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 } -// TempPath indicates an expected call of TempPath. -func (mr *MockFileServiceMockRecorder) TempPath(prefix, pattern any) *gomock.Call { +// 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, "TempPath", reflect.TypeOf((*MockFileService)(nil).TempPath), prefix, pattern) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReadFile", reflect.TypeOf((*MockFileService)(nil).ReadFile), ctx, name) } // WriteFile mocks base method. @@ -116,6 +287,20 @@ func (mr *MockFileServiceMockRecorder) WriteFile(ctx, name, data, opts any) *gom 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() @@ -130,3 +315,42 @@ func (mr *MockFileServiceMockRecorder) WriteTempFile(ctx, prefix, reader, opts a 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 index 881593cfb770..4b2d13a661b7 100644 --- a/cmd/api/src/services/storage/storage.go +++ b/cmd/api/src/services/storage/storage.go @@ -30,7 +30,7 @@ import ( "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 . FileService +//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 @@ -169,7 +169,6 @@ func (s *LocalFileService) ListFiles(ctx context.Context, name string, options L return s.Storage.List(ctx, name, options) } -// TODO MC: is this functionality necessary? func MoveFileBetweenServices( ctx context.Context, sourceService FileService, diff --git a/cmd/api/src/services/upload/upload.go b/cmd/api/src/services/upload/upload.go index d7713bd97b4b..52cc8fed0f3f 100644 --- a/cmd/api/src/services/upload/upload.go +++ b/cmd/api/src/services/upload/upload.go @@ -97,18 +97,18 @@ func WriteAndValidateFile(ctx context.Context, fileService storage.FileService, // Check if validation failed first — the temp file should be cleaned up. if validationErr != nil { - slog.ErrorContext(ctx, "Validation failed", slog.String("tempFileName", tempFileName), attr.Error(validationErr)) + slog.ErrorContext(ctx, "Validation failed", slog.String("temp_file_name", tempFileName), attr.Error(validationErr)) fileService.DeleteFile(ctx, tempFileName) return "", validationErr } // Check if writing failed — the temp file should be cleaned up. if writeErr != nil { - slog.ErrorContext(ctx, "Write failed", slog.String("tempFileName", tempFileName), attr.Error(writeErr)) + slog.ErrorContext(ctx, "Write failed", slog.String("temp_file_name", tempFileName), attr.Error(writeErr)) fileService.DeleteFile(ctx, tempFileName) return "", writeErr } - slog.InfoContext(ctx, "File written and validated", slog.String("tempFileName", tempFileName)) + slog.InfoContext(ctx, "File written and validated", slog.String("temp_file_name", tempFileName)) return tempFileName, nil } diff --git a/cmd/api/src/test/lab/fixtures/api.go b/cmd/api/src/test/lab/fixtures/api.go index 4348e097e10d..ea0c70613248 100644 --- a/cmd/api/src/test/lab/fixtures/api.go +++ b/cmd/api/src/test/lab/fixtures/api.go @@ -62,12 +62,12 @@ func NewCustomApiFixture(cfgFixture *lab.Fixture[config.Configuration]) *lab.Fix 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) { + 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) }, } From 9d4910d7260d4bcd27fb1f1af62d3c0d90ec85f5 Mon Sep 17 00:00:00 2001 From: Michael Cuomo Date: Fri, 15 May 2026 16:58:38 -0400 Subject: [PATCH 17/26] chore: linting cleanup --- cmd/api/src/api/tools/ingest.go | 14 ++++++++-- cmd/api/src/daemons/datapipe/cleanup.go | 14 ++++++++-- cmd/api/src/services/graphify/tasks.go | 36 ------------------------- cmd/api/src/services/upload/upload.go | 15 +++++++++-- 4 files changed, 37 insertions(+), 42 deletions(-) diff --git a/cmd/api/src/api/tools/ingest.go b/cmd/api/src/api/tools/ingest.go index 9a7872505d8a..e22bcd4320d5 100644 --- a/cmd/api/src/api/tools/ingest.go +++ b/cmd/api/src/api/tools/ingest.go @@ -142,7 +142,12 @@ func (s *IngestControl) FetchRetainedIngestFiles(response http.ResponseWriter, r } 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)) + slog.WarnContext( + request.Context(), + "Failed writing retained file to response", + slog.String("path", file.Path), + attr.Error(err), + ) break } } @@ -205,7 +210,12 @@ func (s *IngestControl) DisableIngestFileRetention(response http.ResponseWriter, } 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)) + slog.WarnContext( + ctx, + "Failed removing retained file", + slog.String("path", file.Path), + attr.Error(err), + ) } } }(cleanupCtx) diff --git a/cmd/api/src/daemons/datapipe/cleanup.go b/cmd/api/src/daemons/datapipe/cleanup.go index 115525672b1d..c57ee84810dc 100644 --- a/cmd/api/src/daemons/datapipe/cleanup.go +++ b/cmd/api/src/daemons/datapipe/cleanup.go @@ -104,7 +104,12 @@ func NewOrphanFileSweeper(fileOps FileOperations, tempDirectoryRootPath string, func ClearLocalIngestScratch(ctx context.Context, scratchDirectory string, minimumAge time.Duration) { entries, err := os.ReadDir(scratchDirectory) if err != nil { - slog.ErrorContext(ctx, "Error reading scratch directory", slog.String("scratch_directory", scratchDirectory), attr.Error(err)) + slog.ErrorContext( + ctx, + "Error reading scratch directory", + slog.String("scratch_directory", scratchDirectory), + attr.Error(err), + ) return } @@ -232,7 +237,12 @@ func (s *OrphanFileSweeper) clearStoredIngestFiles(ctx context.Context, ingestFi } if err := ingestFileService.DeleteFile(ctx, logicalPath); err != nil { - slog.WarnContext(ctx, "Failed deleting orphaned ingest file", slog.String("path", logicalPath), attr.Error(err)) + slog.WarnContext( + ctx, + "Failed deleting orphaned ingest file", + slog.String("path", logicalPath), + attr.Error(err), + ) } } } diff --git a/cmd/api/src/services/graphify/tasks.go b/cmd/api/src/services/graphify/tasks.go index a84fd5f6d628..f1f207cb8d7e 100644 --- a/cmd/api/src/services/graphify/tasks.go +++ b/cmd/api/src/services/graphify/tasks.go @@ -59,42 +59,6 @@ type IngestFileData struct { UserDataErrs []string } -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, fileService storage.FileService, task model.IngestTask) ([]IngestFileData, error) { diff --git a/cmd/api/src/services/upload/upload.go b/cmd/api/src/services/upload/upload.go index 52cc8fed0f3f..c515184bd2b5 100644 --- a/cmd/api/src/services/upload/upload.go +++ b/cmd/api/src/services/upload/upload.go @@ -97,14 +97,25 @@ func WriteAndValidateFile(ctx context.Context, fileService storage.FileService, // Check if validation failed first — the temp file should be cleaned up. if validationErr != nil { - slog.ErrorContext(ctx, "Validation failed", slog.String("temp_file_name", tempFileName), attr.Error(validationErr)) + slog.ErrorContext( + ctx, + "Validation failed", + slog.String("temp_file_name", + tempFileName), + attr.Error(validationErr), + ) fileService.DeleteFile(ctx, tempFileName) return "", validationErr } // 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)) + slog.ErrorContext( + ctx, + "Write failed", + slog.String("temp_file_name", tempFileName), + attr.Error(writeErr), + ) fileService.DeleteFile(ctx, tempFileName) return "", writeErr } From 6edc2073a594434695ca2a614bfe41fa506fe48d Mon Sep 17 00:00:00 2001 From: Michael Cuomo Date: Fri, 15 May 2026 17:01:08 -0400 Subject: [PATCH 18/26] chore: pfc --- cmd/api/src/services/graphify/ingest_storage.go | 2 +- cmd/api/src/services/graphify/tasks.go | 3 --- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/cmd/api/src/services/graphify/ingest_storage.go b/cmd/api/src/services/graphify/ingest_storage.go index 56ea5d467631..9c84a8c570c1 100644 --- a/cmd/api/src/services/graphify/ingest_storage.go +++ b/cmd/api/src/services/graphify/ingest_storage.go @@ -4,7 +4,7 @@ // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // -// http://www.apache.org/licenses/LICENSE-2.0 +// http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, diff --git a/cmd/api/src/services/graphify/tasks.go b/cmd/api/src/services/graphify/tasks.go index f1f207cb8d7e..5f83bf2be0be 100644 --- a/cmd/api/src/services/graphify/tasks.go +++ b/cmd/api/src/services/graphify/tasks.go @@ -17,11 +17,9 @@ package graphify import ( - "archive/zip" "context" "errors" "fmt" - "io" "io/fs" "log/slog" "os" @@ -33,7 +31,6 @@ import ( "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" ) From 6944b2c1fad2a30d3b60fc9025845d46f8460859 Mon Sep 17 00:00:00 2001 From: Michael Cuomo Date: Fri, 15 May 2026 17:14:23 -0400 Subject: [PATCH 19/26] chore: pfc --- go.mod | 28 ++++++++++++++-------------- go.sum | 56 ++++++++++++++++++++++++++++---------------------------- 2 files changed, 42 insertions(+), 42 deletions(-) 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= From 4571b732b404f630fc7af54b165828503543598d Mon Sep 17 00:00:00 2001 From: Michael Cuomo Date: Mon, 18 May 2026 08:49:15 -0400 Subject: [PATCH 20/26] chore: comment cleanup --- cmd/api/src/api/v2/model.go | 4 ---- cmd/api/src/bootstrap/util.go | 7 +++--- cmd/api/src/daemons/datapipe/pipeline.go | 2 +- cmd/api/src/services/entrypoint.go | 9 ++++---- cmd/api/src/services/storage/storage.go | 28 ++++++++++++++++++++++++ 5 files changed, 37 insertions(+), 13 deletions(-) diff --git a/cmd/api/src/api/v2/model.go b/cmd/api/src/api/v2/model.go index 93b9a2aff18b..eca05d084d9a 100644 --- a/cmd/api/src/api/v2/model.go +++ b/cmd/api/src/api/v2/model.go @@ -134,10 +134,6 @@ func NewResources( dogtagsService dogtags.Service, openGraphSchemaService OpenGraphSchemaService, ) Resources { - // localStore, err := storage.NewLocalStore("/tmp/") - // if err != nil { - // slog.Error("Error creating local store") - // } return Resources{ Decoder: schema.NewDecoder(), DB: rdms, diff --git a/cmd/api/src/bootstrap/util.go b/cmd/api/src/bootstrap/util.go index f736d48b2625..980100605e0d 100644 --- a/cmd/api/src/bootstrap/util.go +++ b/cmd/api/src/bootstrap/util.go @@ -34,10 +34,9 @@ import ( ) // 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. +// 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 } diff --git a/cmd/api/src/daemons/datapipe/pipeline.go b/cmd/api/src/daemons/datapipe/pipeline.go index cc81261489e9..1782f43c9542 100644 --- a/cmd/api/src/daemons/datapipe/pipeline.go +++ b/cmd/api/src/daemons/datapipe/pipeline.go @@ -61,7 +61,7 @@ func NewPipeline(ctx context.Context, cfg config.Configuration, db database.Data graphdb: graphDB, cache: cache, cfg: cfg, - orphanedFileSweeper: NewOrphanFileSweeper(NewOSFileOperations(), cfg.TempDirectory(), cfg.ScratchDirectory(), "retained"), + orphanedFileSweeper: NewOrphanFileSweeper(NewOSFileOperations(), cfg.TempDirectory(), cfg.ScratchDirectory(), string(storage.FileServiceRetained)), ingestSchema: ingestSchema, fileServiceResolver: fileServiceResolver, jobService: job.NewJobService(ctx, db), diff --git a/cmd/api/src/services/entrypoint.go b/cmd/api/src/services/entrypoint.go index 0cbdef1a114e..fe15ec715315 100644 --- a/cmd/api/src/services/entrypoint.go +++ b/cmd/api/src/services/entrypoint.go @@ -75,16 +75,17 @@ 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. +// 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. + // 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 { diff --git a/cmd/api/src/services/storage/storage.go b/cmd/api/src/services/storage/storage.go index 4b2d13a661b7..99e25d0bddbe 100644 --- a/cmd/api/src/services/storage/storage.go +++ b/cmd/api/src/services/storage/storage.go @@ -94,14 +94,38 @@ type Storage interface { 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) } @@ -193,6 +217,8 @@ func MoveFileBetweenServices( // 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) } @@ -241,6 +267,8 @@ func (s *fileServiceResolver) Resolve(name FileServiceName) (FileService, error) 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) From 01dd7e3efeb538566f6453f9a94d32af7c556bf2 Mon Sep 17 00:00:00 2001 From: Michael Cuomo Date: Mon, 18 May 2026 12:38:30 -0400 Subject: [PATCH 21/26] test: add localstorage tests --- cmd/api/src/services/storage/localstore.go | 19 +- .../src/services/storage/localstore_test.go | 1712 +++++++++++++++++ 2 files changed, 1729 insertions(+), 2 deletions(-) create mode 100644 cmd/api/src/services/storage/localstore_test.go diff --git a/cmd/api/src/services/storage/localstore.go b/cmd/api/src/services/storage/localstore.go index 2e4281cf0d9d..4d7588666104 100644 --- a/cmd/api/src/services/storage/localstore.go +++ b/cmd/api/src/services/storage/localstore.go @@ -232,11 +232,18 @@ func (s *LocalStore) Delete(ctx context.Context, name string) error { if err := ctx.Err(); err != nil { return err } - err := s.root.Remove(name) + + stat, err := s.root.Stat(name) if errors.Is(err, os.ErrNotExist) { return nil } - return err + 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) { @@ -348,6 +355,14 @@ func (s *LocalStore) Move(ctx context.Context, srcName, dstName string, options 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 { 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) + }) + } +} From 698faa673a55b6eb5e6593d40a2221a4595cd337 Mon Sep 17 00:00:00 2001 From: Michael Cuomo Date: Mon, 18 May 2026 14:05:38 -0400 Subject: [PATCH 22/26] test: add storage_test.go --- cmd/api/src/services/storage/storage.go | 8 +- cmd/api/src/services/storage/storage_test.go | 1115 ++++++++++++++++++ 2 files changed, 1122 insertions(+), 1 deletion(-) create mode 100644 cmd/api/src/services/storage/storage_test.go diff --git a/cmd/api/src/services/storage/storage.go b/cmd/api/src/services/storage/storage.go index 99e25d0bddbe..0a5fdbebe71a 100644 --- a/cmd/api/src/services/storage/storage.go +++ b/cmd/api/src/services/storage/storage.go @@ -205,9 +205,15 @@ func MoveFileBetweenServices( if err != nil { return err } - defer sourceFile.Close() 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 } 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..1daaa5b9bbdb --- /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 TestLocalFileService_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 TestLocalFileService_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 TestLocalFileService_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 TestLocalFileService_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 TestLocalFileService_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 TestLocalFileService_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 TestLocalFileService_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 TestLocalFileService_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 { + localFileService, ok := fileService.(*storage.LocalFileService) + require.True(t, ok) + + localStore, ok := localFileService.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) + }) + } +} From 8f8eeb689d99512048524825e15333dcd31d1ce3 Mon Sep 17 00:00:00 2001 From: Michael Cuomo Date: Tue, 19 May 2026 12:38:37 -0400 Subject: [PATCH 23/26] test: add tests for new functionality --- cmd/api/src/api/tools/ingest_test.go | 554 +++++++++++++++ cmd/api/src/api/v2/collectors_test.go | 38 ++ cmd/api/src/api/v2/fileingest_test.go | 24 + cmd/api/src/bootstrap/server_test.go | 35 + cmd/api/src/daemons/datapipe/cleanup.go | 8 +- .../daemons/datapipe/cleanup_internal_test.go | 302 +++++++++ .../datapipe/delete_integration_test.go | 8 +- .../datapipe/pipeline_integration_test.go | 8 +- cmd/api/src/services/entrypoint_test.go | 118 ++++ .../services/graphify/ingest_storage_test.go | 633 ++++++++++++++++++ .../graphify/tasks_integration_test.go | 92 +-- cmd/api/src/services/upload/upload.go | 8 +- cmd/api/src/services/upload/upload_test.go | 123 ++++ 13 files changed, 1894 insertions(+), 57 deletions(-) create mode 100644 cmd/api/src/api/tools/ingest_test.go create mode 100644 cmd/api/src/daemons/datapipe/cleanup_internal_test.go create mode 100644 cmd/api/src/services/entrypoint_test.go create mode 100644 cmd/api/src/services/graphify/ingest_storage_test.go 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_test.go b/cmd/api/src/api/v2/collectors_test.go index a3ca271b11b6..3059626f219d 100644 --- a/cmd/api/src/api/v2/collectors_test.go +++ b/cmd/api/src/api/v2/collectors_test.go @@ -110,6 +110,25 @@ func TestResources_DownloadCollectorByVersion(t *testing.T) { 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, + 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: "Success: download collector not-latest release - OK", buildRequest: func() *http.Request { @@ -262,6 +281,25 @@ func TestResources_DownloadCollectorChecksumByVersion(t *testing.T) { 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, + 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: "Success: download collector not-latest release - OK", buildRequest: func() *http.Request { diff --git a/cmd/api/src/api/v2/fileingest_test.go b/cmd/api/src/api/v2/fileingest_test.go index d773bafe4893..ec9bdf849773 100644 --- a/cmd/api/src/api/v2/fileingest_test.go +++ b/cmd/api/src/api/v2/fileingest_test.go @@ -597,6 +597,30 @@ func TestResources_ProcessIngestTask(t *testing.T) { 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 { 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/daemons/datapipe/cleanup.go b/cmd/api/src/daemons/datapipe/cleanup.go index c57ee84810dc..51df798e8b92 100644 --- a/cmd/api/src/daemons/datapipe/cleanup.go +++ b/cmd/api/src/daemons/datapipe/cleanup.go @@ -181,6 +181,7 @@ func (s *OrphanFileSweeper) addExpectedLocalPath(expectedLocalFiles map[string]s // 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 } @@ -199,7 +200,12 @@ func (s *OrphanFileSweeper) addExpectedStoragePath(expectedFiles map[string]stru return } - expectedFiles[path.Clean(filepath.ToSlash(expectedFileName))] = struct{}{} + 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) { 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/delete_integration_test.go b/cmd/api/src/daemons/datapipe/delete_integration_test.go index 661e3ec9d44b..c4ae3a7ee3f2 100644 --- a/cmd/api/src/daemons/datapipe/delete_integration_test.go +++ b/cmd/api/src/daemons/datapipe/delete_integration_test.go @@ -43,7 +43,7 @@ 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) @@ -94,7 +94,7 @@ 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) @@ -142,7 +142,7 @@ 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) @@ -189,7 +189,7 @@ 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) diff --git a/cmd/api/src/daemons/datapipe/pipeline_integration_test.go b/cmd/api/src/daemons/datapipe/pipeline_integration_test.go index 7c78a4b3b415..d3c9793c446f 100644 --- a/cmd/api/src/daemons/datapipe/pipeline_integration_test.go +++ b/cmd/api/src/daemons/datapipe/pipeline_integration_test.go @@ -45,7 +45,7 @@ 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) @@ -97,7 +97,7 @@ 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) @@ -149,7 +149,7 @@ 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) @@ -211,7 +211,7 @@ func TestPartialIngest(t *testing.T) { ingestCtx := graphify.NewIngestContext(ctx, graphify.WithEndpointResolver(endpoint.NewResolver(testSuite.GraphDB))) fileData, err := testSuite.GraphifyService.ProcessIngestFile(ingestCtx, ingestFileService, model.IngestTask{ - StoredFileName: path.Join(testSuite.WorkDir, "oneGoodOneInvalidRel.json"), + StoredFileName: "oneGoodOneInvalidRel.json", FileType: model.FileTypeJson, }) require.NoError(t, err) 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/graphify/ingest_storage_test.go b/cmd/api/src/services/graphify/ingest_storage_test.go new file mode 100644 index 000000000000..8b4fd32d3567 --- /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 stroage`, + }, + }, + } + + 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 stroage: 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/tasks_integration_test.go b/cmd/api/src/services/graphify/tasks_integration_test.go index fc8d4c49d5f4..1cb03f07ab77 100644 --- a/cmd/api/src/services/graphify/tasks_integration_test.go +++ b/cmd/api/src/services/graphify/tasks_integration_test.go @@ -40,14 +40,14 @@ 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) @@ -86,7 +86,7 @@ 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) @@ -125,19 +125,19 @@ 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) @@ -176,7 +176,7 @@ 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) @@ -215,19 +215,19 @@ 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) @@ -266,7 +266,7 @@ 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) @@ -305,14 +305,14 @@ 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) @@ -351,7 +351,7 @@ 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) diff --git a/cmd/api/src/services/upload/upload.go b/cmd/api/src/services/upload/upload.go index c515184bd2b5..b072ad0069ba 100644 --- a/cmd/api/src/services/upload/upload.go +++ b/cmd/api/src/services/upload/upload.go @@ -104,7 +104,9 @@ func WriteAndValidateFile(ctx context.Context, fileService storage.FileService, tempFileName), attr.Error(validationErr), ) - fileService.DeleteFile(ctx, tempFileName) + if tempFileName != "" { + fileService.DeleteFile(ctx, tempFileName) + } return "", validationErr } @@ -116,7 +118,9 @@ func WriteAndValidateFile(ctx context.Context, fileService storage.FileService, slog.String("temp_file_name", tempFileName), attr.Error(writeErr), ) - fileService.DeleteFile(ctx, tempFileName) + if tempFileName != "" { + fileService.DeleteFile(ctx, tempFileName) + } return "", writeErr } 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 From e1a90867874bab70e85ee32e33301b77bc19c80b Mon Sep 17 00:00:00 2001 From: Michael Cuomo Date: Tue, 19 May 2026 13:18:21 -0400 Subject: [PATCH 24/26] chore: minor cleanup of todos and flatten files for temp --- cmd/api/src/services/graphify/ingest_storage.go | 2 +- cmd/api/src/services/storage/storage.go | 3 +-- cmd/api/src/services/storage/storage_test.go | 4 ++-- cmd/api/src/test/lab/fixtures/api.go | 5 +++-- 4 files changed, 7 insertions(+), 7 deletions(-) diff --git a/cmd/api/src/services/graphify/ingest_storage.go b/cmd/api/src/services/graphify/ingest_storage.go index 9c84a8c570c1..eb7207173d41 100644 --- a/cmd/api/src/services/graphify/ingest_storage.go +++ b/cmd/api/src/services/graphify/ingest_storage.go @@ -114,7 +114,7 @@ func WriteArchiveFileToStorage(ctx context.Context, fileService storage.FileServ ctx, prefix, normalizedFile, - storage.WriteOptions{}, // TODO MC: include write options? + storage.WriteOptions{}, ) if err != nil { return "", fmt.Errorf("write archive file %q to stroage: %w", archiveFile.Name, err) diff --git a/cmd/api/src/services/storage/storage.go b/cmd/api/src/services/storage/storage.go index 0a5fdbebe71a..5fb39b5c0302 100644 --- a/cmd/api/src/services/storage/storage.go +++ b/cmd/api/src/services/storage/storage.go @@ -24,7 +24,6 @@ import ( "errors" "fmt" "io" - "path" "time" "github.com/specterops/bloodhound/cmd/api/src/config" @@ -177,7 +176,7 @@ func (s *LocalFileService) WriteTempFile(ctx context.Context, prefix string, rea return "", err } - tempPath := path.Join(prefix, "tmp-"+id) + tempPath := prefix + "tmp-" + id if err := s.Storage.Put(ctx, tempPath, reader, opts); err != nil { return "", err } diff --git a/cmd/api/src/services/storage/storage_test.go b/cmd/api/src/services/storage/storage_test.go index 1daaa5b9bbdb..49067b52eb06 100644 --- a/cmd/api/src/services/storage/storage_test.go +++ b/cmd/api/src/services/storage/storage_test.go @@ -468,14 +468,14 @@ func TestLocalFileService_WriteTempFile(t *testing.T) { 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) + 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) + tempPath, err := fileService.WriteTempFile(ctx, "prefix_", strings.NewReader("content"), options) // Assert if testCase.expected.errIs != nil { diff --git a/cmd/api/src/test/lab/fixtures/api.go b/cmd/api/src/test/lab/fixtures/api.go index ea0c70613248..df01d5a0ea89 100644 --- a/cmd/api/src/test/lab/fixtures/api.go +++ b/cmd/api/src/test/lab/fixtures/api.go @@ -60,8 +60,9 @@ func NewCustomApiFixture(cfgFixture *lab.Fixture[config.Configuration]) *lab.Fix defer wg.Done() initializer := bootstrap.Initializer[*database.BloodhoundDB, *graph.DatabaseSwitch]{ - Configuration: cfg, - DBConnector: services.ConnectDatabases, + 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 From c85404ef3bde896068f2f6759c861a6391609a8b Mon Sep 17 00:00:00 2001 From: Michael Cuomo Date: Tue, 19 May 2026 13:27:05 -0400 Subject: [PATCH 25/26] chore: rename LocalFileService to StorageFileService and typo fixes --- cmd/api/src/bootstrap/initializer.go | 2 +- .../src/services/graphify/ingest_storage.go | 2 +- .../services/graphify/ingest_storage_test.go | 4 ++-- cmd/api/src/services/storage/storage.go | 22 +++++++++---------- cmd/api/src/services/storage/storage_test.go | 20 ++++++++--------- 5 files changed, 25 insertions(+), 25 deletions(-) diff --git a/cmd/api/src/bootstrap/initializer.go b/cmd/api/src/bootstrap/initializer.go index d6c0a2a4eb92..50b5298ecbb0 100644 --- a/cmd/api/src/bootstrap/initializer.go +++ b/cmd/api/src/bootstrap/initializer.go @@ -33,7 +33,7 @@ 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 RuntimeDependenciesConstructor[DBType database.Database, GraphType graph.Database] func(ctx context.Context, cfg config.Configuration, conections DatabaseConnections[DBType, GraphType]) (RuntimeDependencies, 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 { diff --git a/cmd/api/src/services/graphify/ingest_storage.go b/cmd/api/src/services/graphify/ingest_storage.go index eb7207173d41..4377355d82c6 100644 --- a/cmd/api/src/services/graphify/ingest_storage.go +++ b/cmd/api/src/services/graphify/ingest_storage.go @@ -117,7 +117,7 @@ func WriteArchiveFileToStorage(ctx context.Context, fileService storage.FileServ storage.WriteOptions{}, ) if err != nil { - return "", fmt.Errorf("write archive file %q to stroage: %w", archiveFile.Name, err) + return "", fmt.Errorf("write archive file %q to storage: %w", archiveFile.Name, err) } return extractedPath, nil diff --git a/cmd/api/src/services/graphify/ingest_storage_test.go b/cmd/api/src/services/graphify/ingest_storage_test.go index 8b4fd32d3567..f0995ff25f28 100644 --- a/cmd/api/src/services/graphify/ingest_storage_test.go +++ b/cmd/api/src/services/graphify/ingest_storage_test.go @@ -351,7 +351,7 @@ func TestWriteArchiveFileToStorage(t *testing.T) { writeErr: errWrite, expected: expected{ errIs: errWrite, - errContains: `write archive file "file.json" to stroage`, + errContains: `write archive file "file.json" to storage`, }, }, } @@ -572,7 +572,7 @@ func TestExtractIngestFiles(t *testing.T) { { Name: "one.json", ParentFile: "provided.zip", - Errors: []string{`write archive file "one.json" to stroage: write failed`}, + Errors: []string{`write archive file "one.json" to storage: write failed`}, }, { Name: "two.json", diff --git a/cmd/api/src/services/storage/storage.go b/cmd/api/src/services/storage/storage.go index 5fb39b5c0302..1060888c51ba 100644 --- a/cmd/api/src/services/storage/storage.go +++ b/cmd/api/src/services/storage/storage.go @@ -128,12 +128,12 @@ type FileService interface { ListFiles(ctx context.Context, name string, opts ListOptions) ([]FileInfo, error) } -type LocalFileService struct { +type StorageFileService struct { Storage Storage } -func NewFileService(storage Storage) *LocalFileService { - return &LocalFileService{Storage: storage} +func NewFileService(storage Storage) *StorageFileService { + return &StorageFileService{Storage: storage} } func randomID() (string, error) { @@ -144,11 +144,11 @@ func randomID() (string, error) { return hex.EncodeToString(b[:]), nil } -func (s *LocalFileService) GetFile(ctx context.Context, name string) (io.ReadCloser, FileInfo, error) { +func (s *StorageFileService) GetFile(ctx context.Context, name string) (io.ReadCloser, FileInfo, error) { return s.Storage.Get(ctx, name) } -func (s *LocalFileService) ReadFile(ctx context.Context, name string) ([]byte, error) { +func (s *StorageFileService) ReadFile(ctx context.Context, name string) ([]byte, error) { rc, _, err := s.Storage.Get(ctx, name) if err != nil { return nil, err @@ -158,19 +158,19 @@ func (s *LocalFileService) ReadFile(ctx context.Context, name string) ([]byte, e return io.ReadAll(rc) } -func (s *LocalFileService) WriteFile(ctx context.Context, name string, data []byte, opts WriteOptions) error { +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 *LocalFileService) WriteFileFromReader(ctx context.Context, name string, reader io.Reader, opts WriteOptions) error { +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 *LocalFileService) DeleteFile(ctx context.Context, name string) error { +func (s *StorageFileService) DeleteFile(ctx context.Context, name string) error { return s.Storage.Delete(ctx, name) } -func (s *LocalFileService) WriteTempFile(ctx context.Context, prefix string, reader io.Reader, opts WriteOptions) (string, error) { +func (s *StorageFileService) WriteTempFile(ctx context.Context, prefix string, reader io.Reader, opts WriteOptions) (string, error) { id, err := randomID() if err != nil { return "", err @@ -184,11 +184,11 @@ func (s *LocalFileService) WriteTempFile(ctx context.Context, prefix string, rea return tempPath, nil } -func (s *LocalFileService) MoveFile(ctx context.Context, srcName, dstName string, options WriteOptions) error { +func (s *StorageFileService) MoveFile(ctx context.Context, srcName, dstName string, options WriteOptions) error { return s.Storage.Move(ctx, srcName, dstName, options) } -func (s *LocalFileService) ListFiles(ctx context.Context, name string, options ListOptions) ([]FileInfo, error) { +func (s *StorageFileService) ListFiles(ctx context.Context, name string, options ListOptions) ([]FileInfo, error) { return s.Storage.List(ctx, name, options) } diff --git a/cmd/api/src/services/storage/storage_test.go b/cmd/api/src/services/storage/storage_test.go index 49067b52eb06..f4e2269e57c4 100644 --- a/cmd/api/src/services/storage/storage_test.go +++ b/cmd/api/src/services/storage/storage_test.go @@ -77,7 +77,7 @@ func TestNewFileService(t *testing.T) { require.Same(t, mockStorage, fileService.Storage) } -func TestLocalFileService_GetFile(t *testing.T) { +func TestStorageFileService_GetFile(t *testing.T) { t.Parallel() var ( @@ -148,7 +148,7 @@ func TestLocalFileService_GetFile(t *testing.T) { } } -func TestLocalFileService_ReadFile(t *testing.T) { +func TestStorageFileService_ReadFile(t *testing.T) { t.Parallel() var ( @@ -244,7 +244,7 @@ func TestLocalFileService_ReadFile(t *testing.T) { } } -func TestLocalFileService_WriteFile(t *testing.T) { +func TestStorageFileService_WriteFile(t *testing.T) { t.Parallel() var errPut = errors.New("put failed") @@ -307,7 +307,7 @@ func TestLocalFileService_WriteFile(t *testing.T) { } } -func TestLocalFileService_WriteFileFromReader(t *testing.T) { +func TestStorageFileService_WriteFileFromReader(t *testing.T) { t.Parallel() var errPut = errors.New("put failed") @@ -365,7 +365,7 @@ func TestLocalFileService_WriteFileFromReader(t *testing.T) { } } -func TestLocalFileService_DeleteFile(t *testing.T) { +func TestStorageFileService_DeleteFile(t *testing.T) { t.Parallel() var errDelete = errors.New("delete failed") @@ -421,7 +421,7 @@ func TestLocalFileService_DeleteFile(t *testing.T) { } } -func TestLocalFileService_WriteTempFile(t *testing.T) { +func TestStorageFileService_WriteTempFile(t *testing.T) { t.Parallel() var errPut = errors.New("put failed") @@ -489,7 +489,7 @@ func TestLocalFileService_WriteTempFile(t *testing.T) { } } -func TestLocalFileService_MoveFile(t *testing.T) { +func TestStorageFileService_MoveFile(t *testing.T) { t.Parallel() var errMove = errors.New("move failed") @@ -546,7 +546,7 @@ func TestLocalFileService_MoveFile(t *testing.T) { } } -func TestLocalFileService_ListFiles(t *testing.T) { +func TestStorageFileService_ListFiles(t *testing.T) { t.Parallel() var ( @@ -1042,10 +1042,10 @@ func TestNewDefaultFileServices(t *testing.T) { require.Len(t, fileServices, 4) for _, fileService := range fileServices { - localFileService, ok := fileService.(*storage.LocalFileService) + storageFileService, ok := fileService.(*storage.StorageFileService) require.True(t, ok) - localStore, ok := localFileService.Storage.(*storage.LocalStore) + localStore, ok := storageFileService.Storage.(*storage.LocalStore) require.True(t, ok) require.NoError(t, localStore.Close()) } From cd0604e4dda567a3f8dfe29ed06a1ff58808e8b6 Mon Sep 17 00:00:00 2001 From: Michael Cuomo Date: Tue, 19 May 2026 13:30:31 -0400 Subject: [PATCH 26/26] chore: pfc --- cmd/api/src/services/upload/streamdecoder.go | 18 +++++++----------- packages/go/schemagen/generator/sql.go | 16 ++++++++-------- 2 files changed, 15 insertions(+), 19 deletions(-) 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/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")