lisa/fix-screenshots-only-as-image-types

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))

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(

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()