smartproxy/rust/crates/rustproxy-http/src/template.rs

//! Header template variable expansion.
//!
//! Supports expanding template variables like `{clientIp}`, `{domain}`, etc.
//! in header values before they are applied to requests or responses.

use std::collections::HashMap;
use std::time::{SystemTime, UNIX_EPOCH};

/// Context for template variable expansion.
pub struct RequestContext {
    pub client_ip: String,
    pub domain: String,
    pub port: u16,
    pub path: String,
    pub route_name: String,
    pub connection_id: u64,
}

/// Expand template variables in a header value.
/// Supported variables: {clientIp}, {domain}, {port}, {path}, {routeName}, {connectionId}, {timestamp}
pub fn expand_template(template: &str, ctx: &RequestContext) -> String {
    let timestamp = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs();

    template
        .replace("{clientIp}", &ctx.client_ip)
        .replace("{domain}", &ctx.domain)
        .replace("{port}", &ctx.port.to_string())
        .replace("{path}", &ctx.path)
        .replace("{routeName}", &ctx.route_name)
        .replace("{connectionId}", &ctx.connection_id.to_string())
        .replace("{timestamp}", &timestamp.to_string())
}

/// Expand templates in a map of header key-value pairs.
pub fn expand_headers(
    headers: &HashMap<String, String>,
    ctx: &RequestContext,
) -> HashMap<String, String> {
    headers.iter()
        .map(|(k, v)| (k.clone(), expand_template(v, ctx)))
        .collect()
}

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

    fn test_context() -> RequestContext {
        RequestContext {
            client_ip: "192.168.1.100".to_string(),
            domain: "example.com".to_string(),
            port: 443,
            path: "/api/v1/users".to_string(),
            route_name: "api-route".to_string(),
            connection_id: 42,
        }
    }

    #[test]
    fn test_expand_client_ip() {
        let ctx = test_context();
        assert_eq!(expand_template("{clientIp}", &ctx), "192.168.1.100");
    }

    #[test]
    fn test_expand_domain() {
        let ctx = test_context();
        assert_eq!(expand_template("{domain}", &ctx), "example.com");
    }

    #[test]
    fn test_expand_port() {
        let ctx = test_context();
        assert_eq!(expand_template("{port}", &ctx), "443");
    }

    #[test]
    fn test_expand_path() {
        let ctx = test_context();
        assert_eq!(expand_template("{path}", &ctx), "/api/v1/users");
    }

    #[test]
    fn test_expand_route_name() {
        let ctx = test_context();
        assert_eq!(expand_template("{routeName}", &ctx), "api-route");
    }

    #[test]
    fn test_expand_connection_id() {
        let ctx = test_context();
        assert_eq!(expand_template("{connectionId}", &ctx), "42");
    }

    #[test]
    fn test_expand_timestamp() {
        let ctx = test_context();
        let result = expand_template("{timestamp}", &ctx);
        // Timestamp should be a valid number
        let ts: u64 = result.parse().expect("timestamp should be a number");
        // Should be a reasonable Unix timestamp (after 2020)
        assert!(ts > 1_577_836_800);
    }

    #[test]
    fn test_expand_mixed_template() {
        let ctx = test_context();
        let result = expand_template("client={clientIp}, host={domain}:{port}", &ctx);
        assert_eq!(result, "client=192.168.1.100, host=example.com:443");
    }

    #[test]
    fn test_expand_no_variables() {
        let ctx = test_context();
        assert_eq!(expand_template("plain-value", &ctx), "plain-value");
    }

    #[test]
    fn test_expand_empty_string() {
        let ctx = test_context();
        assert_eq!(expand_template("", &ctx), "");
    }

    #[test]
    fn test_expand_multiple_same_variable() {
        let ctx = test_context();
        let result = expand_template("{clientIp}-{clientIp}", &ctx);
        assert_eq!(result, "192.168.1.100-192.168.1.100");
    }

    #[test]
    fn test_expand_headers_map() {
        let ctx = test_context();
        let mut headers = HashMap::new();
        headers.insert("X-Forwarded-For".to_string(), "{clientIp}".to_string());
        headers.insert("X-Route".to_string(), "{routeName}".to_string());
        headers.insert("X-Static".to_string(), "no-template".to_string());

        let result = expand_headers(&headers, &ctx);
        assert_eq!(result.get("X-Forwarded-For").unwrap(), "192.168.1.100");
        assert_eq!(result.get("X-Route").unwrap(), "api-route");
        assert_eq!(result.get("X-Static").unwrap(), "no-template");
    }

    #[test]
    fn test_expand_all_variables_in_one() {
        let ctx = test_context();
        let template = "{clientIp}|{domain}|{port}|{path}|{routeName}|{connectionId}";
        let result = expand_template(template, &ctx);
        assert_eq!(result, "192.168.1.100|example.com|443|/api/v1/users|api-route|42");
    }

    #[test]
    fn test_expand_unknown_variable_left_as_is() {
        let ctx = test_context();
        let result = expand_template("{unknownVar}", &ctx);
        assert_eq!(result, "{unknownVar}");
    }
}