agentmux_mcp/
self_keys.rs1use std::sync::Mutex;
24use std::time::{Duration, Instant};
25
26use agentmux_common::AUTH_KEY_HEADER;
27use serde::Deserialize;
28
29const REFRESH_AFTER: Duration = Duration::from_secs(20 * 60 * 60);
31
32#[derive(Deserialize, Default, Clone)]
33struct Wire {
34 slug: String,
35 jekt_key: Option<String>,
36 #[serde(default)]
38 jekt_key_expires_at: Option<i64>,
39 lan_key: Option<String>,
40 wan_key: Option<String>,
41 #[serde(default)]
42 uid: String,
43 uid_lan_key: Option<String>,
44}
45
46struct Cached {
47 keys: Wire,
48 fetched_at: Instant,
49}
50
51static CACHE: Mutex<Option<Cached>> = Mutex::new(None);
52
53static LAST_ATTEMPT: Mutex<Option<Instant>> = Mutex::new(None);
57const RETRY_AFTER: Duration = Duration::from_secs(60);
58const FETCH_TIMEOUT: Duration = Duration::from_secs(3);
60
61pub(crate) async fn refresh_if_stale(client: &reqwest::Client, local_url: &str, auth_key: &str) {
65 let fresh = CACHE.lock().unwrap_or_else(|e| e.into_inner()).as_ref().is_some_and(|c| still_good(c, agentmux_common::time::now_secs()));
66 if fresh || local_url.is_empty() || std::env::var("AGENTMUX_AGENT_TOKEN").map_or(true, |t| t.trim().is_empty()) {
67 return;
68 }
69 {
70 let mut last = LAST_ATTEMPT.lock().unwrap_or_else(|e| e.into_inner());
71 if last.is_some_and(|t| t.elapsed() < RETRY_AFTER) {
72 return;
73 }
74 *last = Some(Instant::now());
75 }
76 let url = format!("{}/agentmux/agents/self/keys", local_url.trim_end_matches('/'));
77 let Ok(resp) = client.get(&url).header(AUTH_KEY_HEADER, auth_key).timeout(FETCH_TIMEOUT).send().await else { return };
78 if !resp.status().is_success() {
79 return;
80 }
81 if let Ok(keys) = resp.json::<Wire>().await {
82 *CACHE.lock().unwrap_or_else(|e| e.into_inner()) = Some(Cached { keys, fetched_at: Instant::now() });
83 }
84}
85
86fn still_good(c: &Cached, now: i64) -> bool {
88 c.fetched_at.elapsed() < REFRESH_AFTER && c.keys.jekt_key_expires_at.is_none_or(|expires| now < expires)
89}
90
91fn name_key(pick: impl Fn(&Wire) -> Option<&String>) -> Option<Vec<u8>> {
94 let own = std::env::var("AGENTMUX_AGENT_ID").ok()?;
95 let cache = CACHE.lock().unwrap_or_else(|e| e.into_inner());
96 let cached = cache.as_ref()?;
97 if cached.keys.slug != own {
98 return None;
99 }
100 agentmux_common::jekt_sign::decode_key(pick(&cached.keys)?)
101}
102
103pub(crate) fn jekt_key() -> Option<Vec<u8>> {
105 name_key(|k| k.jekt_key.as_ref()).or_else(|| env_key("AGENTMUX_JEKT_KEY"))
106}
107
108pub(crate) fn lan_key() -> Option<Vec<u8>> {
110 name_key(|k| k.lan_key.as_ref()).or_else(|| env_key("AGENTMUX_LAN_KEY"))
111}
112
113pub(crate) fn wan_key() -> Option<Vec<u8>> {
115 name_key(|k| k.wan_key.as_ref()).or_else(|| env_key("AGENTMUX_WAN_KEY"))
116}
117
118pub(crate) fn uid_signer() -> Option<(String, Vec<u8>)> {
122 let cache = CACHE.lock().unwrap_or_else(|e| e.into_inner());
123 let keys = &cache.as_ref()?.keys;
124 let uid = keys.uid.trim();
125 if uid.is_empty() {
126 return None;
127 }
128 Some((uid.to_string(), agentmux_common::jekt_sign::decode_key(keys.uid_lan_key.as_ref()?)?))
129}
130
131fn env_key(var: &str) -> Option<Vec<u8>> {
132 std::env::var(var).ok().filter(|s| !s.is_empty()).and_then(|b64| agentmux_common::jekt_sign::decode_key(&b64))
133}
134
135#[cfg(test)]
136pub(crate) fn set_for_test(slug: &str, jekt_key: Option<&str>) {
137 *CACHE.lock().unwrap_or_else(|e| e.into_inner()) = Some(Cached {
138 keys: Wire { slug: slug.to_string(), jekt_key: jekt_key.map(str::to_string), ..Default::default() },
139 fetched_at: Instant::now(),
140 });
141}
142
143#[cfg(test)]
144pub(crate) fn clear_for_test() {
145 *CACHE.lock().unwrap_or_else(|e| e.into_inner()) = None;
146}
147
148#[cfg(test)]
149mod tests {
150 use super::*;
151
152 const FETCHED_B64: &str = "BwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwc=";
154 const ENV_B64: &str = "AQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQE=";
155
156 static LOCK: Mutex<()> = Mutex::new(());
158
159 fn with_env<T>(agent_id: &str, env_key: Option<&str>, f: impl FnOnce() -> T) -> T {
160 let _guard = LOCK.lock().unwrap_or_else(|e| e.into_inner());
161 std::env::set_var("AGENTMUX_AGENT_ID", agent_id);
162 match env_key {
163 Some(k) => std::env::set_var("AGENTMUX_JEKT_KEY", k),
164 None => std::env::remove_var("AGENTMUX_JEKT_KEY"),
165 }
166 let out = f();
167 clear_for_test();
168 std::env::remove_var("AGENTMUX_JEKT_KEY");
169 out
170 }
171
172 #[test]
173 fn a_fetched_key_for_exactly_this_name_wins_over_the_env_copy() {
174 let fetched = [7u8; 32];
175 let key = with_env("aria", Some(ENV_B64), || {
176 set_for_test("aria", Some(FETCHED_B64));
177 jekt_key()
178 });
179 assert_eq!(key.as_deref(), Some(&fetched[..]));
180 }
181
182 #[test]
183 fn a_key_served_for_another_name_is_not_used() {
184 let env = [1u8; 32];
187 let key = with_env("aria", Some(ENV_B64), || {
188 set_for_test("aria-2", Some(FETCHED_B64));
189 jekt_key()
190 });
191 assert_eq!(key.as_deref(), Some(&env[..]));
192 }
193
194 #[test]
195 fn nothing_fetched_means_the_env_key_as_before() {
196 let env = [1u8; 32];
197 assert_eq!(with_env("aria", Some(ENV_B64), jekt_key).as_deref(), Some(&env[..]));
198 assert_eq!(with_env("aria", None, jekt_key), None);
199 }
200
201 #[test]
202 fn a_copy_is_good_until_its_key_expires_then_refetched() {
203 let cached = |expires: Option<i64>| Cached {
204 keys: Wire { slug: "aria".into(), jekt_key_expires_at: expires, ..Default::default() },
205 fetched_at: Instant::now(),
206 };
207 assert!(still_good(&cached(Some(1_000)), 999));
208 assert!(!still_good(&cached(Some(1_000)), 1_000), "expired: srv may rotate it now");
209 assert!(still_good(&cached(None), 1_000_000), "no expiry served: the 20 h clock");
210 }
211
212 #[test]
213 fn the_v2_signer_is_only_ever_the_fetched_uid_key() {
214 let _guard = LOCK.lock().unwrap_or_else(|e| e.into_inner());
215 clear_for_test();
216 assert!(
217 uid_signer().is_none(),
218 "no fetch: no v2, whatever the env holds"
219 );
220 let set = |uid: &str, key: Option<&str>| {
221 *CACHE.lock().unwrap_or_else(|e| e.into_inner()) = Some(Cached {
222 keys: Wire {
223 slug: "aria".into(),
224 uid: uid.into(),
225 uid_lan_key: key.map(str::to_string),
226 ..Default::default()
227 },
228 fetched_at: Instant::now(),
229 });
230 };
231 set("uid-1", Some(FETCHED_B64));
232 assert_eq!(uid_signer(), Some(("uid-1".to_string(), vec![7u8; 32])));
233 set("", Some(FETCHED_B64));
234 assert!(uid_signer().is_none(), "a key with no UID to claim");
235 set("uid-1", None);
236 assert!(uid_signer().is_none());
237 clear_for_test();
238 }
239}