From 5b7de043f789b9da4a82108696242024e743de8d Mon Sep 17 00:00:00 2001 From: Gustavo Carvalho Date: Fri, 15 May 2026 11:40:10 -0300 Subject: [PATCH] implement: fix-screenshots-only-as-image-types --- .../api/handler/workflows/step_execution.go | 3 + internal/workflow/step_transition.go | 61 ++++++++++++++++ internal/workflow/step_transition_test.go | 70 +++++++++++++++++++ 3 files changed, 134 insertions(+) diff --git a/internal/api/handler/workflows/step_execution.go b/internal/api/handler/workflows/step_execution.go index d9c1fac3..aefa316c 100644 --- a/internal/api/handler/workflows/step_execution.go +++ b/internal/api/handler/workflows/step_execution.go @@ -333,6 +333,9 @@ func (h *StepExecutionHandler) TransitionStep(ctx echo.Context) error { if errors.Is(err, workflow.ErrInvalidStepTransition) { return ctx.JSON(http.StatusBadRequest, api.NewError(err)) } + if errors.Is(err, workflow.ErrInvalidEvidenceSubmission) { + return ctx.JSON(http.StatusBadRequest, api.NewError(err)) + } // Check if it's a permission error if isPermissionError(err) { return ctx.JSON(http.StatusForbidden, api.NewError(err)) diff --git a/internal/workflow/step_transition.go b/internal/workflow/step_transition.go index fbca8bd3..301cd3f9 100644 --- a/internal/workflow/step_transition.go +++ b/internal/workflow/step_transition.go @@ -31,6 +31,7 @@ type StepTransitionService struct { } var ErrInvalidStepTransition = errors.New("invalid step transition") +var ErrInvalidEvidenceSubmission = errors.New("invalid evidence submission") var errTransitionForbidden = errors.New("forbidden") // WorkflowDefinitionServiceInterface defines the interface for workflow definition operations @@ -293,6 +294,9 @@ func (s *StepTransitionService) validateEvidenceRequirements(stepDef *workflows. // Build a map of submitted evidence types submittedTypes := make(map[string]int) for _, evidence := range submittedEvidence { + if err := validateEvidenceSubmissionFileType(evidence); err != nil { + return err + } submittedTypes[evidence.EvidenceType]++ } @@ -313,6 +317,63 @@ func (s *StepTransitionService) validateEvidenceRequirements(stepDef *workflows. return nil } +func validateEvidenceSubmissionFileType(evidence EvidenceSubmission) error { + if !strings.EqualFold(strings.TrimSpace(evidence.EvidenceType), "screenshot") { + return nil + } + + allowedMediaTypes := map[string]struct{}{ + "image/png": {}, + "image/jpg": {}, + "image/jpeg": {}, + "image/gif": {}, + "image/webp": {}, + } + allowedExtensions := map[string]struct{}{ + ".png": {}, + ".jpg": {}, + ".jpeg": {}, + ".gif": {}, + ".webp": {}, + } + + mediaType := strings.ToLower(strings.TrimSpace(strings.Split(evidence.MediaType, ";")[0])) + if mediaType != "" { + if _, ok := allowedMediaTypes[mediaType]; !ok { + return fmt.Errorf("%w: screenshot evidence file type %q is not supported", ErrInvalidEvidenceSubmission, evidence.MediaType) + } + } + + filename := evidence.Name + if evidence.FilePath != "" { + filename = evidence.FilePath + } + extension := strings.ToLower(strings.TrimSpace(fileExtension(filename))) + if extension != "" { + if _, ok := allowedExtensions[extension]; !ok { + return fmt.Errorf("%w: screenshot evidence file extension %q is not supported", ErrInvalidEvidenceSubmission, extension) + } + } + + if mediaType == "" && extension == "" && evidence.FileContent != "" { + return fmt.Errorf("%w: screenshot evidence requires an image media type or file extension", ErrInvalidEvidenceSubmission) + } + + return nil +} + +func fileExtension(filename string) string { + lastSlash := strings.LastIndexAny(filename, `/\`) + if lastSlash >= 0 { + filename = filename[lastSlash+1:] + } + lastDot := strings.LastIndex(filename, ".") + if lastDot < 0 { + return "" + } + return filename[lastDot:] +} + // storeStepEvidence stores the submitted evidence for a step execution as relational.Evidence // with BackMatter resources for uploaded files and proper labels for the workflow execution stream func (s *StepTransitionService) storeStepEvidence( diff --git a/internal/workflow/step_transition_test.go b/internal/workflow/step_transition_test.go index 05c00b77..9f638667 100644 --- a/internal/workflow/step_transition_test.go +++ b/internal/workflow/step_transition_test.go @@ -216,6 +216,76 @@ func TestVerifyTransitionActorPermission_UnexpectedRoleLookupErrorBubbles(t *tes mockRole.AssertExpectations(t) } +func TestValidateEvidenceRequirements_ScreenshotFileTypes(t *testing.T) { + svc := &StepTransitionService{} + stepDef := &workflows.WorkflowStepDefinition{} + + tests := []struct { + name string + evidence EvidenceSubmission + wantError bool + }{ + { + name: "allows screenshot image media type and extension", + evidence: EvidenceSubmission{ + EvidenceType: "screenshot", + Name: "screen.PNG", + MediaType: "image/png", + FileContent: "ZmFrZQ==", + }, + }, + { + name: "allows screenshot jpg alias", + evidence: EvidenceSubmission{ + EvidenceType: "screenshot", + Name: "screen.jpg", + MediaType: "image/jpg", + FileContent: "ZmFrZQ==", + }, + }, + { + name: "rejects screenshot document media type", + evidence: EvidenceSubmission{ + EvidenceType: "screenshot", + Name: "screen.pdf", + MediaType: "application/pdf", + FileContent: "ZmFrZQ==", + }, + wantError: true, + }, + { + name: "rejects screenshot document extension", + evidence: EvidenceSubmission{ + EvidenceType: "screenshot", + Name: "screen.doc", + MediaType: "image/png", + FileContent: "ZmFrZQ==", + }, + wantError: true, + }, + { + name: "keeps document evidence broad", + evidence: EvidenceSubmission{ + EvidenceType: "document", + Name: "attestation.pdf", + MediaType: "application/pdf", + FileContent: "ZmFrZQ==", + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := svc.validateEvidenceRequirements(stepDef, []EvidenceSubmission{tt.evidence}) + if tt.wantError { + require.ErrorIs(t, err, ErrInvalidEvidenceSubmission) + return + } + require.NoError(t, err) + }) + } +} + func TestTransitionStepStatus_UnexpectedPermissionLookupErrorBubbles(t *testing.T) { stepExecutionID := uuid.New() stepDefID := uuid.New()