agentmux_srv/backend/storage/
agent_uid_keys.rs1use 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#[derive(Debug, Clone, PartialEq, Eq)]
26pub struct UidKeypair {
27 pub public_key: String,
28 pub private_key: String,
29}
30
31#[derive(Debug, Clone, PartialEq, Eq)]
33pub struct AgentUidKeys {
34 pub lan: UidKeypair,
35 pub wan: UidKeypair,
36}
37
38#[derive(Debug, Clone, Copy, PartialEq, Eq)]
40pub enum UidKeyOrigin {
41 Existing,
43 Copied,
45 Fresh(FreshReason),
47}
48
49#[derive(Debug, Clone, Copy, PartialEq, Eq)]
50pub enum FreshReason {
51 NoNameKey,
53 SlugShared,
55 RowUndated,
57 KeyOlderThanRow,
59 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#[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
106struct Row {
108 slug: String,
109 created_at_ms: i64,
110}
111
112impl Store {
113 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 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
160fn 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
184fn 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 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 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
225pub(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;