zesdex/apps/infrastructure/src/auth/oauth_loopback.rs

//! Minimal loopback HTTP server for capturing OAuth authorization-code redirects.

use std::io::{Read, Write};
use std::net::{TcpListener, TcpStream};

/// A single-use HTTP listener on `127.0.0.1` that receives the OAuth
/// `?code=...` redirect and serves back a static confirmation page.
pub struct LoopbackServer {
    listener: TcpListener,
    port: u16,
}

impl LoopbackServer {
    pub fn bind() -> std::io::Result<Self> {
        let listener = TcpListener::bind("127.0.0.1:0")?;
        let port = listener.local_addr()?.port();
        Ok(LoopbackServer { listener, port })
    }

    pub fn redirect_uri(&self) -> String {
        format!("http://127.0.0.1:{}/callback", self.port)
    }

    pub fn wait_for_code(
        &self,
        timeout_ms: u64,
        expected_state: &str,
    ) -> std::io::Result<String> {
        let (mut stream, _) = self.listener.accept()?;
        stream.set_read_timeout(Some(std::time::Duration::from_millis(timeout_ms)))?;
        Self::read_callback(&mut stream, expected_state)
    }

    fn read_callback(
        stream: &mut TcpStream,
        expected_state: &str,
    ) -> std::io::Result<String> {
        let mut buf = [0u8; 4096];
        let n = stream.read(&mut buf)?;
        let request = String::from_utf8_lossy(&buf[..n]);
        let code = Self::extract_code(&request);
        let state = Self::extract_state(&request);
        let state_ok = state.as_deref() == Some(expected_state);
        let response = match (code.as_ref(), state_ok) {
            (Some(_), true) => {
                "HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\n\r\n\
                 Authorization complete. You may close this tab."
            }
            (Some(_), false) => {
                "HTTP/1.1 400 Bad Request\r\nContent-Type: text/plain\r\n\r\n\
                 State mismatch — possible CSRF attack."
            }
            (None, _) => {
                "HTTP/1.1 400 Bad Request\r\nContent-Type: text/plain\r\n\r\n\
                 Missing authorization code."
            }
        };
        if let Err(e) = stream.write_all(response.as_bytes()) {
            tracing::warn!("OAuth loopback write error: {e}");
        }
        if let Err(e) = stream.flush() {
            tracing::warn!("OAuth loopback flush error: {e}");
        }
        if !state_ok {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                "state mismatch",
            ));
        }
        code.ok_or_else(|| {
            std::io::Error::new(std::io::ErrorKind::InvalidData, "code not found in callback")
        })
    }

    fn extract_code(request: &str) -> Option<String> {
        let line = request.lines().next()?;
        let path = line.split(' ').nth(1)?;
        let query = path.split('?').nth(1)?;
        for pair in query.split('&') {
            let mut parts = pair.splitn(2, '=');
            if parts.next()? == "code" {
                return parts.next().map(urlencoding);
            }
        }
        None
    }

    fn extract_state(request: &str) -> Option<String> {
        let line = request.lines().next()?;
        let path = line.split(' ').nth(1)?;
        let query = path.split('?').nth(1)?;
        for pair in query.split('&') {
            let mut parts = pair.splitn(2, '=');
            if parts.next()? == "state" {
                return parts.next().map(urlencoding);
            }
        }
        None
    }
}

/// Percent-decode a string (e.g. `%20` -> space).
fn urlencoding(s: &str) -> String {
    let mut result = String::with_capacity(s.len());
    let mut chars = s.chars();
    while let Some(c) = chars.next() {
        if c == '%' {
            match (
                chars.next().and_then(|c| c.to_digit(16)),
                chars.next().and_then(|c| c.to_digit(16)),
            ) {
                (Some(hi), Some(lo)) => {
                    let byte: u8 = (hi as u8) * 16 + lo as u8;
                    result.push(char::from(byte));
                }
                _ => {
                    result.push('%');
                }
            }
        } else {
            result.push(c);
        }
    }
    result
}