From fd6992ac00e8de1308d20f4cb865cf94ebbd44ad Mon Sep 17 00:00:00 2001 From: Bryan Frimin Date: Mon, 5 May 2025 01:06:59 -0700 Subject: [PATCH] Add filetype validation Signed-off-by: Bryan Frimin --- pkg/filevalidation/validator.go | 186 ++++++++++ pkg/filevalidation/validator_test.go | 506 +++++++++++++++++++++++++++ pkg/probo/evidence_service.go | 45 ++- pkg/probo/organization_service.go | 66 +++- pkg/probo/service.go | 19 +- 5 files changed, 809 insertions(+), 13 deletions(-) create mode 100644 pkg/filevalidation/validator.go create mode 100644 pkg/filevalidation/validator_test.go diff --git a/pkg/filevalidation/validator.go b/pkg/filevalidation/validator.go new file mode 100644 index 000000000..cd166467a --- /dev/null +++ b/pkg/filevalidation/validator.go @@ -0,0 +1,186 @@ +// Copyright (c) 2025 Probo Inc . +// +// Permission to use, copy, modify, and/or distribute this software for any +// purpose with or without fee is hereby granted, provided that the above +// copyright notice and this permission notice appear in all copies. +// +// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +// PERFORMANCE OF THIS SOFTWARE. + +package filevalidation + +import ( + "fmt" + "path/filepath" + "strings" +) + +const ( + // DefaultMaxFileSize defines the default maximum allowed file size (50MB) + DefaultMaxFileSize = 50 * 1024 * 1024 +) + +// File type categories +const ( + CategoryDocument = "document" + CategorySpreadsheet = "spreadsheet" + CategoryPresentation = "presentation" + CategoryText = "text" + CategoryImage = "image" + CategoryData = "data" + CategoryVideo = "video" +) + +// FileType defines a supported file type with its MIME type and extensions +type FileType struct { + MimeType string + Extensions []string + Category string +} + +// FileTypes is a list of all supported file types +var FileTypes = []FileType{ + // Document types + {MimeType: "application/pdf", Extensions: []string{".pdf"}, Category: CategoryDocument}, + {MimeType: "application/msword", Extensions: []string{".doc"}, Category: CategoryDocument}, + {MimeType: "application/vnd.openxmlformats-officedocument.wordprocessingml.document", Extensions: []string{".docx"}, Category: CategoryDocument}, + {MimeType: "application/vnd.oasis.opendocument.text", Extensions: []string{".odt"}, Category: CategoryDocument}, + + // Spreadsheet types + {MimeType: "application/vnd.ms-excel", Extensions: []string{".xls"}, Category: CategorySpreadsheet}, + {MimeType: "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", Extensions: []string{".xlsx"}, Category: CategorySpreadsheet}, + {MimeType: "application/vnd.oasis.opendocument.spreadsheet", Extensions: []string{".ods"}, Category: CategorySpreadsheet}, + + // Presentation types + {MimeType: "application/vnd.ms-powerpoint", Extensions: []string{".ppt"}, Category: CategoryPresentation}, + {MimeType: "application/vnd.openxmlformats-officedocument.presentationml.presentation", Extensions: []string{".pptx"}, Category: CategoryPresentation}, + {MimeType: "application/vnd.oasis.opendocument.presentation", Extensions: []string{".odp"}, Category: CategoryPresentation}, + + // Text types + {MimeType: "text/plain", Extensions: []string{".txt"}, Category: CategoryText}, + {MimeType: "text/x-log", Extensions: []string{".log"}, Category: CategoryText}, + {MimeType: "text/uri-list", Extensions: []string{".uri"}, Category: CategoryText}, + + // Image types + {MimeType: "image/jpeg", Extensions: []string{".jpg", ".jpeg"}, Category: CategoryImage}, + {MimeType: "image/png", Extensions: []string{".png"}, Category: CategoryImage}, + {MimeType: "image/gif", Extensions: []string{".gif"}, Category: CategoryImage}, + {MimeType: "image/svg+xml", Extensions: []string{".svg"}, Category: CategoryImage}, + {MimeType: "image/webp", Extensions: []string{".webp"}, Category: CategoryImage}, + + // Data types + {MimeType: "application/yaml", Extensions: []string{".yaml", ".yml"}, Category: CategoryData}, + {MimeType: "application/json", Extensions: []string{".json"}, Category: CategoryData}, + {MimeType: "text/yaml", Extensions: []string{".yaml", ".yml"}, Category: CategoryData}, + {MimeType: "text/json", Extensions: []string{".json"}, Category: CategoryData}, + + // Video types + {MimeType: "video/mp4", Extensions: []string{".mp4"}, Category: CategoryVideo}, + {MimeType: "video/mpeg", Extensions: []string{".mpeg", ".mpg"}, Category: CategoryVideo}, + {MimeType: "video/quicktime", Extensions: []string{".mov"}, Category: CategoryVideo}, + {MimeType: "video/x-msvideo", Extensions: []string{".avi"}, Category: CategoryVideo}, + {MimeType: "video/webm", Extensions: []string{".webm"}, Category: CategoryVideo}, +} + +// FileValidator is a configurable file validator +type FileValidator struct { + // MaxFileSize specifies the maximum file size in bytes + MaxFileSize int64 + + // AllowedMimeTypes is a map of allowed MIME types + AllowedMimeTypes map[string]bool + + // AllowedExtensions maps file extensions to their expected MIME types + AllowedExtensions map[string][]string + + // Categories is a list of categories to include + Categories []string +} + +// NewValidator creates a new file validator using supported file types +func NewValidator(categories ...string) *FileValidator { + v := &FileValidator{ + MaxFileSize: DefaultMaxFileSize, + AllowedMimeTypes: make(map[string]bool), + AllowedExtensions: make(map[string][]string), + Categories: categories, + } + + if len(categories) == 0 { + for _, fileType := range FileTypes { + v.AllowedMimeTypes[fileType.MimeType] = true + for _, ext := range fileType.Extensions { + if v.AllowedExtensions[ext] == nil { + v.AllowedExtensions[ext] = []string{} + } + v.AllowedExtensions[ext] = append(v.AllowedExtensions[ext], fileType.MimeType) + } + } + return v + } + + categoryMap := make(map[string]bool) + for _, category := range categories { + categoryMap[category] = true + } + + for _, fileType := range FileTypes { + if categoryMap[fileType.Category] { + v.AllowedMimeTypes[fileType.MimeType] = true + for _, ext := range fileType.Extensions { + if v.AllowedExtensions[ext] == nil { + v.AllowedExtensions[ext] = []string{} + } + v.AllowedExtensions[ext] = append(v.AllowedExtensions[ext], fileType.MimeType) + } + } + } + + return v +} + +// WithMaxFileSize sets the maximum file size and returns the validator +func (v *FileValidator) WithMaxFileSize(maxSize int64) *FileValidator { + v.MaxFileSize = maxSize + return v +} + +// Validate validates that the file meets the configured requirements +func (v *FileValidator) Validate(filename string, contentType string, size int64) error { + if size > v.MaxFileSize { + return fmt.Errorf("file size exceeds maximum allowed size of %d bytes", v.MaxFileSize) + } + + if !v.AllowedMimeTypes[contentType] { + return fmt.Errorf("content type %q is not allowed", contentType) + } + + ext := strings.ToLower(filepath.Ext(filename)) + if ext == "" { + return fmt.Errorf("file has no extension") + } + + allowedTypes, extAllowed := v.AllowedExtensions[ext] + if !extAllowed { + return fmt.Errorf("file extension %q is not allowed", ext) + } + + validType := false + for _, allowedType := range allowedTypes { + if contentType == allowedType { + validType = true + break + } + } + + if !validType { + return fmt.Errorf("content type %q does not match extension %q", contentType, ext) + } + + return nil +} diff --git a/pkg/filevalidation/validator_test.go b/pkg/filevalidation/validator_test.go new file mode 100644 index 000000000..45cb35825 --- /dev/null +++ b/pkg/filevalidation/validator_test.go @@ -0,0 +1,506 @@ +// Copyright (c) 2025 Probo Inc . +// +// Permission to use, copy, modify, and/or distribute this software for any +// purpose with or without fee is hereby granted, provided that the above +// copyright notice and this permission notice appear in all copies. +// +// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +// PERFORMANCE OF THIS SOFTWARE. + +package filevalidation + +import ( + "fmt" + "strings" + "testing" +) + +func TestNewValidator(t *testing.T) { + tests := []struct { + name string + categories []string + expectedMimes []string + expectedExts []string + unexpectedMime string + }{ + { + name: "No categories (all types)", + categories: []string{}, + expectedMimes: []string{"application/pdf", "image/jpeg", "text/plain", "video/mp4"}, + expectedExts: []string{".pdf", ".jpg", ".txt", ".mp4"}, + unexpectedMime: "application/octet-stream", + }, + { + name: "Only documents", + categories: []string{CategoryDocument}, + expectedMimes: []string{"application/pdf", "application/msword"}, + expectedExts: []string{".pdf", ".doc", ".docx"}, + unexpectedMime: "image/jpeg", + }, + { + name: "Only images", + categories: []string{CategoryImage}, + expectedMimes: []string{"image/jpeg", "image/png", "image/gif"}, + expectedExts: []string{".jpg", ".png", ".gif"}, + unexpectedMime: "application/pdf", + }, + { + name: "Multiple categories", + categories: []string{CategoryImage, CategoryDocument}, + expectedMimes: []string{"image/jpeg", "application/pdf"}, + expectedExts: []string{".jpg", ".pdf"}, + unexpectedMime: "video/mp4", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + v := NewValidator(tc.categories...) + + // Test expected mime types are allowed + for _, mime := range tc.expectedMimes { + if !v.AllowedMimeTypes[mime] { + t.Errorf("Expected mime type %q to be allowed, but it wasn't", mime) + } + } + + // Test expected extensions are allowed + for _, ext := range tc.expectedExts { + if _, ok := v.AllowedExtensions[ext]; !ok { + t.Errorf("Expected extension %q to be allowed, but it wasn't", ext) + } + } + + // Test unexpected mime type is not allowed + if v.AllowedMimeTypes[tc.unexpectedMime] { + t.Errorf("Unexpected mime type %q should not be allowed, but it was", tc.unexpectedMime) + } + + // Check categories were set correctly + if len(tc.categories) != len(v.Categories) { + t.Errorf("Expected %d categories, got %d", len(tc.categories), len(v.Categories)) + } + + // Default max size should be set + if v.MaxFileSize != DefaultMaxFileSize { + t.Errorf("Expected default max file size %d, got %d", DefaultMaxFileSize, v.MaxFileSize) + } + }) + } +} + +func TestWithMaxFileSize(t *testing.T) { + v := NewValidator() + originalSize := v.MaxFileSize + + // Test changing the max file size + newSize := int64(10 * 1024 * 1024) // 10MB + v = v.WithMaxFileSize(newSize) + + if v.MaxFileSize != newSize { + t.Errorf("Expected max file size to be %d after WithMaxFileSize, got %d", newSize, v.MaxFileSize) + } + + if originalSize == v.MaxFileSize { + t.Errorf("Max file size should have changed from original %d", originalSize) + } +} + +func TestValidate(t *testing.T) { + tests := []struct { + name string + validator *FileValidator + filename string + contentType string + fileSize int64 + shouldError bool + errorMsg string + }{ + { + name: "Valid PDF file", + validator: NewValidator(CategoryDocument), + filename: "test.pdf", + contentType: "application/pdf", + fileSize: 1024 * 1024, // 1MB + shouldError: false, + }, + { + name: "File too large", + validator: NewValidator().WithMaxFileSize(1024 * 1024), // 1MB max + filename: "test.pdf", + contentType: "application/pdf", + fileSize: 2 * 1024 * 1024, // 2MB + shouldError: true, + errorMsg: "file size exceeds maximum allowed size", + }, + { + name: "Disallowed content type", + validator: NewValidator(CategoryDocument), + filename: "test.jpg", + contentType: "image/jpeg", + fileSize: 1024 * 1024, + shouldError: true, + errorMsg: "content type \"image/jpeg\" is not allowed", + }, + { + name: "Missing file extension", + validator: NewValidator(), + filename: "testfile", + contentType: "application/pdf", + fileSize: 1024, + shouldError: true, + errorMsg: "file has no extension", + }, + { + name: "Disallowed file extension", + validator: NewValidator(CategoryDocument), + filename: "test.exe", + contentType: "application/octet-stream", + fileSize: 1024, + shouldError: true, + errorMsg: "content type \"application/octet-stream\" is not allowed", + }, + { + name: "Content type doesn't match extension", + validator: NewValidator(), + filename: "test.pdf", + contentType: "image/jpeg", + fileSize: 1024, + shouldError: true, + errorMsg: "content type \"image/jpeg\" does not match extension \".pdf\"", + }, + { + name: "Valid image file", + validator: NewValidator(CategoryImage), + filename: "test.jpg", + contentType: "image/jpeg", + fileSize: 1024 * 1024, + shouldError: false, + }, + { + name: "Valid with multiple categories", + validator: NewValidator(CategoryImage, CategoryDocument), + filename: "test.jpg", + contentType: "image/jpeg", + fileSize: 1024 * 1024, + shouldError: false, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + err := tc.validator.Validate(tc.filename, tc.contentType, tc.fileSize) + + if tc.shouldError && err == nil { + t.Errorf("Expected error but got none") + } + + if !tc.shouldError && err != nil { + t.Errorf("Expected no error but got: %v", err) + } + + if tc.shouldError && err != nil && tc.errorMsg != "" { + if !contains(err.Error(), tc.errorMsg) { + t.Errorf("Error message '%s' does not contain expected text '%s'", err.Error(), tc.errorMsg) + } + } + }) + } +} + +func TestValidateEdgeCases(t *testing.T) { + tests := []struct { + name string + validator *FileValidator + filename string + contentType string + fileSize int64 + shouldError bool + errorMsg string + }{ + { + name: "Zero file size", + validator: NewValidator(), + filename: "empty.txt", + contentType: "text/plain", + fileSize: 0, + shouldError: false, + }, + { + name: "Exact max file size", + validator: NewValidator().WithMaxFileSize(1024), + filename: "exact.txt", + contentType: "text/plain", + fileSize: 1024, + shouldError: false, + }, + { + name: "File with uppercase extension", + validator: NewValidator(), + filename: "test.PDF", + contentType: "application/pdf", + fileSize: 1024, + shouldError: false, + }, + { + name: "File with multiple extensions", + validator: NewValidator(), + filename: "test.tar.gz", + contentType: "application/gzip", + fileSize: 1024, + shouldError: true, + errorMsg: "content type \"application/gzip\" is not allowed", + }, + { + name: "Empty filename", + validator: NewValidator(), + filename: "", + contentType: "text/plain", + fileSize: 1024, + shouldError: true, + errorMsg: "file has no extension", + }, + { + name: "Empty content type", + validator: NewValidator(), + filename: "test.txt", + contentType: "", + fileSize: 1024, + shouldError: true, + errorMsg: "content type \"\" is not allowed", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + err := tc.validator.Validate(tc.filename, tc.contentType, tc.fileSize) + + if tc.shouldError && err == nil { + t.Errorf("Expected error but got none") + } + + if !tc.shouldError && err != nil { + t.Errorf("Expected no error but got: %v", err) + } + + if tc.shouldError && err != nil && tc.errorMsg != "" { + if !contains(err.Error(), tc.errorMsg) { + t.Errorf("Error message '%s' does not contain expected text '%s'", err.Error(), tc.errorMsg) + } + } + }) + } +} + +// Test each file category individually +func TestFileCategories(t *testing.T) { + categoryTests := []struct { + category string + validExt string + validMimeType string + }{ + {CategoryDocument, ".pdf", "application/pdf"}, + {CategorySpreadsheet, ".xlsx", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"}, + {CategoryPresentation, ".pptx", "application/vnd.openxmlformats-officedocument.presentationml.presentation"}, + {CategoryText, ".txt", "text/plain"}, + {CategoryImage, ".png", "image/png"}, + {CategoryData, ".json", "application/json"}, + {CategoryVideo, ".mp4", "video/mp4"}, + } + + for _, tc := range categoryTests { + t.Run(tc.category, func(t *testing.T) { + v := NewValidator(tc.category) + + // Test that the validator accepts files of this category + err := v.Validate("test"+tc.validExt, tc.validMimeType, 1024) + if err != nil { + t.Errorf("Expected no error for %s file but got: %v", tc.category, err) + } + + // Test that all other categories are rejected + for _, otherTC := range categoryTests { + if otherTC.category == tc.category { + continue + } + + err := v.Validate("test"+otherTC.validExt, otherTC.validMimeType, 1024) + if err == nil { + t.Errorf("Expected error when validating %s file with %s validator, but got none", + otherTC.category, tc.category) + } + } + }) + } +} + +func TestFileTypes(t *testing.T) { + // Check that FileTypes has entries for all defined categories + categories := map[string]bool{ + CategoryDocument: false, + CategorySpreadsheet: false, + CategoryPresentation: false, + CategoryText: false, + CategoryImage: false, + CategoryData: false, + CategoryVideo: false, + } + + for _, ft := range FileTypes { + categories[ft.Category] = true + + // Verify each file type has required fields + if ft.MimeType == "" { + t.Errorf("FileType with category %q has empty MimeType", ft.Category) + } + + if len(ft.Extensions) == 0 { + t.Errorf("FileType with MimeType %q has no extensions", ft.MimeType) + } + } + + // Check if all categories have at least one file type + for category, found := range categories { + if !found { + t.Errorf("No file types defined for category %q", category) + } + } +} + +// Helper function to check if a string contains a substring +func contains(s, substr string) bool { + return strings.Contains(s, substr) +} + +func TestMultipleExtensionsPerMimeType(t *testing.T) { + // Test MIME types that have multiple allowed extensions + for _, fileType := range FileTypes { + if len(fileType.Extensions) > 1 { + v := NewValidator(fileType.Category) + + // All extensions for this MIME type should be valid + for _, ext := range fileType.Extensions { + err := v.Validate("test"+ext, fileType.MimeType, 1024) + if err != nil { + t.Errorf("Extension %q should be valid for MIME type %q but got error: %v", + ext, fileType.MimeType, err) + } + } + } + } +} + +func TestExtensionsWithMultipleMimeTypes(t *testing.T) { + // Create a map of extensions to MIME types + extToMimes := make(map[string][]string) + for _, fileType := range FileTypes { + for _, ext := range fileType.Extensions { + extToMimes[ext] = append(extToMimes[ext], fileType.MimeType) + } + } + + // Find extensions that have multiple MIME types + for ext, mimeTypes := range extToMimes { + if len(mimeTypes) > 1 { + v := NewValidator() + + // All MIME types for this extension should be valid + for _, mimeType := range mimeTypes { + err := v.Validate("test"+ext, mimeType, 1024) + if err != nil { + t.Errorf("MIME type %q should be valid for extension %q but got error: %v", + mimeType, ext, err) + } + } + } + } +} + +// BenchmarkValidate benchmarks the Validate function +func BenchmarkValidate(b *testing.B) { + v := NewValidator() + b.ResetTimer() + + for i := 0; i < b.N; i++ { + _ = v.Validate("test.pdf", "application/pdf", 1024) + } +} + +// BenchmarkNewValidator benchmarks the NewValidator function +func BenchmarkNewValidator(b *testing.B) { + for i := 0; i < b.N; i++ { + _ = NewValidator(CategoryDocument, CategoryImage) + } +} + +// TestUtilFunctions tests utility functions +func TestUtilFunctions(t *testing.T) { + // Test contains function + testCases := []struct { + str string + substr string + expected bool + }{ + {"hello world", "hello", true}, + {"hello world", "world", true}, + {"hello world", "goodbye", false}, + {"hello world", "", true}, // Empty string is always contained + {"", "hello", false}, // Empty string contains nothing except empty string + {"", "", true}, // Empty string contains empty string + } + + for _, tc := range testCases { + t.Run(fmt.Sprintf("contains(%q,%q)", tc.str, tc.substr), func(t *testing.T) { + result := contains(tc.str, tc.substr) + if result != tc.expected { + t.Errorf("contains(%q,%q) = %v; want %v", tc.str, tc.substr, result, tc.expected) + } + }) + } +} + +func TestValidateCustomAllowedExtensions(t *testing.T) { + // Test the file extension not allowed path + v := &FileValidator{ + MaxFileSize: DefaultMaxFileSize, + AllowedMimeTypes: map[string]bool{"application/pdf": true, "image/jpeg": true}, + AllowedExtensions: map[string][]string{ + ".pdf": {"application/pdf"}, + // No .jpg extension here + }, + } + + // This will trigger the "file extension is not allowed" path + err := v.Validate("test.jpg", "image/jpeg", 1024) + if err == nil { + t.Error("Expected error for unregistered extension, got none") + } + if !contains(err.Error(), "file extension \".jpg\" is not allowed") { + t.Errorf("Unexpected error message: %s", err.Error()) + } + + // Test the content type doesn't match extension path + v2 := &FileValidator{ + MaxFileSize: DefaultMaxFileSize, + AllowedMimeTypes: map[string]bool{"application/pdf": true, "image/jpeg": true, "image/png": true}, + AllowedExtensions: map[string][]string{ + ".pdf": {"application/pdf"}, + ".jpg": {"image/jpeg"}, + ".png": {"image/png"}, + }, + } + + // This will trigger the content type mismatch path + err = v2.Validate("test.jpg", "image/png", 1024) + if err == nil { + t.Error("Expected error for content type not matching extension, got none") + } + if !contains(err.Error(), "content type \"image/png\" does not match extension \".jpg\"") { + t.Errorf("Unexpected error message: %s", err.Error()) + } +} diff --git a/pkg/probo/evidence_service.go b/pkg/probo/evidence_service.go index 7af7f4f0a..e0abf646c 100644 --- a/pkg/probo/evidence_service.go +++ b/pkg/probo/evidence_service.go @@ -15,6 +15,7 @@ package probo import ( + "bytes" "context" "fmt" "io" @@ -26,6 +27,7 @@ import ( "github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/service/s3" "github.com/getprobo/probo/pkg/coredata" + "github.com/getprobo/probo/pkg/filevalidation" "github.com/getprobo/probo/pkg/gid" "github.com/getprobo/probo/pkg/page" "go.gearno.de/crypto/uuid" @@ -34,7 +36,8 @@ import ( type ( EvidenceService struct { - svc *TenantService + svc *TenantService + fileValidator *filevalidation.FileValidator } File struct { @@ -153,6 +156,30 @@ func (s EvidenceService) Fulfill( if req.File != nil { evidence.Type = coredata.EvidenceTypeFile + var fileSize int64 + var fileContent io.ReadSeeker + + if seeker, ok := req.File.(io.Seeker); ok { + size, err := seeker.Seek(0, io.SeekEnd) + if err != nil { + return fmt.Errorf("cannot determine file size: %w", err) + } + + _, err = seeker.Seek(0, io.SeekStart) + if err != nil { + return fmt.Errorf("cannot reset file position: %w", err) + } + fileSize = size + fileContent = req.File.(io.ReadSeeker) + } else { + buf, err := io.ReadAll(req.File) + if err != nil { + return fmt.Errorf("cannot read file: %w", err) + } + fileSize = int64(len(buf)) + fileContent = bytes.NewReader(buf) + } + contentType := "application/octet-stream" if req.Filename != nil { evidence.Filename = *req.Filename @@ -161,6 +188,10 @@ func (s EvidenceService) Fulfill( } } + if err := s.fileValidator.Validate(evidence.Filename, contentType, fileSize); err != nil { + return err + } + objectKey, err := uuid.NewV7() if err != nil { return fmt.Errorf("cannot generate object key: %w", err) @@ -169,7 +200,7 @@ func (s EvidenceService) Fulfill( _, err = s.svc.s3.PutObject(ctx, &s3.PutObjectInput{ Bucket: aws.String(s.svc.bucket), Key: aws.String(objectKey.String()), - Body: req.File, + Body: fileContent, ContentType: aws.String(contentType), }) if err != nil { @@ -228,11 +259,14 @@ func (s EvidenceService) UploadTaskEvidence( UpdatedAt: now, } - // TODO validate content type if req.File.ContentType == "" { req.File.ContentType = "application/octet-stream" } + if err := s.fileValidator.Validate(req.File.Filename, req.File.ContentType, req.File.Size); err != nil { + return nil, err + } + objectKey, err := uuid.NewV7() if err != nil { return nil, fmt.Errorf("cannot generate object key: %w", err) @@ -305,11 +339,14 @@ func (s EvidenceService) UploadMeasureEvidence( UpdatedAt: now, } - // TODO validate content type if req.File.ContentType == "" { req.File.ContentType = "application/octet-stream" } + if err := s.fileValidator.Validate(req.File.Filename, req.File.ContentType, req.File.Size); err != nil { + return nil, err + } + objectKey, err := uuid.NewV7() if err != nil { return nil, fmt.Errorf("cannot generate object key: %w", err) diff --git a/pkg/probo/organization_service.go b/pkg/probo/organization_service.go index 6ee00d69d..03a9aa633 100644 --- a/pkg/probo/organization_service.go +++ b/pkg/probo/organization_service.go @@ -15,15 +15,19 @@ package probo import ( + "bytes" "context" "fmt" "io" + "mime" "net/url" + "path/filepath" "time" "github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/service/s3" "github.com/getprobo/probo/pkg/coredata" + "github.com/getprobo/probo/pkg/filevalidation" "github.com/getprobo/probo/pkg/gid" "go.gearno.de/crypto/uuid" "go.gearno.de/kit/pg" @@ -31,7 +35,8 @@ import ( type ( OrganizationService struct { - svc *TenantService + svc *TenantService + fileValidator *filevalidation.FileValidator } CreateOrganizationRequest struct { @@ -39,9 +44,12 @@ type ( } UpdateOrganizationRequest struct { - ID gid.GID - Name *string - File io.Reader + ID gid.GID + Name *string + File io.Reader + Filename string + FileSize int64 + ContentType string } ) @@ -128,10 +136,54 @@ func (s OrganizationService) Update( return fmt.Errorf("cannot generate object key: %w", err) } + var fileSize int64 + var fileContent io.ReadSeeker + filename := req.Filename + contentType := req.ContentType + + if seeker, ok := req.File.(io.Seeker); ok { + if req.FileSize <= 0 { + size, err := seeker.Seek(0, io.SeekEnd) + if err != nil { + return fmt.Errorf("cannot determine file size: %w", err) + } + fileSize = size + + _, err = seeker.Seek(0, io.SeekStart) + if err != nil { + return fmt.Errorf("cannot reset file position: %w", err) + } + } else { + fileSize = req.FileSize + } + fileContent = req.File.(io.ReadSeeker) + } else { + buf, err := io.ReadAll(req.File) + if err != nil { + return fmt.Errorf("cannot read file: %w", err) + } + fileSize = int64(len(buf)) + fileContent = bytes.NewReader(buf) + } + + if contentType == "" { + contentType = "application/octet-stream" + if filename != "" { + if detectedType := mime.TypeByExtension(filepath.Ext(filename)); detectedType != "" { + contentType = detectedType + } + } + } + + if err := s.fileValidator.Validate(filename, contentType, fileSize); err != nil { + return err + } + _, err = s.svc.s3.PutObject(ctx, &s3.PutObjectInput{ - Bucket: aws.String(s.svc.bucket), - Key: aws.String(objectKey.String()), - Body: req.File, + Bucket: aws.String(s.svc.bucket), + Key: aws.String(objectKey.String()), + Body: fileContent, + ContentType: aws.String(contentType), }) if err != nil { diff --git a/pkg/probo/service.go b/pkg/probo/service.go index 4ea174db6..cb3b6dac8 100644 --- a/pkg/probo/service.go +++ b/pkg/probo/service.go @@ -21,6 +21,7 @@ import ( "github.com/aws/aws-sdk-go-v2/service/s3" "github.com/getprobo/probo/pkg/coredata" "github.com/getprobo/probo/pkg/crypto/cipher" + "github.com/getprobo/probo/pkg/filevalidation" "github.com/getprobo/probo/pkg/gid" "go.gearno.de/kit/pg" ) @@ -97,11 +98,25 @@ func (s *Service) WithTenant(tenantID gid.TenantID) *TenantService { tenantService.Frameworks = &FrameworkService{svc: tenantService} tenantService.Measures = &MeasureService{svc: tenantService} tenantService.Tasks = &TaskService{svc: tenantService} - tenantService.Evidences = &EvidenceService{svc: tenantService} + tenantService.Evidences = &EvidenceService{ + svc: tenantService, + fileValidator: filevalidation.NewValidator( + filevalidation.CategoryDocument, + filevalidation.CategorySpreadsheet, + filevalidation.CategoryPresentation, + filevalidation.CategoryImage, + filevalidation.CategoryVideo, + ), + } tenantService.Peoples = &PeopleService{svc: tenantService} tenantService.Vendors = &VendorService{svc: tenantService} tenantService.Policies = &PolicyService{svc: tenantService} - tenantService.Organizations = &OrganizationService{svc: tenantService} + tenantService.Organizations = &OrganizationService{ + svc: tenantService, + fileValidator: filevalidation.NewValidator( + filevalidation.CategoryImage, + ), + } tenantService.Controls = &ControlService{svc: tenantService} tenantService.Risks = &RiskService{svc: tenantService} tenantService.VendorComplianceReports = &VendorComplianceReportService{svc: tenantService}