Merged in feature/preserve-folder-paths (pull request #179)

support folder paths * support folder paths * s3 storage docs
2025-09-05 18:25:31 +00:00
parent 0ad41de26b
commit 730f3ba11a
27 changed files with 8281 additions and 63 deletions
@@ -23,12 +23,6 @@ import (
 	"github.com/stretchr/testify/require"
 )

-// TestConfig embeds objectstore for S3 operations
-type TestConfig struct {
-	serviceconfig.BaseConfig
-	objectstore.ObjectStoreConfig
-}
-
 // TestProcessBatchWork tests the main batch processing work function
 func TestProcessBatchWork(t *testing.T) {
 	if testing.Short() {
@@ -36,7 +30,7 @@ func TestProcessBatchWork(t *testing.T) {
 	}

 	ctx := t.Context()
-	cfg := &TestConfig{}
+	cfg := &BaseConfig{}
 	_ = serviceconfig.InitializeConfig(cfg)
 	test.CreateDB(t, cfg)
 	test.CreateAWSResources(t, cfg)
@@ -115,7 +109,7 @@ func TestProcessSingleBatch(t *testing.T) {
 	}

 	ctx := t.Context()
-	cfg := &TestConfig{}
+	cfg := &BaseConfig{}
 	_ = serviceconfig.InitializeConfig(cfg)
 	test.CreateDB(t, cfg)
 	test.CreateAWSResources(t, cfg)
@@ -191,7 +185,7 @@ func TestProcessSingleBatchWithFailure(t *testing.T) {
 	}

 	ctx := t.Context()
-	cfg := &TestConfig{}
+	cfg := &BaseConfig{}
 	_ = serviceconfig.InitializeConfig(cfg)
 	test.CreateDB(t, cfg)
 	test.CreateAWSResources(t, cfg)
@@ -269,7 +263,7 @@ func TestDownloadZipFromS3(t *testing.T) {
 	}

 	ctx := t.Context()
-	cfg := &TestConfig{}
+	cfg := &BaseConfig{}
 	_ = serviceconfig.InitializeConfig(cfg)
 	test.CreateAWSResources(t, cfg)

@@ -302,7 +296,7 @@ func TestUploadToS3(t *testing.T) {
 	}

 	ctx := t.Context()
-	cfg := &TestConfig{}
+	cfg := &BaseConfig{}
 	_ = serviceconfig.InitializeConfig(cfg)
 	test.CreateAWSResources(t, cfg)

@@ -380,3 +374,422 @@ func createMixedContentZIP(t *testing.T) []byte {

 	return buf.Bytes()
 }
+
+// Helper function to create a ZIP with nested folder structure
+func createNestedFolderZIP(t *testing.T) []byte {
+	var buf bytes.Buffer
+	zipWriter := zip.NewWriter(&buf)
+
+	// Create files in various nested folder structures
+	testFiles := []struct {
+		path    string
+		content string
+	}{
+		{"document1.pdf", "%%PDF-1.4\n%%Root level PDF\n%%%%EOF"},
+		{"foldera/file1.pdf", "%%PDF-1.4\n%%FolderA File1\n%%%%EOF"},
+		{"foldera/file2.pdf", "%%PDF-1.4\n%%FolderA File2\n%%%%EOF"},
+		{"folderb/folderc/file3.pdf", "%%PDF-1.4\n%%Nested File3\n%%%%EOF"},
+		{"folderb/folderc/file4.pdf", "%%PDF-1.4\n%%Nested File4\n%%%%EOF"},
+		{"folderb/file5.pdf", "%%PDF-1.4\n%%FolderB File5\n%%%%EOF"},
+		{"deep/nested/structure/test/file6.pdf", "%%PDF-1.4\n%%Deep nested\n%%%%EOF"},
+		// Add a non-PDF to verify it's skipped
+		{"foldera/readme.txt", "This should be skipped"},
+	}
+
+	for _, tf := range testFiles {
+		fileWriter, err := zipWriter.Create(tf.path)
+		require.NoError(t, err)
+		_, err = fileWriter.Write([]byte(tf.content))
+		require.NoError(t, err)
+	}
+
+	err := zipWriter.Close()
+	require.NoError(t, err)
+
+	return buf.Bytes()
+}
+
+// Helper function to create a ZIP with path traversal attempts
+func createMaliciousPathZIP(t *testing.T) []byte {
+	var buf bytes.Buffer
+	zipWriter := zip.NewWriter(&buf)
+
+	// Create files with various path traversal attempts
+	testFiles := []struct {
+		path    string
+		content string
+	}{
+		{"../escape.pdf", "%%PDF-1.4\n%%Escape attempt\n%%%%EOF"},
+		{"../../etc/passwd", "malicious content"},
+		{"./valid.pdf", "%%PDF-1.4\n%%Valid PDF\n%%%%EOF"},
+		{"folder/../sibling.pdf", "%%PDF-1.4\n%%Sibling\n%%%%EOF"},
+		{"normal/file.pdf", "%%PDF-1.4\n%%Normal\n%%%%EOF"},
+	}
+
+	for _, tf := range testFiles {
+		// Note: zip.Writer doesn't validate paths, so we can create these entries
+		header := &zip.FileHeader{
+			Name:   tf.path,
+			Method: zip.Deflate,
+		}
+		fileWriter, err := zipWriter.CreateHeader(header)
+		require.NoError(t, err)
+		_, err = fileWriter.Write([]byte(tf.content))
+		require.NoError(t, err)
+	}
+
+	err := zipWriter.Close()
+	require.NoError(t, err)
+
+	return buf.Bytes()
+}
+
+// TestSanitizeZipPath tests the path sanitization function
+func TestSanitizeZipPath(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    string
+		expected string
+	}{
+		{"simple file", "document.pdf", "document.pdf"},
+		{"file in folder", "folder/document.pdf", "folder/document.pdf"},
+		{"nested folders", "folder/subfolder/document.pdf", "folder/subfolder/document.pdf"},
+		{"deep nesting", "a/b/c/d/e/f/document.pdf", "a/b/c/d/e/f/document.pdf"},
+		{"dot prefix cleaned", "./document.pdf", "document.pdf"},
+		{"dot in path cleaned", "folder/./document.pdf", "folder/document.pdf"},
+		{"parent traversal rejected", "../document.pdf", ""},
+		{"parent in middle cleaned", "folder/../document.pdf", "document.pdf"},
+		{"absolute path rejected", "/etc/passwd", ""},
+		{"backslash rejected", "folder\\document.pdf", ""},
+		{"double parent rejected", "../../document.pdf", ""},
+		{"complex traversal attempt", "folder/../../etc/passwd", ""},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := sanitizeZipPath(tt.input)
+			assert.Equal(t, tt.expected, result, "sanitizeZipPath(%q) = %q, want %q", tt.input, result, tt.expected)
+		})
+	}
+}
+
+// TestProcessBatchWithNestedFolders tests processing a ZIP with nested folder structure
+func TestProcessBatchWithNestedFolders(t *testing.T) {
+	if testing.Short() {
+		t.SkipNow()
+	}
+
+	ctx := t.Context()
+	cfg := &BaseConfig{}
+	_ = serviceconfig.InitializeConfig(cfg)
+	test.CreateDB(t, cfg)
+	test.CreateAWSResources(t, cfg)
+
+	// Create test client
+	clientID := "test_nested_folders"
+	err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
+		Clientid: clientID,
+		Name:     "Test Nested Folders Client",
+	})
+	require.NoError(t, err)
+
+	// Create services
+	batchService := batch.New(cfg)
+	s3ClientInterface := cfg.GetStoreClient()
+	s3Client, ok := s3ClientInterface.(*s3.Client)
+	require.True(t, ok, "s3Client should be *s3.Client")
+	bucket := cfg.GetBucket()
+
+	// Upload test ZIP with nested folders
+	zipContent := createNestedFolderZIP(t)
+	archiveKey := fmt.Sprintf("test/%s/nested.zip", clientID)
+
+	_, err = s3Client.PutObject(ctx, &s3.PutObjectInput{
+		Bucket: aws.String(bucket),
+		Key:    aws.String(archiveKey),
+		Body:   bytes.NewReader(zipContent),
+	})
+	require.NoError(t, err)
+
+	// Create batch
+	batchID, err := batchService.CreateWithStorage(ctx, clientID, "nested.zip", 7, bucket, archiveKey, int64(len(zipContent)))
+	require.NoError(t, err)
+
+	// Use the real upload handler that creates documents in the database
+	uploadHandler := createDocumentUploadHandler(cfg)
+
+	// Create batch details
+	batchDetails := &batch.BatchUploadDetails{
+		BatchUploadSummary: batch.BatchUploadSummary{
+			ID:               batchID,
+			ClientID:         clientID,
+			OriginalFilename: "nested.zip",
+			Status:           "processing",
+		},
+		ArchiveKey:    archiveKey,
+		ArchiveBucket: bucket,
+	}
+
+	// Process the batch
+	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
+	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails)
+	require.NoError(t, err)
+
+	// Verify batch status
+	batchInfo, err := batchService.Get(ctx, clientID, batchID)
+	require.NoError(t, err)
+	assert.Equal(t, int32(7), batchInfo.ProcessedDocuments)   // 7 PDFs
+	assert.Equal(t, int32(1), batchInfo.InvalidTypeDocuments) // 1 txt file
+	assert.Equal(t, "completed", batchInfo.Status)
+
+	// Manually trigger document initialization for uploaded files since S3 notifications
+	// don't work in test environment. Get all document uploads for this batch.
+	uploads, err := cfg.GetDBQueries().ListDocumentUploads(ctx, clientID)
+	require.NoError(t, err)
+
+	// Process each uploaded file through the document initialization pipeline
+	for _, upload := range uploads {
+		if upload.BatchID != nil && *upload.BatchID == batchID {
+			// Parse the bucket key
+			key, err := objectstore.ParseBucketKey(upload.Key)
+			require.NoError(t, err)
+
+			// Get file hash from S3 object metadata (normally provided by S3 event)
+			headResult, err := s3Client.HeadObject(ctx, &s3.HeadObjectInput{
+				Bucket: aws.String(upload.Bucket),
+				Key:    aws.String(upload.Key),
+			})
+			require.NoError(t, err)
+
+			// Create document directly in database since full pipeline requires additional setup
+			_, err = cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{
+				Clientid: key.ClientID,
+				Hash:     *headResult.ETag, // Use ETag as hash for testing
+				BatchID:  upload.BatchID,
+				Filename: upload.Filename,
+			})
+			require.NoError(t, err)
+		}
+	}
+
+	// Get documents for this specific batch from database
+	batchDocs, err := cfg.GetDBQueries().ListDocumentsByBatch(ctx, &batchID)
+	require.NoError(t, err)
+
+	// Verify we have the expected number of documents
+	assert.Equal(t, 7, len(batchDocs), "Should have exactly 7 documents in database")
+
+	// Verify filenames are preserved with their paths
+	foundFilenames := make(map[string]bool)
+	for _, doc := range batchDocs {
+		if doc.Filename != nil {
+			foundFilenames[*doc.Filename] = true
+		}
+	}
+
+	// Check that we found all expected filenames with their paths
+	expectedFilenames := []string{
+		"document1.pdf",
+		"foldera/file1.pdf",
+		"foldera/file2.pdf",
+		"folderb/folderc/file3.pdf",
+		"folderb/folderc/file4.pdf",
+		"folderb/file5.pdf",
+		"deep/nested/structure/test/file6.pdf",
+	}
+
+	for _, expectedFilename := range expectedFilenames {
+		assert.True(t, foundFilenames[expectedFilename],
+			"Should have preserved filename with path: %s", expectedFilename)
+	}
+}
+
+// TestProcessBatchWithPathTraversal tests that path traversal attempts are handled safely
+func TestProcessBatchWithPathTraversal(t *testing.T) {
+	if testing.Short() {
+		t.SkipNow()
+	}
+
+	ctx := t.Context()
+	cfg := &BaseConfig{}
+	_ = serviceconfig.InitializeConfig(cfg)
+	test.CreateDB(t, cfg)
+	test.CreateAWSResources(t, cfg)
+
+	// Create test client
+	clientID := "test_path_traversal"
+	err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
+		Clientid: clientID,
+		Name:     "Test Path Traversal Client",
+	})
+	require.NoError(t, err)
+
+	// Create services
+	batchService := batch.New(cfg)
+	s3ClientInterface := cfg.GetStoreClient()
+	s3Client, ok := s3ClientInterface.(*s3.Client)
+	require.True(t, ok, "s3Client should be *s3.Client")
+	bucket := cfg.GetBucket()
+
+	// Upload test ZIP with path traversal attempts
+	zipContent := createMaliciousPathZIP(t)
+	archiveKey := fmt.Sprintf("test/%s/malicious.zip", clientID)
+
+	_, err = s3Client.PutObject(ctx, &s3.PutObjectInput{
+		Bucket: aws.String(bucket),
+		Key:    aws.String(archiveKey),
+		Body:   bytes.NewReader(zipContent),
+	})
+	require.NoError(t, err)
+
+	// Create batch
+	batchID, err := batchService.CreateWithStorage(ctx, clientID, "malicious.zip", 5, bucket, archiveKey, int64(len(zipContent)))
+	require.NoError(t, err)
+
+	// Track uploaded documents
+	uploadedDocs := make(map[string][]byte)
+	uploadHandler := func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error {
+		data, err := io.ReadAll(docData)
+		if err != nil {
+			return err
+		}
+		uploadedDocs[filename] = data
+		return nil
+	}
+
+	// Create batch details
+	batchDetails := &batch.BatchUploadDetails{
+		BatchUploadSummary: batch.BatchUploadSummary{
+			ID:               batchID,
+			ClientID:         clientID,
+			OriginalFilename: "malicious.zip",
+			Status:           "processing",
+		},
+		ArchiveKey:    archiveKey,
+		ArchiveBucket: bucket,
+	}
+
+	// Process the batch
+	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
+	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails)
+	require.NoError(t, err)
+
+	// Verify that path traversal attempts were blocked
+	assert.NotContains(t, uploadedDocs, "../escape.pdf", "Should reject parent traversal")
+	assert.NotContains(t, uploadedDocs, "../../etc/passwd", "Should reject multiple parent traversal")
+	assert.NotContains(t, uploadedDocs, "/etc/passwd", "Should reject absolute paths")
+
+	// Verify that safe files were processed correctly
+	assert.Contains(t, uploadedDocs, "valid.pdf", "Should process ./valid.pdf as valid.pdf")
+	assert.Contains(t, uploadedDocs, "sibling.pdf", "Should clean folder/../sibling.pdf to sibling.pdf")
+	assert.Contains(t, uploadedDocs, "normal/file.pdf", "Should process normal paths correctly")
+
+	// Verify only safe PDFs were processed
+	assert.Equal(t, 3, len(uploadedDocs), "Should have processed exactly 3 safe PDFs")
+}
+
+// TestProcessBatchWithDuplicateFilenames tests handling of duplicate filenames in different folders
+func TestProcessBatchWithDuplicateFilenames(t *testing.T) {
+	if testing.Short() {
+		t.SkipNow()
+	}
+
+	ctx := t.Context()
+	cfg := &BaseConfig{}
+	_ = serviceconfig.InitializeConfig(cfg)
+	test.CreateDB(t, cfg)
+	test.CreateAWSResources(t, cfg)
+
+	// Create test client
+	clientID := "test_duplicate_names"
+	err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
+		Clientid: clientID,
+		Name:     "Test Duplicate Names Client",
+	})
+	require.NoError(t, err)
+
+	// Create services
+	batchService := batch.New(cfg)
+	s3ClientInterface := cfg.GetStoreClient()
+	s3Client, ok := s3ClientInterface.(*s3.Client)
+	require.True(t, ok, "s3Client should be *s3.Client")
+	bucket := cfg.GetBucket()
+
+	// Create ZIP with duplicate filenames in different folders
+	var buf bytes.Buffer
+	zipWriter := zip.NewWriter(&buf)
+
+	// Create files with same names in different folders
+	testFiles := []struct {
+		path    string
+		content string
+	}{
+		{"document.pdf", "%%PDF-1.4\n%%Root document\n%%%%EOF"},
+		{"folder1/document.pdf", "%%PDF-1.4\n%%Folder1 document\n%%%%EOF"},
+		{"folder2/document.pdf", "%%PDF-1.4\n%%Folder2 document\n%%%%EOF"},
+		{"folder1/subfolder/document.pdf", "%%PDF-1.4\n%%Subfolder document\n%%%%EOF"},
+	}
+
+	for _, tf := range testFiles {
+		fileWriter, err := zipWriter.Create(tf.path)
+		require.NoError(t, err)
+		_, err = fileWriter.Write([]byte(tf.content))
+		require.NoError(t, err)
+	}
+
+	err = zipWriter.Close()
+	require.NoError(t, err)
+	zipContent := buf.Bytes()
+
+	// Upload the ZIP
+	archiveKey := fmt.Sprintf("test/%s/duplicates.zip", clientID)
+	_, err = s3Client.PutObject(ctx, &s3.PutObjectInput{
+		Bucket: aws.String(bucket),
+		Key:    aws.String(archiveKey),
+		Body:   bytes.NewReader(zipContent),
+	})
+	require.NoError(t, err)
+
+	// Create batch
+	batchID, err := batchService.CreateWithStorage(ctx, clientID, "duplicates.zip", 4, bucket, archiveKey, int64(len(zipContent)))
+	require.NoError(t, err)
+
+	// Track uploaded documents
+	uploadedDocs := make(map[string]string)
+	uploadHandler := func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error {
+		data, err := io.ReadAll(docData)
+		if err != nil {
+			return err
+		}
+		// Extract the content identifier from the PDF
+		content := string(data)
+		uploadedDocs[filename] = content
+		return nil
+	}
+
+	// Create batch details
+	batchDetails := &batch.BatchUploadDetails{
+		BatchUploadSummary: batch.BatchUploadSummary{
+			ID:               batchID,
+			ClientID:         clientID,
+			OriginalFilename: "duplicates.zip",
+			Status:           "processing",
+		},
+		ArchiveKey:    archiveKey,
+		ArchiveBucket: bucket,
+	}
+
+	// Process the batch
+	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
+	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails)
+	require.NoError(t, err)
+
+	// Verify all files were uploaded with correct paths
+	assert.Equal(t, 4, len(uploadedDocs), "Should have uploaded all 4 PDFs")
+
+	// Verify each file has unique path and correct content
+	assert.Contains(t, uploadedDocs["document.pdf"], "Root document")
+	assert.Contains(t, uploadedDocs["folder1/document.pdf"], "Folder1 document")
+	assert.Contains(t, uploadedDocs["folder2/document.pdf"], "Folder2 document")
+	assert.Contains(t, uploadedDocs["folder1/subfolder/document.pdf"], "Subfolder document")
+}