//! Generic, compiled request/response header transformations. No app-specific exceptions.
use anyhow::{bail, ensure, Context, Result};
use http::{HeaderMap, HeaderName, HeaderValue};
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Operation {
    Set,
    Append,
    Remove,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct Rule {
    pub op: Operation,
    pub name: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub value: Option<String>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct HeaderRules {
    #[serde(default)]
    pub request: Vec<Rule>,
    #[serde(default)]
    pub response: Vec<Rule>,
}
#[derive(Debug, Clone)]
pub struct CompiledRule {
    op: Operation,
    name: HeaderName,
    value: Option<HeaderValue>,
}
#[derive(Debug, Clone, Default)]
pub struct CompiledHeaders {
    pub request: Vec<CompiledRule>,
    pub response: Vec<CompiledRule>,
}
impl HeaderRules {
    pub fn is_empty(&self) -> bool {
        self.request.is_empty() && self.response.is_empty()
    }
    pub fn compile(&self) -> Result<CompiledHeaders> {
        Ok(CompiledHeaders {
            request: compile_rules(&self.request)?,
            response: compile_rules(&self.response)?,
        })
    }
}
fn protected(name: &str) -> bool {
    matches!(
        name,
        "host"
            | "connection"
            | "content-length"
            | "transfer-encoding"
            | "te"
            | "trailer"
            | "content-encoding"
            | "content-range"
            | "upgrade"
            | "keep-alive"
            | "proxy-connection"
            | "proxy-authorization"
            | "proxy-authenticate"
            | "forwarded"
            | "x-real-ip"
            | "x-request-id"
            | "x-rust-proxy-backend"
            | "x-rust-proxy-backend-time"
            | "x-rustproxy-backend-time"
    ) || name.starts_with("x-forwarded-")
        || name.starts_with("sec-websocket-")
}
fn compile_rules(rules: &[Rule]) -> Result<Vec<CompiledRule>> {
    ensure!(
        rules.len() <= 32,
        "at most 32 header rules per direction and scope"
    );
    rules
        .iter()
        .map(|rule| {
            ensure!(rule.name.len() <= 128, "header name is too long");
            let name: HeaderName = rule.name.parse().context("invalid header name")?;
            ensure!(
                !protected(name.as_str()),
                "transport/identity header cannot be changed: {}",
                name
            );
            let value = match (&rule.op, &rule.value) {
                (Operation::Remove, None) => None,
                (Operation::Remove, Some(_)) => bail!("remove header rule cannot include a value"),
                (_, Some(value)) => {
                    ensure!(value.len() <= 8192, "header value too long");
                    Some(value.parse().context("invalid header value")?)
                }
                (_, None) => bail!("set/append header rule requires a value"),
            };
            Ok(CompiledRule {
                op: rule.op,
                name,
                value,
            })
        })
        .collect()
}
pub fn apply(headers: &mut HeaderMap, rules: &[CompiledRule]) {
    for rule in rules {
        match (&rule.op, &rule.value) {
            (Operation::Remove, _) => {
                headers.remove(&rule.name);
            }
            (Operation::Set, Some(value)) => {
                headers.insert(rule.name.clone(), value.clone());
            }
            (Operation::Append, Some(value)) => {
                headers.append(rule.name.clone(), value.clone());
            }
            _ => unreachable!("only compiled header rules reach the request path"),
        }
    }
}
#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn transformations_preserve_distinct_set_cookie_fields() {
        let source = r#"{"request":[],"response":[{"op":"remove","name":"server"},{"op":"set","name":"x-test","value":"one"},{"op":"append","name":"x-test","value":"two"},{"op":"append","name":"set-cookie","value":"b=2"}]}"#;
        let rules: HeaderRules = serde_json::from_str(source).unwrap();
        let compiled = rules.compile().unwrap();
        let mut headers = HeaderMap::new();
        headers.insert("server", "upstream".parse().unwrap());
        headers.append("set-cookie", "a=1".parse().unwrap());
        apply(&mut headers, &compiled.response);
        assert!(!headers.contains_key("server"));
        assert_eq!(headers.get_all("x-test").iter().count(), 2);
        assert_eq!(headers.get_all("set-cookie").iter().count(), 2);
    }
    #[test]
    fn transport_identity_and_crlf_are_rejected() {
        for name in [
            "Host",
            "Connection",
            "Content-Length",
            "x-forwarded-for",
            "Upgrade",
            "X-RustProxy-Backend",
            "X-RustProxy-Backend-Time",
            "X-Rust-Proxy-Backend-Time",
        ] {
            let rules = HeaderRules {
                request: vec![Rule {
                    op: Operation::Set,
                    name: name.into(),
                    value: Some("x".into()),
                }],
                response: vec![],
            };
            assert!(rules.compile().is_err(), "{name}");
        }
        let rules = HeaderRules {
            request: vec![Rule {
                op: Operation::Set,
                name: "x-test".into(),
                value: Some("hello\r\ninjected: yes".into()),
            }],
            response: vec![],
        };
        assert!(rules.compile().is_err());
    }
}