Skip to main content

aion_package/
extraction.rs

1//! Explicit inflate budgets for `.aion` archive extraction.
2//!
3//! ZIP entries declare their own compressed/uncompressed sizes, so a hostile
4//! archive can lie and inflate far past any upload ceiling (DEFLATE bombs
5//! reach ~1000:1). Extraction therefore charges every inflated byte against
6//! an [`ExtractionLimits`] budget the caller chooses explicitly — there is
7//! deliberately no `Default`.
8
9use std::io::Read;
10
11use zip::result::ZipError;
12
13use crate::PackageError;
14
15/// Caller-chosen budget on the total inflated size of all archive entries.
16///
17/// Every package-loading entry point requires one; the caller decides whether
18/// the input is trusted ([`ExtractionLimits::unbounded`]) or hostile network
19/// bytes ([`ExtractionLimits::bounded`]).
20#[derive(Clone, Copy, Debug, Eq, PartialEq)]
21pub struct ExtractionLimits {
22    max_inflated_bytes: Option<u64>,
23}
24
25impl ExtractionLimits {
26    /// Caps the running total of inflated bytes across all archive entries
27    /// (manifest included) at `max_inflated_bytes`. Exceeding the budget
28    /// fails extraction loudly with
29    /// [`PackageError::InflatedSizeExceeded`] — entries are never truncated.
30    #[must_use]
31    pub const fn bounded(max_inflated_bytes: u64) -> Self {
32        Self {
33            max_inflated_bytes: Some(max_inflated_bytes),
34        }
35    }
36
37    /// No inflate ceiling.
38    ///
39    /// ONLY for trusted operator-local files (engine startup packages, CLI
40    /// and build tooling reading archives they just wrote, test fixtures) —
41    /// never for network input. A hostile archive under any upload ceiling
42    /// can inflate ~1000:1 and exhaust memory if extracted unbounded.
43    #[must_use]
44    pub const fn unbounded() -> Self {
45        Self {
46            max_inflated_bytes: None,
47        }
48    }
49
50    /// Begins one extraction's running budget.
51    pub(crate) const fn budget(self) -> ExtractionBudget {
52        match self.max_inflated_bytes {
53            Some(limit) => ExtractionBudget::Bounded {
54                limit,
55                remaining: limit,
56            },
57            None => ExtractionBudget::Unbounded,
58        }
59    }
60}
61
62/// Running inflate budget consumed entry by entry during one extraction.
63#[derive(Debug)]
64pub(crate) enum ExtractionBudget {
65    /// Trusted-input extraction with no inflate ceiling.
66    Unbounded,
67    /// Bounded extraction tracking the bytes still admissible.
68    Bounded {
69        /// The caller-configured ceiling, reported on refusal.
70        limit: u64,
71        /// Budget left for the remaining entries.
72        remaining: u64,
73    },
74}
75
76impl ExtractionBudget {
77    /// Reads one archive entry to completion, charging its inflated size
78    /// against the remaining budget.
79    ///
80    /// The reader is wrapped in [`Read::take`] at one byte past the remaining
81    /// budget, so a single entry can never buffer meaningfully past the
82    /// budget before the refusal fires.
83    pub(crate) fn read_entry<R>(&mut self, reader: &mut R) -> Result<Vec<u8>, PackageError>
84    where
85        R: Read,
86    {
87        let mut bytes = Vec::new();
88        match self {
89            Self::Unbounded => {
90                reader
91                    .read_to_end(&mut bytes)
92                    .map_err(|source| PackageError::ArchiveRead(ZipError::Io(source)))?;
93            }
94            Self::Bounded { limit, remaining } => {
95                // One sentinel byte past the budget distinguishes "exactly on
96                // budget" from "would exceed it" without unbounded buffering.
97                let probe = remaining.saturating_add(1);
98                let mut taken = reader.take(probe);
99                taken
100                    .read_to_end(&mut bytes)
101                    .map_err(|source| PackageError::ArchiveRead(ZipError::Io(source)))?;
102                let inflated = probe - taken.limit();
103                if inflated > *remaining {
104                    return Err(PackageError::InflatedSizeExceeded { limit: *limit });
105                }
106                *remaining -= inflated;
107            }
108        }
109        Ok(bytes)
110    }
111}
112
113#[cfg(test)]
114mod tests {
115    use super::ExtractionLimits;
116    use crate::PackageError;
117
118    #[test]
119    fn bounded_budget_charges_across_entries() -> Result<(), PackageError> {
120        let mut budget = ExtractionLimits::bounded(10).budget();
121
122        let first = budget.read_entry(&mut &[0_u8; 6][..])?;
123        assert_eq!(first.len(), 6);
124
125        let second = budget.read_entry(&mut &[0_u8; 4][..])?;
126        assert_eq!(second.len(), 4);
127
128        let result = budget.read_entry(&mut &[0_u8; 1][..]);
129        assert!(matches!(
130            result,
131            Err(PackageError::InflatedSizeExceeded { limit: 10 })
132        ));
133        Ok(())
134    }
135
136    #[test]
137    fn single_entry_past_budget_is_refused_reporting_the_limit() {
138        let mut budget = ExtractionLimits::bounded(4).budget();
139
140        let result = budget.read_entry(&mut &[0_u8; 5][..]);
141
142        assert!(matches!(
143            result,
144            Err(PackageError::InflatedSizeExceeded { limit: 4 })
145        ));
146    }
147
148    #[test]
149    fn unbounded_budget_reads_everything() -> Result<(), PackageError> {
150        let mut budget = ExtractionLimits::unbounded().budget();
151
152        let bytes = budget.read_entry(&mut &[0_u8; 4096][..])?;
153
154        assert_eq!(bytes.len(), 4096);
155        Ok(())
156    }
157}