agentmux_srv/backend/storage/

agent_wan_keys.rs

// Copyright 2026, AgentMux Corp.
// SPDX-License-Identifier: Apache-2.0

//! Per-agent Ed25519 keypair for general agent-to-agent WAN-tier jekt
//! signing. See `db_agent_wan_keys` in migrations.rs (OBJECT_SCHEMA_VERSION
//! v36).
//!
//! Mirrors `agent_lan_keys.rs`: one row per locally-registered agent_id,
//! minted once (first call to `agent_wan_key_ensure`), with the PRIVATE half
//! leaving this srv instance only by being injected into that ONE agent's own
//! MCP server process env (`AGENTMUX_WAN_KEY`) — never into any other agent's
//! env, never returned over any RPC. The public half is not secret: it is what
//! gets published to muxbus's agent-ownership registry so a WAN counterpart
//! can verify this agent's outgoing signatures.
//!
//! **Why this is a separate table from `db_agent_lan_keys` rather than a reuse
//! of the same keypair.** The two tiers need different key *lifecycles*, and
//! that difference is structural, not stylistic:
//!
//!  - A LAN public key is permanently trust-on-first-use pinned by every peer
//!    that has ever seen it (`db_lan_peer_pubkey_pins`, no expiry, no re-pin
//!    path), which is exactly why `SPEC_JEKT_HOST_KEY_TTL_ROTATION_2026_09_14.md`
//!    §4 declined to rotate LAN keys — rotation would turn every pinned peer
//!    into a permanent false "active forgery" alarm.
//!  - A WAN public key is served from an authoritative registry the receiver
//!    queries, so rotation is a republish rather than a spoofing alarm, and
//!    the spec intends WAN keys to rotate.
//!
//! One keypair cannot have both lifecycles. Keeping them separate also means
//! a compromise of one tier's key does not hand over the other's.
//!
//! **Scope note:** minting and signing ship ahead of verification on purpose.
//! Verifying a WAN signature requires resolving `(sender_account,
//! source_agent)` to exactly one key, which waits on relay-side work tracked
//! in the private cloud repo. An agent
//! only receives a key when it is spawned, so minting now is what makes
//! verification meaningful later instead of applying to almost nobody.

use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
use rusqlite::params;

use super::error::StoreError;
use super::store::Store;

use agentmux_common::time::now_secs;

use super::agent_lan_keys::random_seed_bytes;

/// An agent's WAN signing keypair, both halves base64-encoded for
/// storage/transport, plus the version the publish/rotation phases will use.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AgentWanKeypair {
    pub public_key: String,
    pub private_key: String,
    /// Monotonic version of this keypair, starting at 1. Reserved for the
    /// version-synchronised publish (§3.2) and rotation (§3.6) phases —
    /// nothing reads it yet, and it is stored now only so adopting those
    /// phases needs no second schema change.
    pub key_version: i64,
}

impl Store {
    /// Load this agent's WAN keypair if one has already been minted, without
    /// creating one.
    pub fn agent_wan_key_load(&self, agent_id: &str) -> Result<Option<AgentWanKeypair>, StoreError> {
        let key = agent_id.to_lowercase();
        let conn = self.conn.lock().unwrap();
        let mut stmt = conn.prepare(
            "SELECT public_key, private_key, key_version FROM db_agent_wan_keys WHERE agent_id = ?1",
        )?;
        match stmt.query_row(params![key], |row| {
            Ok(AgentWanKeypair {
                public_key: row.get(0)?,
                private_key: row.get(1)?,
                key_version: row.get(2)?,
            })
        }) {
            Ok(v) => Ok(Some(v)),
            Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
            Err(e) => Err(e.into()),
        }
    }

    /// Return this agent's WAN keypair, minting and persisting a fresh random
    /// one on first use. Race-safe under concurrent first-use — same
    /// `INSERT OR IGNORE` + re-read pattern as `agent_lan_key_ensure`: two
    /// concurrent callers for the same never-before-seen agent_id agree on one
    /// keypair instead of each minting and using a different one.
    ///
    /// Note there is deliberately no TTL/rotation branch here yet, unlike
    /// `agent_jekt_key_ensure`. Rotation needs the continuity chain described
    /// in the spec's §3.6 (a new key signed by its predecessor) to be
    /// distinguishable from a registry rewrite; rotating before that exists
    /// would trade one gap for a worse one.
    pub fn agent_wan_key_ensure(&self, agent_id: &str) -> Result<AgentWanKeypair, StoreError> {
        let key = agent_id.to_lowercase();
        if let Some(existing) = self.agent_wan_key_load(&key)? {
            return Ok(existing);
        }
        let conn = self.conn.lock().unwrap();
        let seed = random_seed_bytes();
        let (public_bytes, private_bytes) = agentmux_common::jekt_sign::generate_wan_keypair(seed);
        let public_key = BASE64.encode(public_bytes);
        let private_key = BASE64.encode(private_bytes);
        conn.execute(
            "INSERT OR IGNORE INTO db_agent_wan_keys (agent_id, public_key, private_key, key_version, created_at) \
             VALUES (?1, ?2, ?3, 1, ?4)",
            params![key, public_key, private_key, now_secs()],
        )?;
        // Re-read regardless of whether our insert won the race — single
        // source of truth for "what keypair does this agent have," not an
        // assumption that our own insert succeeded.
        let mut stmt = conn.prepare(
            "SELECT public_key, private_key, key_version FROM db_agent_wan_keys WHERE agent_id = ?1",
        )?;
        stmt.query_row(params![key], |row| {
            Ok(AgentWanKeypair {
                public_key: row.get(0)?,
                private_key: row.get(1)?,
                key_version: row.get(2)?,
            })
        })
        .map_err(Into::into)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn object_store() -> Store {
        Store::open(std::path::Path::new(":memory:")).unwrap()
    }

    #[test]
    fn load_returns_none_when_unminted() {
        let store = object_store();
        assert!(store.agent_wan_key_load("agentx").unwrap().is_none());
    }

    #[test]
    fn ensure_mints_a_valid_base64_keypair() {
        let store = object_store();
        let keypair = store.agent_wan_key_ensure("agentx").unwrap();
        let pub_bytes = BASE64.decode(&keypair.public_key).unwrap();
        let priv_bytes = BASE64.decode(&keypair.private_key).unwrap();
        assert_eq!(pub_bytes.len(), 32);
        assert_eq!(priv_bytes.len(), 32);
        assert_ne!(keypair.public_key, keypair.private_key);
        assert_eq!(keypair.key_version, 1);
    }

    #[test]
    fn ensure_is_idempotent_same_agent_gets_same_keypair() {
        let store = object_store();
        let first = store.agent_wan_key_ensure("AgentX").unwrap();
        let second = store.agent_wan_key_ensure("agentx").unwrap(); // case-insensitive
        assert_eq!(first, second);
    }

    #[test]
    fn different_agents_get_different_keypairs() {
        let store = object_store();
        let a = store.agent_wan_key_ensure("agentx").unwrap();
        let b = store.agent_wan_key_ensure("agenty").unwrap();
        assert_ne!(a.public_key, b.public_key);
        assert_ne!(a.private_key, b.private_key);
    }

    #[test]
    fn ensure_then_load_round_trips() {
        let store = object_store();
        let minted = store.agent_wan_key_ensure("agentx").unwrap();
        let loaded = store.agent_wan_key_load("agentx").unwrap().unwrap();
        assert_eq!(minted, loaded);
    }

    #[test]
    fn a_minted_keypair_can_sign_and_verify_a_real_jekt() {
        // End-to-end sanity through the actual jekt_sign functions, not just
        // storage round-tripping — proves the stored encoding is exactly what
        // generate_wan_keypair/sign_wan_jekt/verify_wan_jekt expect.
        let store = object_store();
        let keypair = store.agent_wan_key_ensure("agentx").unwrap();
        let private = BASE64.decode(&keypair.private_key).unwrap();
        let public = BASE64.decode(&keypair.public_key).unwrap();
        let sig = agentmux_common::jekt_sign::sign_wan_jekt(
            &private, "msg-1", "agentx", "narko", "stable", "agenty", 1_000, "hello",
        )
        .unwrap();
        assert!(agentmux_common::jekt_sign::verify_wan_jekt(
            &public, "msg-1", "agentx", "narko", "stable", "agenty", 1_000, "hello", &sig,
        ));
    }

    #[test]
    fn an_agents_wan_and_lan_keypairs_are_independent() {
        // The tables must not share key material — see this module's doc
        // comment for why their lifecycles differ. A shared keypair would also
        // make the domain separation in jekt_sign the only thing standing
        // between a LAN signature and a WAN trust claim.
        let store = object_store();
        let wan = store.agent_wan_key_ensure("agentx").unwrap();
        let lan = store.agent_lan_key_ensure("agentx").unwrap();
        assert_ne!(wan.public_key, lan.public_key);
        assert_ne!(wan.private_key, lan.private_key);
    }
}