9434f4c8e6
Add MaxArtifactFiles and MaxArtifactTotalBytes to SandboxConfig: - Default MaxArtifactFiles: 10,000 (configurable via SecurityDefaults) - Default MaxArtifactTotalBytes: 100GB (configurable via SecurityDefaults) - ApplySecurityDefaults() sets defaults if not specified Enforce caps in scanArtifacts() during directory walk: - Returns error immediately when MaxArtifactFiles exceeded - Returns error immediately when MaxArtifactTotalBytes exceeded - Prevents resource exhaustion attacks from malicious artifact trees Update all call sites to pass SandboxConfig for cap enforcement: - Native bridge libs updated to pass caps argument - Benchmark tests updated with nil caps (unlimited for benchmarks) - Unit tests updated with nil caps Closes: artifact ingestion caps items from security plan
68 lines
1.6 KiB
Go
68 lines
1.6 KiB
Go
package worker_test
|
|
|
|
import (
|
|
"os"
|
|
"path/filepath"
|
|
"testing"
|
|
|
|
"github.com/jfraeys/fetch_ml/internal/worker"
|
|
)
|
|
|
|
func TestScanArtifacts_SkipsKnownPathsAndLogs(t *testing.T) {
|
|
runDir := t.TempDir()
|
|
|
|
mustWrite := func(rel string, data []byte) {
|
|
p := filepath.Join(runDir, rel)
|
|
if err := os.MkdirAll(filepath.Dir(p), 0750); err != nil {
|
|
t.Fatalf("mkdir: %v", err)
|
|
}
|
|
if err := os.WriteFile(p, data, 0600); err != nil {
|
|
t.Fatalf("write file: %v", err)
|
|
}
|
|
}
|
|
|
|
mustWrite("run_manifest.json", []byte("{}"))
|
|
mustWrite("output.log", []byte("log"))
|
|
mustWrite("code/ignored.txt", []byte("ignore"))
|
|
mustWrite("snapshot/ignored.bin", []byte("ignore"))
|
|
|
|
mustWrite("results/metrics.jsonl", []byte("m"))
|
|
mustWrite("checkpoints/best.pt", []byte("checkpoint"))
|
|
mustWrite("plots/loss.png", []byte("png"))
|
|
|
|
art, err := worker.ScanArtifacts(runDir, false, nil)
|
|
if err != nil {
|
|
t.Fatalf("scanArtifacts: %v", err)
|
|
}
|
|
if art == nil {
|
|
t.Fatalf("expected artifacts")
|
|
}
|
|
|
|
paths := make([]string, 0, len(art.Files))
|
|
var total int64
|
|
for _, f := range art.Files {
|
|
paths = append(paths, f.Path)
|
|
total += f.SizeBytes
|
|
}
|
|
|
|
want := []string{
|
|
"checkpoints/best.pt",
|
|
"plots/loss.png",
|
|
"results/metrics.jsonl",
|
|
}
|
|
if len(paths) != len(want) {
|
|
t.Fatalf("expected %d files, got %d: %v", len(want), len(paths), paths)
|
|
}
|
|
for i := range want {
|
|
if paths[i] != want[i] {
|
|
t.Fatalf("expected paths[%d]=%q, got %q", i, want[i], paths[i])
|
|
}
|
|
}
|
|
|
|
if art.TotalSizeBytes != total {
|
|
t.Fatalf("expected total_size_bytes=%d, got %d", total, art.TotalSizeBytes)
|
|
}
|
|
if art.DiscoveryTime.IsZero() {
|
|
t.Fatalf("expected discovery_time")
|
|
}
|
|
}
|