fn0 0.7.0

FaaS platform powered by wasmtime
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

use bytes::Bytes;
use http_body_util::BodyExt;
use http_body_util::combinators::UnsyncBoxBody;
use hyper::HeaderMap;
use std::sync::{Arc, OnceLock};
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
use wasmtime_wasi_http::p3::bindings::http::types::ErrorCode;

pub(crate) const QUEUE_BODY_LIMIT: usize = 128 * 1024;
pub(crate) const VAULT_BODY_LIMIT: usize = 16 * 1024;
pub(crate) const STATIC_PAGE_CACHE_BODY_LIMIT: usize = 256 * 1024;
pub(crate) const SINGLETON_CONNECT_BODY_LIMIT: usize = 256 * 1024;
pub(crate) const SINGLETON_ACTIVATION_BODY_LIMIT: usize = 16 * 1024;
pub(crate) const OTLP_BODY_LIMIT: usize = 8 * 1024 * 1024;
const AGGREGATE_BODY_BUDGET_CHUNK_SIZE: usize = 64 * 1024;
const AGGREGATE_BODY_BUDGET_SIZE: usize = 32 * 1024 * 1024;

#[derive(Debug)]
pub(crate) enum BodyLimitError {
    TooLarge,
    Body(ErrorCode),
}

pub(crate) struct LimitedBody {
    pub(crate) bytes: Bytes,
    pub(crate) permit: OwnedSemaphorePermit,
}

pub(crate) fn declared_content_length_exceeds_limit(headers: &HeaderMap, limit: usize) -> bool {
    headers
        .get(hyper::header::CONTENT_LENGTH)
        .and_then(|value| value.to_str().ok())
        .and_then(|value| value.parse::<u64>().ok())
        .is_some_and(|length| length > limit as u64)
}

pub(crate) async fn collect_body_limited(
    mut body: UnsyncBoxBody<Bytes, ErrorCode>,
    limit: usize,
) -> Result<LimitedBody, BodyLimitError> {
    let permit_count = limit.div_ceil(AGGREGATE_BODY_BUDGET_CHUNK_SIZE) as u32;
    let permit = aggregate_body_budget()
        .acquire_many_owned(permit_count)
        .await
        .map_err(|error| {
            BodyLimitError::Body(ErrorCode::InternalError(Some(format!(
                "request body budget unavailable: {error}"
            ))))
        })?;
    let mut collected = Vec::new();
    while let Some(frame) = body.frame().await {
        let frame = frame.map_err(BodyLimitError::Body)?;
        let Some(data) = frame.data_ref() else {
            continue;
        };
        if data.len() > limit.saturating_sub(collected.len()) {
            return Err(BodyLimitError::TooLarge);
        }
        collected.try_reserve_exact(data.len()).map_err(|error| {
            BodyLimitError::Body(ErrorCode::InternalError(Some(format!(
                "request body buffer allocation failed: {error}"
            ))))
        })?;
        collected.extend_from_slice(data);
    }
    Ok(LimitedBody {
        bytes: Bytes::from(collected),
        permit,
    })
}

fn aggregate_body_budget() -> Arc<Semaphore> {
    static BUDGET: OnceLock<Arc<Semaphore>> = OnceLock::new();
    BUDGET
        .get_or_init(|| {
            Arc::new(Semaphore::new(
                AGGREGATE_BODY_BUDGET_SIZE / AGGREGATE_BODY_BUDGET_CHUNK_SIZE,
            ))
        })
        .clone()
}

#[cfg(test)]
mod tests {
    use super::*;
    use futures::stream;
    use http_body::Frame;
    use http_body_util::{BodyExt, StreamBody};

    fn body(chunks: Vec<Bytes>) -> UnsyncBoxBody<Bytes, ErrorCode> {
        StreamBody::new(stream::iter(
            chunks
                .into_iter()
                .map(|chunk| Ok::<_, std::convert::Infallible>(Frame::data(chunk))),
        ))
        .map_err(|never| match never {})
        .boxed_unsync()
    }

    #[tokio::test]
    async fn accepts_body_at_limit_across_chunks() {
        let result = collect_body_limited(
            body(vec![Bytes::from_static(b"abc"), Bytes::from_static(b"def")]),
            6,
        )
        .await
        .expect("body within limit");
        assert_eq!(result.bytes, Bytes::from_static(b"abcdef"));
    }

    #[tokio::test]
    async fn refuses_body_after_crossing_limit() {
        let result = collect_body_limited(
            body(vec![Bytes::from_static(b"abc"), Bytes::from_static(b"def")]),
            5,
        )
        .await;
        assert!(matches!(result, Err(BodyLimitError::TooLarge)));
    }

    #[test]
    fn refuses_declared_length_above_limit() {
        let headers = HeaderMap::from_iter([(hyper::header::CONTENT_LENGTH, "7".parse().unwrap())]);
        assert!(declared_content_length_exceeds_limit(&headers, 6));
        assert!(!declared_content_length_exceeds_limit(&headers, 7));
    }
}