agent_core/core/auth/
callback.rs

1use std::{collections::HashMap, sync::Arc};
2use tokio::sync::{oneshot, Mutex};
3
4use super::{CallbackResult, CALLBACK_HOST};
5
6pub(crate) const SUCCESS_HTML: &str = "<!doctype html><title>Login Successful</title><h1>Authentication successful</h1><p>You can close this window.</p>";
7pub(crate) const ERROR_HTML: &str = "<!doctype html><title>Login Failed</title><h1>Authentication failed</h1><p>Return to the terminal.</p>";
8
9#[derive(Debug, Clone, PartialEq, Eq)]
10pub enum CallbackOutcome {
11    Authorized(CallbackResult),
12    Denied {
13        error: String,
14        description: Option<String>,
15    },
16    Invalid,
17}
18
19pub struct CallbackServerHandle {
20    shutdown: Option<oneshot::Sender<()>>,
21    task: Option<tokio::task::JoinHandle<()>>,
22}
23impl CallbackServerHandle {
24    pub async fn shutdown(mut self) {
25        if let Some(tx) = self.shutdown.take() {
26            let _ = tx.send(());
27        }
28        if let Some(task) = self.task.take() {
29            let _ = tokio::time::timeout(std::time::Duration::from_secs(2), task).await;
30        }
31    }
32}
33impl Drop for CallbackServerHandle {
34    fn drop(&mut self) {
35        if let Some(tx) = self.shutdown.take() {
36            let _ = tx.send(());
37        }
38        if let Some(task) = self.task.take() {
39            task.abort();
40        }
41    }
42}
43
44pub async fn start_callback_server(
45    expected_state: String,
46    port: u16,
47) -> Result<(oneshot::Receiver<CallbackOutcome>, CallbackServerHandle), String> {
48    start_callback_server_at(expected_state, CALLBACK_HOST, port, "/callback").await
49}
50
51pub async fn start_callback_server_at(
52    expected_state: String,
53    host: &str,
54    port: u16,
55    path: &str,
56) -> Result<(oneshot::Receiver<CallbackOutcome>, CallbackServerHandle), String> {
57    if !path.starts_with('/') || path.contains("..") {
58        return Err("invalid callback path".into());
59    }
60    let host_ip: std::net::IpAddr = host
61        .parse()
62        .map_err(|_| "callback host must be a loopback IP address")?;
63    if !host_ip.is_loopback() {
64        return Err("callback host must be loopback".into());
65    }
66    let (tx, rx) = oneshot::channel();
67    let tx = Arc::new(Mutex::new(Some(tx)));
68    let handler = {
69        let tx = tx.clone();
70        move |axum::extract::Query(query): axum::extract::Query<HashMap<String, String>>| {
71            let tx = tx.clone();
72            let expected = expected_state.clone();
73            async move {
74                let state_matches = query.get("state") == Some(&expected);
75                let outcome = if let Some(error) = query.get("error") {
76                    if state_matches {
77                        CallbackOutcome::Denied {
78                            error: sanitize(error),
79                            description: query.get("error_description").map(|s| sanitize(s)),
80                        }
81                    } else {
82                        CallbackOutcome::Invalid
83                    }
84                } else if let (Some(code), Some(state)) = (query.get("code"), query.get("state")) {
85                    if state == &expected && !code.trim().is_empty() {
86                        CallbackOutcome::Authorized(CallbackResult {
87                            code: code.clone(),
88                            state: state.clone(),
89                        })
90                    } else {
91                        CallbackOutcome::Invalid
92                    }
93                } else {
94                    CallbackOutcome::Invalid
95                };
96                if let Some(sender) = tx.lock().await.take() {
97                    let _ = sender.send(outcome.clone());
98                }
99                axum::response::Html(if matches!(outcome, CallbackOutcome::Authorized(_)) {
100                    SUCCESS_HTML
101                } else {
102                    ERROR_HTML
103                })
104            }
105        }
106    };
107    let app = axum::Router::new().route(path, axum::routing::get(handler));
108    let addr = format!("{host}:{port}");
109    let listener = tokio::net::TcpListener::bind(&addr)
110        .await
111        .map_err(|e| format!("Failed to bind callback server on {addr}: {e}"))?;
112    let (shutdown, stop) = oneshot::channel();
113    let task = tokio::spawn(async move {
114        let _ = axum::serve(listener, app)
115            .with_graceful_shutdown(async {
116                let _ = stop.await;
117            })
118            .await;
119    });
120    Ok((
121        rx,
122        CallbackServerHandle {
123            shutdown: Some(shutdown),
124            task: Some(task),
125        },
126    ))
127}
128
129fn sanitize(value: &str) -> String {
130    value
131        .chars()
132        .filter(|c| c.is_ascii_alphanumeric() || matches!(c, ' ' | '_' | '-' | '.'))
133        .take(160)
134        .collect()
135}
136
137#[cfg(test)]
138mod tests {
139    use super::*;
140    #[tokio::test]
141    async fn empty_http_code_is_invalid_and_listener_is_released() {
142        let probe = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
143        let port = probe.local_addr().unwrap().port();
144        drop(probe);
145        let (rx, handle) = start_callback_server("expected".into(), port)
146            .await
147            .unwrap();
148        let response = reqwest::get(format!(
149            "http://127.0.0.1:{port}/callback?code=&state=expected"
150        ))
151        .await
152        .unwrap();
153        assert!(!response.text().await.unwrap().contains("successful</h1>"));
154        assert_eq!(rx.await.unwrap(), CallbackOutcome::Invalid);
155        handle.shutdown().await;
156        assert!(tokio::net::TcpListener::bind(("127.0.0.1", port))
157            .await
158            .is_ok());
159    }
160
161    #[test]
162    fn sanitizes_provider_errors() {
163        assert_eq!(sanitize("denied\nsecret=?"), "deniedsecret");
164    }
165}
agent_core/core/auth/callback.rs

agent_core/core/auth/
callback.rs
