1use std::collections::HashMap;
26use std::sync::{Arc, Mutex, OnceLock};
27use std::time::{Duration, Instant};
28
29use crate::backend::storage::store::Store;
30use crate::muxbus::cloud_subscriber::load_valid_token;
31
32const CREDENTIAL_HTTP_TIMEOUT: Duration = Duration::from_secs(10);
40
41const CREDENTIAL_RETRY_COOLDOWN: Duration = Duration::from_secs(60);
52
53static CREDENTIAL_COOLDOWN: OnceLock<Mutex<HashMap<String, Instant>>> = OnceLock::new();
54
55fn credential_cooldown() -> &'static Mutex<HashMap<String, Instant>> {
56 CREDENTIAL_COOLDOWN.get_or_init(|| Mutex::new(HashMap::new()))
57}
58
59fn credential_recently_failed(agent_id: &str) -> bool {
60 let map = credential_cooldown().lock().unwrap();
61 map.get(agent_id)
62 .is_some_and(|t| t.elapsed() < CREDENTIAL_RETRY_COOLDOWN)
63}
64
65fn record_credential_failure(agent_id: &str) {
66 credential_cooldown()
67 .lock()
68 .unwrap()
69 .insert(agent_id.to_string(), Instant::now());
70}
71
72#[derive(Debug)]
76enum CredentialFailure {
77 Unrecognized(String),
81 Other(String),
84}
85
86impl std::fmt::Display for CredentialFailure {
87 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
88 match self {
89 CredentialFailure::Unrecognized(m) | CredentialFailure::Other(m) => f.write_str(m),
90 }
91 }
92}
93
94fn token_rejection_is_unrecognized(status: u16, body: &str) -> bool {
98 matches!(status, 400 | 401)
99 && serde_json::from_str::<serde_json::Value>(body)
100 .ok()
101 .and_then(|v| v.get("error").and_then(|e| e.as_str()).map(|e| e == "invalid_client"))
102 .unwrap_or(false)
103}
104
105fn provision_rejection_is_unrecognized(status: u16) -> bool {
109 matches!(status, 401 | 403)
110}
111
112fn drop_unrecognized_credential(agent_id: &str, failure: &CredentialFailure, mstore: &Arc<Store>) {
116 if !matches!(failure, CredentialFailure::Unrecognized(_)) {
117 return;
118 }
119 match mstore.agent_credential_delete(agent_id) {
120 Ok(()) => tracing::info!(
121 agent_id = %agent_id,
122 "muxbus: cloud no longer recognizes this agent's credential — dropped it; the next use re-provisions",
123 ),
124 Err(e) => tracing::warn!(
125 agent_id = %agent_id, error = %e,
126 "muxbus: failed to drop an unrecognized per-agent credential",
127 ),
128 }
129}
130
131pub async fn ensure_agent_credential(
136 agent_id: &str,
137 mstore: &Arc<Store>,
138 http: &reqwest::Client,
139) -> Option<String> {
140 ensure_agent_credential_at(
141 agent_id,
142 mstore,
143 http,
144 &crate::muxbus::relay::rest_base_url(),
145 user_token(mstore),
146 )
147 .await
148}
149
150async fn ensure_agent_credential_at(
153 agent_id: &str,
154 mstore: &Arc<Store>,
155 http: &reqwest::Client,
156 base: &str,
157 user_token: impl std::future::Future<Output = Result<String, String>>,
158) -> Option<String> {
159 let key = agent_id.to_lowercase();
160
161 if credential_recently_failed(&key) {
165 return None;
166 }
167
168 let creds = mstore.agent_credential_load(&key).ok().flatten();
169
170 let creds = match creds {
171 Some(c) if !c.client_id.is_empty() => c,
172 _ => {
173 if let Err(e) = provision_agent_client(&key, mstore, http, base, user_token).await {
175 tracing::warn!(
176 agent_id = %key, error = %e,
177 "muxbus: agent credential provisioning failed, backing off {}s",
178 CREDENTIAL_RETRY_COOLDOWN.as_secs(),
179 );
180 drop_unrecognized_credential(&key, &e, mstore);
181 record_credential_failure(&key);
182 return None;
183 }
184 mstore.agent_credential_load(&key).ok().flatten()?
185 }
186 };
187
188 if creds.is_valid() {
189 return Some(creds.access_token);
190 }
191
192 match fetch_m2m_token(&key, &creds.client_id, &creds.client_secret, &creds.token_endpoint, mstore, http).await {
193 Ok(access_token) => Some(access_token),
194 Err(e) => {
195 tracing::warn!(
196 agent_id = %key, error = %e,
197 "muxbus: agent m2m token fetch failed, backing off {}s",
198 CREDENTIAL_RETRY_COOLDOWN.as_secs(),
199 );
200 drop_unrecognized_credential(&key, &e, mstore);
201 record_credential_failure(&key);
202 None
203 }
204 }
205}
206
207pub fn invalidate_cached_token(agent_id: &str, mstore: &Arc<Store>) {
215 if let Err(e) = mstore.agent_credential_invalidate_token(&agent_id.to_lowercase()) {
216 tracing::warn!(
217 agent_id = %agent_id, error = %e,
218 "muxbus: failed to invalidate stale per-agent credential",
219 );
220 }
221}
222
223pub fn invalidate_binding_mismatched_credential(agent_id: &str, mstore: &Arc<Store>) {
244 let key = agent_id.to_lowercase();
245 invalidate_cached_token(&key, mstore);
246 record_credential_failure(&key);
247}
248
249#[cfg(test)]
250mod tests {
251 use super::*;
252
253 #[test]
259 fn an_untouched_agent_has_no_recorded_failure() {
260 assert!(!credential_recently_failed("test-agent-credentials-untouched"));
261 }
262
263 #[test]
264 fn record_credential_failure_is_visible_to_credential_recently_failed() {
265 let agent_id = "test-agent-credentials-record-failure";
266 assert!(!credential_recently_failed(agent_id));
267 record_credential_failure(agent_id);
268 assert!(credential_recently_failed(agent_id));
269 }
270
271 #[test]
278 fn invalidate_binding_mismatched_credential_starts_a_cooldown() {
279 let mstore = Arc::new(crate::backend::storage::store::Store::open_in_memory().unwrap());
280 let agent_id = "test-agent-credentials-binding-mismatch";
281 assert!(!credential_recently_failed(agent_id));
282
283 invalidate_binding_mismatched_credential(agent_id, &mstore);
284
285 assert!(
286 credential_recently_failed(agent_id),
287 "a 403 binding mismatch must cool down the per-agent pipeline, not just clear the cached token"
288 );
289 }
290
291 #[test]
296 fn invalidate_cached_token_alone_does_not_start_a_cooldown() {
297 let mstore = Arc::new(crate::backend::storage::store::Store::open_in_memory().unwrap());
298 let agent_id = "test-agent-credentials-plain-401";
299
300 invalidate_cached_token(agent_id, &mstore);
301
302 assert!(!credential_recently_failed(agent_id));
303 }
304
305 #[test]
309 fn only_invalid_client_on_400_or_401_means_the_client_is_unrecognized() {
310 let invalid_client = r#"{"error":"invalid_client"}"#;
311 assert!(token_rejection_is_unrecognized(400, invalid_client));
312 assert!(token_rejection_is_unrecognized(401, invalid_client));
313 assert!(!token_rejection_is_unrecognized(500, invalid_client));
314 assert!(!token_rejection_is_unrecognized(429, invalid_client));
315 assert!(!token_rejection_is_unrecognized(400, r#"{"error":"invalid_request"}"#));
316 assert!(!token_rejection_is_unrecognized(400, "<html>bad gateway</html>"));
317
318 assert!(provision_rejection_is_unrecognized(401));
319 assert!(provision_rejection_is_unrecognized(403));
320 for status in [400, 404, 429, 500, 503] {
321 assert!(!provision_rejection_is_unrecognized(status), "{status}");
322 }
323 }
324
325 type Route = (&'static str, u16, serde_json::Value);
326
327 async fn fake_cloud(routes: impl FnOnce(&str) -> Vec<Route>) -> (String, Arc<Mutex<Vec<String>>>) {
331 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
332 let base = format!("http://{}", listener.local_addr().unwrap());
333 let routes = routes(&base);
334 let seen: Arc<Mutex<Vec<String>>> = Arc::default();
335 let app = axum::Router::new().fallback({
336 let seen = Arc::clone(&seen);
337 move |req: axum::extract::Request| {
338 let (seen, routes) = (Arc::clone(&seen), routes.clone());
339 async move {
340 let line = format!("{} {}", req.method(), req.uri().path());
341 seen.lock().unwrap().push(line.clone());
342 let (status, body) = routes
343 .iter()
344 .find(|(route, _, _)| *route == line)
345 .map(|(_, status, body)| (*status, body.clone()))
346 .unwrap_or((404, serde_json::Value::Null));
347 (axum::http::StatusCode::from_u16(status).unwrap(), axum::Json(body))
348 }
349 }
350 });
351 tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
352 (base, seen)
353 }
354
355 fn reprovisioning(base: &str) -> Vec<Route> {
357 vec![
358 (
359 "POST /agents/provision",
360 200,
361 serde_json::json!({
362 "client_id": "new-client",
363 "client_secret": "new-secret",
364 "token_endpoint": format!("{base}/new/oauth2/token"),
365 }),
366 ),
367 ("POST /new/oauth2/token", 200, serde_json::json!({ "access_token": "fresh", "expires_in": 3600 })),
368 ]
369 }
370
371 fn unique_agent() -> String {
372 format!("test-agent-credentials-{}", uuid::Uuid::new_v4())
373 }
374
375 fn provisioned_store(agent: &str, token_endpoint: &str) -> Arc<Store> {
377 let mstore = Arc::new(Store::open_in_memory().unwrap());
378 mstore.agent_credential_save(agent, "old-client", "old-secret", token_endpoint).unwrap();
379 mstore
380 }
381
382 async fn ensure(agent: &str, mstore: &Arc<Store>, base: &str) -> Option<String> {
383 let user_token = async { Ok("user-token".to_string()) };
384 ensure_agent_credential_at(agent, mstore, &reqwest::Client::new(), base, user_token).await
385 }
386
387 fn clear_cooldown(agent: &str) {
388 credential_cooldown().lock().unwrap().remove(agent);
389 }
390
391 #[tokio::test]
392 async fn invalid_client_drops_the_row_and_the_next_use_re_provisions() {
393 for status in [400, 401] {
394 let agent = unique_agent();
395 let (base, seen) = fake_cloud(|base| {
396 let mut routes = reprovisioning(base);
397 routes.push(("POST /old/oauth2/token", status, serde_json::json!({ "error": "invalid_client" })));
398 routes
399 })
400 .await;
401 let mstore = provisioned_store(&agent, &format!("{base}/old/oauth2/token"));
402
403 assert_eq!(ensure(&agent, &mstore, &base).await, None);
404 assert!(mstore.agent_credential_load(&agent).unwrap().is_none(), "{status}: row dropped");
405 assert!(credential_recently_failed(&agent), "{status}: the cooldown still applies");
406
407 clear_cooldown(&agent);
409 assert_eq!(ensure(&agent, &mstore, &base).await.as_deref(), Some("fresh"));
410 assert_eq!(mstore.agent_credential_load(&agent).unwrap().unwrap().client_id, "new-client");
411 assert_eq!(
412 *seen.lock().unwrap(),
413 ["POST /old/oauth2/token", "POST /agents/provision", "POST /new/oauth2/token"]
414 );
415 }
416 }
417
418 #[tokio::test]
419 async fn a_401_or_403_from_provisioning_drops_the_agents_row() {
420 for status in [401, 403] {
421 let agent = unique_agent();
422 let (base, _) = fake_cloud(|_| vec![("POST /agents/provision", status, serde_json::json!({ "error": "forbidden" }))]).await;
423 let mstore = Arc::new(Store::open_in_memory().unwrap());
425 mstore.agent_credential_save(&agent, "", "", "").unwrap();
426
427 assert_eq!(ensure(&agent, &mstore, &base).await, None);
428 assert!(mstore.agent_credential_load(&agent).unwrap().is_none(), "{status}: row dropped");
429 assert!(credential_recently_failed(&agent), "{status}: the cooldown still applies");
430 }
431 }
432
433 #[tokio::test]
434 async fn transient_failures_keep_the_row() {
435 for (status, body) in [
437 (500, serde_json::json!({ "error": "invalid_client" })),
438 (503, serde_json::Value::Null),
439 (429, serde_json::json!({ "error": "slow_down" })),
440 (400, serde_json::json!({ "error": "invalid_request" })),
441 ] {
442 let agent = unique_agent();
443 let (base, _) = fake_cloud(|_| vec![("POST /old/oauth2/token", status, body)]).await;
444 let mstore = provisioned_store(&agent, &format!("{base}/old/oauth2/token"));
445
446 assert_eq!(ensure(&agent, &mstore, &base).await, None);
447 assert_eq!(
448 mstore.agent_credential_load(&agent).unwrap().unwrap().client_id,
449 "old-client",
450 "{status}: row kept"
451 );
452 assert!(credential_recently_failed(&agent));
453 }
454
455 let closed = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
457 let dead = format!("http://{}/oauth2/token", closed.local_addr().unwrap());
458 drop(closed);
459 let agent = unique_agent();
460 let mstore = provisioned_store(&agent, &dead);
461 assert_eq!(ensure(&agent, &mstore, "http://127.0.0.1:9").await, None);
462 assert_eq!(mstore.agent_credential_load(&agent).unwrap().unwrap().client_id, "old-client");
463 }
464}
465
466async fn user_token(mstore: &Arc<Store>) -> Result<String, String> {
468 let scheduler = crate::broker::get_global()
475 .ok_or_else(|| "muxbus refresh scheduler not initialized yet".to_string())?;
476 load_valid_token(mstore, &scheduler)
477 .await
478 .ok_or_else(|| "no valid user-level muxbus login to provision from".to_string())
479}
480
481async fn provision_agent_client(
485 agent_id: &str,
486 mstore: &Arc<Store>,
487 http: &reqwest::Client,
488 base: &str,
489 user_token: impl std::future::Future<Output = Result<String, String>>,
490) -> Result<(), CredentialFailure> {
491 let user_token = user_token.await.map_err(CredentialFailure::Other)?;
492
493 #[derive(serde::Deserialize)]
494 struct ProvisionResp {
495 client_id: String,
496 client_secret: String,
497 token_endpoint: String,
498 }
499
500 let url = format!("{}/agents/provision", base);
501 let resp = http
507 .post(&url)
508 .header("Authorization", format!("Bearer {}", user_token))
509 .json(&serde_json::json!({ "agent_id": agent_id }))
510 .timeout(CREDENTIAL_HTTP_TIMEOUT)
511 .send()
512 .await
513 .map_err(|e| CredentialFailure::Other(format!("provision request failed: {e}")))?;
514
515 if !resp.status().is_success() {
516 let status = resp.status().as_u16();
517 let body = resp.text().await.unwrap_or_default();
518 let msg = format!("provision failed ({status}): {body}");
519 return Err(if provision_rejection_is_unrecognized(status) {
520 CredentialFailure::Unrecognized(msg)
521 } else {
522 CredentialFailure::Other(msg)
523 });
524 }
525
526 let parsed: ProvisionResp = resp
527 .json()
528 .await
529 .map_err(|e| CredentialFailure::Other(format!("provision response parse failed: {e}")))?;
530
531 mstore
532 .agent_credential_save(agent_id, &parsed.client_id, &parsed.client_secret, &parsed.token_endpoint)
533 .map_err(|e| CredentialFailure::Other(format!("failed to save agent credential: {e}")))?;
534
535 tracing::info!(agent_id = %agent_id, "muxbus: provisioned per-agent credential");
536 Ok(())
537}
538
539async fn fetch_m2m_token(
543 agent_id: &str,
544 client_id: &str,
545 client_secret: &str,
546 token_endpoint: &str,
547 mstore: &Arc<Store>,
548 http: &reqwest::Client,
549) -> Result<String, CredentialFailure> {
550 if client_id.is_empty() || token_endpoint.is_empty() {
551 return Err(CredentialFailure::Other(
552 "agent credential missing client_id/token_endpoint".to_string(),
553 ));
554 }
555
556 let params = [
557 ("grant_type", "client_credentials"),
558 ("client_id", client_id),
559 ("client_secret", client_secret),
560 ];
561 let resp = http
562 .post(token_endpoint)
563 .form(¶ms)
564 .timeout(CREDENTIAL_HTTP_TIMEOUT)
565 .send()
566 .await
567 .map_err(|e| CredentialFailure::Other(format!("m2m token request failed: {e}")))?;
568
569 if !resp.status().is_success() {
570 let status = resp.status().as_u16();
571 let body = resp.text().await.unwrap_or_default();
572 let msg = format!("m2m token request failed ({status}): {body}");
573 return Err(if token_rejection_is_unrecognized(status, &body) {
574 CredentialFailure::Unrecognized(msg)
575 } else {
576 CredentialFailure::Other(msg)
577 });
578 }
579
580 let json: serde_json::Value = resp
581 .json()
582 .await
583 .map_err(|e| CredentialFailure::Other(format!("m2m token response parse failed: {e}")))?;
584
585 let access_token = json["access_token"].as_str().unwrap_or("").to_string();
586 if access_token.is_empty() {
587 return Err(CredentialFailure::Other("m2m token response missing access_token".to_string()));
588 }
589 let expires_in = json["expires_in"].as_i64().unwrap_or(3600);
590 let expires_at = (agentmux_common::time::now_secs())
591 + expires_in;
592
593 if let Err(e) = mstore.agent_credential_save_token(agent_id, &access_token, expires_at) {
594 tracing::warn!(agent_id = %agent_id, error = %e, "muxbus: failed to cache m2m token (will re-fetch next time)");
595 }
596
597 Ok(access_token)
598}