Skip to main content

agentmux_srv/backend/storage/
agent_uid_keys.rs

1// Copyright 2026, AgentMux Corp.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Identity M4d-2: each agent's LAN and WAN signing keypairs filed under its
5//! UID (`db_agents.id`), beside the name-keyed tables they are copied from
6//! (`SPEC_AGENT_IDENTITY_CARRIED_NOT_DERIVED_2026_09_23.md` §6.5.10).
7//!
8//! A name-keyed key is copied to a UID only on evidence that the agent owns
9//! it: the name is the row's slug, no other row holds that slug (folded),
10//! the key is no older than the row, and the name is not a deleted agent's
11//! (`db_agent_key_tombstones`). Otherwise the UID gets a fresh keypair. The
12//! v1 signatures keep using the name-keyed keys until M5; a UID key signs
13//! only the v2 signatures (M4d-6), so neither outcome changes what peers
14//! verify today.
15
16use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
17use rusqlite::{params, OptionalExtension};
18
19use super::agent_lan_keys::random_seed_bytes;
20use super::error::StoreError;
21use super::store::Store;
22use agentmux_common::time::now_secs;
23
24/// One UID-keyed keypair, both halves base64.
25#[derive(Debug, Clone, PartialEq, Eq)]
26pub struct UidKeypair {
27    pub public_key: String,
28    pub private_key: String,
29}
30
31/// An agent's UID-keyed LAN and WAN keypairs.
32#[derive(Debug, Clone, PartialEq, Eq)]
33pub struct AgentUidKeys {
34    pub lan: UidKeypair,
35    pub wan: UidKeypair,
36}
37
38/// How a UID key came to be, for the `m4d.key_*` counters.
39#[derive(Debug, Clone, Copy, PartialEq, Eq)]
40pub enum UidKeyOrigin {
41    /// It was already there.
42    Existing,
43    /// Copied from the name-keyed key, on ownership evidence.
44    Copied,
45    /// Minted, and why the name-keyed key wasn't copied.
46    Fresh(FreshReason),
47}
48
49#[derive(Debug, Clone, Copy, PartialEq, Eq)]
50pub enum FreshReason {
51    /// The row has no slug, or no name-keyed key exists under it.
52    NoNameKey,
53    /// Another row holds the same slug, folded.
54    SlugShared,
55    /// The row's `created_at` is 0 (a legacy or test row): no age to compare.
56    RowUndated,
57    /// The name-keyed key predates the row: another agent's.
58    KeyOlderThanRow,
59    /// The name is a deleted agent's.
60    Tombstoned,
61}
62
63impl UidKeyOrigin {
64    fn counter(self) -> Option<&'static str> {
65        match self {
66            UidKeyOrigin::Existing => None,
67            UidKeyOrigin::Copied => Some("m4d.key_copied"),
68            UidKeyOrigin::Fresh(FreshReason::NoNameKey) => Some("m4d.key_fresh.no_name_key"),
69            UidKeyOrigin::Fresh(FreshReason::SlugShared) => Some("m4d.key_fresh.slug_shared"),
70            UidKeyOrigin::Fresh(FreshReason::RowUndated) => Some("m4d.key_fresh.row_undated"),
71            UidKeyOrigin::Fresh(FreshReason::KeyOlderThanRow) => Some("m4d.key_fresh.key_older_than_row"),
72            UidKeyOrigin::Fresh(FreshReason::Tombstoned) => Some("m4d.key_fresh.tombstoned"),
73        }
74    }
75}
76
77/// The two key tiers, each with its legacy name-keyed table and its UID table.
78#[derive(Clone, Copy)]
79enum Tier {
80    Lan,
81    Wan,
82}
83
84impl Tier {
85    fn name_table(self) -> &'static str {
86        match self {
87            Tier::Lan => "db_agent_lan_keys",
88            Tier::Wan => "db_agent_wan_keys",
89        }
90    }
91    fn uid_table(self) -> &'static str {
92        match self {
93            Tier::Lan => "db_agent_lan_keys_by_uid",
94            Tier::Wan => "db_agent_wan_keys_by_uid",
95        }
96    }
97    fn generate(self) -> ([u8; 32], [u8; 32]) {
98        let seed = random_seed_bytes();
99        match self {
100            Tier::Lan => agentmux_common::jekt_sign::generate_lan_keypair(seed),
101            Tier::Wan => agentmux_common::jekt_sign::generate_wan_keypair(seed),
102        }
103    }
104}
105
106/// The UID's row, as the evidence test needs it.
107struct Row {
108    slug: String,
109    created_at_ms: i64,
110}
111
112impl Store {
113    /// The UID's LAN and WAN keypairs, copied or minted on first use (§6.5.10
114    /// M4d-2). One critical section: the row is re-read inside the lock, so a
115    /// row purged since the caller read it gets nothing (`Ok(None)`) and no
116    /// private key is left without an owner.
117    pub fn agent_uid_keys_ensure(&self, uid: &str) -> Result<Option<(AgentUidKeys, [UidKeyOrigin; 2])>, StoreError> {
118        let uid = uid.trim();
119        if uid.is_empty() {
120            return Ok(None);
121        }
122        let conn = self.conn.lock().unwrap();
123        let row = conn
124            .query_row(
125                "SELECT slug, created_at FROM db_agents WHERE id = ?1",
126                params![uid],
127                |r| Ok(Row { slug: r.get::<_, Option<String>>(0)?.unwrap_or_default(), created_at_ms: r.get::<_, Option<i64>>(1)?.unwrap_or(0) }),
128            )
129            .optional()?;
130        let Some(row) = row else { return Ok(None) };
131        let (lan, lan_origin) = ensure_tier(&conn, Tier::Lan, uid, &row)?;
132        let (wan, wan_origin) = ensure_tier(&conn, Tier::Wan, uid, &row)?;
133        drop(conn);
134        for origin in [lan_origin, wan_origin] {
135            if let Some(counter) = origin.counter() {
136                crate::backend::agent_resolve::record_uid_fallback(counter);
137            }
138        }
139        Ok(Some((AgentUidKeys { lan, wan }, [lan_origin, wan_origin])))
140    }
141
142    /// The UID's LAN public key, if it has one; never creates one.
143    pub fn agent_uid_lan_public_key_load(&self, uid: &str) -> Result<Option<String>, StoreError> {
144        let conn = self.conn.lock().unwrap();
145        conn.query_row("SELECT public_key FROM db_agent_lan_keys_by_uid WHERE uid = ?1", params![uid.trim()], |r| r.get(0))
146            .optional()
147            .map_err(Into::into)
148    }
149}
150
151fn load_uid_key(conn: &rusqlite::Connection, tier: Tier, uid: &str) -> rusqlite::Result<Option<UidKeypair>> {
152    conn.query_row(
153        &format!("SELECT public_key, private_key FROM {} WHERE uid = ?1", tier.uid_table()),
154        params![uid],
155        |r| Ok(UidKeypair { public_key: r.get(0)?, private_key: r.get(1)? }),
156    )
157    .optional()
158}
159
160/// One tier: the existing UID key, else a copy of the name-keyed key when the
161/// evidence holds, else a fresh one. Runs under the caller's lock.
162fn ensure_tier(conn: &rusqlite::Connection, tier: Tier, uid: &str, row: &Row) -> Result<(UidKeypair, UidKeyOrigin), StoreError> {
163    if let Some(existing) = load_uid_key(conn, tier, uid)? {
164        return Ok((existing, UidKeyOrigin::Existing));
165    }
166    let (pair, copied_from, origin) = match copyable_name_key(conn, tier, row)? {
167        Ok((pair, name)) => (pair, name, UidKeyOrigin::Copied),
168        Err(reason) => {
169            let (public, private) = tier.generate();
170            (UidKeypair { public_key: BASE64.encode(public), private_key: BASE64.encode(private) }, String::new(), UidKeyOrigin::Fresh(reason))
171        }
172    };
173    conn.execute(
174        &format!(
175            "INSERT OR IGNORE INTO {} (uid, public_key, private_key, created_at, copied_from) VALUES (?1, ?2, ?3, ?4, ?5)",
176            tier.uid_table()
177        ),
178        params![uid, pair.public_key, pair.private_key, now_secs(), copied_from],
179    )?;
180    let stored = load_uid_key(conn, tier, uid)?.ok_or(StoreError::NotFound)?;
181    Ok((stored, origin))
182}
183
184/// The name-keyed key the UID may take, and its name, or why not.
185fn copyable_name_key(conn: &rusqlite::Connection, tier: Tier, row: &Row) -> Result<Result<(UidKeypair, String), FreshReason>, StoreError> {
186    let slug = row.slug.trim().to_lowercase();
187    if slug.is_empty() {
188        return Ok(Err(FreshReason::NoNameKey));
189    }
190    let key = conn
191        .query_row(
192            &format!("SELECT public_key, private_key, created_at FROM {} WHERE agent_id = ?1", tier.name_table()),
193            params![slug],
194            |r| Ok((UidKeypair { public_key: r.get(0)?, private_key: r.get(1)? }, r.get::<_, i64>(2)?)),
195        )
196        .optional()?;
197    let Some((pair, key_created_secs)) = key else { return Ok(Err(FreshReason::NoNameKey)) };
198    // Slug uniqueness is case-sensitive, the key tables are not: "Aria" and
199    // "aria" share one key row, so either holding it disqualifies both.
200    // Folded in Rust as the key tables fold, never with SQLite's ASCII-only
201    // lower(): "Ä" and "ä" would each count as the only holder and both
202    // take one private key.
203    let holders = slug_holders_folded(conn, &slug)?;
204    if holders != 1 {
205        return Ok(Err(FreshReason::SlugShared));
206    }
207    if row.created_at_ms <= 0 {
208        return Ok(Err(FreshReason::RowUndated));
209    }
210    // Units differ: keys in seconds, rows in milliseconds. A key minted in the
211    // same second as its row passes.
212    if key_created_secs < row.created_at_ms / 1000 {
213        return Ok(Err(FreshReason::KeyOlderThanRow));
214    }
215    let tombstoned: bool = conn
216        .query_row("SELECT 1 FROM db_agent_key_tombstones WHERE name = ?1", params![slug], |_| Ok(true))
217        .optional()?
218        .unwrap_or(false);
219    if tombstoned {
220        return Ok(Err(FreshReason::Tombstoned));
221    }
222    Ok(Ok((pair, slug)))
223}
224
225/// How many rows hold `folded` as their slug, trimmed and folded with Rust's
226/// Unicode-aware `to_lowercase` (the key tables' folding).
227pub(super) fn slug_holders_folded(conn: &rusqlite::Connection, folded: &str) -> Result<usize, StoreError> {
228    let mut stmt = conn.prepare("SELECT slug FROM db_agents WHERE slug IS NOT NULL AND slug <> ''")?;
229    let slugs = stmt.query_map([], |r| r.get::<_, String>(0))?;
230    let mut n = 0;
231    for slug in slugs {
232        if slug?.trim().to_lowercase() == folded {
233            n += 1;
234        }
235    }
236    Ok(n)
237}
238
239#[cfg(test)]
240mod tests;