agentmux_srv/server/
caller.rs1use axum::body::Body;
22use axum::extract::State;
23use axum::http::Request;
24use axum::middleware::Next;
25use axum::response::Response;
26
27use super::{AppState, ReactiveAuthVia};
28
29pub(crate) const AGENT_TOKEN_HEADER: &str = "X-Agent-Token";
32
33#[derive(Debug, Clone, PartialEq, Eq)]
36pub(crate) enum Caller {
37 Agent { uid: String },
39 Unattributed,
42}
43
44impl Caller {
45 pub(crate) fn uid(&self) -> Option<&str> {
47 match self {
48 Caller::Agent { uid } => Some(uid),
49 Caller::Unattributed => None,
50 }
51 }
52}
53
54pub(crate) fn attributed_uid(caller: Option<&Caller>) -> String {
58 caller
59 .and_then(Caller::uid)
60 .map(str::to_string)
61 .unwrap_or_default()
62}
63
64pub(crate) enum Lookup {
69 Index(Option<String>),
70 NoIndex,
73}
74
75pub(crate) fn derive_caller(
76 token: Option<&str>,
77 full_key: bool,
78 lookup: impl FnOnce(&str) -> Lookup,
79) -> Caller {
80 let Some(token) = token.map(str::trim).filter(|t| !t.is_empty()) else {
81 return Caller::Unattributed;
82 };
83 if !full_key {
84 crate::backend::agent_resolve::record_uid_fallback("m4.token_on_lan_key");
85 return Caller::Unattributed;
86 }
87 match lookup(token) {
88 Lookup::Index(Some(uid)) => Caller::Agent { uid },
89 Lookup::Index(None) => {
90 crate::backend::agent_resolve::record_uid_fallback("m4.token_unknown");
91 Caller::Unattributed
92 }
93 Lookup::NoIndex => {
94 crate::backend::agent_resolve::record_uid_fallback("m4.token_index_unavailable");
95 Caller::Unattributed
96 }
97 }
98}
99
100pub(crate) async fn caller_middleware(
108 State(state): State<AppState>,
109 mut req: Request<Body>,
110 next: Next,
111) -> Response {
112 let full_key = req.extensions().get::<ReactiveAuthVia>() == Some(&ReactiveAuthVia::FullAuthKey);
113 let pinned = req
117 .extensions()
118 .get::<crate::backend::container_credential::ContainerGrant>()
119 .map(|g| g.agent_token.clone());
120 let header = if req.uri().path() == "/ws" || req.method() == axum::http::Method::OPTIONS {
121 None
122 } else {
123 req.headers()
124 .get(AGENT_TOKEN_HEADER)
125 .and_then(|v| v.to_str().ok())
126 };
127 let token = match &pinned {
128 Some(agent_token) => agent_token.as_deref(),
129 None => header,
130 };
131 let caller = derive_caller(token, full_key, |t| match state.mstore.token_index() {
132 Some(index) => Lookup::Index(index.uid_for(t)),
133 None => Lookup::NoIndex,
134 });
135 req.extensions_mut().insert(caller);
136 next.run(req).await
137}
138
139pub(crate) async fn handle_identity_fallbacks(
147 State(state): State<AppState>,
148 caller: Option<axum::Extension<Caller>>,
149) -> axum::Json<serde_json::Value> {
150 let counters: serde_json::Map<String, serde_json::Value> =
151 crate::backend::agent_resolve::uid_fallback_counts()
152 .into_iter()
153 .map(|(site, n)| (site.to_string(), serde_json::json!(n)))
154 .collect();
155 let registrations = state.reactive_handler.list_agents();
158 let mstore = state.mstore.clone();
159 let gauges = tokio::task::spawn_blocking(move || {
161 crate::backend::identity_spawn::live_gauges(
162 ®istrations,
163 |block| crate::backend::blockcontroller::get_controller(block).is_some(),
164 |block| crate::backend::identity_spawn::store_block_has_row(&mstore, block),
165 )
166 })
167 .await
168 .ok();
169 let (tokenless_or_unknown, unidentified) = match gauges {
170 Some((t, u)) => (serde_json::json!(t), serde_json::json!(u)),
171 None => (serde_json::Value::Null, serde_json::Value::Null),
172 };
173 let caller_uid = caller.as_ref().and_then(|c| c.0.uid().map(str::to_string));
174 axum::Json(serde_json::json!({
175 "boot_id": state.boot_id.to_string(),
179 "caller": caller_uid,
180 "counters": counters,
181 "gauges": {
182 "live.tokenless_or_unknown": tokenless_or_unknown,
183 "live.unidentified": unidentified,
184 },
185 }))
186}
187
188#[cfg(test)]
189mod tests {
190 use super::*;
191
192 fn index(token: &str) -> Lookup {
193 Lookup::Index((token == "tok-y").then(|| "uid-y".to_string()))
194 }
195
196 #[test]
197 fn a_known_token_on_the_full_key_is_the_agent() {
198 assert_eq!(
199 derive_caller(Some("tok-y"), true, index),
200 Caller::Agent {
201 uid: "uid-y".into()
202 }
203 );
204 }
205
206 #[test]
207 fn no_token_is_unattributed() {
208 assert_eq!(derive_caller(None, true, index), Caller::Unattributed);
209 assert_eq!(derive_caller(Some(" "), true, index), Caller::Unattributed);
210 }
211
212 #[test]
215 fn an_unknown_token_is_unattributed_not_refused() {
216 assert_eq!(
217 derive_caller(Some("tok-gone"), true, index),
218 Caller::Unattributed
219 );
220 }
221
222 #[test]
223 fn a_missing_index_is_unattributed() {
224 assert_eq!(
225 derive_caller(Some("tok-y"), true, |_| Lookup::NoIndex),
226 Caller::Unattributed
227 );
228 }
229
230 #[test]
232 fn a_token_on_a_lan_key_request_is_ignored() {
233 assert_eq!(
234 derive_caller(Some("tok-y"), false, index),
235 Caller::Unattributed
236 );
237 }
238}