Skip to content

First pass at a simple integration test #1353

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 42 additions & 14 deletions cmd/api/src/api/v2/file_uploads.go
Original file line number Diff line number Diff line change
Expand Up @@ -131,23 +131,44 @@ func (s Resources) ProcessFileUpload(response http.ResponseWriter, request *http
defer request.Body.Close()
}

if !IsValidContentTypeForUpload(request.Header) {
contentType, err := ParseUploadContentType(request.Header)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd like us to go back to using a pipeline approach (if/else chaining) for this set of requirements, primarily to keep it in line with other endpoints. It does come with the upside of not requiring any explicit returns.

if err != nil {
api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusBadRequest, "Content type must be application/json or application/zip", request), response)
} else if fileUploadJobID, err := strconv.Atoi(fileUploadJobIdString); err != nil {
return
}

fileUploadJobID, err := strconv.Atoi(fileUploadJobIdString)
if err != nil {
api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusBadRequest, api.ErrorResponseDetailsIDMalformed, request), response)
} else if ingestJob, err := ingest.GetIngestJobByID(request.Context(), s.DB, int64(fileUploadJobID)); err != nil {
return
}

ingestJob, err := ingest.GetIngestJobByID(request.Context(), s.DB, int64(fileUploadJobID))
if err != nil {
api.HandleDatabaseError(request, response, err)
} else if fileName, fileType, err := ingest.SaveIngestFile(s.Config.TempDirectory(), request); errors.Is(err, ingest.ErrInvalidJSON) {
return
}

fileName, fileType, err := ingest.SaveIngestFile(s.Config.TempDirectory(), contentType, request.Body)
if errors.Is(err, ingest.ErrInvalidJSON) {
api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusBadRequest, fmt.Sprintf("Error saving ingest file: %v", err), request), response)
return
} 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 = ingest.CreateIngestTask(request.Context(), s.DB, fileName, fileType, requestId, int64(fileUploadJobID)); err != nil {
return
}

if _, err = ingest.CreateIngestTask(request.Context(), s.DB, fileName, fileType, requestId, int64(fileUploadJobID)); err != nil {
api.HandleDatabaseError(request, response, err)
} else if err = ingest.TouchIngestJobLastIngest(request.Context(), s.DB, ingestJob); err != nil {
return
}

if err = ingest.TouchIngestJobLastIngest(request.Context(), s.DB, ingestJob); err != nil {
api.HandleDatabaseError(request, response, err)
} else {
response.WriteHeader(http.StatusAccepted)
return
}

response.WriteHeader(http.StatusAccepted)
}

func (s Resources) EndFileUploadJob(response http.ResponseWriter, request *http.Request) {
Expand All @@ -172,13 +193,20 @@ func (s Resources) ListAcceptedFileUploadTypes(response http.ResponseWriter, req
api.WriteBasicResponse(request.Context(), ingestModel.AllowedFileUploadTypes, http.StatusOK, response)
}

func IsValidContentTypeForUpload(header http.Header) bool {
func ParseUploadContentType(header http.Header) (string, error) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This should probably be private, I'm not sure why the previous was public tbh, but we can fix that here.

rawValue := header.Get(headers.ContentType.String())
if rawValue == "" {
return false
} else if parsed, _, err := mime.ParseMediaType(rawValue); err != nil {
return false
} else {
return slices.Contains(ingestModel.AllowedFileUploadTypes, parsed)
return "", fmt.Errorf("missing Content-Type header")
}

parsed, _, err := mime.ParseMediaType(rawValue)
if err != nil {
return "", fmt.Errorf("invalid Content-Type format: %w", err)
}

if !slices.Contains(ingestModel.AllowedFileUploadTypes, parsed) {
return "", fmt.Errorf("unsupported Content-Type: %s", parsed)
}

return parsed, nil
Comment on lines +199 to +211
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Returning this to a pipeline style is preferred right now. There are no current edge cases in this logic

}
15 changes: 8 additions & 7 deletions cmd/api/src/services/ingest/job.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,16 +23,14 @@ import (
"fmt"
"io"
"log/slog"
"net/http"
"os"
"slices"
"time"

"github.com/specterops/bloodhound/bomenc"
"github.com/specterops/bloodhound/headers"
"github.com/specterops/bloodhound/mediatypes"
"github.com/specterops/bloodhound/src/model"
"github.com/specterops/bloodhound/src/model/ingest"
"github.com/specterops/bloodhound/src/utils"
)

const jobActivityTimeout = time.Minute * 20
Expand Down Expand Up @@ -74,6 +72,7 @@ func GetAllIngestJobs(ctx context.Context, db IngestData, skip int, limit int, o
return db.GetAllIngestJobs(ctx, skip, limit, order, filter)
}

// TODO: Integration test here
func StartIngestJob(ctx context.Context, db IngestData, user model.User) (model.IngestJob, error) {
job := model.IngestJob{
UserID: user.ID,
Expand Down Expand Up @@ -104,16 +103,17 @@ func WriteAndValidateJSON(src io.Reader, dst io.Writer) error {
return err
}

func SaveIngestFile(location string, request *http.Request) (string, model.FileType, error) {
fileData := request.Body
// TODO: Integration test here
func SaveIngestFile(location string, contentType string, fileData io.ReadCloser) (string, model.FileType, error) {
tempFile, err := os.CreateTemp(location, "bh")
if err != nil {
return "", model.FileTypeJson, fmt.Errorf("error creating ingest file: %w", err)
}

if utils.HeaderMatches(request.Header, headers.ContentType.String(), mediatypes.ApplicationJson.String()) {
if contentType == mediatypes.ApplicationJson.String() {
return tempFile.Name(), model.FileTypeJson, WriteAndValidateFile(fileData, tempFile, WriteAndValidateJSON)
} else if utils.HeaderMatches(request.Header, headers.ContentType.String(), ingest.AllowedZipFileUploadTypes...) {

} else if slices.Contains(ingest.AllowedZipFileUploadTypes, contentType) {
return tempFile.Name(), model.FileTypeZip, WriteAndValidateFile(fileData, tempFile, WriteAndValidateZip)
} else {
// We should never get here since this is checked a level above
Expand Down Expand Up @@ -144,6 +144,7 @@ func TouchIngestJobLastIngest(ctx context.Context, db IngestData, job model.Inge
return db.UpdateIngestJob(ctx, job)
}

// TODO: Integration test here
func EndIngestJob(ctx context.Context, db IngestData, job model.IngestJob) error {
job.Status = model.JobStatusIngesting

Expand Down
37 changes: 37 additions & 0 deletions cmd/api/src/services/ingest/job_integration_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
package ingest

import (
"bytes"
"io"
"os"
"testing"

"github.com/specterops/bloodhound/src/model"
)

func TestIntegration_SaveIngestFile_JSON(t *testing.T) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's name this TestSaveIngestFile and use subtests to split by JSON/zip/etc

// TODO: Make this
body := io.NopCloser(bytes.NewBufferString(`{"meta":{"methods":46067,"type":"computers","count":0,"vers ion":6},"data":[]}`))

// TODO: Make this take it a FS?
location, fileType, err := SaveIngestFile("/tmp", "application/json", body)
if err != nil {
t.Fatalf("SaveIngestFile failed: %v", err)
}

if fileType != model.FileTypeJson {
t.Errorf("expected fileType %v, got %v", model.FileTypeJson, fileType)
}

info, statErr := os.Stat(location)
if statErr != nil {
t.Fatalf("output file not found: %v", statErr)
}

if info.Size() == 0 {
t.Error("file saved is empty")
}
Comment on lines +17 to +33
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We have access to testify for require and assert assertions, it makes the test code a bit more concise and easier to reason about. I'd recommend using require.NoError for your error conditions that should never fail because they'd cause the remaining assertions to autofail, and appropriate assertions like assert.Equal to check size. See https://pkg.go.dev/github.com/stretchr/testify/assert#pkg-functions for more docs on what assertions exist. The more specific the better, since it acts as documentation for what specifically you're asserting.


// TODO: Make this unnecessary
os.Remove(location)
}
Loading