lmrc-proxy 0.3.16

//! Proxy middleware
//!
//! Axum middleware for subdomain extraction and proxy routing.

use axum::{
    extract::{Request, State},
    http::HeaderMap,
    middleware::Next,
    response::Response,
};
use std::sync::Arc;

use crate::{client::{proxy_request, ProxyConfig}, error::ProxyError, routing::{extract_subdomain, RouteResolver}};

/// Subdomain marker type for request extensions
#[derive(Debug, Clone)]
pub struct Subdomain(pub String);

/// Extract subdomain from Host header and attach to request extensions
///
/// ## Example
///
/// ```rust,ignore
/// use axum::{Router, middleware};
/// use lmrc_proxy::middleware::subdomain_extractor;
///
/// let app = Router::new()
///     .layer(middleware::from_fn(subdomain_extractor));
/// ```
pub async fn subdomain_extractor(
    headers: HeaderMap,
    mut request: Request,
    next: Next,
) -> Response {
    let subdomain = headers
        .get("host")
        .and_then(|host| host.to_str().ok())
        .and_then(extract_subdomain);

    if let Some(subdomain) = subdomain {
        request.extensions_mut().insert(Subdomain(subdomain));
    }

    next.run(request).await
}

/// Proxy middleware that routes based on subdomain
///
/// ## Example
///
/// ```rust,ignore
/// use axum::{Router, middleware};
/// use lmrc_proxy::{middleware::subdomain_proxy, routing::StaticRouteResolver, ProxyConfig};
/// use std::sync::Arc;
///
/// let resolver = Arc::new(
///     StaticRouteResolver::new()
///         .add_route("api", "http://api-service:8080")
/// );
///
/// let app = Router::new()
///     .layer(middleware::from_fn_with_state(
///         (resolver, ProxyConfig::default()),
///         subdomain_proxy
///     ));
/// ```
pub async fn subdomain_proxy<R>(
    State((resolver, config)): State<(Arc<R>, ProxyConfig)>,
    request: Request,
) -> Result<Response, ProxyError>
where
    R: RouteResolver,
{
    // Extract subdomain from extensions (set by subdomain_extractor middleware)
    let subdomain = request
        .extensions()
        .get::<Subdomain>()
        .map(|s| s.0.clone());

    let subdomain = subdomain.ok_or_else(|| ProxyError::RouteNotFound("No subdomain found".to_string()))?;

    // Resolve backend URL
    let backend_url = resolver
        .resolve(&subdomain)
        .await
        .ok_or_else(|| ProxyError::RouteNotFound(format!("No route for subdomain: {}", subdomain)))?;

    // Proxy request
    proxy_request(request, &backend_url, config).await
}

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

    #[test]
    fn test_subdomain_type() {
        let subdomain = Subdomain("api".to_string());
        assert_eq!(subdomain.0, "api");
    }
}