Skip to main content

aether_auth/
browser.rs

1use crate::error::OAuthError;
2use crate::handler::OAuthHandler;
3use futures::future::BoxFuture;
4use std::process::Command;
5use std::time::Duration;
6use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
7use tokio::net::TcpListener;
8use tokio::time::timeout;
9
10/// Default `OAuthHandler` that opens the system browser and listens
11/// for the OAuth callback on a dynamically-assigned local port.
12pub struct BrowserOAuthHandler {
13    listener: TcpListener,
14    redirect_uri: String,
15}
16
17impl BrowserOAuthHandler {
18    pub fn new() -> Result<Self, std::io::Error> {
19        let std_listener = std::net::TcpListener::bind("127.0.0.1:0")?;
20        let port = std_listener.local_addr()?.port();
21        std_listener.set_nonblocking(true)?;
22        let listener = TcpListener::from_std(std_listener)?;
23        Ok(Self { listener, redirect_uri: format!("http://127.0.0.1:{port}/oauth2callback") })
24    }
25
26    /// Create a handler bound to a specific port with a custom redirect URI.
27    pub fn with_redirect_uri(redirect_uri: impl Into<String>, port: u16) -> Result<Self, std::io::Error> {
28        let std_listener = std::net::TcpListener::bind(format!("127.0.0.1:{port}"))?;
29        std_listener.set_nonblocking(true)?;
30        let listener = TcpListener::from_std(std_listener)?;
31        Ok(Self { listener, redirect_uri: redirect_uri.into() })
32    }
33}
34
35impl OAuthHandler for BrowserOAuthHandler {
36    fn redirect_uri(&self) -> &str {
37        &self.redirect_uri
38    }
39
40    fn authorize(&self, auth_url: &str) -> BoxFuture<'_, Result<String, OAuthError>> {
41        let auth_url = auth_url.to_string();
42        Box::pin(async move {
43            if let Err(error) = open_browser(&auth_url) {
44                tracing::warn!("Failed to open browser: {error}");
45            }
46            accept_oauth_callback(&self.listener).await
47        })
48    }
49}
50
51/// Accept an OAuth callback and return its absolute URL.
52pub async fn accept_oauth_callback(listener: &TcpListener) -> Result<String, OAuthError> {
53    loop {
54        let (mut socket, _) = listener.accept().await?;
55        let request_line = {
56            let mut reader = BufReader::new(&mut socket);
57            let mut line = String::new();
58            let bytes_read =
59                timeout(Duration::from_secs(2), reader.read_line(&mut line)).await.ok().and_then(Result::ok);
60            let Some(1..) = bytes_read else { continue };
61            line
62        };
63
64        match callback_url(&request_line) {
65            Ok(callback_url) => {
66                let _ = socket.write_all(success_response().as_bytes()).await;
67                return Ok(callback_url);
68            }
69            Err(error) if request_line.contains('?') => return Err(error),
70            Err(_) => {}
71        }
72    }
73}
74
75/// Start a local callback server and return the OAuth callback URL.
76pub async fn wait_for_callback(port: u16) -> Result<String, OAuthError> {
77    let listener = TcpListener::bind(format!("127.0.0.1:{port}")).await?;
78    accept_oauth_callback(&listener).await
79}
80
81/// Open a URL in the default browser.
82pub fn open_browser(url: &str) -> Result<(), OAuthError> {
83    #[cfg(target_os = "macos")]
84    {
85        Command::new("open").arg(url).spawn().map_err(std::io::Error::other)?;
86    }
87
88    #[cfg(target_os = "linux")]
89    {
90        Command::new("xdg-open").arg(url).spawn().map_err(std::io::Error::other)?;
91    }
92
93    #[cfg(target_os = "windows")]
94    {
95        Command::new("cmd").args(["/C", "start", url]).spawn().map_err(std::io::Error::other)?;
96    }
97
98    Ok(())
99}
100
101fn callback_url(request_line: &str) -> Result<String, OAuthError> {
102    let path = request_line
103        .split_whitespace()
104        .nth(1)
105        .ok_or_else(|| OAuthError::InvalidCallback("Invalid HTTP request format".to_string()))?;
106    let url = url::Url::parse(&format!("http://localhost{path}"))
107        .map_err(|error| OAuthError::InvalidCallback(format!("Invalid callback URL: {error}")))?;
108    if url.query().is_none() {
109        return Err(OAuthError::InvalidCallback("No query parameters in callback".to_string()));
110    }
111    if url.query_pairs().any(|(key, _)| key == "error") {
112        return Err(OAuthError::InvalidCallback("OAuth authorization failed".to_string()));
113    }
114    Ok(url.to_string())
115}
116
117fn success_response() -> String {
118    let body = include_str!("oauth_success.html");
119    format!(
120        "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
121        body.len(),
122        body
123    )
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129
130    #[test]
131    fn callback_request_preserves_encoded_parameters() {
132        let callback = callback_url(
133            "GET /oauth2callback?code=caf%C3%A9&state=test+state&iss=https%3A%2F%2Fauth.example.com HTTP/1.1\r\n",
134        )
135        .unwrap();
136        let url = url::Url::parse(&callback).unwrap();
137        let params = url.query_pairs().collect::<std::collections::HashMap<_, _>>();
138        assert_eq!(params.get("code").map(AsRef::as_ref), Some("café"));
139        assert_eq!(params.get("state").map(AsRef::as_ref), Some("test state"));
140        assert_eq!(params.get("iss").map(AsRef::as_ref), Some("https://auth.example.com"));
141    }
142
143    #[test]
144    fn callback_error_is_sanitized() {
145        let error =
146            callback_url("GET /oauth2callback?error=access_denied&error_description=attacker+controlled HTTP/1.1\r\n")
147                .unwrap_err()
148                .to_string();
149        assert!(error.contains("OAuth authorization failed"));
150        assert!(!error.contains("attacker"));
151    }
152
153    #[tokio::test]
154    async fn callback_listener_skips_stale_requests() {
155        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
156        let port = listener.local_addr().unwrap().port();
157        let handle = tokio::spawn(async move { accept_oauth_callback(&listener).await });
158
159        let mut stale = tokio::net::TcpStream::connect(("127.0.0.1", port)).await.unwrap();
160        stale.write_all(b"GET /favicon.ico HTTP/1.1\r\n").await.unwrap();
161        let mut callback = tokio::net::TcpStream::connect(("127.0.0.1", port)).await.unwrap();
162        callback.write_all(b"GET /?code=abc&state=xyz HTTP/1.1\r\n").await.unwrap();
163
164        assert!(handle.await.unwrap().unwrap().contains("code=abc&state=xyz"));
165    }
166
167    #[tokio::test]
168    async fn custom_redirect_uri_is_retained() {
169        let handler = BrowserOAuthHandler::with_redirect_uri("http://localhost:9999/callback", 0).unwrap();
170        assert_eq!(handler.redirect_uri(), "http://localhost:9999/callback");
171    }
172}