zesdex/crates/zesdex-iam/src/infrastructure/oauth_loopback.rs

//! Minimal loopback HTTP server for capturing OAuth authorization-code redirects.
//!
//! Ported from `zesdex-backend::service::oauth::loopback` to centralise OAuth
//! primitives in the `zesdex-iam` crate.
//!
//! # Flow
//!
//! 1. [`LoopbackServer::bind`] — bind to `127.0.0.1:0` (OS-assigned port).
//! 2. [`redirect_uri`](LoopbackServer::redirect_uri) — caller gets the full
//!    `http://127.0.0.1:<port>/callback` URI to pass to `start_flow`.
//! 3. [`wait_for_code`](LoopbackServer::wait_for_code) — block until browser
//!    redirect hits the loopback → parse `?code=` and `?state=` from the HTTP
//!    request line → validate state → respond with 200/400 → return the code.
//!
//! # Components
//!
//! - `LoopbackServer` — single-use TCP listener for one OAuth callback
//! - `wait_for_code` / `read_callback` / `extract_code` / `extract_state`
//! - `urlencoding` — minimal percent-decoder for query parameters
use std::io::{Read, Write};
use std::net::{TcpListener, TcpStream};
use tracing;

/// 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 {
    /// The bound TCP listener (accepts one connection per `wait_for_code` call).
    listener: TcpListener,
    /// The OS-assigned port number.
    port: u16,
}

impl LoopbackServer {
    /// Bind to an OS-assigned free port on localhost.
    ///
    /// Return: `Err` if the loopback interface can't be bound.
    pub fn bind() -> std::io::Result<Self> {
        let listener = TcpListener::bind("127.0.0.1:0")?;
        let port = listener.local_addr()?.port();
        tracing::debug!(port, "loopback server bound");
        Ok(LoopbackServer { listener, port })
    }

    /// The redirect URI to hand to the OAuth authorization endpoint.
    pub fn redirect_uri(&self) -> String {
        format!("http://127.0.0.1:{}/callback", self.port)
    }

    /// Block until one HTTP request arrives, then extract the `code` query param
    /// and validate that the `state` param matches the expected value.
    ///
    /// Flow: accept one connection → apply read timeout → parse request line
    /// → verify state matches → respond 200/400 depending on whether the code
    /// was found and state matched.
    ///
    /// Return: `Err(InvalidData)` if no `code` param is present or the state
    /// doesn't match `expected_state`.
    pub fn wait_for_code(&self, timeout_ms: u64, expected_state: &str) -> std::io::Result<String> {
        tracing::debug!(port = self.port, timeout_ms, "waiting for OAuth callback");
        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)
    }

    /// Read and parse a single HTTP callback request off `stream`, replying with a status page.
    ///
    /// Why: writes the HTTP response before returning so the browser tab
    /// shows a result regardless of whether the code was found.
    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\nAuthorization complete. You may close this tab.",
            (Some(_), false) => "HTTP/1.1 400 Bad Request\r\nContent-Type: text/plain\r\n\r\nState mismatch — possible CSRF attack.",
            (None, _) => "HTTP/1.1 400 Bad Request\r\nContent-Type: text/plain\r\n\r\nMissing authorization code.",
        };
        let _ = stream.write_all(response.as_bytes());
        let _ = stream.flush();
        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",
            )
        })
    }

    /// Extract and percent-decode the `code` query parameter from an HTTP request line.
    ///
    /// Return: `None` if the request is malformed or has no `code` param.
    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
    }

    /// Extract the `state` query parameter from an HTTP request line.
    ///
    /// Return: `None` if the request is malformed or has no `state` param.
    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).
///
/// Why: invalid escape sequences (missing/non-hex digits) are passed through
/// literally as `%` rather than erroring, since this only handles a redirect
/// query param, not untrusted binary data.
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)) => result.push(char::from((hi * 16 + lo) as u8)),
                _ => {
                    result.push('%');
                }
            }
        } else {
            result.push(c);
        }
    }
    result
}