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}