wasm_runner_sdk/

router.rs

//! HTTP router for dispatching requests to handlers.
//!
//! Provides a simple, ergonomic router with support for path parameters.
//!
//! # Example
//!
//! ```ignore
//! use wasm_runner_sdk::prelude::*;
//!
//! fn get_user(Param(id): Param<String>) -> impl IntoResponse {
//!     Json(serde_json::json!({ "userId": id }))
//! }
//!
//! fn create_user(Json(body): Json<CreateUser>) -> impl IntoResponse {
//!     (StatusCode::CREATED, Json(body))
//! }
//!
//! let router = Router::new()
//!     .get("/users/:id", get_user)
//!     .post("/users", create_user)
//!     .not_found(|| StatusCode::NOT_FOUND);
//!
//! #[unsafe(no_mangle)]
//! pub extern "C" fn _start() {
//!     router.run();
//! }
//! ```

use crate::extract::{ExtractError, FromRequest};
use crate::handler::Handler;
use crate::request::{Method, Request};
use crate::response::{IntoResponse, Response, StatusCode};
use std::collections::HashMap;

/// A path parameter extracted from the URL.
///
/// # Example
///
/// ```ignore
/// // For route "/users/:id" and request "/users/123"
/// fn handler(Param(id): Param<String>) -> impl IntoResponse {
///     format!("User ID: {}", id)
/// }
/// ```
#[derive(Debug, Clone)]
pub struct Param<T>(pub T);

/// Multiple path parameters extracted from the URL.
///
/// # Example
///
/// ```ignore
/// // For route "/users/:user_id/posts/:post_id"
/// fn handler(Params(params): Params) -> impl IntoResponse {
///     let user_id = params.get("user_id").unwrap();
///     let post_id = params.get("post_id").unwrap();
///     format!("User: {}, Post: {}", user_id, post_id)
/// }
/// ```
#[derive(Debug, Clone, Default)]
pub struct Params(pub HashMap<String, String>);

impl Params {
    /// Gets a parameter by name.
    pub fn get(&self, name: &str) -> Option<&str> {
        self.0.get(name).map(|s| s.as_str())
    }

    /// Gets a parameter by name, parsing it to the desired type.
    pub fn get_parse<T: std::str::FromStr>(&self, name: &str) -> Option<T> {
        self.0.get(name).and_then(|s| s.parse().ok())
    }
}

// Thread-local storage for route parameters during request handling
std::thread_local! {
    static CURRENT_PARAMS: std::cell::RefCell<HashMap<String, String>> = std::cell::RefCell::new(HashMap::new());
}

fn set_current_params(params: HashMap<String, String>) {
    CURRENT_PARAMS.with(|p| {
        *p.borrow_mut() = params;
    });
}

fn get_current_params() -> HashMap<String, String> {
    CURRENT_PARAMS.with(|p| p.borrow().clone())
}

impl FromRequest for Params {
    type Error = ExtractError;

    fn from_request(_req: &mut Request) -> Result<Self, Self::Error> {
        Ok(Params(get_current_params()))
    }
}

impl FromRequest for Param<String> {
    type Error = ExtractError;

    fn from_request(_req: &mut Request) -> Result<Self, Self::Error> {
        let params = get_current_params();
        // Return the first parameter value
        params
            .into_values()
            .next()
            .map(Param)
            .ok_or_else(|| ExtractError::bad_request("No path parameter found"))
    }
}

impl FromRequest for Param<i64> {
    type Error = ExtractError;

    fn from_request(_req: &mut Request) -> Result<Self, Self::Error> {
        let params = get_current_params();
        params
            .into_values()
            .next()
            .ok_or_else(|| ExtractError::bad_request("No path parameter found"))?
            .parse()
            .map(Param)
            .map_err(|_| ExtractError::bad_request("Invalid integer parameter"))
    }
}

impl FromRequest for Param<u64> {
    type Error = ExtractError;

    fn from_request(_req: &mut Request) -> Result<Self, Self::Error> {
        let params = get_current_params();
        params
            .into_values()
            .next()
            .ok_or_else(|| ExtractError::bad_request("No path parameter found"))?
            .parse()
            .map(Param)
            .map_err(|_| ExtractError::bad_request("Invalid unsigned integer parameter"))
    }
}

/// A boxed handler function that can be stored in the router.
type BoxedHandler = Box<dyn Fn(&mut Request) -> Response>;

/// A route pattern with parameter extraction.
#[derive(Clone)]
struct RoutePattern {
    segments: Vec<PatternSegment>,
}

#[derive(Clone)]
enum PatternSegment {
    /// Literal segment that must match exactly.
    Literal(String),
    /// Parameter segment that captures a value.
    Param(String),
    /// Wildcard that matches any remaining path.
    Wildcard,
}

impl RoutePattern {
    fn parse(pattern: &str) -> Self {
        let segments = pattern
            .split('/')
            .filter(|s| !s.is_empty())
            .map(|s| {
                if s == "*" {
                    PatternSegment::Wildcard
                } else if let Some(name) = s.strip_prefix(':') {
                    PatternSegment::Param(name.to_string())
                } else {
                    PatternSegment::Literal(s.to_string())
                }
            })
            .collect();

        RoutePattern { segments }
    }

    fn matches(&self, path: &str) -> Option<HashMap<String, String>> {
        let path_segments: Vec<&str> = path.split('/').filter(|s| !s.is_empty()).collect();
        let mut params = HashMap::new();
        let mut path_idx = 0;

        for (pattern_idx, segment) in self.segments.iter().enumerate() {
            match segment {
                PatternSegment::Literal(lit) => {
                    if path_idx >= path_segments.len() || path_segments[path_idx] != lit {
                        return None;
                    }
                    path_idx += 1;
                }
                PatternSegment::Param(name) => {
                    if path_idx >= path_segments.len() {
                        return None;
                    }
                    params.insert(name.clone(), path_segments[path_idx].to_string());
                    path_idx += 1;
                }
                PatternSegment::Wildcard => {
                    // Wildcard matches rest of path - must be last segment
                    if pattern_idx == self.segments.len() - 1 {
                        return Some(params);
                    }
                    return None;
                }
            }
        }

        // All path segments must be consumed (unless wildcard)
        if path_idx == path_segments.len() {
            Some(params)
        } else {
            None
        }
    }
}

/// A route entry combining method, pattern, and handler.
struct Route {
    method: Option<Method>, // None means any method
    pattern: RoutePattern,
    handler: BoxedHandler,
}

/// HTTP router for dispatching requests to handlers.
///
/// # Example
///
/// ```ignore
/// let router = Router::new()
///     .get("/", || "Hello, World!")
///     .get("/users", list_users)
///     .get("/users/:id", get_user)
///     .post("/users", create_user)
///     .put("/users/:id", update_user)
///     .delete("/users/:id", delete_user);
/// ```
pub struct Router {
    routes: Vec<Route>,
    not_found_handler: Option<BoxedHandler>,
    method_not_allowed_handler: Option<BoxedHandler>,
}

impl Default for Router {
    fn default() -> Self {
        Self::new()
    }
}

impl Router {
    /// Creates a new empty router.
    pub fn new() -> Self {
        Router {
            routes: Vec::new(),
            not_found_handler: None,
            method_not_allowed_handler: None,
        }
    }

    /// Adds a route that matches any HTTP method.
    pub fn route<H, Args>(self, pattern: &str, handler: H) -> Self
    where
        H: Handler<Args> + 'static,
        H::Output: IntoResponse,
    {
        self.add_route(None, pattern, handler)
    }

    /// Adds a GET route.
    pub fn get<H, Args>(self, pattern: &str, handler: H) -> Self
    where
        H: Handler<Args> + 'static,
        H::Output: IntoResponse,
    {
        self.add_route(Some(Method::Get), pattern, handler)
    }

    /// Adds a POST route.
    pub fn post<H, Args>(self, pattern: &str, handler: H) -> Self
    where
        H: Handler<Args> + 'static,
        H::Output: IntoResponse,
    {
        self.add_route(Some(Method::Post), pattern, handler)
    }

    /// Adds a PUT route.
    pub fn put<H, Args>(self, pattern: &str, handler: H) -> Self
    where
        H: Handler<Args> + 'static,
        H::Output: IntoResponse,
    {
        self.add_route(Some(Method::Put), pattern, handler)
    }

    /// Adds a DELETE route.
    pub fn delete<H, Args>(self, pattern: &str, handler: H) -> Self
    where
        H: Handler<Args> + 'static,
        H::Output: IntoResponse,
    {
        self.add_route(Some(Method::Delete), pattern, handler)
    }

    /// Adds a PATCH route.
    pub fn patch<H, Args>(self, pattern: &str, handler: H) -> Self
    where
        H: Handler<Args> + 'static,
        H::Output: IntoResponse,
    {
        self.add_route(Some(Method::Patch), pattern, handler)
    }

    /// Adds a HEAD route.
    pub fn head<H, Args>(self, pattern: &str, handler: H) -> Self
    where
        H: Handler<Args> + 'static,
        H::Output: IntoResponse,
    {
        self.add_route(Some(Method::Head), pattern, handler)
    }

    /// Adds an OPTIONS route.
    pub fn options<H, Args>(self, pattern: &str, handler: H) -> Self
    where
        H: Handler<Args> + 'static,
        H::Output: IntoResponse,
    {
        self.add_route(Some(Method::Options), pattern, handler)
    }

    /// Sets a custom 404 Not Found handler.
    pub fn not_found<H, Args>(mut self, handler: H) -> Self
    where
        H: Handler<Args> + 'static,
        H::Output: IntoResponse,
    {
        let handler_cell = std::cell::RefCell::new(Some(handler));
        self.not_found_handler = Some(Box::new(move |req| {
            if let Some(h) = handler_cell.borrow_mut().take() {
                h.call(req).into_response()
            } else {
                Response::new(StatusCode::INTERNAL_SERVER_ERROR)
                    .with_text("Handler already consumed")
            }
        }));
        self
    }

    /// Sets a custom 405 Method Not Allowed handler.
    pub fn method_not_allowed<H, Args>(mut self, handler: H) -> Self
    where
        H: Handler<Args> + 'static,
        H::Output: IntoResponse,
    {
        let handler_cell = std::cell::RefCell::new(Some(handler));
        self.method_not_allowed_handler = Some(Box::new(move |req| {
            if let Some(h) = handler_cell.borrow_mut().take() {
                h.call(req).into_response()
            } else {
                Response::new(StatusCode::INTERNAL_SERVER_ERROR)
                    .with_text("Handler already consumed")
            }
        }));
        self
    }

    /// Nests another router under a path prefix.
    ///
    /// # Example
    ///
    /// ```ignore
    /// let api_router = Router::new()
    ///     .get("/users", list_users)
    ///     .get("/users/:id", get_user);
    ///
    /// let main_router = Router::new()
    ///     .get("/", home)
    ///     .nest("/api", api_router);
    /// ```
    pub fn nest(mut self, prefix: &str, nested: Router) -> Self {
        let prefix = prefix.trim_end_matches('/');

        for route in nested.routes {
            // Combine prefix with route pattern
            let combined_pattern = if route.pattern.segments.is_empty() {
                prefix.to_string()
            } else {
                let pattern_str: String = route
                    .pattern
                    .segments
                    .iter()
                    .map(|seg| match seg {
                        PatternSegment::Literal(s) => format!("/{}", s),
                        PatternSegment::Param(s) => format!("/:{}", s),
                        PatternSegment::Wildcard => "/*".to_string(),
                    })
                    .collect();
                format!("{}{}", prefix, pattern_str)
            };

            self.routes.push(Route {
                method: route.method,
                pattern: RoutePattern::parse(&combined_pattern),
                handler: route.handler,
            });
        }

        self
    }

    fn add_route<H, Args>(mut self, method: Option<Method>, pattern: &str, handler: H) -> Self
    where
        H: Handler<Args> + 'static,
        H::Output: IntoResponse,
    {
        // We need to store the handler in a way that can be called multiple times.
        // Since Handler::call takes self by value, we need a different approach.
        // We'll use a Cell to allow taking ownership each time.
        let handler_cell = std::cell::RefCell::new(Some(handler));

        let boxed: BoxedHandler = Box::new(move |req| {
            // This is a limitation - each handler can only be called once per route registration.
            // For a real router, we'd need handlers to be Clone or use Fn instead of FnOnce.
            // For now, we'll panic if called twice (which shouldn't happen in normal use).
            if let Some(h) = handler_cell.borrow_mut().take() {
                h.call(req).into_response()
            } else {
                // Handler was already consumed - this shouldn't happen with proper usage
                Response::new(StatusCode::INTERNAL_SERVER_ERROR)
                    .with_text("Handler already consumed")
            }
        });

        self.routes.push(Route {
            method,
            pattern: RoutePattern::parse(pattern),
            handler: boxed,
        });

        self
    }

    /// Handles a request and returns a response.
    pub fn handle(&self, req: &mut Request) -> Response {
        let method = req.method();
        let path = req.path().to_string();

        // Track if we found a matching path (for 405 vs 404)
        let mut path_matched = false;

        for route in &self.routes {
            if let Some(params) = route.pattern.matches(&path) {
                path_matched = true;

                // Check method
                if let Some(route_method) = route.method {
                    if route_method != method {
                        continue;
                    }
                }

                // Set params for extractors
                set_current_params(params);

                // Call handler
                return (route.handler)(req);
            }
        }

        // No matching route found
        if path_matched {
            // Path matched but method didn't - 405
            if let Some(ref handler) = self.method_not_allowed_handler {
                return handler(req);
            }
            Response::new(StatusCode::METHOD_NOT_ALLOWED)
                .with_text("Method Not Allowed")
        } else {
            // No path match - 404
            if let Some(ref handler) = self.not_found_handler {
                return handler(req);
            }
            Response::new(StatusCode::NOT_FOUND).with_text("Not Found")
        }
    }

    /// Runs the router, handling the current request.
    ///
    /// This is the main entry point for using the router.
    ///
    /// # Example
    ///
    /// ```ignore
    /// #[unsafe(no_mangle)]
    /// pub extern "C" fn _start() {
    ///     Router::new()
    ///         .get("/", || "Hello!")
    ///         .run();
    /// }
    /// ```
    pub fn run(self) {
        let mut request = Request::new();
        let method = request.method();
        let path = request.path().to_string();
        crate::log_info!("Handling request {} {}", method, path);
        let response = self.handle(&mut request);
        response.send();
    }
}