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
10pub 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 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
51pub 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
75pub 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
81pub 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}