A failed strict decode charges the page's decode budget for what it decoded

crates/cos/src/budget.rs

@@ -12,6 +12,15 @@ pub const MAX_PAGE_DECODED: usize = 256 << 20;
 /// Form XObjects one page's content walk may enter, each entry counted again.
 pub const MAX_PAGE_FORMS: usize = 50_000;
 
+/// A stream decode within a byte bound: the data, or the error and the bytes decoded before it.
+type Decoder = fn(&Stream, usize) -> Result<Vec<u8>, (FilterError, usize)>;
+
+/// A tolerant decode recovers corrupt data as partial output, so its only errors are a spent
+/// limit and an unsupported filter, which is refused before any decoding.
+fn tolerant(stream: &Stream, max: usize) -> Result<Vec<u8>, (FilterError, usize)> {
+    stream.decoded_within_typed(max).map_err(|e| (e, 0))
+}
+
 /// What one operation may still decode and enter across all the streams it reads.
 ///
 /// [`MAX_DECODED`](crate::MAX_DECODED) bounds one decode. It does not bound a page whose
@@ -53,13 +62,13 @@ 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<Vec<u8>, CosError> {
-        self.charge(stream, usize::MAX, Stream::decoded_within_typed)
+        self.charge(stream, usize::MAX, tolerant)
     }
 
     /// Like [`DecodeBudget::decode`], for a stream with a tighter bound of its own: decoding
     /// past `max` bytes is an error that charges `max`, the work it took to find out.
     pub fn decode_within(&mut self, stream: &Stream, max: usize) -> Result<Vec<u8>, CosError> {
-        self.charge(stream, max, Stream::decoded_within_typed)
+        self.charge(stream, max, tolerant)
     }
 
     /// `stream` decoded like [`Stream::decoded_strict`] (any corruption is an error), charged to
@@ -68,7 +77,10 @@ impl DecodeBudget {
         self.charge(stream, usize::MAX, Stream::decoded_strict_within_typed)
     }
 
-    fn charge(&mut self, stream: &Stream, max: usize, decode: fn(&Stream, usize) -> Result<Vec<u8>, FilterError>) -> Result<Vec<u8>, CosError> {
+    /// Decodes `stream` within `max` and the allowance, charging what it decoded: a corrupt
+    /// stream pays for the bytes decoded before the error, so referencing it again and again
+    /// spends the allowance like any other stream.
+    fn charge(&mut self, stream: &Stream, max: usize, decode: Decoder) -> Result<Vec<u8>, CosError> {
         if self.bytes == 0 && !stream.raw.is_empty() {
             return Err(CosError::Filter(FilterError::LimitExceeded(self.limit).to_string()));
         }
@@ -76,22 +88,25 @@ impl DecodeBudget {
             let result = decode(stream, max);
             let spent = match &result {
                 Ok(data) => data.len(),
-                Err(FilterError::LimitExceeded(_)) => max,
-                Err(_) => 0,
+                Err((FilterError::LimitExceeded(_), _)) => max,
+                Err((_, spent)) => *spent,
             };
             self.bytes = self.bytes.saturating_sub(spent);
-            return result.map_err(|e| CosError::Filter(e.to_string()));
+            return result.map_err(|(e, _)| CosError::Filter(e.to_string()));
         }
         match decode(stream, self.bytes) {
             Ok(data) => {
                 self.bytes = self.bytes.saturating_sub(data.len());
                 Ok(data)
             }
-            Err(FilterError::LimitExceeded(_)) => {
+            Err((FilterError::LimitExceeded(_), _)) => {
                 self.bytes = 0;
                 Err(CosError::Filter(FilterError::LimitExceeded(self.limit).to_string()))
             }
-            Err(e) => Err(CosError::Filter(e.to_string())),
+            Err((e, spent)) => {
+                self.bytes = self.bytes.saturating_sub(spent);
+                Err(CosError::Filter(e.to_string()))
+            }
         }
     }
 
@@ -169,4 +184,21 @@ mod tests {
         assert!(budget.first_visit(bomb));
         assert!(!budget.first_visit(bomb));
     }
+
+    /// A corrupt stream fails a strict decode after decoding part of it: each reference pays
+    /// for that part, so references to it can't decode without limit.
+    #[test]
+    fn failed_strict_decodes_charge_what_they_decoded() {
+        let whole = Stream::flate(Dict::new(), &vec![b' '; 1 << 20]);
+        let cut = whole.raw.len() / 2;
+        let corrupt = Stream::from_raw(whole.dict.clone(), whole.raw.get(..cut).unwrap_or_default().to_vec());
+        let partial = corrupt.decoded().map(|d| d.len()).unwrap_or(0);
+        assert!(partial > 0, "a cut flate stream decodes in part");
+        let mut budget = DecodeBudget::new(10 << 20, 1);
+        assert!(budget.decode_strict(&corrupt).is_err());
+        assert_eq!(budget.remaining(), (10 << 20) - partial);
+        let mut budget = DecodeBudget::new(10 << 20, 1);
+        assert!(budget.decode_within(&corrupt, 1 << 30).is_ok());
+        assert_eq!(budget.remaining(), (10 << 20) - partial);
+    }
 }

crates/cos/src/object.rs

@@ -299,17 +299,18 @@ impl Stream {
 
     /// Decoded data, bounded by [`MAX_DECODED`]; any corruption is an error.
     pub fn decoded_strict(&self) -> Result<Vec<u8>, CosError> {
-        self.decoded_strict_within_typed(MAX_DECODED).map_err(|e| CosError::Filter(e.to_string()))
+        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<Vec<u8>, pdfcraft_filters::FilterError> {
+    /// [`Stream::decoded_strict`] bounded by `max`, with the filter error kept typed and the
+    /// bytes decoded before it.
+    pub(crate) fn decoded_strict_within_typed(&self, max: usize) -> Result<Vec<u8>, (pdfcraft_filters::FilterError, usize)> {
         let max = max.min(MAX_DECODED);
-        let chain = self.readable_filters()?;
+        let chain = self.readable_filters().map_err(|e| (e, 0))?;
         if chain.is_empty() {
-            return self.raw_within_typed(max);
+            return self.raw_within_typed(max).map_err(|e| (e, 0));
         }
-        pdfcraft_filters::decode(&chain, &self.raw, max)
+        pdfcraft_filters::decode_counted(&chain, &self.raw, max)
     }
 
     /// Data after the first filter alone, bounded by `max`: the rest of the chain stays encoded.

crates/filters/src/lib.rs

@@ -140,6 +140,12 @@ pub(crate) type Step = Result<Vec<u8>, Failure>;
 /// [`FilterError::Unsupported`] (callers keep such data encoded), as does a chain longer than
 /// [`MAX_CHAIN`]. `max_output` bounds every intermediate buffer.
 pub fn decode(chain: &[(Filter, Params)], data: &[u8], max_output: usize) -> Result<Vec<u8>, FilterError> {
+    decode_counted(chain, data, max_output).map_err(|(e, _)| e)
+}
+
+/// Like [`decode`], but a failure also gives the bytes decoded before it, across the stages
+/// that ran: the work a caller charging an allowance spent finding out.
+pub fn decode_counted(chain: &[(Filter, Params)], data: &[u8], max_output: usize) -> Result<Vec<u8>, (FilterError, usize)> {
     run(chain, data, max_output, true).map(|(v, _)| v)
 }
 
@@ -147,22 +153,24 @@ pub fn decode(chain: &[(Filter, Params)], data: &[u8], max_output: usize) -> Res
 /// reports `partial = true` (the behaviour viewers expect). The remaining filters of the
 /// chain still run on the partial data. Limit and unsupported errors are still errors.
 pub fn decode_tolerant(chain: &[(Filter, Params)], data: &[u8], max_output: usize) -> Result<(Vec<u8>, bool), FilterError> {
-    run(chain, data, max_output, false)
+    run(chain, data, max_output, false).map_err(|(e, _)| e)
 }
 
 /// Longest filter chain decoded. Real streams use one to three filters; every stage may produce
 /// up to the output limit, so the chain length multiplies the work.
 pub const MAX_CHAIN: usize = 8;
 
-fn run(chain: &[(Filter, Params)], data: &[u8], max: usize, strict: bool) -> Result<(Vec<u8>, bool), FilterError> {
+/// The decoded data and whether it is partial, or the error and the bytes decoded before it.
+fn run(chain: &[(Filter, Params)], data: &[u8], max: usize, strict: bool) -> Result<(Vec<u8>, bool), (FilterError, usize)> {
     if chain.len() > MAX_CHAIN {
-        return Err(FilterError::Unsupported(format!("a chain of {} filters", chain.len())));
+        return Err((FilterError::Unsupported(format!("a chain of {} filters", chain.len())), 0));
     }
     if let Some((f, _)) = chain.iter().find(|(f, _)| !f.is_decodable()) {
-        return Err(FilterError::Unsupported(f.name().to_owned()));
+        return Err((FilterError::Unsupported(f.name().to_owned()), 0));
     }
     let mut cur: Cow<[u8]> = Cow::Borrowed(data);
     let mut partial = false;
+    let mut spent = 0usize;
     for (filter, params) in chain {
         #[cfg(test)]
         STAGE_BYTES.with(|c| c.set(c.get() + cur.len()));
@@ -172,8 +180,9 @@ fn run(chain: &[(Filter, Params)], data: &[u8], max: usize, strict: bool) -> Res
                 partial = true;
                 out
             }
-            Err(f) => return Err(f.error),
+            Err(f) => return Err((f.error, spent.saturating_add(f.partial.len()))),
         });
+        spent = spent.saturating_add(cur.len());
     }
     Ok((cur.into_owned(), partial))
 }