Building Custom Guards

use chio_wasm_guards::{
    host::create_shared_engine,
    runtime::guard::WasmGuard,
    runtime::wasmtime_backend::{load_guards_from_policy, PolicyCustomGuards},
};
use serde_json::{json, Value};
use std::collections::HashMap;

pub struct Prepared {
    pub guards: Vec<WasmGuard>,
    pub context: Value,
    pub policy_bytes: Vec<u8>,
    pub module_bytes: Vec<u8>,
}

pub fn load() -> anyhow::Result<Prepared> {
    let policy_bytes = chio_example_runtime::files::read_bounded(
        std::path::Path::new("guard-policy.json"),
        65_536,
    )?;
    let value: Value = serde_json::from_slice(&policy_bytes)?;
    anyhow::ensure!(
        value
            .as_object()
            .is_some_and(|v| v.len() == 1 && v.contains_key("modules")),
        "Expected only a modules field"
    );
    let module_values = value["modules"]
        .as_array()
        .ok_or_else(|| anyhow::anyhow!("modules must be an array"))?;
    for entry in module_values {
        let fields = entry
            .as_object()
            .ok_or_else(|| anyhow::anyhow!("Each module must be an object"))?;
        anyhow::ensure!(
            fields.keys().all(|key| [
                "name",
                "version",
                "module_path",
                "capabilities",
                "config",
                "fuel_limit",
                "priority",
                "advisory",
                "signer_public_key",
                "allow_unsigned"
            ]
            .contains(&key.as_str())),
            "Unknown guard module field"
        );
    }
    let policy: PolicyCustomGuards = serde_json::from_value(value)?;
    anyhow::ensure!(
        policy.modules.len() == 1,
        "This application loads one word-limit guard"
    );
    let spec = &policy.modules[0];
    anyhow::ensure!(
        spec.name == "word-limit" && !spec.advisory,
        "The word-limit guard must be blocking"
    );
    anyhow::ensure!(
        spec.config
            .as_object()
            .is_some_and(|v| v.len() == 1 && v.contains_key("max_words")),
        "Expected only max_words in guard config"
    );
    let path = spec
        .module
        .path()
        .ok_or_else(|| anyhow::anyhow!("Supply module_path"))?;
    let module_bytes =
        chio_example_runtime::files::read_bounded(std::path::Path::new(path), 4 * 1024 * 1024)?;
    let module_sha256 = chio_core::sha256_hex(&module_bytes);
    let handles = load_guards_from_policy(
        &policy,
        &HashMap::<String, String>::new(),
        &["chio.get_config".into()],
        create_shared_engine()?,
    )?;
    anyhow::ensure!(
        handles[0].guard().manifest_sha256().as_deref() == Some(module_sha256.as_str()),
        "Module changed while loading; retry with a stable file"
    );
    let context = json!({"policy": policy, "module_sha256": module_sha256});
    Ok(Prepared {
        guards: handles.into_iter().map(|h| h.into_guard()).collect(),
        context,
        policy_bytes,
        module_bytes,
    })
}

use chio_core::{
    capability::{
        scope::{ChioScope, Operation, ToolGrant},
        token::CapabilityToken,
    },
    crypto::Keypair,
};
use chio_kernel::{
    ChioKernel, KernelConfig, DEFAULT_CHECKPOINT_BATCH_SIZE, DEFAULT_MAX_STREAM_DURATION_SECS,
    DEFAULT_MAX_STREAM_TOTAL_BYTES,
};
use std::path::Path;

pub struct Host {
    pub kernel: ChioKernel,
    pub capability: CapabilityToken,
}

pub fn boot(directory: &Path, prepared: crate::loader::Prepared) -> anyhow::Result<Host> {
    let authority_bytes =
        chio_example_runtime::files::read_bounded(Path::new("authority.json"), 65_536)?;
    let config: serde_json::Value = serde_json::from_slice(&authority_bytes)?;
    let fields = config
        .as_object()
        .ok_or_else(|| anyhow::anyhow!("Authority must be an object"))?;
    anyhow::ensure!(
        fields.keys().all(
            |key| ["schema", "tool", "max_invocations", "lifetime_seconds"].contains(&key.as_str())
        ),
        "Unknown authority field"
    );
    anyhow::ensure!(
        config["schema"] == "document-tools.authority.v1",
        "Unknown authority schema"
    );
    anyhow::ensure!(
        config["tool"] == crate::service::TOOL,
        "This host grants only count_words"
    );
    let call_limit = config["max_invocations"]
        .as_u64()
        .filter(|v| (1..=50).contains(v))
        .ok_or_else(|| anyhow::anyhow!("max_invocations must be an integer from 1 to 50"))?
        as u32;
    let lifetime = config["lifetime_seconds"]
        .as_u64()
        .filter(|v| (1..=3600).contains(v))
        .ok_or_else(|| anyhow::anyhow!("lifetime_seconds must be an integer from 1 to 3600"))?;
    // A new application run gets a fresh session. Refuse to reuse its directory.
    std::fs::create_dir(directory)?;
    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        std::fs::set_permissions(directory, std::fs::Permissions::from_mode(0o700))?;
    }
    std::fs::write(directory.join("authority.json"), &authority_bytes)?;
    let authority =
        chio_control_plane::DurableAdmissionRuntime::open(&directory.join("admission.db"))?;
    let keypair = authority.kernel_keypair();
    // The host chooses this key before any receipt exists.
    std::fs::write(
        directory.join("kernel-public-key.txt"),
        keypair.public_key().to_hex(),
    )?;
    let mut kernel = ChioKernel::new(KernelConfig {
        keypair,
        ca_public_keys: Vec::new(),
        max_delegation_depth: 3,
        policy_hash: chio_core::sha256_hex(&serde_json::to_vec(&serde_json::json!({
            "authority": config, "guard": prepared.context
        }))?),
        allow_sampling: false,
        allow_sampling_tool_use: false,
        allow_elicitation: false,
        max_stream_duration_secs: DEFAULT_MAX_STREAM_DURATION_SECS,
        max_stream_total_bytes: DEFAULT_MAX_STREAM_TOTAL_BYTES,
        require_web3_evidence: false,
        allow_ephemeral_receipt_log: true,
        allow_ephemeral_revocation_store: true,
        checkpoint_batch_size: DEFAULT_CHECKPOINT_BATCH_SIZE,
        retention_config: None,
        memory_budget: chio_kernel::MemoryBudgetConfig::defaults(),
        deadlines: chio_kernel::HotPathDeadlineConfig::default(),
    });
    kernel.set_receipt_store(Box::new(chio_store_sqlite::SqliteReceiptStore::open(
        directory.join("receipts.db"),
    )?))?;
    authority.attach(&mut kernel)?;
    std::fs::write(directory.join("guard-policy.json"), &prepared.policy_bytes)?;
    std::fs::write(directory.join("word_limit.wasm"), &prepared.module_bytes)?;
    std::fs::write(
        directory.join("guard-context.json"),
        serde_json::to_vec_pretty(&prepared.context)?,
    )?;
    for guard in prepared.guards {
        kernel.add_guard(Box::new(guard));
    }
    let service_key = Keypair::generate();
    let manifest = crate::service::manifest(service_key.public_key().to_hex());
    let signed = chio_manifest::sign_manifest(&manifest, &service_key)?;
    chio_manifest::verify_manifest(&signed, &service_key.public_key())?;
    std::fs::write(
        directory.join("manifest.json"),
        serde_json::to_vec_pretty(&signed)?,
    )?;
    kernel.register_tool_server(Box::new(crate::service::DocumentTools));
    let agent = Keypair::generate();
    let capability = kernel.issue_capability(
        &agent.public_key(),
        ChioScope {
            grants: vec![ToolGrant {
                server_id: crate::service::SERVER.into(),
                tool_name: crate::service::TOOL.into(),
                operations: vec![Operation::Invoke],
                constraints: Vec::new(),
                max_invocations: Some(call_limit),
                max_cost_per_invocation: None,
                max_total_cost: None,
                dpop_required: None,
            }],
            ..ChioScope::default()
        },
        lifetime,
    )?;
    std::fs::write(
        directory.join("capability.json"),
        serde_json::to_vec_pretty(&capability)?,
    )?;
    Ok(Host { kernel, capability })
}

mod host;
mod loader;
mod service;

use chio_kernel::{ToolCallOutput, ToolCallRequest};
use serde_json::json;
use std::{io::Read, path::Path};
use uuid::Uuid;

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter("chio_wasm_guards=warn")
        .with_writer(std::io::stderr)
        .without_time()
        .with_ansi(false)
        .with_target(false)
        .init();
    let paths: Vec<String> = std::env::args().skip(1).collect();
    anyhow::ensure!(
        !paths.is_empty(),
        "usage: guarded-documents FILE [FILE ...]"
    );
    let mut documents = Vec::new();
    for path in paths {
        let mut bytes = Vec::new();
        std::fs::File::open(&path)?
            .take((service::TEXT_LIMIT + 1) as u64)
            .read_to_end(&mut bytes)?;
        anyhow::ensure!(
            bytes.len() <= service::TEXT_LIMIT,
            "Document exceeds 65536 UTF-8 bytes: {path}"
        );
        documents.push((path, String::from_utf8(bytes)?));
    }
    let prepared = loader::load()?;
    let directory = Path::new("runs").join(Uuid::new_v4().to_string());
    let mut parents = std::fs::DirBuilder::new();
    parents.recursive(true);
    #[cfg(unix)]
    {
        use std::os::unix::fs::DirBuilderExt;
        parents.mode(0o700);
    }
    parents.create("runs")?;
    let host = host::boot(&directory, prepared)?;
    println!("Evidence: {}", directory.display());
    let mut refused = false;
    for (index, (path, text)) in documents.into_iter().enumerate() {
        let response = host
            .kernel
            .evaluate_tool_call(&ToolCallRequest {
                request_id: Uuid::new_v4().to_string(),
                capability: host.capability.clone(),
                server_id: service::SERVER.into(),
                tool_name: service::TOOL.into(),
                agent_id: host.capability.subject.to_hex(),
                arguments: json!({"text": text}),
                dpop_proof: None,
                execution_nonce: None,
                governed_intent: None,
                approval_token: None,
                approval_tokens: Vec::new(),
                threshold_approval_proposal: None,
                supplemental_authorization: None,
                model_metadata: None,
                federated_origin_kernel_id: None,
            })
            .await?;
        let receipt_file = format!("receipt-{}.json", index + 1);
        std::fs::write(
            directory.join(&receipt_file),
            serde_json::to_vec_pretty(&response.receipt)?,
        )?;
        if response.receipt.is_allowed() {
            let Some(ToolCallOutput::Value(value)) = response.output else {
                anyhow::bail!("Expected a count result");
            };
            println!("{path}: ALLOW {value}");
        } else {
            refused = true;
            println!(
                "{path}: DENY {}",
                response.reason.as_deref().unwrap_or("see signed receipt")
            );
        }
        for evidence in &response.receipt.evidence {
            if !evidence.verdict {
                if let Some(details) = &evidence.details {
                    println!("Guard {}: {details}", evidence.guard_name);
                }
            }
        }
        println!("Receipt: {} ({receipt_file})", response.receipt.id);
    }
    std::fs::write(
        directory.join("execution-count.json"),
        serde_json::to_vec_pretty(&json!({"tool_invocations": service::invocation_count()}))?,
    )?;
    if refused {
        std::process::exit(2);
    }
    Ok(())
}

use chio_kernel::{KernelError, NestedFlowBridge, ToolServerConnection};
use chio_manifest::{ToolDefinition, ToolManifest};
use serde_json::{json, Value};

pub const SERVER: &str = "documents";
pub const TOOL: &str = "count_words";
pub const TEXT_LIMIT: usize = 65_536;

pub struct DocumentTools;
static INVOCATIONS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
pub fn invocation_count() -> usize {
    INVOCATIONS.load(std::sync::atomic::Ordering::Relaxed)
}

fn count(tool: &str, arguments: &Value) -> Result<Value, KernelError> {
    if tool != TOOL {
        return Err(KernelError::ToolNotRegistered(tool.to_owned()));
    }
    let text = arguments
        .get("text")
        .and_then(Value::as_str)
        .ok_or_else(|| KernelError::RequestIncomplete("text must be a string".into()))?;
    if text.len() > TEXT_LIMIT {
        return Err(KernelError::RequestIncomplete(
            "text exceeds 65536 UTF-8 bytes".into(),
        ));
    }
    Ok(json!({"words": text.split_whitespace().count(), "bytes": text.len()}))
}

#[async_trait::async_trait]
impl ToolServerConnection for DocumentTools {
    fn server_id(&self) -> &str {
        SERVER
    }
    fn tool_names(&self) -> Vec<String> {
        vec![TOOL.into()]
    }

    async fn invoke(
        &self,
        tool: &str,
        arguments: Value,
        _nested: Option<&mut dyn NestedFlowBridge>,
    ) -> Result<Value, KernelError> {
        INVOCATIONS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        count(tool, &arguments)
    }
}

pub fn manifest(public_key: String) -> ToolManifest {
    ToolManifest {
        schema: "chio.manifest.v1".into(),
        server_id: SERVER.into(),
        name: "Document tools".into(),
        version: "1.0.0".into(),
        description: Some("Count words in caller-provided text".into()),
        tools: vec![ToolDefinition {
            name: TOOL.into(),
            description: "Count whitespace-separated words and UTF-8 bytes".into(),
            input_schema: json!({
                "type": "object", "properties": {"text": {"type": "string"}},
                "required": ["text"]
            }),
            output_schema: None,
            pricing: None,
            has_side_effects: false,
            latency_hint: None,
        }],
        server_tools: Vec::new(),
        required_permissions: None,
        public_key,
    }
}

[package]
name = "guarded-documents"
version = "0.1.0"
edition = "2021"
default-run = "guarded-documents"

[workspace]

[dependencies]
chio-example-runtime = { path = "runtime" }
chio-core = { package = "chio-core-types", git = "https://github.com/backbay-labs/chio", rev = "1346840e1192b5d6baf581c447cd37550e5c0144" }
chio-kernel = { git = "https://github.com/backbay-labs/chio", rev = "1346840e1192b5d6baf581c447cd37550e5c0144" }
chio-manifest = { git = "https://github.com/backbay-labs/chio", rev = "1346840e1192b5d6baf581c447cd37550e5c0144" }
chio-control-plane = { git = "https://github.com/backbay-labs/chio", rev = "1346840e1192b5d6baf581c447cd37550e5c0144" }
chio-store-sqlite = { git = "https://github.com/backbay-labs/chio", rev = "1346840e1192b5d6baf581c447cd37550e5c0144" }
anyhow = "1"
async-trait = "0.1"
serde_json = "1"
chio-wasm-guards = { git = "https://github.com/backbay-labs/chio", rev = "1346840e1192b5d6baf581c447cd37550e5c0144" }
tokio = { version = "1", features = ["macros", "rt-multi-thread"] }
uuid = { version = "1", features = ["v4"] }
tracing-subscriber = { version = "0.3", features = ["env-filter"] }

[package]
name = "word-limit"
version = "0.1.0"
edition = "2021"

[workspace]

[lib]
crate-type = ["cdylib"]

[dependencies]
chio-guard-sdk = { git = "https://github.com/backbay-labs/chio", rev = "1346840e1192b5d6baf581c447cd37550e5c0144" }
chio-guard-sdk-macros = { git = "https://github.com/backbay-labs/chio", rev = "1346840e1192b5d6baf581c447cd37550e5c0144" }