# Security Fixes Design - Agent 122

**Date**: 2025-10-14
**Agent**: Agent 122
**Priority**: HIGH (Critical security issues blocking production deployment)

---

## Executive Summary

This document outlines the implementation design for three critical security vulnerabilities identified in Agent 108's security audit:

1. **SEC-001**: Missing checkpoint cryptographic signatures
2. **SEC-002**: No model poisoning detection
3. **SEC-003**: Insufficient prediction sanity checks

**Implementation Timeline**: 2 weeks
**Files Modified**: ~8 files
**New Files Created**: ~4 files
**Test Coverage**: 100% for security features

---

## Architecture Overview

```
┌─────────────────────────────────────────────────────────────┐
│                   Security Layer Architecture                │
└─────────────────────────────────────────────────────────────┘

┌─────────────────┐      ┌──────────────────┐      ┌─────────────────┐
│  Checkpoint     │      │   Prediction     │      │    Ensemble     │
│  Signature      │──────│   Validator      │──────│    Anomaly      │
│  System         │      │                  │      │    Detector     │
└────────┬────────┘      └────────┬─────────┘      └────────┬────────┘
         │                        │                         │
         ▼                        ▼                         ▼
┌─────────────────┐      ┌──────────────────┐      ┌─────────────────┐
│  Vault          │      │  Statistical     │      │  Temporal       │
│  Key Storage    │      │  Bounds Check    │      │  Pattern        │
│  (HMAC keys)    │      │  (Z-score)       │      │  Analysis       │
└─────────────────┘      └──────────────────┘      └─────────────────┘
         │                        │                         │
         └────────────────────────┴─────────────────────────┘
                                  │
                                  ▼
                    ┌──────────────────────────┐
                    │   Security Audit Logger   │
                    │  (ml_security_events)     │
                    └──────────────────────────┘
```

---

## 1. Checkpoint Signature System (SEC-001)

### Design Goals

- **Authenticity**: Verify checkpoint origin using HMAC-SHA256 signatures
- **Integrity**: Detect tampering via signature validation
- **Key Rotation**: Support quarterly key rotation with versioned keys
- **Performance**: <100μs signature overhead per checkpoint operation

### Implementation Components

#### 1.1 Extended CheckpointMetadata

```rust
pub struct CheckpointMetadata {
    // ... existing fields ...

    // Security fields (new)
    pub signature: Option<String>,           // HMAC-SHA256 signature (hex-encoded)
    pub signature_algorithm: String,         // "HMAC-SHA256"
    pub signing_key_id: String,              // Key rotation support (e.g., "2024-Q4")
    pub signed_at: Option<DateTime<Utc>>,    // Signature timestamp
}
```

#### 1.2 CheckpointSigner

**File**: `ml/src/checkpoint/signer.rs` (NEW)

```rust
pub struct CheckpointSigner {
    vault_client: Arc<VaultClient>,
    key_cache: RwLock<HashMap<String, Vec<u8>>>,  // key_id -> secret_key
    cache_ttl: Duration,  // 5 minutes
}

impl CheckpointSigner {
    /// Sign checkpoint data with HMAC-SHA256
    pub async fn sign_checkpoint(
        &self,
        data: &[u8],
        model_type: ModelType,
    ) -> Result<SignatureInfo, MLError>;

    /// Verify checkpoint signature
    pub async fn verify_signature(
        &self,
        data: &[u8],
        signature: &str,
        key_id: &str,
    ) -> Result<(), MLError>;

    /// Rotate signing keys (quarterly)
    pub async fn rotate_keys(&self) -> Result<(), MLError>;
}
```

#### 1.3 Vault Key Storage

**Key Path Structure**:
```
secret/foxhunt/ml/checkpoint-signing/
  ├── DQN/
  │   ├── 2024-Q4  (32-byte HMAC key)
  │   └── 2025-Q1  (next rotation)
  ├── PPO/
  ├── MAMBA-2/
  └── TFT/
```

**Key Generation**:
```bash
# Generate 256-bit HMAC key
openssl rand -hex 32 > /tmp/hmac_key.txt

# Store in Vault
vault kv put secret/foxhunt/ml/checkpoint-signing/DQN/2024-Q4 \
  key_value=@/tmp/hmac_key.txt
```

### Integration Points

1. **CheckpointManager::save_checkpoint()**: Sign before saving
2. **CheckpointManager::load_checkpoint()**: Verify after loading
3. **ValidationManager::validate_checkpoint()**: Add signature validation step

### Performance Impact

- **Signing**: ~50μs (HMAC-SHA256 computation)
- **Verification**: ~50μs
- **Key cache hit**: <1μs (no Vault call)
- **Key cache miss**: ~10ms (Vault HTTP request)

---

## 2. Prediction Validator (SEC-002)

### Design Goals

- **Outlier Detection**: Identify predictions >3 standard deviations from mean
- **Range Validation**: Enforce [-1.0, 1.0] bounds
- **Rate Limiting**: Detect >10% extreme predictions as poisoning signal
- **Performance**: <10μs validation overhead per prediction

### Implementation Components

#### 2.1 PredictionValidator

**File**: `ml/src/security/prediction_validator.rs` (NEW)

```rust
pub struct PredictionValidator {
    // Historical statistics (rolling window)
    prediction_stats: RwLock<PredictionStatistics>,

    // Configuration
    z_score_threshold: f64,           // Default: 3.0 (99.7% of normal distribution)
    confidence_threshold: f64,         // Default: 0.5
    extreme_prediction_threshold: f64, // Default: 0.9
    max_extreme_rate: f64,             // Default: 0.1 (10%)

    // Temporal tracking
    extreme_prediction_window: VecDeque<Instant>,
    window_duration: Duration,         // Default: 60 seconds
}

pub struct PredictionStatistics {
    mean: f64,
    std_dev: f64,
    sample_count: u64,
    last_updated: Instant,
}

pub struct ValidatedPrediction {
    pub value: f64,
    pub is_outlier: bool,
    pub z_score: f64,
    pub should_override: bool,  // Use ensemble fallback
    pub validation_flags: Vec<ValidationFlag>,
}

pub enum ValidationFlag {
    OutOfBounds { value: f64 },
    Outlier { z_score: f64 },
    LowConfidence { confidence: f64 },
    ExtremeRateLimit { rate: f64 },
}
```

#### 2.2 Validation Logic

```rust
impl PredictionValidator {
    /// Validate prediction with multi-layer checks
    pub fn validate(
        &self,
        prediction: f64,
        confidence: f64,
        model_id: &str,
    ) -> Result<ValidatedPrediction, MLError> {
        let mut flags = Vec::new();

        // Layer 1: Range check
        if prediction < -1.0 || prediction > 1.0 {
            flags.push(ValidationFlag::OutOfBounds { value: prediction });
            return Err(MLError::ValidationError { ... });
        }

        // Layer 2: Z-score outlier detection
        let stats = self.prediction_stats.read().unwrap();
        let z_score = (prediction - stats.mean) / stats.std_dev;

        if z_score.abs() > self.z_score_threshold {
            flags.push(ValidationFlag::Outlier { z_score });
            // Log security event
            self.log_security_event(SecurityEvent::OutlierDetected { ... });
        }

        // Layer 3: Confidence sanity check
        if confidence < self.confidence_threshold {
            flags.push(ValidationFlag::LowConfidence { confidence });
        }

        // Layer 4: Rate limit extreme predictions
        if prediction.abs() > self.extreme_prediction_threshold {
            let extreme_rate = self.calculate_extreme_rate();
            if extreme_rate > self.max_extreme_rate {
                flags.push(ValidationFlag::ExtremeRateLimit { rate: extreme_rate });
                return Err(MLError::ValidationError { ... });
            }
        }

        Ok(ValidatedPrediction {
            value: prediction,
            is_outlier: !flags.is_empty(),
            z_score,
            should_override: z_score.abs() > self.z_score_threshold,
            validation_flags: flags,
        })
    }

    /// Update statistics with new prediction (exponential moving average)
    pub fn update_statistics(&self, prediction: f64) {
        let mut stats = self.prediction_stats.write().unwrap();

        // Exponential moving average
        let alpha = 0.05; // Weight for new samples
        stats.mean = (1.0 - alpha) * stats.mean + alpha * prediction;

        // Update standard deviation (Welford's online algorithm)
        stats.sample_count += 1;
        // ... (standard deviation update logic)
    }
}
```

### Integration Points

1. **InferenceEngine::process_onnx_inference()**: Validate before returning result
2. **EnsembleCoordinator::aggregate_predictions()**: Validate before aggregation
3. **Database logging**: Store `is_outlier` flag in `ensemble_predictions.metadata`

### Statistical Initialization

**Bootstrap Phase** (first 1000 predictions):
- Use conservative defaults: mean=0.0, std_dev=0.3
- Gradually refine statistics as data accumulates
- After 1000 samples, switch to exponential moving average

---

## 3. Ensemble Anomaly Detector (SEC-003)

### Design Goals

- **Temporal Pattern Detection**: Identify sudden signal shifts (>50%)
- **Coordinated Attack Detection**: Flag when >80% models predict extreme values
- **Model Drift Detection**: Track per-model behavioral changes
- **Performance**: <20μs overhead per ensemble decision

### Implementation Components

#### 3.1 EnsembleAnomalyDetector

**File**: `ml/src/security/anomaly_detector.rs` (NEW)

```rust
pub struct EnsembleAnomalyDetector {
    // Historical ensemble signal distribution
    signal_history: RwLock<VecDeque<f64>>,
    window_size: usize,  // Default: 100

    // Per-model tracking
    model_signal_history: RwLock<HashMap<String, VecDeque<f64>>>,
    model_window_size: usize,  // Default: 50

    // Anomaly thresholds
    sudden_shift_threshold: f64,        // Default: 0.5 (50% signal change)
    coordinated_attack_threshold: f64,  // Default: 0.8 (80% models agree on extreme)
    drift_threshold: f64,               // Default: 0.7
}

pub struct AnomalyReport {
    pub has_anomalies: bool,
    pub anomalies: Vec<Anomaly>,
    pub timestamp: DateTime<Utc>,
    pub severity: AnomalySeverity,
}

pub enum Anomaly {
    SuddenShift {
        previous: f64,
        current: f64,
        magnitude: f64,
    },
    CoordinatedAttack {
        extreme_ratio: f64,
        affected_models: Vec<String>,
    },
    ModelDrift {
        model_id: String,
        historical_mean: f64,
        current_signal: f64,
        drift_magnitude: f64,
    },
}

pub enum AnomalySeverity {
    Low,      // Single outlier
    Medium,   // Multiple outliers or moderate drift
    High,     // Coordinated attack or extreme shift
    Critical, // System-wide compromise suspected
}
```

#### 3.2 Detection Logic

```rust
impl EnsembleAnomalyDetector {
    /// Detect anomalies in ensemble prediction
    pub fn detect_anomaly(
        &self,
        decision: &EnsembleDecision,
    ) -> AnomalyReport {
        let mut anomalies = Vec::new();

        // Detection 1: Sudden signal shift
        let signal_history = self.signal_history.read().unwrap();
        if let Some(prev_signal) = signal_history.back() {
            let shift = (decision.signal - prev_signal).abs();
            if shift > self.sudden_shift_threshold {
                anomalies.push(Anomaly::SuddenShift {
                    previous: *prev_signal,
                    current: decision.signal,
                    magnitude: shift,
                });
            }
        }

        // Detection 2: Coordinated attack (all models extreme)
        let extreme_count = decision.model_votes.values()
            .filter(|v| v.signal.abs() > 0.9)
            .count();
        let extreme_ratio = extreme_count as f64 / decision.model_votes.len() as f64;

        if extreme_ratio > self.coordinated_attack_threshold {
            anomalies.push(Anomaly::CoordinatedAttack {
                extreme_ratio,
                affected_models: decision.model_votes.keys().cloned().collect(),
            });
        }

        // Detection 3: Per-model behavioral drift
        let model_history = self.model_signal_history.read().unwrap();
        for (model_id, vote) in &decision.model_votes {
            if let Some(history) = model_history.get(model_id) {
                let mean = history.iter().sum::<f64>() / history.len() as f64;
                let drift = (vote.signal - mean).abs();

                if drift > self.drift_threshold {
                    anomalies.push(Anomaly::ModelDrift {
                        model_id: model_id.clone(),
                        historical_mean: mean,
                        current_signal: vote.signal,
                        drift_magnitude: drift,
                    });
                }
            }
        }

        // Determine severity
        let severity = self.calculate_severity(&anomalies);

        // Log to security audit system
        if severity >= AnomalySeverity::High {
            self.log_security_event(SecurityEvent::EnsembleAnomaly { ... });
        }

        AnomalyReport {
            has_anomalies: !anomalies.is_empty(),
            anomalies,
            timestamp: Utc::now(),
            severity,
        }
    }

    /// Update history with new decision
    pub fn update_history(&self, decision: &EnsembleDecision) {
        // Update ensemble signal history
        let mut signal_history = self.signal_history.write().unwrap();
        signal_history.push_back(decision.signal);
        if signal_history.len() > self.window_size {
            signal_history.pop_front();
        }

        // Update per-model history
        let mut model_history = self.model_signal_history.write().unwrap();
        for (model_id, vote) in &decision.model_votes {
            model_history
                .entry(model_id.clone())
                .or_insert_with(VecDeque::new)
                .push_back(vote.signal);

            let history = model_history.get_mut(model_id).unwrap();
            if history.len() > self.model_window_size {
                history.pop_front();
            }
        }
    }
}
```

### Integration Points

1. **EnsembleCoordinator::aggregate_predictions()**: Detect anomalies after aggregation
2. **Automated rollback**: Trigger model rollback if severity is Critical
3. **Monitoring**: Export anomaly metrics to Prometheus

---

## 4. Security Event Logging

### Database Schema

**File**: `migrations/024_ml_security_events.sql` (NEW)

```sql
-- ML security events table
CREATE TABLE ml_security_events (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    timestamp TIMESTAMPTZ NOT NULL DEFAULT NOW(),

    -- Event classification
    event_type VARCHAR(50) NOT NULL,  -- signature_failure, outlier_detected, etc.
    severity VARCHAR(20) NOT NULL,    -- low, medium, high, critical

    -- Context
    model_id VARCHAR(50),
    checkpoint_id VARCHAR(255),
    prediction_id UUID REFERENCES ensemble_predictions(id),

    -- Event details
    description TEXT NOT NULL,
    metadata JSONB,

    -- Response
    action_taken VARCHAR(100),  -- rejected, flagged, alerted, rollback

    CONSTRAINT chk_severity CHECK (severity IN ('low', 'medium', 'high', 'critical'))
);

-- Indexes for fast querying
CREATE INDEX idx_ml_security_events_timestamp ON ml_security_events (timestamp DESC);
CREATE INDEX idx_ml_security_events_severity ON ml_security_events (severity)
    WHERE severity IN ('high', 'critical');
CREATE INDEX idx_ml_security_events_type ON ml_security_events (event_type);
CREATE INDEX idx_ml_security_events_model ON ml_security_events (model_id)
    WHERE model_id IS NOT NULL;

-- TimescaleDB hypertable for time-series data
SELECT create_hypertable('ml_security_events', 'timestamp');

-- Retention policy: Keep high/critical events for 1 year, others for 90 days
SELECT add_retention_policy('ml_security_events', INTERVAL '90 days');
```

### Event Types

```rust
pub enum SecurityEventType {
    // Checkpoint security
    CheckpointSignatureFailure,
    CheckpointSignatureMissing,
    CheckpointTamperingDetected,

    // Prediction validation
    PredictionOutlierDetected,
    PredictionOutOfBounds,
    ExtremeRateExceeded,

    // Ensemble anomalies
    EnsembleSuddenShift,
    CoordinatedAttackSuspected,
    ModelBehavioralDrift,

    // System events
    AutomaticRollback,
    ManualIntervention,
}
```

---

## 5. Testing Strategy

### Unit Tests

**File**: `ml/src/checkpoint/signer_tests.rs`
- Test HMAC signature generation
- Test signature verification (valid/invalid)
- Test key rotation
- Test Vault integration (mocked)

**File**: `ml/src/security/prediction_validator_tests.rs`
- Test range validation
- Test Z-score outlier detection
- Test extreme rate limiting
- Test statistics update

**File**: `ml/src/security/anomaly_detector_tests.rs`
- Test sudden shift detection
- Test coordinated attack detection
- Test model drift detection
- Test severity calculation

### Integration Tests

**File**: `ml/tests/security_integration_test.rs`
- End-to-end checkpoint signing workflow
- Prediction validation with database logging
- Ensemble anomaly detection with rollback
- Security event logging

### Adversarial Tests

**File**: `ml/tests/adversarial_tests.rs`
- Tampered checkpoint rejection
- Poisoned model detection (synthetic adversarial predictions)
- Coordinated attack simulation
- Performance under attack

---

## 6. Performance Benchmarks

### Target Performance

| Operation | Target | Expected |
|-----------|--------|----------|
| Checkpoint signing | <100μs | ~50μs |
| Checkpoint verification | <100μs | ~50μs |
| Prediction validation | <10μs | ~5μs |
| Anomaly detection | <20μs | ~15μs |
| Security event logging | <500μs | ~300μs (async) |

### Benchmark Suite

**File**: `ml/benches/security_benchmarks.rs`
```rust
#[bench]
fn bench_checkpoint_signing(b: &mut Bencher) { ... }

#[bench]
fn bench_prediction_validation(b: &mut Bencher) { ... }

#[bench]
fn bench_anomaly_detection(b: &mut Bencher) { ... }
```

---

## 7. Monitoring & Alerting

### Prometheus Metrics

```rust
// Checkpoint security metrics
checkpoint_signature_failures_total
checkpoint_signature_verification_duration_seconds

// Prediction validation metrics
prediction_outliers_total
prediction_out_of_bounds_total
prediction_extreme_rate

// Ensemble anomaly metrics
ensemble_anomalies_total{severity="high|critical"}
ensemble_sudden_shifts_total
model_drift_events_total
```

### Alerting Rules

**File**: `monitoring/prometheus/alerts/ml_security_alerts.yml`

```yaml
groups:
  - name: ml_security
    interval: 30s
    rules:
      - alert: CheckpointSignatureFailure
        expr: rate(checkpoint_signature_failures_total[5m]) > 0
        for: 1m
        severity: critical
        annotations:
          summary: "Checkpoint signature verification failed"

      - alert: HighPredictionOutlierRate
        expr: prediction_extreme_rate > 0.1
        for: 5m
        severity: high
        annotations:
          summary: "More than 10% of predictions are outliers - possible model poisoning"

      - alert: EnsembleCoordinatedAttack
        expr: ensemble_anomalies_total{severity="critical"} > 0
        for: 1m
        severity: critical
        annotations:
          summary: "Coordinated attack detected - multiple models producing extreme predictions"
```

---

## 8. Implementation Timeline

### Week 1: Core Implementation

**Day 1-2**: Checkpoint signature system
- Extend CheckpointMetadata
- Implement CheckpointSigner
- Vault integration for key management
- Unit tests

**Day 3-4**: Prediction validator
- Implement PredictionValidator
- Statistical bounds checking
- Rate limiting logic
- Unit tests

**Day 5**: Ensemble anomaly detector
- Implement EnsembleAnomalyDetector
- Temporal pattern detection
- Unit tests

### Week 2: Integration & Testing

**Day 6-7**: Integration
- Integrate signer into CheckpointManager
- Integrate validator into InferenceEngine
- Integrate detector into EnsembleCoordinator

**Day 8**: Security event logging
- Database migration
- Logger implementation
- Async logging integration

**Day 9**: Testing
- Integration tests
- Adversarial tests
- E2E security tests

**Day 10**: Documentation & deployment
- Update security documentation
- Performance benchmarks
- Monitoring setup
- Final validation

---

## 9. Rollout Plan

### Phase 1: Staging Deployment (Day 11-12)
1. Deploy to staging environment
2. Run 24-hour security validation
3. Monitor false positive rate
4. Tune thresholds

### Phase 2: Production Deployment (Day 13-14)
1. Enable checkpoint signing (all new checkpoints)
2. Enable prediction validation (monitoring only)
3. Enable anomaly detection (alerting enabled)
4. Monitor for 7 days

### Phase 3: Full Enforcement (Day 21)
1. Reject unsigned checkpoints
2. Reject invalid predictions
3. Automatic rollback on critical anomalies

---

## 10. Success Criteria

- ✅ All critical security issues (SEC-001, SEC-002, SEC-003) resolved
- ✅ 100% test coverage for security features
- ✅ Performance targets met (<100μs overhead)
- ✅ Zero false positives in 7-day production trial
- ✅ Monitoring and alerting operational
- ✅ Security event logging complete
- ✅ Documentation updated

---

## 11. Future Enhancements

### Q1 2026
- Add Ed25519 public key signatures (stronger than HMAC)
- Implement checkpoint certificate chains
- Add hardware security module (HSM) support

### Q2 2026
- Machine learning-based anomaly detection (LSTM autoencoder)
- Adversarial training for robust models
- Automatic retraining on poisoning detection

---

**Design Status**: ✅ Complete
**Next Step**: Begin implementation (Week 1, Day 1)
**Owner**: Agent 122
**Reviewers**: Security Team, ML Team Lead