oxicode/
oauth_listener.rs

1//! Single-shot HTTP listener for OAuth `authorization_code` callbacks.
2//! Binds an ephemeral 127.0.0.1 port, accepts one connection, parses the
3//! `GET <path>?<query>`, returns the `code` and `state`.
4
5use std::time::Duration;
6use tokio::io::{AsyncReadExt, AsyncWriteExt};
7use tokio::net::TcpListener;
8
9#[derive(Debug, Clone)]
10pub struct CallbackReceived {
11    pub code: String,
12    pub state: String,
13}
14
15#[derive(Debug, thiserror::Error)]
16pub enum CallbackError {
17    #[error("timeout waiting for OAuth callback")]
18    Timeout,
19    #[error("invalid request: {0}")]
20    BadRequest(String),
21    #[error("state mismatch (expected {expected:?})")]
22    StateMismatch { expected: String },
23    #[error("missing `code` in callback")]
24    MissingCode,
25    #[error("path mismatch (expected {expected:?})")]
26    PathMismatch { expected: String },
27}
28
29pub async fn await_callback(
30    listener: TcpListener,
31    expected_state: String,
32    expected_path: String,
33    timeout: Duration,
34) -> Result<CallbackReceived, CallbackError> {
35    implement(listener, expected_state, expected_path, timeout).await
36}
37
38async fn implement(
39    listener: TcpListener,
40    expected_state: String,
41    expected_path: String,
42    timeout: Duration,
43) -> Result<CallbackReceived, CallbackError> {
44    let accept = async {
45        let (stream, _addr) = listener
46            .accept()
47            .await
48            .map_err(|e| CallbackError::BadRequest(e.to_string()))?;
49        Ok::<_, CallbackError>(stream)
50    };
51    let timeout_fut = tokio::time::sleep(timeout);
52    tokio::pin!(timeout_fut);
53    let mut stream = tokio::select! {
54        biased;
55        _ = &mut timeout_fut => return Err(CallbackError::Timeout),
56        s = accept => s?,
57    };
58
59    let mut header_buf = Vec::with_capacity(512);
60    let mut tmp = [0u8; 1024];
61    let header_end = loop {
62        let n = tokio::time::timeout(Duration::from_secs(5), stream.read(&mut tmp))
63            .await
64            .map_err(|_| CallbackError::BadRequest("header read timeout".into()))?
65            .map_err(|e| CallbackError::BadRequest(e.to_string()))?;
66        if n == 0 {
67            return Err(CallbackError::BadRequest("empty request".into()));
68        }
69        header_buf.extend_from_slice(&tmp[..n]);
70        if header_buf.len() > 8192 {
71            return Err(CallbackError::BadRequest("headers too large".into()));
72        }
73        if let Some(pos) = find_header_end(&header_buf) {
74            break pos;
75        }
76    };
77
78    let mut headers = [httparse::EMPTY_HEADER; 16];
79    let mut req = httparse::Request::new(&mut headers);
80    let parsed = req
81        .parse(&header_buf)
82        .map_err(|e| CallbackError::BadRequest(format!("header parse: {e}")))?;
83    if !parsed.is_complete() {
84        return Err(CallbackError::BadRequest("incomplete headers".into()));
85    }
86
87    let method = req
88        .method
89        .ok_or_else(|| CallbackError::BadRequest("missing method".into()))?;
90    if method != "GET" {
91        return Err(CallbackError::BadRequest(format!(
92            "expected GET, got {method}"
93        )));
94    }
95    let path_full = req
96        .path
97        .ok_or_else(|| CallbackError::BadRequest("missing path".into()))?;
98
99    let (path, query) = match path_full.split_once('?') {
100        Some((p, q)) => (p, q),
101        None => (path_full, ""),
102    };
103    if path != expected_path {
104        return Err(CallbackError::PathMismatch {
105            expected: expected_path,
106        });
107    }
108
109    let mut code: Option<String> = None;
110    let mut state: Option<String> = None;
111    for (k, v) in url::form_urlencoded::parse(query.as_bytes()) {
112        match k.as_ref() {
113            "code" => code = Some(v.into_owned()),
114            "state" => state = Some(v.into_owned()),
115            _ => {}
116        }
117    }
118    let _ = header_end; // keep header_end alive for clarity; full body is discarded.
119
120    let Some(received_state) = state else {
121        return Err(CallbackError::BadRequest(
122            "missing `state` in callback".into(),
123        ));
124    };
125    if received_state != expected_state {
126        return Err(CallbackError::StateMismatch {
127            expected: expected_state,
128        });
129    }
130    let Some(code) = code else {
131        return Err(CallbackError::MissingCode);
132    };
133
134    let body =
135        "<!DOCTYPE html><html><body>Login complete. You may close this window.</body></html>";
136    let response = format!(
137        "HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=utf-8\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
138        body.len(),
139        body,
140    );
141    let _ = stream.write_all(response.as_bytes()).await;
142    let _ = stream.shutdown().await;
143
144    Ok(CallbackReceived {
145        code,
146        state: received_state,
147    })
148}
149
150fn find_header_end(buf: &[u8]) -> Option<usize> {
151    buf.windows(4).position(|w| w == b"\r\n\r\n")
152}
153#[cfg(test)]
154mod tests {
155    use super::*;
156    use std::time::Duration;
157    use tokio::io::AsyncWriteExt;
158    use tokio::net::TcpStream;
159
160    async fn drive_callback(
161        request: &str,
162        expected_state: &str,
163        expected_path: &str,
164        timeout: Duration,
165    ) -> Result<CallbackReceived, CallbackError> {
166        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
167        let port = listener.local_addr().unwrap().port();
168        let task = tokio::spawn(await_callback(
169            listener,
170            expected_state.to_string(),
171            expected_path.to_string(),
172            timeout,
173        ));
174        tokio::time::sleep(Duration::from_millis(20)).await;
175        let mut conn = TcpStream::connect(("127.0.0.1", port)).await.unwrap();
176        conn.write_all(request.as_bytes()).await.unwrap();
177        conn.flush().await.unwrap();
178        task.await.unwrap()
179    }
180
181    #[tokio::test]
182    async fn parses_valid_callback() {
183        let req = "GET /callback?code=abc&state=ST HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n";
184        let got = drive_callback(req, "ST", "/callback", Duration::from_secs(2))
185            .await
186            .unwrap();
187        assert_eq!(got.code, "abc");
188        assert_eq!(got.state, "ST");
189    }
190
191    #[tokio::test]
192    async fn rejects_state_mismatch() {
193        let req = "GET /callback?code=abc&state=OTHER HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n";
194        let err = drive_callback(req, "ST", "/callback", Duration::from_secs(2))
195            .await
196            .unwrap_err();
197        assert!(matches!(err, CallbackError::StateMismatch { .. }));
198    }
199
200    #[tokio::test]
201    async fn rejects_missing_code() {
202        let req = "GET /callback?state=ST HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n";
203        let err = drive_callback(req, "ST", "/callback", Duration::from_secs(2))
204            .await
205            .unwrap_err();
206        assert!(matches!(err, CallbackError::MissingCode));
207    }
208
209    #[tokio::test]
210    async fn rejects_missing_state_returns_bad_request() {
211        let req = "GET /callback?code=abc HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n";
212        let err = drive_callback(req, "ST", "/callback", Duration::from_secs(2))
213            .await
214            .unwrap_err();
215        match &err {
216            CallbackError::BadRequest(msg) => assert!(
217                msg.contains("state"),
218                "expected BadRequest mentioning `state`, got {msg:?}"
219            ),
220            other => panic!("expected BadRequest, got {other:?}"),
221        }
222    }
223
224    #[tokio::test]
225    async fn rejects_path_mismatch() {
226        let req = "GET /wrong?code=abc&state=ST HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n";
227        let err = drive_callback(req, "ST", "/callback", Duration::from_secs(2))
228            .await
229            .unwrap_err();
230        assert!(matches!(err, CallbackError::PathMismatch { .. }));
231    }
232
233    #[tokio::test]
234    async fn timeout_when_no_connection() {
235        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
236        let err = await_callback(
237            listener,
238            "ST".into(),
239            "/callback".into(),
240            Duration::from_millis(100),
241        )
242        .await
243        .unwrap_err();
244        assert!(matches!(err, CallbackError::Timeout));
245    }
246}
oxicode/oauth_listener.rs

oxicode/
oauth_listener.rs
