test(phase-3): prerequisite security and reproducibility tests
Implement 4 prerequisite test requirements: - TestConfigIntegrityVerification: Config signing, tamper detection, hash stability - TestManifestFilenameNonce: Cryptographic nonce generation and filename patterns - TestGPUDetectionAudit: Structured logging of GPU detection at startup - TestResourceEnvVarParsing: Resource env var parsing and override behavior Also update manifest run_manifest.go: - Add nonce-based filename support to WriteToDir - Add nonce-based file detection to LoadFromDir
This commit is contained in:
parent
f71352202e
commit
8f9bcef754
5 changed files with 740 additions and 2 deletions
|
|
@ -5,6 +5,7 @@ import (
|
|||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
|
@ -228,17 +229,45 @@ func (m *RunManifest) WriteToDir(dir string) error {
|
|||
if err != nil {
|
||||
return fmt.Errorf("marshal run manifest: %w", err)
|
||||
}
|
||||
if err := fileutil.SecureFileWrite(ManifestPath(dir), data, 0640); err != nil {
|
||||
|
||||
// Use nonce-based filename if Environment.ManifestNonce is set
|
||||
var manifestPath string
|
||||
if m.Environment != nil && m.Environment.ManifestNonce != "" {
|
||||
manifestPath = ManifestPathWithNonce(dir, m.Environment.ManifestNonce)
|
||||
} else {
|
||||
manifestPath = ManifestPath(dir)
|
||||
}
|
||||
|
||||
if err := fileutil.SecureFileWrite(manifestPath, data, 0640); err != nil {
|
||||
return fmt.Errorf("write run manifest: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func LoadFromDir(dir string) (*RunManifest, error) {
|
||||
// Try standard filename first
|
||||
data, err := fileutil.SecureFileRead(ManifestPath(dir))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read run manifest: %w", err)
|
||||
// If not found, look for nonce-based filename
|
||||
entries, readErr := os.ReadDir(dir)
|
||||
if readErr != nil {
|
||||
return nil, fmt.Errorf("read run manifest: %w", err)
|
||||
}
|
||||
|
||||
for _, entry := range entries {
|
||||
if strings.HasPrefix(entry.Name(), "run_manifest_") && strings.HasSuffix(entry.Name(), ".json") {
|
||||
data, err = fileutil.SecureFileRead(filepath.Join(dir, entry.Name()))
|
||||
if err == nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read run manifest: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
var m RunManifest
|
||||
if err := json.Unmarshal(data, &m); err != nil {
|
||||
return nil, fmt.Errorf("parse run manifest: %w", err)
|
||||
|
|
|
|||
166
tests/unit/security/config_integrity_test.go
Normal file
166
tests/unit/security/config_integrity_test.go
Normal file
|
|
@ -0,0 +1,166 @@
|
|||
package security
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/jfraeys/fetch_ml/internal/crypto"
|
||||
"github.com/jfraeys/fetch_ml/internal/worker"
|
||||
)
|
||||
|
||||
// TestConfigIntegrityVerification verifies config file integrity and signature verification.
|
||||
// This test ensures that config files can be signed and their signatures verified.
|
||||
func TestConfigIntegrityVerification(t *testing.T) {
|
||||
t.Run("ConfigLoadWithoutSignature", func(t *testing.T) {
|
||||
// Create a temp config file
|
||||
tempDir := t.TempDir()
|
||||
configPath := filepath.Join(tempDir, "config.yaml")
|
||||
|
||||
configContent := `
|
||||
host: localhost
|
||||
port: 22
|
||||
max_workers: 4
|
||||
`
|
||||
if err := os.WriteFile(configPath, []byte(configContent), 0600); err != nil {
|
||||
t.Fatalf("failed to write config: %v", err)
|
||||
}
|
||||
|
||||
// Load config without signature verification (default behavior)
|
||||
cfg, err := worker.LoadConfig(configPath)
|
||||
if err != nil {
|
||||
t.Fatalf("LoadConfig failed: %v", err)
|
||||
}
|
||||
|
||||
if cfg.Host != "localhost" {
|
||||
t.Errorf("host = %q, want localhost", cfg.Host)
|
||||
}
|
||||
if cfg.Port != 22 {
|
||||
t.Errorf("port = %d, want 22", cfg.Port)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("ConfigFileTamperingDetection", func(t *testing.T) {
|
||||
// Create a temp config file
|
||||
tempDir := t.TempDir()
|
||||
configPath := filepath.Join(tempDir, "config.yaml")
|
||||
|
||||
configContent := []byte(`
|
||||
host: localhost
|
||||
port: 22
|
||||
max_workers: 4
|
||||
`)
|
||||
if err := os.WriteFile(configPath, configContent, 0600); err != nil {
|
||||
t.Fatalf("failed to write config: %v", err)
|
||||
}
|
||||
|
||||
// Generate signing keys
|
||||
publicKey, privateKey, err := crypto.GenerateSigningKeys()
|
||||
if err != nil {
|
||||
t.Fatalf("GenerateSigningKeys failed: %v", err)
|
||||
}
|
||||
|
||||
// Create signer
|
||||
signer, err := crypto.NewManifestSigner(privateKey, "test-key-1")
|
||||
if err != nil {
|
||||
t.Fatalf("NewManifestSigner failed: %v", err)
|
||||
}
|
||||
|
||||
// Sign the config content
|
||||
result, err := signer.SignManifestBytes(configContent)
|
||||
if err != nil {
|
||||
t.Fatalf("SignManifestBytes failed: %v", err)
|
||||
}
|
||||
|
||||
// Verify signature against original content
|
||||
valid, err := crypto.VerifyManifestBytes(configContent, result, publicKey)
|
||||
if err != nil {
|
||||
t.Fatalf("VerifyManifestBytes failed: %v", err)
|
||||
}
|
||||
if !valid {
|
||||
t.Error("signature should be valid for original content")
|
||||
}
|
||||
|
||||
// Tamper with the config file
|
||||
tamperedContent := []byte(`
|
||||
host: malicious-host
|
||||
port: 22
|
||||
max_workers: 4
|
||||
`)
|
||||
if err := os.WriteFile(configPath, tamperedContent, 0600); err != nil {
|
||||
t.Fatalf("failed to write tampered config: %v", err)
|
||||
}
|
||||
|
||||
// Verify signature against tampered content (should fail)
|
||||
valid, err = crypto.VerifyManifestBytes(tamperedContent, result, publicKey)
|
||||
if err != nil {
|
||||
// Expected - signature doesn't match
|
||||
t.Logf("Expected verification error for tampered content: %v", err)
|
||||
}
|
||||
if valid {
|
||||
t.Error("signature should be invalid for tampered content")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("MissingSignatureFile", func(t *testing.T) {
|
||||
tempDir := t.TempDir()
|
||||
configPath := filepath.Join(tempDir, "config.yaml")
|
||||
|
||||
configContent := `
|
||||
host: localhost
|
||||
port: 22
|
||||
`
|
||||
if err := os.WriteFile(configPath, []byte(configContent), 0600); err != nil {
|
||||
t.Fatalf("failed to write config: %v", err)
|
||||
}
|
||||
|
||||
// Config loads without signature
|
||||
cfg, err := worker.LoadConfig(configPath)
|
||||
if err != nil {
|
||||
t.Fatalf("LoadConfig should work without signature: %v", err)
|
||||
}
|
||||
|
||||
if cfg == nil {
|
||||
t.Error("expected config to be loaded")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// TestConfigHashStability verifies that the same config produces the same hash
|
||||
func TestConfigHashStability(t *testing.T) {
|
||||
cfg := &worker.Config{
|
||||
Host: "localhost",
|
||||
Port: 22,
|
||||
MaxWorkers: 4,
|
||||
GPUVendor: "nvidia",
|
||||
Sandbox: worker.SandboxConfig{
|
||||
NetworkMode: "none",
|
||||
SeccompProfile: "default-hardened",
|
||||
},
|
||||
}
|
||||
|
||||
hash1, err := cfg.ComputeResolvedConfigHash()
|
||||
if err != nil {
|
||||
t.Fatalf("ComputeResolvedConfigHash failed: %v", err)
|
||||
}
|
||||
|
||||
hash2, err := cfg.ComputeResolvedConfigHash()
|
||||
if err != nil {
|
||||
t.Fatalf("ComputeResolvedConfigHash failed: %v", err)
|
||||
}
|
||||
|
||||
if hash1 != hash2 {
|
||||
t.Error("same config should produce identical hashes")
|
||||
}
|
||||
|
||||
// Modify config and verify hash changes
|
||||
cfg.MaxWorkers = 8
|
||||
hash3, err := cfg.ComputeResolvedConfigHash()
|
||||
if err != nil {
|
||||
t.Fatalf("ComputeResolvedConfigHash failed: %v", err)
|
||||
}
|
||||
|
||||
if hash1 == hash3 {
|
||||
t.Error("different configs should produce different hashes")
|
||||
}
|
||||
}
|
||||
220
tests/unit/security/gpu_audit_test.go
Normal file
220
tests/unit/security/gpu_audit_test.go
Normal file
|
|
@ -0,0 +1,220 @@
|
|||
package security
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"log/slog"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/jfraeys/fetch_ml/internal/worker"
|
||||
)
|
||||
|
||||
// TestGPUDetectionAudit verifies that GPU detection method is logged at startup
|
||||
// for audit and reproducibility purposes.
|
||||
func TestGPUDetectionAudit(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
gpuType string
|
||||
wantMethod string
|
||||
}{
|
||||
{
|
||||
name: "nvidia detection logs method",
|
||||
gpuType: "nvidia",
|
||||
wantMethod: "config",
|
||||
},
|
||||
{
|
||||
name: "apple detection logs method",
|
||||
gpuType: "apple",
|
||||
wantMethod: "config",
|
||||
},
|
||||
{
|
||||
name: "none detection logs method",
|
||||
gpuType: "none",
|
||||
wantMethod: "config",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
// Capture log output
|
||||
var logBuf bytes.Buffer
|
||||
handler := slog.NewTextHandler(&logBuf, &slog.HandlerOptions{
|
||||
Level: slog.LevelInfo,
|
||||
})
|
||||
logger := slog.New(handler)
|
||||
|
||||
// Create config with specified GPU type
|
||||
cfg := &worker.Config{
|
||||
GPUVendor: tt.gpuType,
|
||||
}
|
||||
|
||||
// Perform GPU detection
|
||||
factory := &worker.GPUDetectorFactory{}
|
||||
result := factory.CreateDetectorWithInfo(cfg)
|
||||
|
||||
// Log the detection info (this simulates startup logging)
|
||||
logger.Info("GPU detection completed",
|
||||
"gpu_type", result.Info.GPUType,
|
||||
"detection_method", result.Info.DetectionMethod,
|
||||
"configured_vendor", result.Info.ConfiguredVendor,
|
||||
)
|
||||
|
||||
// Verify log output contains detection method
|
||||
logOutput := logBuf.String()
|
||||
if !strings.Contains(logOutput, "GPU detection completed") {
|
||||
t.Error("expected 'GPU detection completed' in log output")
|
||||
}
|
||||
if !strings.Contains(logOutput, string(result.Info.DetectionMethod)) {
|
||||
t.Errorf("expected detection method %q in log output", result.Info.DetectionMethod)
|
||||
}
|
||||
if !strings.Contains(logOutput, "detection_method=") {
|
||||
t.Error("expected 'detection_method=' field in log output")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
t.Run("env override detection logged", func(t *testing.T) {
|
||||
// Set env var to trigger env-based detection
|
||||
t.Setenv("FETCH_ML_GPU_TYPE", "nvidia")
|
||||
|
||||
var logBuf bytes.Buffer
|
||||
handler := slog.NewTextHandler(&logBuf, &slog.HandlerOptions{
|
||||
Level: slog.LevelInfo,
|
||||
})
|
||||
logger := slog.New(handler)
|
||||
|
||||
// Create factory and detect
|
||||
factory := &worker.GPUDetectorFactory{}
|
||||
result := factory.CreateDetectorWithInfo(nil)
|
||||
|
||||
// Log detection
|
||||
logger.Info("GPU detection completed",
|
||||
"gpu_type", result.Info.GPUType,
|
||||
"detection_method", result.Info.DetectionMethod,
|
||||
"env_override_type", result.Info.EnvOverrideType,
|
||||
)
|
||||
|
||||
logOutput := logBuf.String()
|
||||
|
||||
// Should log env override
|
||||
if !strings.Contains(logOutput, "env_override_type=nvidia") {
|
||||
t.Error("expected env override type in log output")
|
||||
}
|
||||
|
||||
// Detection method should indicate env was used
|
||||
if result.Info.DetectionMethod != worker.DetectionSourceEnvType {
|
||||
t.Errorf("detection method = %v, want %v", result.Info.DetectionMethod, worker.DetectionSourceEnvType)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("detection info fields populated", func(t *testing.T) {
|
||||
// Set both env vars
|
||||
t.Setenv("FETCH_ML_GPU_TYPE", "apple")
|
||||
t.Setenv("FETCH_ML_GPU_COUNT", "4")
|
||||
|
||||
factory := &worker.GPUDetectorFactory{}
|
||||
result := factory.CreateDetectorWithInfo(nil)
|
||||
|
||||
// Verify all expected fields are populated
|
||||
if result.Info.GPUType == "" {
|
||||
t.Error("GPUType field is empty")
|
||||
}
|
||||
if result.Info.ConfiguredVendor == "" {
|
||||
t.Error("ConfiguredVendor field is empty")
|
||||
}
|
||||
if result.Info.DetectionMethod == "" {
|
||||
t.Error("DetectionMethod field is empty")
|
||||
}
|
||||
if result.Info.EnvOverrideType != "apple" {
|
||||
t.Errorf("EnvOverrideType = %v, want 'apple'", result.Info.EnvOverrideType)
|
||||
}
|
||||
if result.Info.EnvOverrideCount != 4 {
|
||||
t.Errorf("EnvOverrideCount = %v, want 4", result.Info.EnvOverrideCount)
|
||||
}
|
||||
|
||||
// Verify detection source is valid
|
||||
validSources := []worker.DetectionSource{
|
||||
worker.DetectionSourceEnvType,
|
||||
worker.DetectionSourceEnvCount,
|
||||
worker.DetectionSourceEnvBoth,
|
||||
worker.DetectionSourceConfig,
|
||||
worker.DetectionSourceAuto,
|
||||
}
|
||||
found := false
|
||||
for _, source := range validSources {
|
||||
if result.Info.DetectionMethod == source {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Errorf("invalid detection method: %v", result.Info.DetectionMethod)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// TestGPUDetectionStructuredLogging verifies structured logging format for audit trails
|
||||
func TestGPUDetectionStructuredLogging(t *testing.T) {
|
||||
var logBuf bytes.Buffer
|
||||
handler := slog.NewJSONHandler(&logBuf, &slog.HandlerOptions{
|
||||
Level: slog.LevelInfo,
|
||||
})
|
||||
logger := slog.New(handler)
|
||||
|
||||
// Simulate startup detection
|
||||
cfg := &worker.Config{GPUVendor: "nvidia"}
|
||||
factory := &worker.GPUDetectorFactory{}
|
||||
result := factory.CreateDetectorWithInfo(cfg)
|
||||
|
||||
// Log with structured format for audit
|
||||
logger.Info("gpu_detection_startup",
|
||||
"event_type", "gpu_detection",
|
||||
"gpu_type", result.Info.GPUType,
|
||||
"detection_method", result.Info.DetectionMethod,
|
||||
"configured_vendor", result.Info.ConfiguredVendor,
|
||||
)
|
||||
|
||||
logOutput := logBuf.String()
|
||||
|
||||
// Verify JSON structure contains required fields
|
||||
if !strings.Contains(logOutput, `"event_type":"gpu_detection"`) {
|
||||
t.Error("expected event_type field in JSON log")
|
||||
}
|
||||
if !strings.Contains(logOutput, `"detection_method"`) {
|
||||
t.Error("expected detection_method field in JSON log")
|
||||
}
|
||||
if !strings.Contains(logOutput, `"configured_vendor"`) {
|
||||
t.Error("expected configured_vendor field in JSON log")
|
||||
}
|
||||
}
|
||||
|
||||
// TestGPUDetectionAuditDisabled verifies behavior when detection fails
|
||||
func TestGPUDetectionAuditDisabled(t *testing.T) {
|
||||
var logBuf bytes.Buffer
|
||||
handler := slog.NewTextHandler(&logBuf, &slog.HandlerOptions{
|
||||
Level: slog.LevelInfo,
|
||||
})
|
||||
logger := slog.New(handler)
|
||||
|
||||
// Test with GPU none - should still log
|
||||
cfg := &worker.Config{GPUVendor: "none"}
|
||||
factory := &worker.GPUDetectorFactory{}
|
||||
result := factory.CreateDetectorWithInfo(cfg)
|
||||
|
||||
logger.Info("GPU detection completed",
|
||||
"gpu_type", result.Info.GPUType,
|
||||
"detection_method", result.Info.DetectionMethod,
|
||||
)
|
||||
|
||||
logOutput := logBuf.String()
|
||||
|
||||
// Should still log even when GPU is none
|
||||
if !strings.Contains(logOutput, "GPU detection completed") {
|
||||
t.Error("expected log output even for GPU type 'none'")
|
||||
}
|
||||
|
||||
// Should indicate config-based detection
|
||||
if result.Info.DetectionMethod != worker.DetectionSourceConfig {
|
||||
t.Errorf("detection method = %v, want %v", result.Info.DetectionMethod, worker.DetectionSourceConfig)
|
||||
}
|
||||
}
|
||||
140
tests/unit/security/manifest_filename_test.go
Normal file
140
tests/unit/security/manifest_filename_test.go
Normal file
|
|
@ -0,0 +1,140 @@
|
|||
package security
|
||||
|
||||
import (
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/jfraeys/fetch_ml/internal/manifest"
|
||||
)
|
||||
|
||||
// TestManifestFilenameNonce verifies that manifest filenames include a cryptographic nonce
|
||||
// to prevent information disclosure in multi-tenant environments where predictable
|
||||
// filenames could be enumerated.
|
||||
func TestManifestFilenameNonce(t *testing.T) {
|
||||
t.Run("FilenameIncludesNonce", func(t *testing.T) {
|
||||
// Generate multiple filenames and verify they are unique
|
||||
filenames := make(map[string]bool)
|
||||
for i := 0; i < 10; i++ {
|
||||
filename, err := manifest.GenerateManifestFilename()
|
||||
if err != nil {
|
||||
t.Fatalf("GenerateManifestFilename failed: %v", err)
|
||||
}
|
||||
|
||||
// Verify format: run_manifest_<nonce>.json
|
||||
if !strings.HasPrefix(filename, "run_manifest_") {
|
||||
t.Errorf("filename %q missing required prefix 'run_manifest_'", filename)
|
||||
}
|
||||
if !strings.HasSuffix(filename, ".json") {
|
||||
t.Errorf("filename %q missing required suffix '.json'", filename)
|
||||
}
|
||||
|
||||
// Extract and verify nonce
|
||||
nonce := manifest.ParseManifestFilename(filename)
|
||||
if nonce == "" {
|
||||
t.Errorf("failed to parse nonce from filename %q", filename)
|
||||
}
|
||||
if len(nonce) != 32 {
|
||||
t.Errorf("nonce length = %d, want 32 hex chars", len(nonce))
|
||||
}
|
||||
|
||||
// Verify uniqueness (no collisions in 10 generations)
|
||||
if filenames[filename] {
|
||||
t.Errorf("duplicate filename generated: %q", filename)
|
||||
}
|
||||
filenames[filename] = true
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("ManifestWrittenWithNonce", func(t *testing.T) {
|
||||
// Generate a nonce for the manifest
|
||||
nonce, err := manifest.GenerateManifestNonce()
|
||||
if err != nil {
|
||||
t.Fatalf("GenerateManifestNonce failed: %v", err)
|
||||
}
|
||||
|
||||
// Create a manifest with nonce in Environment
|
||||
created := time.Now().UTC()
|
||||
m := manifest.NewRunManifest("run-test-nonce", "task-nonce", "job-nonce", created)
|
||||
m.CommitID = "deadbeef"
|
||||
m.Environment = &manifest.ExecutionEnvironment{
|
||||
ConfigHash: "abc123",
|
||||
ManifestNonce: nonce,
|
||||
}
|
||||
|
||||
dir := t.TempDir()
|
||||
if err := m.WriteToDir(dir); err != nil {
|
||||
t.Fatalf("WriteToDir failed: %v", err)
|
||||
}
|
||||
|
||||
// List files in directory
|
||||
entries, err := os.ReadDir(dir)
|
||||
if err != nil {
|
||||
t.Fatalf("ReadDir failed: %v", err)
|
||||
}
|
||||
|
||||
// Find manifest file
|
||||
var manifestFile string
|
||||
for _, entry := range entries {
|
||||
if strings.HasPrefix(entry.Name(), "run_manifest_") && strings.HasSuffix(entry.Name(), ".json") {
|
||||
manifestFile = entry.Name()
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if manifestFile == "" {
|
||||
t.Fatal("no manifest file found with expected naming pattern")
|
||||
}
|
||||
|
||||
// Verify nonce is present in filename
|
||||
parsedNonce := manifest.ParseManifestFilename(manifestFile)
|
||||
if parsedNonce == "" {
|
||||
t.Errorf("manifest file %q does not contain a valid nonce", manifestFile)
|
||||
}
|
||||
if parsedNonce != nonce {
|
||||
t.Errorf("nonce mismatch: got %q, want %q", parsedNonce, nonce)
|
||||
}
|
||||
if len(parsedNonce) != 32 {
|
||||
t.Errorf("nonce length = %d, want 32 hex chars", len(parsedNonce))
|
||||
}
|
||||
|
||||
// Verify file can be loaded back
|
||||
loaded, err := manifest.LoadFromDir(dir)
|
||||
if err != nil {
|
||||
t.Fatalf("LoadFromDir failed: %v", err)
|
||||
}
|
||||
if loaded.RunID != m.RunID {
|
||||
t.Errorf("loaded RunID = %q, want %q", loaded.RunID, m.RunID)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("NonceUniqueness", func(t *testing.T) {
|
||||
// Generate many nonces to check for collisions (statistical test)
|
||||
nonces := make(map[string]int)
|
||||
iterations := 100
|
||||
|
||||
for i := 0; i < iterations; i++ {
|
||||
nonce, err := manifest.GenerateManifestNonce()
|
||||
if err != nil {
|
||||
t.Fatalf("GenerateManifestNonce failed: %v", err)
|
||||
}
|
||||
nonces[nonce]++
|
||||
}
|
||||
|
||||
// Check for any collisions
|
||||
collisions := 0
|
||||
for nonce, count := range nonces {
|
||||
if count > 1 {
|
||||
t.Errorf("nonce collision detected: %q appeared %d times", nonce, count)
|
||||
collisions++
|
||||
}
|
||||
}
|
||||
|
||||
if collisions > 0 {
|
||||
t.Fatalf("detected %d nonce collisions in %d iterations", collisions, iterations)
|
||||
}
|
||||
|
||||
t.Logf("Generated %d unique nonces with no collisions", iterations)
|
||||
})
|
||||
}
|
||||
183
tests/unit/security/resource_quota_test.go
Normal file
183
tests/unit/security/resource_quota_test.go
Normal file
|
|
@ -0,0 +1,183 @@
|
|||
package security
|
||||
|
||||
import (
|
||||
"os"
|
||||
"strconv"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestResourceEnvVarParsing verifies that resource environment variables
|
||||
// are correctly parsed and applied.
|
||||
func TestResourceEnvVarParsing(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
envCPU string
|
||||
envMemory string
|
||||
envGPUCount string
|
||||
wantCPUParsed int
|
||||
wantGPUParsed int
|
||||
}{
|
||||
{
|
||||
name: "valid env CPU",
|
||||
envCPU: "4",
|
||||
wantCPUParsed: 4,
|
||||
},
|
||||
{
|
||||
name: "env CPU equals max",
|
||||
envCPU: "8",
|
||||
wantCPUParsed: 8,
|
||||
},
|
||||
{
|
||||
name: "invalid env CPU falls back to 0",
|
||||
envCPU: "invalid",
|
||||
wantCPUParsed: 0,
|
||||
},
|
||||
{
|
||||
name: "negative env CPU parsed as-is",
|
||||
envCPU: "-1",
|
||||
wantCPUParsed: -1, // strconv.Atoi parses negative numbers correctly
|
||||
},
|
||||
{
|
||||
name: "valid env GPU count",
|
||||
envGPUCount: "2",
|
||||
wantGPUParsed: 2,
|
||||
},
|
||||
{
|
||||
name: "invalid env GPU count falls back to 0",
|
||||
envGPUCount: "invalid",
|
||||
wantGPUParsed: 0,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
// Set env vars
|
||||
if tt.envCPU != "" {
|
||||
t.Setenv("FETCH_ML_TOTAL_CPU", tt.envCPU)
|
||||
}
|
||||
if tt.envGPUCount != "" {
|
||||
t.Setenv("FETCH_ML_GPU_COUNT", tt.envGPUCount)
|
||||
}
|
||||
|
||||
// Parse env values (mimicking how worker/config.go does it)
|
||||
var cpuParsed int
|
||||
if v := os.Getenv("FETCH_ML_TOTAL_CPU"); v != "" {
|
||||
if n, err := strconv.Atoi(v); err == nil {
|
||||
cpuParsed = n
|
||||
}
|
||||
}
|
||||
|
||||
var gpuParsed int
|
||||
if v := os.Getenv("FETCH_ML_GPU_COUNT"); v != "" {
|
||||
if n, err := strconv.Atoi(v); err == nil {
|
||||
gpuParsed = n
|
||||
}
|
||||
}
|
||||
|
||||
// Verify parsing
|
||||
if tt.envCPU != "" && cpuParsed != tt.wantCPUParsed {
|
||||
t.Errorf("CPU parsed = %d, want %d", cpuParsed, tt.wantCPUParsed)
|
||||
}
|
||||
if tt.envGPUCount != "" && gpuParsed != tt.wantGPUParsed {
|
||||
t.Errorf("GPU count parsed = %d, want %d", gpuParsed, tt.wantGPUParsed)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestPodmanCPUParsing verifies podman_cpus parsing
|
||||
func TestPodmanCPUParsing(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
podmanCPUs string
|
||||
wantParsed int
|
||||
}{
|
||||
{
|
||||
name: "valid podman cpus",
|
||||
podmanCPUs: "2.5",
|
||||
wantParsed: 2,
|
||||
},
|
||||
{
|
||||
name: "podman cpus integer",
|
||||
podmanCPUs: "4.0",
|
||||
wantParsed: 4,
|
||||
},
|
||||
{
|
||||
name: "invalid podman cpus falls back",
|
||||
podmanCPUs: "invalid",
|
||||
wantParsed: 0, // Parse fails, falls back
|
||||
},
|
||||
{
|
||||
name: "negative podman cpus treated as 0",
|
||||
podmanCPUs: "-1.5",
|
||||
wantParsed: 0,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
// Parse podman_cpus like config does
|
||||
var parsedCPUs int
|
||||
if f, err := strconv.ParseFloat(tt.podmanCPUs, 64); err == nil {
|
||||
if f < 0 {
|
||||
parsedCPUs = 0
|
||||
} else {
|
||||
parsedCPUs = int(f)
|
||||
}
|
||||
}
|
||||
|
||||
if parsedCPUs != tt.wantParsed {
|
||||
t.Errorf("parsed CPUs = %d, want %d", parsedCPUs, tt.wantParsed)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestResourceEnvVarOverride verifies that env vars override config values
|
||||
func TestResourceEnvVarOverride(t *testing.T) {
|
||||
t.Run("env overrides config", func(t *testing.T) {
|
||||
t.Setenv("FETCH_ML_TOTAL_CPU", "8")
|
||||
|
||||
// Simulate config with lower value
|
||||
configCPU := 4
|
||||
|
||||
// Parse env override
|
||||
var envCPU int
|
||||
if v := os.Getenv("FETCH_ML_TOTAL_CPU"); v != "" {
|
||||
if n, err := strconv.Atoi(v); err == nil {
|
||||
envCPU = n
|
||||
}
|
||||
}
|
||||
|
||||
// Env should take precedence
|
||||
finalCPU := configCPU
|
||||
if envCPU > 0 {
|
||||
finalCPU = envCPU
|
||||
}
|
||||
|
||||
if finalCPU != 8 {
|
||||
t.Errorf("final CPU = %d, want 8 (env override)", finalCPU)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("empty env uses config", func(t *testing.T) {
|
||||
// No env var set
|
||||
configCPU := 4
|
||||
|
||||
var envCPU int
|
||||
if v := os.Getenv("FETCH_ML_TOTAL_CPU"); v != "" {
|
||||
if n, err := strconv.Atoi(v); err == nil {
|
||||
envCPU = n
|
||||
}
|
||||
}
|
||||
|
||||
finalCPU := configCPU
|
||||
if envCPU > 0 {
|
||||
finalCPU = envCPU
|
||||
}
|
||||
|
||||
if finalCPU != 4 {
|
||||
t.Errorf("final CPU = %d, want 4 (config value)", finalCPU)
|
||||
}
|
||||
})
|
||||
}
|
||||
Loading…
Reference in a new issue