aiegis_harness/

main.rs

// SPDX-License-Identifier: Apache-2.0 OR MIT
//! AiEGIS Harness reference daemon — CLI + HTTP server.
//!
//! Ports the HTTP surface of `harness.py` (`_HarnessHandler` + `main`) onto
//! axum + tokio. Endpoints + response shapes are byte-compatible with the
//! Python reference; an agent integrating against either binary can swap
//! `--harness-url` without code change.

mod pack_fetcher;
mod upstream;

use axum::{
    extract::State,
    http::{HeaderMap, StatusCode},
    response::{IntoResponse, Response},
    routing::{get, post},
    Json, Router,
};
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
use clap::Parser;
use harness_core::{
    audit::{AuditEntry, AuditSink, SqliteAuditLog},
    eval::evaluate_packs,
    load_pack, PolicyPack, DAEMON_VERSION, DEFAULT_PACK_ISSUER_PUBKEY_HEX, SPEC_VERSION,
};
use pack_fetcher::RemotePack;
use serde_json::{json, Value};
use std::collections::{HashMap, VecDeque};
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use std::time::{Instant, SystemTime, UNIX_EPOCH};
use upstream::UpstreamClient;

/// CLI arguments. Mirror of the Python ref's `argparse` setup.
#[derive(Parser, Debug)]
#[command(
    name = "aiegis-harness",
    about = "AiEGIS Harness reference daemon (L2 / proxy mode).",
    version
)]
struct Args {
    /// Local HTTP port to listen on (default 8080).
    #[arg(long, default_value_t = 8080)]
    port: u16,

    /// Path to a policy-pack JSON. Pass multiple times for multiple packs.
    #[arg(long = "policy-pack", value_name = "PATH")]
    policy_pack: Vec<PathBuf>,

    /// SQLite file for the append-only audit log.
    /// Day 1: accepted for CLI compatibility; the in-memory sink is used until
    /// the rusqlite backend lands in Day 2.
    #[arg(long, default_value = "harness_audit.db")]
    audit_db: PathBuf,

    /// (Reserved for v0.2) per-agent actions/minute cap.
    #[arg(long, default_value_t = 60)]
    rate_cap: u32,

    /// Verbose logging (DEBUG level).
    #[arg(long, short)]
    verbose: bool,

    /// Remote policy-pack source (directory URL containing index.json).
    /// When set, the daemon fetches Nel-style signed packs from
    /// `<URL>/index.json` on startup, verifies the Ed25519 signature over
    /// sha256(tarball) against `--issuer-pubkey-hex`, and lists them in
    /// /health. Loaded packs are INVENTORY-ONLY (visible + sig-verified
    /// + cached) — they are NOT executed by the daemon's evaluator
    /// because Nel publishes full OPA Rego v1, which the reference
    /// subset evaluator does not interpret. See pack_fetcher.rs for the
    /// needs-iteration note. Local --policy-pack JSON files still drive
    /// the evaluator.
    #[arg(long, value_name = "URL")]
    pack_source: Option<String>,

    /// Issuer Ed25519 public key (hex-encoded 32-byte raw) used to verify
    /// remotely-fetched packs. Defaults to the live aiegis.ie harness
    /// issuer key empirically captured on 2026-05-25.
    #[arg(long, default_value = DEFAULT_PACK_ISSUER_PUBKEY_HEX)]
    issuer_pubkey_hex: String,

    /// Upstream `/v1/harness/receipt` URL (Velo-style). When set + the
    /// local decision is ALLOW/WARN, the daemon POSTs the same payload
    /// upstream, captures the returned `rid`, and surfaces it in both
    /// the response body and the audit ledger. If upstream is
    /// unreachable / 5xx, the daemon honestly falls back to its local
    /// decision and flags `upstream_error: true` in the response.
    #[arg(long, value_name = "URL")]
    upstream_protect: Option<String>,
}

/// Per-agent sliding-window action timestamps for the rate-limit pack.
#[derive(Default)]
struct RateState {
    buckets: Mutex<HashMap<String, VecDeque<f64>>>,
}

impl RateState {
    fn observe(&self, agent_did: &str) -> usize {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map(|d| d.as_secs_f64())
            .unwrap_or(0.0);
        let mut buckets = self.buckets.lock().unwrap();
        let bucket = buckets.entry(agent_did.to_string()).or_default();
        while let Some(&front) = bucket.front() {
            if front < now - 60.0 {
                bucket.pop_front();
            } else {
                break;
            }
        }
        bucket.push_back(now);
        bucket.len()
    }
}

#[derive(Clone)]
struct AppState {
    packs: Arc<Vec<PolicyPack>>,
    remote_packs: Arc<Vec<RemotePack>>,
    pack_source_fingerprint: Arc<String>,
    audit: Arc<dyn AuditSink>,
    rate: Arc<RateState>,
    upstream: Option<UpstreamClient>,
    receipt_counter: Arc<std::sync::atomic::AtomicU64>,
}

#[tokio::main]
async fn main() -> std::process::ExitCode {
    let args = Args::parse();

    let level = if args.verbose { "debug" } else { "info" };
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new(level)),
        )
        .with_target(false)
        .init();

    if args.policy_pack.is_empty() && args.pack_source.is_none() {
        tracing::warn!("No --policy-pack or --pack-source supplied; daemon will ALLOW everything.");
    }

    // Remote pack fetch (inventory-only — see pack_fetcher module docs).
    let remote_packs: Vec<RemotePack> = if let Some(base) = args.pack_source.as_deref() {
        match pack_fetcher::fetch_all(base, &args.issuer_pubkey_hex).await {
            Ok(v) => {
                tracing::info!(
                    "pack_source={} verified_packs={} (inventory-only; reference daemon does not execute OPA Rego v1)",
                    base,
                    v.len()
                );
                v
            }
            Err(e) => {
                tracing::error!("pack_source_failed url={} err={}", base, e);
                return std::process::ExitCode::from(2);
            }
        }
    } else {
        Vec::new()
    };
    let pack_source_fingerprint = pack_fetcher::fingerprint(&remote_packs);

    // PACK-ID DEDUP (Audit-B SEV-LOW closure 2026-05-25): refuse to load two
    // packs with the same pack_id. Spec's rollback-protection design depends
    // on uniqueness; without dedup, second-loaded-of-same-id silently joins
    // the eval set and pack_ids returned by /health are ambiguous.
    let mut packs: Vec<PolicyPack> = Vec::new();
    let mut seen_pack_ids: std::collections::HashSet<String> = std::collections::HashSet::new();
    for path in &args.policy_pack {
        match load_pack(path) {
            Ok(pack) => {
                if !seen_pack_ids.insert(pack.pack_id.clone()) {
                    tracing::error!(
                        "pack_id_collision: refusing to load duplicate pack_id={} from {} (already loaded earlier)",
                        pack.pack_id,
                        path.display()
                    );
                    return std::process::ExitCode::from(2);
                }
                if pack
                    .signature
                    .value
                    .as_deref()
                    .map(|v| v == "REFERENCE_PACK_NOT_SIGNED_LOCAL_DEV_ONLY")
                    .unwrap_or(false)
                {
                    tracing::warn!(
                        "pack={} ships UNSIGNED — local dev only; prod must sign + verify against issuer JWKS",
                        pack.pack_id
                    );
                }
                tracing::info!(
                    "loaded pack={} rules={}",
                    pack.pack_id,
                    pack.rules.len()
                );
                packs.push(pack);
            }
            Err(e) => {
                tracing::error!("pack_load_failed path={} err={}", path.display(), e);
                return std::process::ExitCode::from(2);
            }
        }
    }

    // SQLite-backed audit ledger with append-only triggers (BEFORE DELETE +
    // BEFORE UPDATE both RAISE(ABORT)). Same DDL as the Python reference.
    let audit: Arc<dyn AuditSink> = match SqliteAuditLog::open(&args.audit_db) {
        Ok(s) => {
            tracing::info!("audit ledger ready at {}", args.audit_db.display());
            Arc::new(s)
        }
        Err(e) => {
            tracing::error!(
                "audit_db_open_failed path={} err={}",
                args.audit_db.display(),
                e
            );
            return std::process::ExitCode::from(2);
        }
    };

    let upstream = args.upstream_protect.as_ref().map(|u| {
        tracing::info!("upstream_protect_url={}", u);
        UpstreamClient::new(u.clone())
    });

    let state = AppState {
        packs: Arc::new(packs),
        remote_packs: Arc::new(remote_packs),
        pack_source_fingerprint: Arc::new(pack_source_fingerprint),
        audit,
        rate: Arc::new(RateState::default()),
        upstream,
        receipt_counter: Arc::new(std::sync::atomic::AtomicU64::new(0)),
    };

    let app = Router::new()
        .route("/health", get(health))
        .route("/api/protect/health", get(health))
        .route("/api/protect", post(protect))
        .fallback(not_found)
        .with_state(state);

    let addr = format!("127.0.0.1:{}", args.port);
    let listener = match tokio::net::TcpListener::bind(&addr).await {
        Ok(l) => l,
        Err(e) => {
            tracing::error!("bind failed on {}: {}", addr, e);
            return std::process::ExitCode::from(2);
        }
    };
    tracing::info!(
        "AiEGIS Harness daemon listening on http://{}  (POST /api/protect)",
        addr
    );

    // Ctrl-C support so the smoke test's kill is clean.
    let shutdown = async {
        let _ = tokio::signal::ctrl_c().await;
        tracing::info!("shutting down");
    };

    if let Err(e) = axum::serve(listener, app).with_graceful_shutdown(shutdown).await {
        tracing::error!("server error: {}", e);
        return std::process::ExitCode::from(1);
    }
    std::process::ExitCode::SUCCESS
}

async fn health(State(state): State<AppState>) -> impl IntoResponse {
    let receipt_count = state
        .receipt_counter
        .load(std::sync::atomic::Ordering::Relaxed);
    let body = json!({
        "status": "ready",
        "version": DAEMON_VERSION,
        "spec_version": SPEC_VERSION,
        "packs_loaded": state.packs.len(),
        "pack_ids": state.packs.iter().map(|p| &p.pack_id).collect::<Vec<_>>(),
        "remote_packs_loaded": state.remote_packs.len(),
        "remote_pack_ids": state.remote_packs.iter()
            .map(|p| format!("{}@{}", p.name, p.version)).collect::<Vec<_>>(),
        "pack_source_fingerprint": state.pack_source_fingerprint.as_str(),
        "remote_packs_inventory_only": true,
        "upstream_protect_configured": state.upstream.is_some(),
        "receipts_minted": receipt_count,
    });
    (StatusCode::OK, json_with_version_header(body))
}

async fn not_found() -> impl IntoResponse {
    (
        StatusCode::NOT_FOUND,
        json_with_version_header(json!({"error": "not_found"})),
    )
}

async fn protect(
    State(state): State<AppState>,
    headers: HeaderMap,
    body: axum::body::Bytes,
) -> Response {
    let t0 = Instant::now();

    let payload: Value = if body.is_empty() {
        json!({})
    } else {
        match serde_json::from_slice::<Value>(&body) {
            Ok(v) => v,
            Err(_) => {
                let err_body = json!({
                    "decision": "DENY",
                    "reason": "malformed_json",
                    "layer": "L?",
                    "decision_ms": 0,
                });
                return (StatusCode::BAD_REQUEST, json_with_version_header(err_body))
                    .into_response();
            }
        }
    };

    let agent_tag = headers
        .get("X-AEGIS-Tag")
        .and_then(|h| h.to_str().ok())
        .unwrap_or("")
        .to_string();

    // Lift agent_did from JWT-sub if a 3-part dot token is present, without
    // verifying the signature (matches the Python reference's behaviour).
    let agent_did = extract_agent_did(&agent_tag);

    // Sliding-window counter for rate-limit pack.
    let actions = state.rate.observe(agent_did.as_deref().unwrap_or("anonymous"));
    let daemon_state = json!({"agent_actions_last_minute": actions});

    let result = evaluate_packs(&state.packs, &payload, &daemon_state);

    // Upstream receipt mint — only when --upstream-protect is configured
    // AND the local decision is non-DENY (we never invite an upstream to
    // override our DENY; per spec, DENY is terminal at the daemon).
    let (receipt_id, upstream_error) = if let Some(client) = &state.upstream {
        if result.decision != "DENY" {
            match client.post(&payload, &agent_tag).await {
                Ok(rcp) => {
                    state
                        .receipt_counter
                        .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                    (Some(rcp.rid), false)
                }
                Err(e) => {
                    tracing::warn!("upstream_post_failed err={}", e);
                    (None, true)
                }
            }
        } else {
            (None, false)
        }
    } else {
        (None, false)
    };

    let decision_ms = t0.elapsed().as_millis() as u64;

    let entry = AuditEntry::build_with_receipt(
        agent_did.clone(),
        payload.get("action").and_then(|v| v.as_str()).map(String::from),
        payload.get("target").and_then(|v| v.as_str()).map(String::from),
        result.decision.to_string(),
        result.deciding_pack.clone(),
        result.deciding_rule.clone(),
        result.deciding_layer.clone(),
        result.reason.clone(),
        decision_ms,
        receipt_id.clone(),
        upstream_error,
    );
    // Fail-closed on ledger write failure: a caller seeing ALLOW with no
    // ledger row is un-governable. Switch to DENY 503 if the ledger refuses
    // the row (schema-tamper, file unlinked, disk full, trigger-bypass).
    // Audit Agent D 2026-05-25.
    if let Err(ledger_err) = state.audit.try_append(entry) {
        tracing::error!("ledger_append_failed_failing_closed: {}", ledger_err);
        let resp = json!({
            "decision": "DENY",
            "reason": "ledger_unavailable",
            "layer": "L0-AuditLedger",
            "deciding_pack": null,
            "deciding_rule": null,
            "decision_ms": decision_ms,
            "agent_did": agent_did,
            "receipt_id": null,
            "upstream_error": upstream_error,
            "ledger_error": ledger_err,
        });
        return (StatusCode::SERVICE_UNAVAILABLE, json_with_version_header(resp))
            .into_response();
    }

    let status = if result.decision == "ALLOW" {
        StatusCode::OK
    } else {
        StatusCode::UNAUTHORIZED
    };
    let resp = json!({
        "decision": result.decision,
        "reason": result.reason,
        "layer": result.deciding_layer,
        "deciding_pack": result.deciding_pack,
        "deciding_rule": result.deciding_rule,
        "decision_ms": decision_ms,
        "agent_did": agent_did,
        "receipt_id": receipt_id,
        "upstream_error": upstream_error,
    });
    (status, json_with_version_header(resp)).into_response()
}

fn extract_agent_did(tag: &str) -> Option<String> {
    // RAV audit 2026-05-25: bound the inputs we touch from an untrusted
    // X-AEGIS-Tag header.
    //   - reject obviously-absurd-length tags (default JWTs are <2 KB; we
    //     allow 8 KB headroom). This prevents an attacker from spamming
    //     multi-MB X-AEGIS-Tag values that allocate a base64 buffer per
    //     request.
    //   - cap the returned `agent_did` (sub claim) at 256 chars. The
    //     value is used as a HashMap key in `RateState::buckets` AND
    //     persisted to the audit DB; unbounded values let an
    //     unauthenticated attacker grow daemon RSS by spamming requests
    //     with distinct large `sub` values.
    const MAX_TAG_BYTES: usize = 8 * 1024;
    const MAX_DID_CHARS: usize = 256;
    if tag.len() > MAX_TAG_BYTES {
        return None;
    }
    if tag.matches('.').count() != 2 {
        return None;
    }
    let parts: Vec<&str> = tag.split('.').collect();
    let mut payload_seg = parts[1].to_string();
    while payload_seg.len() % 4 != 0 {
        payload_seg.push('=');
    }
    // URL-safe base64, accept both padded and unpadded forms.
    let decoded = URL_SAFE_NO_PAD
        .decode(payload_seg.trim_end_matches('='))
        .ok()?;
    let claims: Value = serde_json::from_slice(&decoded).ok()?;
    let sub = claims.get("sub").and_then(|v| v.as_str())?;
    if sub.chars().count() > MAX_DID_CHARS {
        return None;
    }
    Some(sub.to_string())
}

/// Helper: render JSON with the standard X-AEGIS-Harness-Version header.
fn json_with_version_header(body: Value) -> Response {
    let mut resp = Json(body).into_response();
    if let Ok(val) = axum::http::HeaderValue::from_str(DAEMON_VERSION) {
        resp.headers_mut().insert("X-AEGIS-Harness-Version", val);
    }
    resp
}