sachy.dev
A pit full of rusty nails
1use std::sync::Arc;
2
3use hashbrown::HashMap;
4use hyper::{HeaderMap, body::Bytes, header::ACCEPT_ENCODING};
5use nailbox::try_arc_within;
6use nailconfig::{DropBehavior, NailConfig, RateLimitingConfig};
7use rapidhash::fast::RandomState;
8
9pub type SpicyPayloads = HashMap<SpicyPayloadKind, Bytes, RandomState>;
10
11static GZIP: &str = "gzip";
12static BROTLI: &str = "br";
13
14#[derive(Debug, PartialEq, Eq, Clone, Copy, Hash)]
15#[repr(u32)]
16pub enum SpicyPayloadKind {
17 Gz,
18 Brotli,
19}
20
21impl SpicyPayloadKind {
22 pub fn as_str(&self) -> &'static str {
23 match self {
24 SpicyPayloadKind::Gz => GZIP,
25 SpicyPayloadKind::Brotli => BROTLI,
26 }
27 }
28}
29
30impl SpicyPayloadKind {
31 fn file_kind(file: impl AsRef<str>) -> Option<SpicyPayloadKind> {
32 let file = file.as_ref();
33
34 if file.ends_with(".gz") {
35 Some(Self::Gz)
36 } else if file.ends_with(".br") {
37 Some(Self::Brotli)
38 } else {
39 None
40 }
41 }
42
43 pub fn accepts_encoding(header: &HeaderMap) -> Option<SpicyPayloadKind> {
44 header
45 .get(ACCEPT_ENCODING)
46 .and_then(|header| header.to_str().ok())
47 .and_then(|header| {
48 header
49 .contains(BROTLI)
50 .then_some(Self::Brotli)
51 .or_else(|| header.contains(GZIP).then_some(Self::Gz))
52 })
53 }
54}
55
56pub fn get_spicy_payload(config: &NailConfig) -> Option<Arc<SpicyPayloads>> {
57 match &config.rate_limiting {
58 RateLimitingConfig::HardLimit {
59 drop_behavior: DropBehavior::Spicy { payload },
60 ..
61 }
62 | RateLimitingConfig::SoftWithHardLimit {
63 drop_behavior: DropBehavior::Spicy { payload },
64 ..
65 } => try_arc_within(|| {
66 payload
67 .iter()
68 .filter_map(|file| SpicyPayloadKind::file_kind(file).map(|kind| (kind, file)))
69 .map(|(kind, file)| {
70 Ok((
71 kind,
72 std::fs::read(file)
73 .inspect_err(|err| {
74 tracing::error!("Failed to load spicy payload: {err}")
75 })
76 .map(Bytes::from)?,
77 ))
78 })
79 .collect::<std::io::Result<_>>()
80 })
81 .ok(),
82 _ => None,
83 }
84}