mirror of
https://github.com/ivuorinen/gibidify.git
synced 2026-01-26 03:24:05 +00:00
feat(security): improve security features, fixes
This commit is contained in:
@@ -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)
|
||||
|
||||
423
fileproc/resource_monitor.go
Normal file
423
fileproc/resource_monitor.go
Normal 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)
|
||||
}
|
||||
377
fileproc/resource_monitor_test.go
Normal file
377
fileproc/resource_monitor_test.go
Normal 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()
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user