yt/auth/
loopback.rs

1use std::collections::HashMap;
2use std::time::Duration;
3
4use miette::IntoDiagnostic as _;
5use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
6use tokio::net::{TcpListener, TcpStream};
7use tokio::sync::mpsc;
8use url::Url;
9
10use crate::error::{Result, YtError};
11
12pub const CALLBACK_PATH: &str = "/callback";
13const READ_TIMEOUT: Duration = Duration::from_secs(10);
14
15/// A one-shot listener on 127.0.0.1 for the browser's redirect after sign-in.
16pub struct Callback {
17    listener: TcpListener,
18    port: u16,
19}
20
21pub fn redirect_uri(port: u16) -> String {
22    format!("http://127.0.0.1:{port}{CALLBACK_PATH}")
23}
24
25impl Callback {
26    pub async fn bind() -> Result<Self> {
27        let listener = TcpListener::bind(("127.0.0.1", 0))
28            .await
29            .into_diagnostic()?;
30        let port = listener.local_addr().into_diagnostic()?.port();
31        Ok(Self { listener, port })
32    }
33
34    pub fn redirect_uri(&self) -> String {
35        redirect_uri(self.port)
36    }
37
38    /// Waits for the redirect and returns the authorization code.
39    pub async fn wait(self, state: &str, timeout: Duration) -> Result<String> {
40        tokio::time::timeout(timeout, self.accept(state))
41            .await
42            .map_err(|_| {
43                miette::Report::from(YtError::OAuth {
44                    message: "the browser did not come back".into(),
45                    help: Some(
46                        "run `yt auth login --no-browser` to finish by pasting the address instead"
47                            .into(),
48                    ),
49                })
50            })?
51    }
52
53    async fn accept(&self, state: &str) -> Result<String> {
54        // Each connection is read on its own task, so an idle one (a browser preconnect) cannot hold up the callback.
55        let (requests, mut callbacks) = mpsc::channel(8);
56        loop {
57            tokio::select! {
58                accepted = self.listener.accept() => {
59                    let (stream, _) = accepted.into_diagnostic()?;
60                    tokio::spawn(read_callback(stream, requests.clone()));
61                }
62                Some((mut stream, params)) = callbacks.recv() => {
63                    if !state_matches(&params, state) {
64                        respond(&mut stream, 400, "This address is not for the running sign-in.").await;
65                        continue;
66                    }
67                    let result = code_from_params(&params, state);
68                    match &result {
69                        Ok(_) => {
70                            respond(
71                                &mut stream,
72                                200,
73                                "Signed in. You can close this tab and return to the terminal.",
74                            )
75                            .await
76                        }
77                        Err(_) => {
78                            respond(
79                                &mut stream,
80                                400,
81                                "Sign-in failed. Return to the terminal for details.",
82                            )
83                            .await
84                        }
85                    }
86                    return result;
87                }
88            }
89        }
90    }
91}
92
93type Request = (TcpStream, HashMap<String, String>);
94
95async fn read_callback(mut stream: TcpStream, callbacks: mpsc::Sender<Request>) {
96    let Ok(Some(target)) = tokio::time::timeout(READ_TIMEOUT, read_target(&mut stream)).await
97    else {
98        return;
99    };
100    match callback_params(&target) {
101        Some(params) => {
102            let _ = callbacks.send((stream, params)).await;
103        }
104        None => respond(&mut stream, 404, "Not found.").await,
105    }
106}
107
108/// The code from the address of the page the browser landed on, pasted by the user (SSH or headless sign-in).
109pub fn code_from_pasted(pasted: &str, state: &str) -> Result<String> {
110    let url = Url::parse(pasted.trim()).map_err(|_| YtError::Usage {
111        message: "that is not the address of the page the browser landed on".into(),
112    })?;
113    let params: HashMap<String, String> = url.query_pairs().into_owned().collect();
114    code_from_params(&params, state)
115}
116
117fn state_matches(params: &HashMap<String, String>, state: &str) -> bool {
118    params.get("state").map(String::as_str) == Some(state)
119}
120
121/// Checks `state` before anything else, so a response that is not for this sign-in cannot show its text.
122fn code_from_params(params: &HashMap<String, String>, state: &str) -> Result<String> {
123    if !state_matches(params, state) {
124        return Err(YtError::OAuth {
125            message: "the sign-in response did not match this request".into(),
126            help: Some("start `yt auth login` again".into()),
127        }
128        .into());
129    }
130    if let Some(error) = params.get("error") {
131        let message = params
132            .get("error_description")
133            .filter(|text| !text.is_empty())
134            .cloned()
135            .unwrap_or_else(|| error.clone());
136        return Err(YtError::OAuth {
137            message,
138            help: crate::auth::oauth::help_for(error),
139        }
140        .into());
141    }
142    params.get("code").cloned().ok_or_else(|| {
143        YtError::OAuth {
144            message: "the sign-in response had no code".into(),
145            help: None,
146        }
147        .into()
148    })
149}
150
151fn callback_params(target: &str) -> Option<HashMap<String, String>> {
152    let url = Url::parse(&format!("http://127.0.0.1{target}")).ok()?;
153    (url.path() == CALLBACK_PATH).then(|| url.query_pairs().into_owned().collect())
154}
155
156async fn read_target(stream: &mut TcpStream) -> Option<String> {
157    let mut buffer = Vec::with_capacity(1024);
158    let mut chunk = [0u8; 1024];
159    while !buffer.windows(4).any(|window| window == b"\r\n\r\n") && buffer.len() < 16 * 1024 {
160        let read = stream.read(&mut chunk).await.ok()?;
161        if read == 0 {
162            break;
163        }
164        buffer.extend_from_slice(&chunk[..read]);
165    }
166    let head = String::from_utf8_lossy(&buffer);
167    let mut parts = head.lines().next()?.split_whitespace();
168    if parts.next()? != "GET" {
169        return None;
170    }
171    parts.next().map(str::to_string)
172}
173
174async fn respond(stream: &mut TcpStream, status: u16, message: &str) {
175    let reason = match status {
176        200 => "OK",
177        400 => "Bad Request",
178        _ => "Not Found",
179    };
180    let body = format!("<!doctype html><meta charset=\"utf-8\"><title>yt</title><p>{message}</p>");
181    let response = format!(
182        "HTTP/1.1 {status} {reason}\r\nContent-Type: text/html; charset=utf-8\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
183        body.len()
184    );
185    let _ = stream.write_all(response.as_bytes()).await;
186    let _ = stream.shutdown().await;
187}
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192
193    async fn send(port: u16, target: &str) -> String {
194        let mut stream = TcpStream::connect(("127.0.0.1", port)).await.unwrap();
195        stream
196            .write_all(format!("GET {target} HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n").as_bytes())
197            .await
198            .unwrap();
199        let mut response = String::new();
200        stream.read_to_string(&mut response).await.unwrap();
201        response
202    }
203
204    async fn bound() -> (Callback, u16) {
205        let callback = Callback::bind().await.unwrap();
206        let port = callback.port;
207        (callback, port)
208    }
209
210    #[tokio::test]
211    async fn returns_the_code_and_ignores_other_paths() {
212        let (callback, port) = bound().await;
213        assert!(callback.redirect_uri().starts_with("http://127.0.0.1:"));
214        let browser = tokio::spawn(async move {
215            let favicon = send(port, "/favicon.ico").await;
216            let page = send(port, "/callback?code=c0de&state=s1").await;
217            (favicon, page)
218        });
219        let code = callback.wait("s1", Duration::from_secs(5)).await.unwrap();
220        assert_eq!(code, "c0de");
221        let (favicon, page) = browser.await.unwrap();
222        assert!(favicon.starts_with("HTTP/1.1 404"));
223        assert!(page.starts_with("HTTP/1.1 200") && page.contains("You can close this tab"));
224    }
225
226    #[tokio::test]
227    async fn requests_for_another_state_are_refused_and_the_wait_goes_on() {
228        let (callback, port) = bound().await;
229        let browser = tokio::spawn(async move {
230            let forged = send(port, "/callback?code=f0rged&state=forged").await;
231            let injected = send(port, "/callback?error=x&error_description=Injected").await;
232            let page = send(port, "/callback?code=c0de&state=s1").await;
233            (forged, injected, page)
234        });
235        let code = callback.wait("s1", Duration::from_secs(5)).await.unwrap();
236        assert_eq!(code, "c0de");
237        let (forged, injected, page) = browser.await.unwrap();
238        assert!(forged.starts_with("HTTP/1.1 400"), "{forged}");
239        assert!(injected.starts_with("HTTP/1.1 400"), "{injected}");
240        assert!(page.starts_with("HTTP/1.1 200"), "{page}");
241    }
242
243    #[tokio::test]
244    async fn an_oauth_error_in_the_redirect_is_reported() {
245        let (callback, port) = bound().await;
246        tokio::spawn(async move {
247            send(
248                port,
249                "/callback?error=access_denied&error_description=User%20denied&state=s1",
250            )
251            .await
252        });
253        let error = callback
254            .wait("s1", Duration::from_secs(5))
255            .await
256            .unwrap_err();
257        assert!(format!("{error:?}").contains("User denied"));
258    }
259
260    #[tokio::test]
261    async fn waiting_times_out_with_a_hint() {
262        let (callback, _) = bound().await;
263        let error = callback
264            .wait("s1", Duration::from_millis(50))
265            .await
266            .unwrap_err();
267        assert!(format!("{error:?}").contains("--no-browser"));
268    }
269
270    #[test]
271    fn pasted_addresses() {
272        assert_eq!(
273            code_from_pasted("http://127.0.0.1:8765/callback?code=c0de&state=s1", "s1").unwrap(),
274            "c0de"
275        );
276        assert!(
277            code_from_pasted("http://127.0.0.1:8765/callback?code=c0de&state=s2", "s1").is_err()
278        );
279        assert!(code_from_pasted("not an address", "s1").is_err());
280        let error = code_from_pasted(
281            "http://127.0.0.1:8765/callback?error=x&error_description=Injected",
282            "s1",
283        )
284        .unwrap_err();
285        assert!(error.to_string().contains("did not match"), "{error}");
286    }
287
288    #[tokio::test]
289    async fn an_idle_connection_does_not_block_the_callback() {
290        let (callback, port) = bound().await;
291        let idle = TcpStream::connect(("127.0.0.1", port)).await.unwrap();
292        tokio::spawn(async move { send(port, "/callback?code=c0de&state=s1").await });
293        let code = callback.wait("s1", Duration::from_secs(2)).await.unwrap();
294        assert_eq!(code, "c0de");
295        drop(idle);
296    }
297}
yt/auth/loopback.rs

yt/auth/
loopback.rs
