lmrc-proxy 0.3.16

//! Routing and subdomain handling
//!
//! Provides traits and utilities for routing requests based on subdomains or other criteria.

use async_trait::async_trait;
use std::collections::HashMap;

/// Route resolver trait
///
/// Implement this trait to provide custom routing logic.
///
/// ## Example
///
/// ```rust
/// use lmrc_proxy::routing::RouteResolver;
/// use async_trait::async_trait;
/// use std::collections::HashMap;
///
/// struct StaticRouter {
///     routes: HashMap<String, String>,
/// }
///
/// #[async_trait]
/// impl RouteResolver for StaticRouter {
///     async fn resolve(&self, key: &str) -> Option<String> {
///         self.routes.get(key).cloned()
///     }
/// }
/// ```
#[async_trait]
pub trait RouteResolver: Send + Sync {
    /// Resolve a routing key (e.g., subdomain) to a backend URL
    async fn resolve(&self, key: &str) -> Option<String>;
}

/// Static route resolver using a HashMap
pub struct StaticRouteResolver {
    routes: HashMap<String, String>,
}

impl StaticRouteResolver {
    /// Create a new static route resolver
    pub fn new() -> Self {
        Self {
            routes: HashMap::new(),
        }
    }

    /// Add a route
    pub fn add_route(mut self, key: impl Into<String>, backend_url: impl Into<String>) -> Self {
        self.routes.insert(key.into(), backend_url.into());
        self
    }

    /// Add routes from a HashMap
    pub fn with_routes(mut self, routes: HashMap<String, String>) -> Self {
        self.routes = routes;
        self
    }
}

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

#[async_trait]
impl RouteResolver for StaticRouteResolver {
    async fn resolve(&self, key: &str) -> Option<String> {
        self.routes.get(key).cloned()
    }
}

/// Extract subdomain from Host header
///
/// ## Examples
///
/// ```rust
/// use lmrc_proxy::routing::extract_subdomain;
///
/// assert_eq!(extract_subdomain("api.example.com"), Some("api".to_string()));
/// assert_eq!(extract_subdomain("infra.example.com:8080"), Some("infra".to_string()));
/// assert_eq!(extract_subdomain("example.com"), None);
/// assert_eq!(extract_subdomain("localhost"), Some("infra".to_string())); // Development default
/// ```
pub fn extract_subdomain(host: &str) -> Option<String> {
    // Remove port if present
    let host = host.split(':').next().unwrap_or(host);

    // Split by dots
    let parts: Vec<&str> = host.split('.').collect();

    // If we have at least 3 parts (subdomain.domain.tld), extract subdomain
    if parts.len() >= 3 {
        Some(parts[0].to_string())
    } else if parts.len() == 2 {
        // Could be domain.tld or localhost:port
        None
    } else if parts.len() == 1 && parts[0] == "localhost" {
        // Development: treat localhost as infra subdomain
        Some("infra".to_string())
    } else {
        None
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_extract_subdomain() {
        assert_eq!(
            extract_subdomain("infra.example.com"),
            Some("infra".to_string())
        );
        assert_eq!(
            extract_subdomain("api.example.com"),
            Some("api".to_string())
        );
        assert_eq!(extract_subdomain("example.com"), None);
        assert_eq!(extract_subdomain("localhost"), Some("infra".to_string()));
        assert_eq!(
            extract_subdomain("infra.example.com:8080"),
            Some("infra".to_string())
        );
        assert_eq!(
            extract_subdomain("my-service.example.com"),
            Some("my-service".to_string())
        );
    }

    #[tokio::test]
    async fn test_static_route_resolver() {
        let resolver = StaticRouteResolver::new()
            .add_route("api", "http://api-backend:8080")
            .add_route("admin", "http://admin-backend:9000");

        assert_eq!(
            resolver.resolve("api").await,
            Some("http://api-backend:8080".to_string())
        );
        assert_eq!(
            resolver.resolve("admin").await,
            Some("http://admin-backend:9000".to_string())
        );
        assert_eq!(resolver.resolve("unknown").await, None);
    }
}