diff --git a/crates/cos/src/budget.rs b/crates/cos/src/budget.rs index 35319f13b..cf1c3fc2a 100644 --- a/crates/cos/src/budget.rs +++ b/crates/cos/src/budget.rs @@ -53,10 +53,20 @@ impl DecodeBudget { /// `stream` decoded like [`Stream::decoded`], charged to the allowance. Fails, without /// decoding, once the allowance is spent. pub fn decode(&mut self, stream: &Stream) -> Result, CosError> { + self.charge(stream, Stream::decoded_within_typed) + } + + /// `stream` decoded like [`Stream::decoded_strict`] (any corruption is an error), charged to + /// the allowance. + pub fn decode_strict(&mut self, stream: &Stream) -> Result, CosError> { + self.charge(stream, Stream::decoded_strict_within_typed) + } + + fn charge(&mut self, stream: &Stream, decode: fn(&Stream, usize) -> Result, FilterError>) -> Result, CosError> { if self.bytes == 0 && !stream.raw.is_empty() { return Err(CosError::Filter(FilterError::LimitExceeded(self.limit).to_string())); } - match stream.decoded_within_typed(self.bytes) { + match decode(stream, self.bytes) { Ok(data) => { self.bytes = self.bytes.saturating_sub(data.len()); Ok(data) @@ -127,6 +137,10 @@ mod tests { let mut budget = DecodeBudget::new((1 << 20) - 1, 1); assert!(budget.decode_contents(&doc, &contents).iter().all(Result::is_err)); assert_eq!(budget.remaining(), 0); + let mut budget = DecodeBudget::new((1 << 20) - 1, 1); + let Object::Stream(stream) = &*doc.get(bomb) else { panic!("not a stream") }; + assert!(budget.decode_strict(stream).is_err()); + assert_eq!(budget.remaining(), 0); assert!(budget.enter_form()); assert!(!budget.enter_form()); diff --git a/crates/cos/src/object.rs b/crates/cos/src/object.rs index fe33388be..9a07d4822 100644 --- a/crates/cos/src/object.rs +++ b/crates/cos/src/object.rs @@ -223,21 +223,24 @@ impl Stream { let max = max.min(MAX_DECODED); let chain = self.filters(); if chain.is_empty() { - if self.raw.len() > max { - return Err(pdfcraft_filters::FilterError::LimitExceeded(max)); - } - return Ok(self.raw.to_vec()); + return self.raw_within_typed(max); } pdfcraft_filters::decode_tolerant(&chain, &self.raw, max).map(|(v, _)| v) } /// Decoded data, bounded by [`MAX_DECODED`]; any corruption is an error. pub fn decoded_strict(&self) -> Result, CosError> { + self.decoded_strict_within_typed(MAX_DECODED).map_err(|e| CosError::Filter(e.to_string())) + } + + /// [`Stream::decoded_strict`] bounded by `max`, with the filter error kept typed. + pub(crate) fn decoded_strict_within_typed(&self, max: usize) -> Result, pdfcraft_filters::FilterError> { + let max = max.min(MAX_DECODED); let chain = self.filters(); if chain.is_empty() { - return self.raw_within(MAX_DECODED); + return self.raw_within_typed(max); } - pdfcraft_filters::decode(&chain, &self.raw, MAX_DECODED).map_err(|e| CosError::Filter(e.to_string())) + pdfcraft_filters::decode(&chain, &self.raw, max) } /// Data after the first filter alone, bounded by `max`: the rest of the chain stays encoded. @@ -252,11 +255,15 @@ impl Stream { } fn raw_within(&self, max: usize) -> Result, CosError> { + self.raw_within_typed(max).map_err(|e| CosError::Filter(e.to_string())) + } + + fn raw_within_typed(&self, max: usize) -> Result, pdfcraft_filters::FilterError> { let max = max.min(MAX_DECODED); // An empty decoder chain still produces an owned output buffer: check its limit // before cloning, including streams whose /Crypt filter was applied on load. if self.raw.len() > max { - return Err(CosError::Filter(pdfcraft_filters::FilterError::LimitExceeded(max).to_string())); + return Err(pdfcraft_filters::FilterError::LimitExceeded(max)); } Ok(self.raw.to_vec()) }