Jay Brown
989 lines
30 KiB
Go
989 lines
30 KiB
Go
package api
|
|
|
|
import (
|
|
"archive/zip"
|
|
"bytes"
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"fmt"
|
|
"io"
|
|
"log/slog"
|
|
"os"
|
|
"testing"
|
|
|
|
"queryorchestration/internal/database/repository"
|
|
"queryorchestration/internal/document/batch"
|
|
"queryorchestration/internal/serviceconfig"
|
|
"queryorchestration/internal/serviceconfig/objectstore"
|
|
"queryorchestration/internal/test"
|
|
|
|
"github.com/aws/aws-sdk-go-v2/aws"
|
|
"github.com/aws/aws-sdk-go-v2/service/s3"
|
|
"github.com/google/uuid"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
// TestProcessBatchWork tests the main batch processing work function
|
|
func TestProcessBatchWork(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 with root folder
|
|
clientID := "test_batch_worker"
|
|
test.CreateTestClient(t, cfg, clientID, "Test Batch Worker Client")
|
|
|
|
// Create a test batch service
|
|
batchService := batch.New(cfg)
|
|
|
|
// Upload a test ZIP file to S3
|
|
zipContent := createTestZIPForWorker(t, 2)
|
|
s3ClientInterface := cfg.GetStoreClient()
|
|
s3Client, ok := s3ClientInterface.(*s3.Client)
|
|
require.True(t, ok, "s3Client should be *s3.Client")
|
|
bucket := cfg.GetBucket()
|
|
archiveKey := fmt.Sprintf("test/%s/batch.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 a batch record in the database
|
|
batchID, err := batchService.CreateWithStorage(ctx, clientID, "batch.zip", 2, 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 worker config
|
|
workerConfig := map[string]any{
|
|
ConfigKeyBatchService: batchService,
|
|
ConfigKeyS3Client: s3Client,
|
|
ConfigKeyBucket: bucket,
|
|
ConfigKeyUploadHandler: uploadHandler,
|
|
}
|
|
|
|
// Execute the batch work
|
|
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
|
err = processBatchWork(ctx, logger, workerConfig)
|
|
require.NoError(t, err)
|
|
|
|
// Verify that documents were uploaded
|
|
assert.Len(t, uploadedDocs, 2)
|
|
assert.Contains(t, uploadedDocs, "document_1.pdf")
|
|
assert.Contains(t, uploadedDocs, "document_2.pdf")
|
|
|
|
// Verify batch status was updated
|
|
batchInfo, err := batchService.Get(ctx, clientID, batchID)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "completed", batchInfo.Status)
|
|
assert.Equal(t, int32(2), batchInfo.ProcessedDocuments)
|
|
assert.Equal(t, int32(0), batchInfo.FailedDocuments)
|
|
}
|
|
|
|
// TestProcessSingleBatch tests processing a single batch
|
|
func TestProcessSingleBatch(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 with root folder
|
|
clientID := "test_single_batch"
|
|
test.CreateTestClient(t, cfg, clientID, "Test Single Batch Client")
|
|
|
|
// 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 mixed content
|
|
zipContent := createMixedContentZIP(t)
|
|
archiveKey := fmt.Sprintf("test/%s/mixed.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, "mixed.zip", 3, bucket, archiveKey, int64(len(zipContent)))
|
|
require.NoError(t, err)
|
|
|
|
// Mock upload handler
|
|
uploadCount := 0
|
|
uploadHandler := func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error {
|
|
uploadCount++
|
|
return nil
|
|
}
|
|
|
|
// Create batch details
|
|
batchDetails := &batch.BatchUploadDetails{
|
|
BatchUploadSummary: batch.BatchUploadSummary{
|
|
ID: batchID,
|
|
ClientID: clientID,
|
|
OriginalFilename: "mixed.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 only PDFs were uploaded
|
|
assert.Equal(t, 1, uploadCount) // Only 1 PDF file
|
|
|
|
// Verify batch status
|
|
batchInfo, err := batchService.Get(ctx, clientID, batchID)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int32(1), batchInfo.ProcessedDocuments)
|
|
assert.Equal(t, int32(2), batchInfo.InvalidTypeDocuments) // 1 txt + 1 docx
|
|
}
|
|
|
|
// TestProcessSingleBatchWithFailure tests handling of upload failures
|
|
func TestProcessSingleBatchWithFailure(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 with root folder
|
|
clientID := "test_batch_failure"
|
|
test.CreateTestClient(t, cfg, clientID, "Test Batch Failure Client")
|
|
|
|
// 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
|
|
zipContent := createTestZIPForWorker(t, 3)
|
|
archiveKey := fmt.Sprintf("test/%s/fail.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, "fail.zip", 3, bucket, archiveKey, int64(len(zipContent)))
|
|
require.NoError(t, err)
|
|
|
|
// Upload handler that fails on second document
|
|
uploadCount := 0
|
|
uploadHandler := func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error {
|
|
uploadCount++
|
|
if uploadCount == 2 {
|
|
return fmt.Errorf("simulated upload failure")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Create batch details
|
|
batchDetails := &batch.BatchUploadDetails{
|
|
BatchUploadSummary: batch.BatchUploadSummary{
|
|
ID: batchID,
|
|
ClientID: clientID,
|
|
OriginalFilename: "fail.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 results
|
|
batchInfo, err := batchService.Get(ctx, clientID, batchID)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int32(2), batchInfo.ProcessedDocuments) // 1st and 3rd succeeded
|
|
assert.Equal(t, int32(1), batchInfo.FailedDocuments) // 2nd failed
|
|
assert.Equal(t, "completed", batchInfo.Status) // Status is still completed since some succeeded
|
|
assert.Contains(t, batchInfo.FailedFilenames, "document_2.pdf")
|
|
}
|
|
|
|
// TestDownloadZipFromS3 tests downloading a ZIP file from S3
|
|
func TestDownloadZipFromS3(t *testing.T) {
|
|
if testing.Short() {
|
|
t.SkipNow()
|
|
}
|
|
|
|
ctx := t.Context()
|
|
cfg := &BaseConfig{}
|
|
_ = serviceconfig.InitializeConfig(cfg)
|
|
test.CreateAWSResources(t, cfg)
|
|
|
|
s3ClientInterface := cfg.GetStoreClient()
|
|
s3Client, ok := s3ClientInterface.(*s3.Client)
|
|
require.True(t, ok, "s3Client should be *s3.Client")
|
|
bucket := cfg.GetBucket()
|
|
|
|
// Upload test data
|
|
testData := []byte("test zip content")
|
|
key := "test/download.zip"
|
|
|
|
_, err := s3Client.PutObject(ctx, &s3.PutObjectInput{
|
|
Bucket: aws.String(bucket),
|
|
Key: aws.String(key),
|
|
Body: bytes.NewReader(testData),
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Download the data
|
|
downloaded, err := downloadZipFromS3(ctx, s3Client, bucket, key)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, testData, downloaded)
|
|
}
|
|
|
|
// TestUploadToS3 tests uploading data to S3
|
|
func TestUploadToS3(t *testing.T) {
|
|
if testing.Short() {
|
|
t.SkipNow()
|
|
}
|
|
|
|
ctx := t.Context()
|
|
cfg := &BaseConfig{}
|
|
_ = serviceconfig.InitializeConfig(cfg)
|
|
test.CreateAWSResources(t, cfg)
|
|
|
|
s3ClientInterface := cfg.GetStoreClient()
|
|
s3Client, ok := s3ClientInterface.(*s3.Client)
|
|
require.True(t, ok, "s3Client should be *s3.Client")
|
|
bucket := cfg.GetBucket()
|
|
|
|
// Upload test data
|
|
testData := []byte("test upload content")
|
|
key := "test/upload.txt"
|
|
|
|
err := uploadToS3(ctx, s3Client, bucket, key, testData)
|
|
require.NoError(t, err)
|
|
|
|
// Verify the upload
|
|
result, err := s3Client.GetObject(ctx, &s3.GetObjectInput{
|
|
Bucket: aws.String(bucket),
|
|
Key: aws.String(key),
|
|
})
|
|
require.NoError(t, err)
|
|
defer result.Body.Close()
|
|
|
|
downloaded, err := io.ReadAll(result.Body)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, testData, downloaded)
|
|
}
|
|
|
|
// Helper function to create a test ZIP file with PDFs
|
|
func createTestZIPForWorker(t *testing.T, numFiles int) []byte {
|
|
var buf bytes.Buffer
|
|
zipWriter := zip.NewWriter(&buf)
|
|
|
|
for i := 1; i <= numFiles; i++ {
|
|
filename := fmt.Sprintf("document_%d.pdf", i)
|
|
fileWriter, err := zipWriter.Create(filename)
|
|
require.NoError(t, err)
|
|
|
|
pdfContent := fmt.Sprintf("%%PDF-1.4\n%%Test content for %s\n%%%%EOF", filename)
|
|
_, err = fileWriter.Write([]byte(pdfContent))
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
err := zipWriter.Close()
|
|
require.NoError(t, err)
|
|
|
|
return buf.Bytes()
|
|
}
|
|
|
|
// Helper function to create a ZIP with mixed content types
|
|
func createMixedContentZIP(t *testing.T) []byte {
|
|
var buf bytes.Buffer
|
|
zipWriter := zip.NewWriter(&buf)
|
|
|
|
// Add a PDF
|
|
pdfWriter, err := zipWriter.Create("document.pdf")
|
|
require.NoError(t, err)
|
|
_, err = pdfWriter.Write([]byte("%%PDF-1.4\n%%Test PDF\n%%%%EOF"))
|
|
require.NoError(t, err)
|
|
|
|
// Add a TXT file
|
|
txtWriter, err := zipWriter.Create("notes.txt")
|
|
require.NoError(t, err)
|
|
_, err = txtWriter.Write([]byte("This is a text file"))
|
|
require.NoError(t, err)
|
|
|
|
// Add a DOCX file
|
|
docxWriter, err := zipWriter.Create("report.docx")
|
|
require.NoError(t, err)
|
|
_, err = docxWriter.Write([]byte("PK")) // Simplified DOCX header
|
|
require.NoError(t, err)
|
|
|
|
err = zipWriter.Close()
|
|
require.NoError(t, err)
|
|
|
|
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 with root folder
|
|
clientID := "test_nested_folders"
|
|
test.CreateTestClient(t, cfg, clientID, "Test Nested Folders Client")
|
|
|
|
// 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 with root folder
|
|
clientID := "test_path_traversal"
|
|
test.CreateTestClient(t, cfg, clientID, "Test Path Traversal Client")
|
|
|
|
// 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")
|
|
}
|
|
|
|
// TestProcessZipFile_RecordsOutcomes tests that processZipFile records per-file
|
|
// outcomes for each extracted file.
|
|
func TestProcessZipFile_RecordsOutcomes(t *testing.T) {
|
|
if testing.Short() {
|
|
t.SkipNow()
|
|
}
|
|
|
|
ctx := t.Context()
|
|
cfg := &BaseConfig{}
|
|
_ = serviceconfig.InitializeConfig(cfg)
|
|
test.CreateDB(t, cfg)
|
|
test.CreateAWSResources(t, cfg)
|
|
|
|
clientID := "test_outcomes_record"
|
|
test.CreateTestClient(t, cfg, clientID, "Test Outcomes Record Client")
|
|
|
|
batchService := batch.New(cfg)
|
|
s3ClientInterface := cfg.GetStoreClient()
|
|
s3Client, ok := s3ClientInterface.(*s3.Client)
|
|
require.True(t, ok)
|
|
bucket := cfg.GetBucket()
|
|
|
|
// Create ZIP with 2 PDFs
|
|
zipContent := createTestZIPForWorker(t, 2)
|
|
archiveKey := fmt.Sprintf("test/%s/outcomes.zip", clientID)
|
|
|
|
_, err := s3Client.PutObject(ctx, &s3.PutObjectInput{
|
|
Bucket: aws.String(bucket),
|
|
Key: aws.String(archiveKey),
|
|
Body: bytes.NewReader(zipContent),
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
batchID, err := batchService.CreateWithStorage(ctx, clientID, "outcomes.zip", 2, bucket, archiveKey, int64(len(zipContent)))
|
|
require.NoError(t, err)
|
|
|
|
uploadHandler := func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error {
|
|
_, _ = io.ReadAll(docData)
|
|
return nil
|
|
}
|
|
|
|
workerConfig := map[string]any{
|
|
ConfigKeyBatchService: batchService,
|
|
ConfigKeyS3Client: s3Client,
|
|
ConfigKeyBucket: bucket,
|
|
ConfigKeyUploadHandler: uploadHandler,
|
|
}
|
|
|
|
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
|
err = processBatchWork(ctx, logger, workerConfig)
|
|
require.NoError(t, err)
|
|
|
|
// Verify outcomes were recorded for each file
|
|
outcomes, err := batchService.ListOutcomes(ctx, batchID)
|
|
require.NoError(t, err)
|
|
assert.Len(t, outcomes, 2)
|
|
|
|
for _, o := range outcomes {
|
|
assert.Equal(t, "submitted", o.Outcome, "each PDF should be recorded as submitted")
|
|
assert.Contains(t, o.Filename, ".pdf")
|
|
}
|
|
}
|
|
|
|
// TestProcessZipFile_DetectsDuplicates verifies that when a file in the batch
|
|
// has the same content hash as an existing document in the system, it is
|
|
// recorded as a duplicate with the existing document's ID.
|
|
func TestProcessZipFile_DetectsDuplicates(t *testing.T) {
|
|
if testing.Short() {
|
|
t.SkipNow()
|
|
}
|
|
|
|
ctx := t.Context()
|
|
cfg := &BaseConfig{}
|
|
_ = serviceconfig.InitializeConfig(cfg)
|
|
test.CreateDB(t, cfg)
|
|
test.CreateAWSResources(t, cfg)
|
|
|
|
clientID := "test_outcomes_dup"
|
|
test.CreateTestClient(t, cfg, clientID, "Test Outcomes Dup Client")
|
|
|
|
batchService := batch.New(cfg)
|
|
s3ClientInterface := cfg.GetStoreClient()
|
|
s3Client, ok := s3ClientInterface.(*s3.Client)
|
|
require.True(t, ok)
|
|
bucket := cfg.GetBucket()
|
|
|
|
// Pre-create an existing document with a known hash
|
|
knownContent := []byte("%%PDF-1.4\n%%Already existing content\n%%%%EOF")
|
|
knownHash := sha256Hash(knownContent)
|
|
existingDocID, err := cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{
|
|
Clientid: clientID,
|
|
Hash: knownHash,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create a ZIP with one file that has the same content as the existing doc
|
|
var buf bytes.Buffer
|
|
zipWriter := zip.NewWriter(&buf)
|
|
w, err := zipWriter.Create("duplicate_of_existing.pdf")
|
|
require.NoError(t, err)
|
|
_, err = w.Write(knownContent)
|
|
require.NoError(t, err)
|
|
require.NoError(t, zipWriter.Close())
|
|
zipContent := buf.Bytes()
|
|
|
|
archiveKey := fmt.Sprintf("test/%s/dup.zip", clientID)
|
|
_, err = s3Client.PutObject(ctx, &s3.PutObjectInput{
|
|
Bucket: aws.String(bucket),
|
|
Key: aws.String(archiveKey),
|
|
Body: bytes.NewReader(zipContent),
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
batchID, err := batchService.CreateWithStorage(ctx, clientID, "dup.zip", 1, bucket, archiveKey, int64(len(zipContent)))
|
|
require.NoError(t, err)
|
|
|
|
uploadHandler := func(ctx context.Context, cID string, docData io.Reader, filename string, bID *uuid.UUID) error {
|
|
_, _ = io.ReadAll(docData)
|
|
return nil
|
|
}
|
|
|
|
workerConfig := map[string]any{
|
|
ConfigKeyBatchService: batchService,
|
|
ConfigKeyS3Client: s3Client,
|
|
ConfigKeyBucket: bucket,
|
|
ConfigKeyUploadHandler: uploadHandler,
|
|
}
|
|
|
|
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
|
err = processBatchWork(ctx, logger, workerConfig)
|
|
require.NoError(t, err)
|
|
|
|
// Verify outcome is "duplicate" with the existing document's ID
|
|
outcomes, err := batchService.ListOutcomes(ctx, batchID)
|
|
require.NoError(t, err)
|
|
require.Len(t, outcomes, 1)
|
|
|
|
assert.Equal(t, "duplicate", outcomes[0].Outcome)
|
|
require.NotNil(t, outcomes[0].DocumentID)
|
|
assert.Equal(t, existingDocID, *outcomes[0].DocumentID)
|
|
}
|
|
|
|
// TestProcessZipFile_MixedOutcomes tests a ZIP with PDF + non-PDF + paths to
|
|
// verify all outcome types (submitted, invalid_type) are correctly recorded.
|
|
func TestProcessZipFile_MixedOutcomes(t *testing.T) {
|
|
if testing.Short() {
|
|
t.SkipNow()
|
|
}
|
|
|
|
ctx := t.Context()
|
|
cfg := &BaseConfig{}
|
|
_ = serviceconfig.InitializeConfig(cfg)
|
|
test.CreateDB(t, cfg)
|
|
test.CreateAWSResources(t, cfg)
|
|
|
|
clientID := "test_outcomes_mixed"
|
|
test.CreateTestClient(t, cfg, clientID, "Test Outcomes Mixed Client")
|
|
|
|
batchService := batch.New(cfg)
|
|
s3ClientInterface := cfg.GetStoreClient()
|
|
s3Client, ok := s3ClientInterface.(*s3.Client)
|
|
require.True(t, ok)
|
|
bucket := cfg.GetBucket()
|
|
|
|
// Create mixed ZIP
|
|
zipContent := createMixedContentZIP(t)
|
|
archiveKey := fmt.Sprintf("test/%s/mixed_outcomes.zip", clientID)
|
|
|
|
_, err := s3Client.PutObject(ctx, &s3.PutObjectInput{
|
|
Bucket: aws.String(bucket),
|
|
Key: aws.String(archiveKey),
|
|
Body: bytes.NewReader(zipContent),
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
batchID, err := batchService.CreateWithStorage(ctx, clientID, "mixed_outcomes.zip", 3, bucket, archiveKey, int64(len(zipContent)))
|
|
require.NoError(t, err)
|
|
|
|
uploadHandler := func(ctx context.Context, cID string, docData io.Reader, filename string, bID *uuid.UUID) error {
|
|
_, _ = io.ReadAll(docData)
|
|
return nil
|
|
}
|
|
|
|
workerConfig := map[string]any{
|
|
ConfigKeyBatchService: batchService,
|
|
ConfigKeyS3Client: s3Client,
|
|
ConfigKeyBucket: bucket,
|
|
ConfigKeyUploadHandler: uploadHandler,
|
|
}
|
|
|
|
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
|
err = processBatchWork(ctx, logger, workerConfig)
|
|
require.NoError(t, err)
|
|
|
|
// Verify outcomes
|
|
outcomes, err := batchService.ListOutcomes(ctx, batchID)
|
|
require.NoError(t, err)
|
|
assert.Len(t, outcomes, 3) // 1 PDF + 1 txt + 1 docx
|
|
|
|
outcomeMap := map[string]string{}
|
|
for _, o := range outcomes {
|
|
outcomeMap[o.Filename] = o.Outcome
|
|
}
|
|
|
|
assert.Equal(t, "submitted", outcomeMap["document.pdf"])
|
|
assert.Equal(t, "invalid_type", outcomeMap["notes.txt"])
|
|
assert.Equal(t, "invalid_type", outcomeMap["report.docx"])
|
|
}
|
|
|
|
// sha256Hash computes the hex-encoded SHA-256 hash of data.
|
|
func sha256Hash(data []byte) string {
|
|
h := sha256.Sum256(data)
|
|
return hex.EncodeToString(h[:])
|
|
}
|
|
|
|
// 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 with root folder
|
|
clientID := "test_duplicate_names"
|
|
test.CreateTestClient(t, cfg, clientID, "Test Duplicate Names Client")
|
|
|
|
// 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")
|
|
}
|