proksi 0.9.0

A batteries-included reverse proxy with automatic HTTPS using Cloudflare Pingora and Let's Encrypt.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

//! `body_limit { max = "1MiB" }`: answers `413` to requests with a larger body

use anyhow::Result;
use async_trait::async_trait;
use bytes::Bytes;
use http::{header, StatusCode};
use pingora::{http::ResponseHeader, proxy::Session, ErrorType::HTTPStatus};

use crate::{config::RoutePlugin, proxy_server::https_proxy::RouterContext};

use super::Middleware;

const INVALID_MAX: &str =
    "body_limit: `max` must be a size such as \"512KiB\", \"1MiB\" or \"10MB\"";

/// Requests with a larger `Content-Length` are answered right away. Bodies without one
/// (chunked) are counted as they're proxied, and the request ends once it's over `max`.
pub(super) struct BodyLimit {
    max: u64,
}

impl BodyLimit {
    pub(super) fn new(plugin: &RoutePlugin) -> Result<Self, String> {
        let max = plugin
            .config
            .as_ref()
            .and_then(|config| config.get("max"))
            .and_then(serde_json::Value::as_str)
            .and_then(parse_size)
            .ok_or(INVALID_MAX)?;
        Ok(Self { max })
    }
}

#[async_trait]
impl Middleware for BodyLimit {
    async fn request(&self, session: &mut Session, _ctx: &mut RouterContext) -> Result<bool> {
        let content_length = session
            .req_header()
            .headers
            .get(header::CONTENT_LENGTH)
            .and_then(|value| value.to_str().ok())
            .and_then(|value| value.parse::<u64>().ok());
        if content_length.is_none_or(|length| length <= self.max) {
            return Ok(false);
        }

        let mut response = ResponseHeader::build_no_case(StatusCode::PAYLOAD_TOO_LARGE, Some(2))?;
        response.insert_header(header::CONTENT_LENGTH, "0")?;
        response.insert_header(header::CONNECTION, "close")?;
        session.set_keepalive(None);
        session
            .write_response_header(Box::new(response), true)
            .await?;
        Ok(true)
    }

    fn request_body(&self, body: Option<&Bytes>, ctx: &mut RouterContext) -> pingora::Result<()> {
        ctx.request_body_bytes += body.map_or(0, |body| body.len() as u64);
        if ctx.request_body_bytes > self.max {
            return Err(pingora::Error::explain(
                HTTPStatus(413),
                "request body larger than body_limit",
            ));
        }
        Ok(())
    }
}

/// A size like `512KiB`, `1MiB` or `10MB`
fn parse_size(value: &str) -> Option<u64> {
    let split = value.find(|c: char| !c.is_ascii_digit())?;
    let (amount, unit) = value.split_at(split);
    let multiplier: u64 = match unit {
        "B" => 1,
        "KB" => 1000,
        "MB" => 1000 * 1000,
        "GB" => 1000 * 1000 * 1000,
        "KiB" => 1024,
        "MiB" => 1024 * 1024,
        "GiB" => 1024 * 1024 * 1024,
        _ => return None,
    };
    amount
        .parse::<u64>()
        .ok()
        .and_then(|amount| amount.checked_mul(multiplier))
        .filter(|size| *size > 0)
}

#[cfg(test)]
mod tests {
    use crate::config::RoutePlugin;

    use super::{parse_size, BodyLimit};

    fn plugin(config: serde_json::Value) -> RoutePlugin {
        serde_json::from_value(serde_json::json!({ "name": "body_limit", "config": config }))
            .unwrap()
    }

    #[test]
    fn test_parse_size() {
        assert_eq!(parse_size("100B"), Some(100));
        assert_eq!(parse_size("10KB"), Some(10_000));
        assert_eq!(parse_size("512KiB"), Some(512 * 1024));
        assert_eq!(parse_size("2MB"), Some(2_000_000));
        assert_eq!(parse_size("1MiB"), Some(1024 * 1024));
        assert_eq!(parse_size("1GiB"), Some(1024 * 1024 * 1024));
        for invalid in ["", "1", "MiB", "0B", "1.5MiB", "-1KB", "1TB", "1 MiB"] {
            assert_eq!(parse_size(invalid), None, "{invalid}");
        }
    }

    #[test]
    fn test_max() {
        let limit = BodyLimit::new(&plugin(serde_json::json!({ "max": "1KiB" }))).unwrap();
        assert_eq!(limit.max, 1024);
    }

    #[test]
    fn test_invalid_settings() {
        for config in [
            serde_json::json!({}),
            serde_json::json!({ "max": "1TB" }),
            serde_json::json!({ "max": 1024 }),
        ] {
            assert_eq!(
                BodyLimit::new(&plugin(config)).err().as_deref(),
                Some("body_limit: `max` must be a size such as \"512KiB\", \"1MiB\" or \"10MB\"")
            );
        }
    }
}