feat(security): improve security features, fixes

This commit is contained in:
2025-07-19 01:37:52 +03:00
parent e35126856d
commit b369d317b1
19 changed files with 2266 additions and 92 deletions

View File

@@ -2,11 +2,15 @@
package fileproc
import (
"context"
"fmt"
"io"
"os"
"path/filepath"
"strings"
"time"
"github.com/sirupsen/logrus"
"github.com/ivuorinen/gibidify/config"
"github.com/ivuorinen/gibidify/utils"
@@ -31,15 +35,26 @@ type WriteRequest struct {
// FileProcessor handles file processing operations.
type FileProcessor struct {
rootPath string
sizeLimit int64
rootPath string
sizeLimit int64
resourceMonitor *ResourceMonitor
}
// NewFileProcessor creates a new file processor.
func NewFileProcessor(rootPath string) *FileProcessor {
return &FileProcessor{
rootPath: rootPath,
sizeLimit: config.GetFileSizeLimit(),
rootPath: rootPath,
sizeLimit: config.GetFileSizeLimit(),
resourceMonitor: NewResourceMonitor(),
}
}
// NewFileProcessorWithMonitor creates a new file processor with a shared resource monitor.
func NewFileProcessorWithMonitor(rootPath string, monitor *ResourceMonitor) *FileProcessor {
return &FileProcessor{
rootPath: rootPath,
sizeLimit: config.GetFileSizeLimit(),
resourceMonitor: monitor,
}
}
@@ -47,30 +62,92 @@ func NewFileProcessor(rootPath string) *FileProcessor {
// It automatically chooses between loading the entire file or streaming based on file size.
func ProcessFile(filePath string, outCh chan<- WriteRequest, rootPath string) {
processor := NewFileProcessor(rootPath)
processor.Process(filePath, outCh)
ctx := context.Background()
processor.ProcessWithContext(ctx, filePath, outCh)
}
// ProcessFileWithMonitor processes a file using a shared resource monitor.
func ProcessFileWithMonitor(ctx context.Context, filePath string, outCh chan<- WriteRequest, rootPath string, monitor *ResourceMonitor) {
processor := NewFileProcessorWithMonitor(rootPath, monitor)
processor.ProcessWithContext(ctx, filePath, outCh)
}
// Process handles file processing with the configured settings.
func (p *FileProcessor) Process(filePath string, outCh chan<- WriteRequest) {
// Validate file
fileInfo, err := p.validateFile(filePath)
ctx := context.Background()
p.ProcessWithContext(ctx, filePath, outCh)
}
// ProcessWithContext handles file processing with context and resource monitoring.
func (p *FileProcessor) ProcessWithContext(ctx context.Context, filePath string, outCh chan<- WriteRequest) {
// Create file processing context with timeout
fileCtx, fileCancel := p.resourceMonitor.CreateFileProcessingContext(ctx)
defer fileCancel()
// Wait for rate limiting
if err := p.resourceMonitor.WaitForRateLimit(fileCtx); err != nil {
if err == context.DeadlineExceeded {
utils.LogErrorf(
utils.NewStructuredError(utils.ErrorTypeValidation, utils.CodeResourceLimitTimeout, "file processing timeout during rate limiting", filePath, nil),
"File processing timeout during rate limiting: %s", filePath,
)
}
return
}
// Validate file and check resource limits
fileInfo, err := p.validateFileWithLimits(fileCtx, filePath)
if err != nil {
return // Error already logged
}
// Acquire read slot for concurrent processing
if err := p.resourceMonitor.AcquireReadSlot(fileCtx); err != nil {
if err == context.DeadlineExceeded {
utils.LogErrorf(
utils.NewStructuredError(utils.ErrorTypeValidation, utils.CodeResourceLimitTimeout, "file processing timeout waiting for read slot", filePath, nil),
"File processing timeout waiting for read slot: %s", filePath,
)
}
return
}
defer p.resourceMonitor.ReleaseReadSlot()
// Check hard memory limits before processing
if err := p.resourceMonitor.CheckHardMemoryLimit(); err != nil {
utils.LogErrorf(err, "Hard memory limit check failed for file: %s", filePath)
return
}
// Get relative path
relPath := p.getRelativePath(filePath)
// Process file with timeout
processStart := time.Now()
defer func() {
// Record successful processing
p.resourceMonitor.RecordFileProcessed(fileInfo.Size())
logrus.Debugf("File processed in %v: %s", time.Since(processStart), filePath)
}()
// Choose processing strategy based on file size
if fileInfo.Size() <= StreamThreshold {
p.processInMemory(filePath, relPath, outCh)
p.processInMemoryWithContext(fileCtx, filePath, relPath, outCh)
} else {
p.processStreaming(filePath, relPath, outCh)
p.processStreamingWithContext(fileCtx, filePath, relPath, outCh)
}
}
// validateFile checks if the file can be processed.
func (p *FileProcessor) validateFile(filePath string) (os.FileInfo, error) {
// validateFileWithLimits checks if the file can be processed with resource limits.
func (p *FileProcessor) validateFileWithLimits(ctx context.Context, filePath string) (os.FileInfo, error) {
// Check context cancellation
select {
case <-ctx.Done():
return nil, ctx.Err()
default:
}
fileInfo, err := os.Stat(filePath)
if err != nil {
structErr := utils.WrapError(err, utils.ErrorTypeFileSystem, utils.CodeFSAccess, "failed to stat file").WithFilePath(filePath)
@@ -78,19 +155,31 @@ func (p *FileProcessor) validateFile(filePath string) (os.FileInfo, error) {
return nil, err
}
// Check size limit
// Check traditional size limit
if fileInfo.Size() > p.sizeLimit {
context := map[string]interface{}{
"file_size": fileInfo.Size(),
"size_limit": p.sizeLimit,
}
utils.LogErrorf(
utils.NewStructuredError(
utils.ErrorTypeValidation,
utils.CodeValidationSize,
fmt.Sprintf("file size (%d bytes) exceeds limit (%d bytes)", fileInfo.Size(), p.sizeLimit),
).WithFilePath(filePath).WithContext("file_size", fileInfo.Size()).WithContext("size_limit", p.sizeLimit),
filePath,
context,
),
"Skipping large file %s", filePath,
)
return nil, fmt.Errorf("file too large")
}
// Check resource limits
if err := p.resourceMonitor.ValidateFileProcessing(filePath, fileInfo.Size()); err != nil {
utils.LogErrorf(err, "Resource limit validation failed for file: %s", filePath)
return nil, err
}
return fileInfo, nil
}
@@ -103,8 +192,20 @@ func (p *FileProcessor) getRelativePath(filePath string) string {
return relPath
}
// processInMemory loads the entire file into memory (for small files).
func (p *FileProcessor) processInMemory(filePath, relPath string, outCh chan<- WriteRequest) {
// processInMemoryWithContext loads the entire file into memory with context awareness.
func (p *FileProcessor) processInMemoryWithContext(ctx context.Context, filePath, relPath string, outCh chan<- WriteRequest) {
// Check context before reading
select {
case <-ctx.Done():
utils.LogErrorf(
utils.NewStructuredError(utils.ErrorTypeValidation, utils.CodeResourceLimitTimeout, "file processing cancelled", filePath, nil),
"File processing cancelled: %s", filePath,
)
return
default:
}
content, err := os.ReadFile(filePath) // #nosec G304 - filePath is validated by walker
if err != nil {
structErr := utils.WrapError(err, utils.ErrorTypeProcessing, utils.CodeProcessingFileRead, "failed to read file").WithFilePath(filePath)
@@ -112,30 +213,79 @@ func (p *FileProcessor) processInMemory(filePath, relPath string, outCh chan<- W
return
}
outCh <- WriteRequest{
// Check context again after reading
select {
case <-ctx.Done():
utils.LogErrorf(
utils.NewStructuredError(utils.ErrorTypeValidation, utils.CodeResourceLimitTimeout, "file processing cancelled after read", filePath, nil),
"File processing cancelled after read: %s", filePath,
)
return
default:
}
// Try to send the result, but respect context cancellation
select {
case <-ctx.Done():
utils.LogErrorf(
utils.NewStructuredError(utils.ErrorTypeValidation, utils.CodeResourceLimitTimeout, "file processing cancelled before output", filePath, nil),
"File processing cancelled before output: %s", filePath,
)
return
case outCh <- WriteRequest{
Path: relPath,
Content: p.formatContent(relPath, string(content)),
IsStream: false,
}:
}
}
// processStreaming creates a streaming reader for large files.
func (p *FileProcessor) processStreaming(filePath, relPath string, outCh chan<- WriteRequest) {
reader := p.createStreamReader(filePath, relPath)
// processStreamingWithContext creates a streaming reader for large files with context awareness.
func (p *FileProcessor) processStreamingWithContext(ctx context.Context, filePath, relPath string, outCh chan<- WriteRequest) {
// Check context before creating reader
select {
case <-ctx.Done():
utils.LogErrorf(
utils.NewStructuredError(utils.ErrorTypeValidation, utils.CodeResourceLimitTimeout, "streaming processing cancelled", filePath, nil),
"Streaming processing cancelled: %s", filePath,
)
return
default:
}
reader := p.createStreamReaderWithContext(ctx, filePath, relPath)
if reader == nil {
return // Error already logged
}
outCh <- WriteRequest{
// Try to send the result, but respect context cancellation
select {
case <-ctx.Done():
utils.LogErrorf(
utils.NewStructuredError(utils.ErrorTypeValidation, utils.CodeResourceLimitTimeout, "streaming processing cancelled before output", filePath, nil),
"Streaming processing cancelled before output: %s", filePath,
)
return
case outCh <- WriteRequest{
Path: relPath,
Content: "", // Empty since content is in Reader
IsStream: true,
Reader: reader,
}:
}
}
// createStreamReader creates a reader that combines header and file content.
func (p *FileProcessor) createStreamReader(filePath, relPath string) io.Reader {
// createStreamReaderWithContext creates a reader that combines header and file content with context awareness.
func (p *FileProcessor) createStreamReaderWithContext(ctx context.Context, filePath, relPath string) io.Reader {
// Check context before opening file
select {
case <-ctx.Done():
return nil
default:
}
file, err := os.Open(filePath) // #nosec G304 - filePath is validated by walker
if err != nil {
structErr := utils.WrapError(err, utils.ErrorTypeProcessing, utils.CodeProcessingFileRead, "failed to open file for streaming").WithFilePath(filePath)

View File

@@ -0,0 +1,423 @@
// Package fileproc provides resource monitoring and limit enforcement for security.
package fileproc
import (
"context"
"runtime"
"sync"
"sync/atomic"
"time"
"github.com/sirupsen/logrus"
"github.com/ivuorinen/gibidify/config"
"github.com/ivuorinen/gibidify/utils"
)
// ResourceMonitor monitors resource usage and enforces limits to prevent DoS attacks.
type ResourceMonitor struct {
enabled bool
maxFiles int
maxTotalSize int64
fileProcessingTimeout time.Duration
overallTimeout time.Duration
maxConcurrentReads int
rateLimitFilesPerSec int
hardMemoryLimitMB int
enableGracefulDegr bool
enableResourceMon bool
// Current state tracking
filesProcessed int64
totalSizeProcessed int64
concurrentReads int64
startTime time.Time
lastRateLimitCheck time.Time
hardMemoryLimitBytes int64
// Rate limiting
rateLimiter *time.Ticker
rateLimitChan chan struct{}
// Synchronization
mu sync.RWMutex
violationLogged map[string]bool
degradationActive bool
emergencyStopRequested bool
}
// ResourceMetrics holds comprehensive resource usage metrics.
type ResourceMetrics struct {
FilesProcessed int64 `json:"files_processed"`
TotalSizeProcessed int64 `json:"total_size_processed"`
ConcurrentReads int64 `json:"concurrent_reads"`
ProcessingDuration time.Duration `json:"processing_duration"`
AverageFileSize float64 `json:"average_file_size"`
ProcessingRate float64 `json:"processing_rate_files_per_sec"`
MemoryUsageMB int64 `json:"memory_usage_mb"`
MaxMemoryUsageMB int64 `json:"max_memory_usage_mb"`
ViolationsDetected []string `json:"violations_detected"`
DegradationActive bool `json:"degradation_active"`
EmergencyStopActive bool `json:"emergency_stop_active"`
LastUpdated time.Time `json:"last_updated"`
}
// ResourceViolation represents a detected resource limit violation.
type ResourceViolation struct {
Type string `json:"type"`
Message string `json:"message"`
Current interface{} `json:"current"`
Limit interface{} `json:"limit"`
Timestamp time.Time `json:"timestamp"`
Context map[string]interface{} `json:"context"`
}
// NewResourceMonitor creates a new resource monitor with configuration.
func NewResourceMonitor() *ResourceMonitor {
rm := &ResourceMonitor{
enabled: config.GetResourceLimitsEnabled(),
maxFiles: config.GetMaxFiles(),
maxTotalSize: config.GetMaxTotalSize(),
fileProcessingTimeout: time.Duration(config.GetFileProcessingTimeoutSec()) * time.Second,
overallTimeout: time.Duration(config.GetOverallTimeoutSec()) * time.Second,
maxConcurrentReads: config.GetMaxConcurrentReads(),
rateLimitFilesPerSec: config.GetRateLimitFilesPerSec(),
hardMemoryLimitMB: config.GetHardMemoryLimitMB(),
enableGracefulDegr: config.GetEnableGracefulDegradation(),
enableResourceMon: config.GetEnableResourceMonitoring(),
startTime: time.Now(),
lastRateLimitCheck: time.Now(),
violationLogged: make(map[string]bool),
hardMemoryLimitBytes: int64(config.GetHardMemoryLimitMB()) * 1024 * 1024,
}
// Initialize rate limiter if rate limiting is enabled
if rm.enabled && rm.rateLimitFilesPerSec > 0 {
interval := time.Second / time.Duration(rm.rateLimitFilesPerSec)
rm.rateLimiter = time.NewTicker(interval)
rm.rateLimitChan = make(chan struct{}, rm.rateLimitFilesPerSec)
// Pre-fill the rate limit channel
for i := 0; i < rm.rateLimitFilesPerSec; i++ {
select {
case rm.rateLimitChan <- struct{}{}:
default:
goto rateLimitFull
}
}
rateLimitFull:
// Start rate limiter refill goroutine
go rm.rateLimiterRefill()
}
return rm
}
// ValidateFileProcessing checks if a file can be processed based on resource limits.
func (rm *ResourceMonitor) ValidateFileProcessing(filePath string, fileSize int64) error {
if !rm.enabled {
return nil
}
rm.mu.RLock()
defer rm.mu.RUnlock()
// Check if emergency stop is active
if rm.emergencyStopRequested {
return utils.NewStructuredError(
utils.ErrorTypeValidation,
utils.CodeResourceLimitMemory,
"processing stopped due to emergency memory condition",
filePath,
map[string]interface{}{
"emergency_stop_active": true,
},
)
}
// Check file count limit
currentFiles := atomic.LoadInt64(&rm.filesProcessed)
if int(currentFiles) >= rm.maxFiles {
return utils.NewStructuredError(
utils.ErrorTypeValidation,
utils.CodeResourceLimitFiles,
"maximum file count limit exceeded",
filePath,
map[string]interface{}{
"current_files": currentFiles,
"max_files": rm.maxFiles,
},
)
}
// Check total size limit
currentTotalSize := atomic.LoadInt64(&rm.totalSizeProcessed)
if currentTotalSize+fileSize > rm.maxTotalSize {
return utils.NewStructuredError(
utils.ErrorTypeValidation,
utils.CodeResourceLimitTotalSize,
"maximum total size limit would be exceeded",
filePath,
map[string]interface{}{
"current_total_size": currentTotalSize,
"file_size": fileSize,
"max_total_size": rm.maxTotalSize,
},
)
}
// Check overall timeout
if time.Since(rm.startTime) > rm.overallTimeout {
return utils.NewStructuredError(
utils.ErrorTypeValidation,
utils.CodeResourceLimitTimeout,
"overall processing timeout exceeded",
filePath,
map[string]interface{}{
"processing_duration": time.Since(rm.startTime),
"overall_timeout": rm.overallTimeout,
},
)
}
return nil
}
// AcquireReadSlot attempts to acquire a slot for concurrent file reading.
func (rm *ResourceMonitor) AcquireReadSlot(ctx context.Context) error {
if !rm.enabled {
return nil
}
// Wait for available read slot
for {
currentReads := atomic.LoadInt64(&rm.concurrentReads)
if currentReads < int64(rm.maxConcurrentReads) {
if atomic.CompareAndSwapInt64(&rm.concurrentReads, currentReads, currentReads+1) {
break
}
// CAS failed, retry
continue
}
// Wait and retry
select {
case <-ctx.Done():
return ctx.Err()
case <-time.After(time.Millisecond):
// Continue loop
}
}
return nil
}
// ReleaseReadSlot releases a concurrent reading slot.
func (rm *ResourceMonitor) ReleaseReadSlot() {
if rm.enabled {
atomic.AddInt64(&rm.concurrentReads, -1)
}
}
// WaitForRateLimit waits for rate limiting if enabled.
func (rm *ResourceMonitor) WaitForRateLimit(ctx context.Context) error {
if !rm.enabled || rm.rateLimitFilesPerSec <= 0 {
return nil
}
select {
case <-ctx.Done():
return ctx.Err()
case <-rm.rateLimitChan:
return nil
case <-time.After(time.Second): // Fallback timeout
logrus.Warn("Rate limiting timeout exceeded, continuing without rate limit")
return nil
}
}
// CheckHardMemoryLimit checks if hard memory limit is exceeded and takes action.
func (rm *ResourceMonitor) CheckHardMemoryLimit() error {
if !rm.enabled || rm.hardMemoryLimitMB <= 0 {
return nil
}
var m runtime.MemStats
runtime.ReadMemStats(&m)
currentMemory := int64(m.Alloc)
if currentMemory > rm.hardMemoryLimitBytes {
rm.mu.Lock()
defer rm.mu.Unlock()
// Log violation if not already logged
violationKey := "hard_memory_limit"
if !rm.violationLogged[violationKey] {
logrus.Errorf("Hard memory limit exceeded: %dMB > %dMB",
currentMemory/1024/1024, rm.hardMemoryLimitMB)
rm.violationLogged[violationKey] = true
}
if rm.enableGracefulDegr {
// Force garbage collection
runtime.GC()
// Check again after GC
runtime.ReadMemStats(&m)
currentMemory = int64(m.Alloc)
if currentMemory > rm.hardMemoryLimitBytes {
// Still over limit, activate emergency stop
rm.emergencyStopRequested = true
return utils.NewStructuredError(
utils.ErrorTypeValidation,
utils.CodeResourceLimitMemory,
"hard memory limit exceeded, emergency stop activated",
"",
map[string]interface{}{
"current_memory_mb": currentMemory / 1024 / 1024,
"limit_mb": rm.hardMemoryLimitMB,
"emergency_stop": true,
},
)
} else {
// Memory freed by GC, continue with degradation
rm.degradationActive = true
logrus.Info("Memory freed by garbage collection, continuing with degradation mode")
}
} else {
// No graceful degradation, hard stop
return utils.NewStructuredError(
utils.ErrorTypeValidation,
utils.CodeResourceLimitMemory,
"hard memory limit exceeded",
"",
map[string]interface{}{
"current_memory_mb": currentMemory / 1024 / 1024,
"limit_mb": rm.hardMemoryLimitMB,
},
)
}
}
return nil
}
// RecordFileProcessed records that a file has been successfully processed.
func (rm *ResourceMonitor) RecordFileProcessed(fileSize int64) {
if rm.enabled {
atomic.AddInt64(&rm.filesProcessed, 1)
atomic.AddInt64(&rm.totalSizeProcessed, fileSize)
}
}
// GetMetrics returns current resource usage metrics.
func (rm *ResourceMonitor) GetMetrics() ResourceMetrics {
if !rm.enableResourceMon {
return ResourceMetrics{}
}
rm.mu.RLock()
defer rm.mu.RUnlock()
var m runtime.MemStats
runtime.ReadMemStats(&m)
filesProcessed := atomic.LoadInt64(&rm.filesProcessed)
totalSize := atomic.LoadInt64(&rm.totalSizeProcessed)
duration := time.Since(rm.startTime)
avgFileSize := float64(0)
if filesProcessed > 0 {
avgFileSize = float64(totalSize) / float64(filesProcessed)
}
processingRate := float64(0)
if duration.Seconds() > 0 {
processingRate = float64(filesProcessed) / duration.Seconds()
}
// Collect violations
violations := make([]string, 0, len(rm.violationLogged))
for violation := range rm.violationLogged {
violations = append(violations, violation)
}
return ResourceMetrics{
FilesProcessed: filesProcessed,
TotalSizeProcessed: totalSize,
ConcurrentReads: atomic.LoadInt64(&rm.concurrentReads),
ProcessingDuration: duration,
AverageFileSize: avgFileSize,
ProcessingRate: processingRate,
MemoryUsageMB: int64(m.Alloc) / 1024 / 1024,
MaxMemoryUsageMB: int64(rm.hardMemoryLimitMB),
ViolationsDetected: violations,
DegradationActive: rm.degradationActive,
EmergencyStopActive: rm.emergencyStopRequested,
LastUpdated: time.Now(),
}
}
// IsEmergencyStopActive returns whether emergency stop is active.
func (rm *ResourceMonitor) IsEmergencyStopActive() bool {
rm.mu.RLock()
defer rm.mu.RUnlock()
return rm.emergencyStopRequested
}
// IsDegradationActive returns whether degradation mode is active.
func (rm *ResourceMonitor) IsDegradationActive() bool {
rm.mu.RLock()
defer rm.mu.RUnlock()
return rm.degradationActive
}
// LogResourceInfo logs current resource limit configuration.
func (rm *ResourceMonitor) LogResourceInfo() {
if rm.enabled {
logrus.Infof("Resource limits enabled: maxFiles=%d, maxTotalSize=%dMB, fileTimeout=%ds, overallTimeout=%ds",
rm.maxFiles, rm.maxTotalSize/1024/1024, int(rm.fileProcessingTimeout.Seconds()), int(rm.overallTimeout.Seconds()))
logrus.Infof("Resource limits: maxConcurrentReads=%d, rateLimitFPS=%d, hardMemoryMB=%d",
rm.maxConcurrentReads, rm.rateLimitFilesPerSec, rm.hardMemoryLimitMB)
logrus.Infof("Resource features: gracefulDegradation=%v, monitoring=%v",
rm.enableGracefulDegr, rm.enableResourceMon)
} else {
logrus.Info("Resource limits disabled")
}
}
// Close cleans up the resource monitor.
func (rm *ResourceMonitor) Close() {
if rm.rateLimiter != nil {
rm.rateLimiter.Stop()
}
}
// rateLimiterRefill refills the rate limiting channel periodically.
func (rm *ResourceMonitor) rateLimiterRefill() {
for range rm.rateLimiter.C {
select {
case rm.rateLimitChan <- struct{}{}:
default:
// Channel is full, skip
}
}
}
// CreateFileProcessingContext creates a context with file processing timeout.
func (rm *ResourceMonitor) CreateFileProcessingContext(parent context.Context) (context.Context, context.CancelFunc) {
if !rm.enabled || rm.fileProcessingTimeout <= 0 {
return parent, func() {}
}
return context.WithTimeout(parent, rm.fileProcessingTimeout)
}
// CreateOverallProcessingContext creates a context with overall processing timeout.
func (rm *ResourceMonitor) CreateOverallProcessingContext(parent context.Context) (context.Context, context.CancelFunc) {
if !rm.enabled || rm.overallTimeout <= 0 {
return parent, func() {}
}
return context.WithTimeout(parent, rm.overallTimeout)
}

View File

@@ -0,0 +1,377 @@
// Package fileproc provides tests for resource monitoring functionality.
package fileproc
import (
"context"
"os"
"path/filepath"
"testing"
"time"
"github.com/spf13/viper"
"github.com/ivuorinen/gibidify/config"
"github.com/ivuorinen/gibidify/testutil"
"github.com/ivuorinen/gibidify/utils"
)
func TestResourceMonitor_NewResourceMonitor(t *testing.T) {
// Reset viper for clean test state
testutil.ResetViperConfig(t, "")
rm := NewResourceMonitor()
if rm == nil {
t.Fatal("NewResourceMonitor() returned nil")
}
// Test default values are set correctly
if !rm.enabled {
t.Error("Expected resource monitor to be enabled by default")
}
if rm.maxFiles != config.DefaultMaxFiles {
t.Errorf("Expected maxFiles to be %d, got %d", config.DefaultMaxFiles, rm.maxFiles)
}
if rm.maxTotalSize != config.DefaultMaxTotalSize {
t.Errorf("Expected maxTotalSize to be %d, got %d", config.DefaultMaxTotalSize, rm.maxTotalSize)
}
if rm.fileProcessingTimeout != time.Duration(config.DefaultFileProcessingTimeoutSec)*time.Second {
t.Errorf("Expected fileProcessingTimeout to be %v, got %v",
time.Duration(config.DefaultFileProcessingTimeoutSec)*time.Second, rm.fileProcessingTimeout)
}
// Clean up
rm.Close()
}
func TestResourceMonitor_DisabledResourceLimits(t *testing.T) {
// Reset viper for clean test state
testutil.ResetViperConfig(t, "")
// Set resource limits disabled
viper.Set("resourceLimits.enabled", false)
rm := NewResourceMonitor()
defer rm.Close()
// Test that validation passes when disabled
err := rm.ValidateFileProcessing("/tmp/test.txt", 1000)
if err != nil {
t.Errorf("Expected no error when resource limits disabled, got %v", err)
}
// Test that read slot acquisition works when disabled
ctx := context.Background()
err = rm.AcquireReadSlot(ctx)
if err != nil {
t.Errorf("Expected no error when acquiring read slot with disabled limits, got %v", err)
}
rm.ReleaseReadSlot()
// Test that rate limiting is bypassed when disabled
err = rm.WaitForRateLimit(ctx)
if err != nil {
t.Errorf("Expected no error when rate limiting disabled, got %v", err)
}
}
func TestResourceMonitor_FileCountLimit(t *testing.T) {
testutil.ResetViperConfig(t, "")
// Set a very low file count limit for testing
viper.Set("resourceLimits.enabled", true)
viper.Set("resourceLimits.maxFiles", 2)
rm := NewResourceMonitor()
defer rm.Close()
// First file should pass
err := rm.ValidateFileProcessing("/tmp/file1.txt", 100)
if err != nil {
t.Errorf("Expected no error for first file, got %v", err)
}
rm.RecordFileProcessed(100)
// Second file should pass
err = rm.ValidateFileProcessing("/tmp/file2.txt", 100)
if err != nil {
t.Errorf("Expected no error for second file, got %v", err)
}
rm.RecordFileProcessed(100)
// Third file should fail
err = rm.ValidateFileProcessing("/tmp/file3.txt", 100)
if err == nil {
t.Error("Expected error for third file (exceeds limit), got nil")
}
// Verify it's the correct error type
structErr, ok := err.(*utils.StructuredError)
if !ok {
t.Errorf("Expected StructuredError, got %T", err)
} else if structErr.Code != utils.CodeResourceLimitFiles {
t.Errorf("Expected error code %s, got %s", utils.CodeResourceLimitFiles, structErr.Code)
}
}
func TestResourceMonitor_TotalSizeLimit(t *testing.T) {
testutil.ResetViperConfig(t, "")
// Set a low total size limit for testing (1KB)
viper.Set("resourceLimits.enabled", true)
viper.Set("resourceLimits.maxTotalSize", 1024)
rm := NewResourceMonitor()
defer rm.Close()
// First small file should pass
err := rm.ValidateFileProcessing("/tmp/small.txt", 500)
if err != nil {
t.Errorf("Expected no error for small file, got %v", err)
}
rm.RecordFileProcessed(500)
// Second small file should pass
err = rm.ValidateFileProcessing("/tmp/small2.txt", 400)
if err != nil {
t.Errorf("Expected no error for second small file, got %v", err)
}
rm.RecordFileProcessed(400)
// Large file that would exceed limit should fail
err = rm.ValidateFileProcessing("/tmp/large.txt", 200)
if err == nil {
t.Error("Expected error for file that would exceed size limit, got nil")
}
// Verify it's the correct error type
structErr, ok := err.(*utils.StructuredError)
if !ok {
t.Errorf("Expected StructuredError, got %T", err)
} else if structErr.Code != utils.CodeResourceLimitTotalSize {
t.Errorf("Expected error code %s, got %s", utils.CodeResourceLimitTotalSize, structErr.Code)
}
}
func TestResourceMonitor_ConcurrentReadsLimit(t *testing.T) {
testutil.ResetViperConfig(t, "")
// Set a low concurrent reads limit for testing
viper.Set("resourceLimits.enabled", true)
viper.Set("resourceLimits.maxConcurrentReads", 2)
rm := NewResourceMonitor()
defer rm.Close()
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
defer cancel()
// First read slot should succeed
err := rm.AcquireReadSlot(ctx)
if err != nil {
t.Errorf("Expected no error for first read slot, got %v", err)
}
// Second read slot should succeed
err = rm.AcquireReadSlot(ctx)
if err != nil {
t.Errorf("Expected no error for second read slot, got %v", err)
}
// Third read slot should timeout (context deadline exceeded)
err = rm.AcquireReadSlot(ctx)
if err == nil {
t.Error("Expected timeout error for third read slot, got nil")
}
// Release one slot and try again
rm.ReleaseReadSlot()
// Create new context for the next attempt
ctx2, cancel2 := context.WithTimeout(context.Background(), 100*time.Millisecond)
defer cancel2()
err = rm.AcquireReadSlot(ctx2)
if err != nil {
t.Errorf("Expected no error after releasing a slot, got %v", err)
}
// Clean up remaining slots
rm.ReleaseReadSlot()
rm.ReleaseReadSlot()
}
func TestResourceMonitor_TimeoutContexts(t *testing.T) {
testutil.ResetViperConfig(t, "")
// Set short timeouts for testing
viper.Set("resourceLimits.enabled", true)
viper.Set("resourceLimits.fileProcessingTimeoutSec", 1) // 1 second
viper.Set("resourceLimits.overallTimeoutSec", 2) // 2 seconds
rm := NewResourceMonitor()
defer rm.Close()
parentCtx := context.Background()
// Test file processing context
fileCtx, fileCancel := rm.CreateFileProcessingContext(parentCtx)
defer fileCancel()
deadline, ok := fileCtx.Deadline()
if !ok {
t.Error("Expected file processing context to have a deadline")
} else if time.Until(deadline) > time.Second+100*time.Millisecond {
t.Error("File processing timeout appears to be too long")
}
// Test overall processing context
overallCtx, overallCancel := rm.CreateOverallProcessingContext(parentCtx)
defer overallCancel()
deadline, ok = overallCtx.Deadline()
if !ok {
t.Error("Expected overall processing context to have a deadline")
} else if time.Until(deadline) > 2*time.Second+100*time.Millisecond {
t.Error("Overall processing timeout appears to be too long")
}
}
func TestResourceMonitor_RateLimiting(t *testing.T) {
testutil.ResetViperConfig(t, "")
// Enable rate limiting with a low rate for testing
viper.Set("resourceLimits.enabled", true)
viper.Set("resourceLimits.rateLimitFilesPerSec", 5) // 5 files per second
rm := NewResourceMonitor()
defer rm.Close()
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
defer cancel()
// First few requests should succeed quickly
start := time.Now()
for i := 0; i < 3; i++ {
err := rm.WaitForRateLimit(ctx)
if err != nil {
t.Errorf("Expected no error for rate limit wait %d, got %v", i, err)
}
}
// Should have taken some time due to rate limiting
duration := time.Since(start)
if duration < 200*time.Millisecond {
t.Logf("Rate limiting may not be working as expected, took only %v", duration)
}
}
func TestResourceMonitor_Metrics(t *testing.T) {
testutil.ResetViperConfig(t, "")
viper.Set("resourceLimits.enabled", true)
viper.Set("resourceLimits.enableResourceMonitoring", true)
rm := NewResourceMonitor()
defer rm.Close()
// Process some files to generate metrics
rm.RecordFileProcessed(1000)
rm.RecordFileProcessed(2000)
rm.RecordFileProcessed(500)
metrics := rm.GetMetrics()
// Verify metrics
if metrics.FilesProcessed != 3 {
t.Errorf("Expected 3 files processed, got %d", metrics.FilesProcessed)
}
if metrics.TotalSizeProcessed != 3500 {
t.Errorf("Expected total size 3500, got %d", metrics.TotalSizeProcessed)
}
expectedAvgSize := float64(3500) / float64(3)
if metrics.AverageFileSize != expectedAvgSize {
t.Errorf("Expected average file size %.2f, got %.2f", expectedAvgSize, metrics.AverageFileSize)
}
if metrics.ProcessingRate <= 0 {
t.Error("Expected positive processing rate")
}
if !metrics.LastUpdated.After(time.Now().Add(-time.Second)) {
t.Error("Expected recent LastUpdated timestamp")
}
}
func TestResourceMonitor_Integration(t *testing.T) {
// Create temporary test directory
tempDir := t.TempDir()
// Create test files
testFiles := []string{"test1.txt", "test2.txt", "test3.txt"}
for _, filename := range testFiles {
testutil.CreateTestFile(t, tempDir, filename, []byte("test content"))
}
testutil.ResetViperConfig(t, "")
// Configure resource limits
viper.Set("resourceLimits.enabled", true)
viper.Set("resourceLimits.maxFiles", 5)
viper.Set("resourceLimits.maxTotalSize", 1024*1024) // 1MB
viper.Set("resourceLimits.fileProcessingTimeoutSec", 10)
viper.Set("resourceLimits.maxConcurrentReads", 3)
rm := NewResourceMonitor()
defer rm.Close()
ctx := context.Background()
// Test file processing workflow
for _, filename := range testFiles {
filePath := filepath.Join(tempDir, filename)
fileInfo, err := os.Stat(filePath)
if err != nil {
t.Fatalf("Failed to stat test file %s: %v", filePath, err)
}
// Validate file can be processed
err = rm.ValidateFileProcessing(filePath, fileInfo.Size())
if err != nil {
t.Errorf("Failed to validate file %s: %v", filePath, err)
continue
}
// Acquire read slot
err = rm.AcquireReadSlot(ctx)
if err != nil {
t.Errorf("Failed to acquire read slot for %s: %v", filePath, err)
continue
}
// Check memory limits
err = rm.CheckHardMemoryLimit()
if err != nil {
t.Errorf("Memory limit check failed for %s: %v", filePath, err)
}
// Record processing
rm.RecordFileProcessed(fileInfo.Size())
// Release read slot
rm.ReleaseReadSlot()
}
// Verify final metrics
metrics := rm.GetMetrics()
if metrics.FilesProcessed != int64(len(testFiles)) {
t.Errorf("Expected %d files processed, got %d", len(testFiles), metrics.FilesProcessed)
}
// Test resource limit logging
rm.LogResourceInfo()
}

View File

@@ -18,11 +18,16 @@ func StartWriter(outFile *os.File, writeCh <-chan WriteRequest, done chan<- stru
case "yaml":
startYAMLWriter(outFile, writeCh, done, prefix, suffix)
default:
context := map[string]interface{}{
"format": format,
}
err := utils.NewStructuredError(
utils.ErrorTypeValidation,
utils.CodeValidationFormat,
fmt.Sprintf("unsupported format: %s", format),
).WithContext("format", format)
"",
context,
)
utils.LogError("Failed to encode output", err)
close(done)
}