RustProxy/src/headers.rs

//! 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>,
    wire_only: bool,
}
#[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, Direction::Request)?,
            response: compile_rules(&self.response, Direction::Response)?,
        })
    }
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum Direction {
    Request,
    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-rustproxy-backend"
            | "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], direction: Direction) -> 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,
                wire_only: direction == Direction::Response
                    && rule.name.eq_ignore_ascii_case("age"),
            })
        })
        .collect()
}
fn apply_rule(headers: &mut HeaderMap, rule: &CompiledRule) {
    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"),
    }
}
pub fn apply(headers: &mut HeaderMap, rules: &[CompiledRule]) {
    for rule in rules {
        // Response Age rules are a wire-only presentation override. Applying
        // them here would let a configured value enter cache freshness/storage.
        if !rule.wire_only {
            apply_rule(headers, rule);
        }
    }
}
pub fn apply_wire_age(headers: &mut HeaderMap, rules: &[CompiledRule]) {
    for rule in rules {
        if rule.wire_only {
            apply_rule(headers, rule);
        }
    }
}
#[cfg(test)]
include!("headers/tests.rs");