A decode charges the page's decode budget for every filter stage's output

crates/cos/src/budget.rs

@@ -12,14 +12,9 @@ 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))
-}
+/// A stream decode within a byte bound, with the bytes its filter stages decoded: the data or
+/// the error, each with that count.
+type Decoder = fn(&Stream, usize) -> Result<(Vec<u8>, usize), (FilterError, usize)>;
 
 /// What one operation may still decode and enter across all the streams it reads.
 ///
@@ -62,24 +57,25 @@ 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, tolerant)
+        self.charge(stream, usize::MAX, Stream::decoded_within_counted)
     }
 
     /// 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, tolerant)
+        self.charge(stream, max, Stream::decoded_within_counted)
     }
 
     /// `stream` decoded like [`Stream::decoded_strict`] (any corruption is an error), charged to
     /// the allowance.
     pub fn decode_strict(&mut self, stream: &Stream) -> Result<Vec<u8>, CosError> {
-        self.charge(stream, usize::MAX, Stream::decoded_strict_within_typed)
+        self.charge(stream, usize::MAX, Stream::decoded_strict_within_counted)
     }
 
-    /// 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.
+    /// Decodes `stream` within `max` and the allowance, charging what every filter stage
+    /// decoded: a stage's output counts even when a later stage shrinks it, and a corrupt stream
+    /// pays for the bytes decoded before the error, so referencing a stream again and again
+    /// spends the allowance by the work each decode does.
     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()));
@@ -87,16 +83,16 @@ impl DecodeBudget {
         if max < self.bytes {
             let result = decode(stream, max);
             let spent = match &result {
-                Ok(data) => data.len(),
+                Ok((_, spent)) => *spent,
                 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(|(data, _)| data).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, spent)) => {
+                self.bytes = self.bytes.saturating_sub(spent);
                 Ok(data)
             }
             Err((FilterError::LimitExceeded(_), _)) => {
@@ -201,4 +197,24 @@ mod tests {
         assert!(budget.decode_within(&corrupt, 1 << 30).is_ok());
         assert_eq!(budget.remaining(), (10 << 20) - partial);
     }
+
+    /// A chain whose later stage shrinks the data pays for every stage's output: a Flate stage
+    /// that inflates to hex text, which ASCIIHex then halves, charges both.
+    #[test]
+    fn decodes_charge_every_stage_of_a_chain() {
+        let hex: Vec<u8> = std::iter::repeat_n(b"ab".as_slice(), 1 << 19).flatten().copied().collect();
+        let flate = Stream::flate(Dict::new(), &hex);
+        let mut dict = Dict::new();
+        dict.set(b"Filter".to_vec(), Object::Array(vec![Object::name("FlateDecode"), Object::name("ASCIIHexDecode")]));
+        let chained = Stream::from_raw(dict, flate.raw.to_vec());
+        assert_eq!(chained.decoded().map(|d| d.len()).ok(), Some(1 << 19));
+        let charged = (1 << 20) + (1 << 19);
+        let mut budget = DecodeBudget::new(10 << 20, 1);
+        assert_eq!(budget.decode(&chained).map(|d| d.len()).ok(), Some(1 << 19));
+        assert_eq!(budget.remaining(), (10 << 20) - charged);
+        assert!(budget.decode_within(&chained, 2 << 20).is_ok());
+        assert_eq!(budget.remaining(), (10 << 20) - 2 * charged);
+        assert!(budget.decode_strict(&chained).is_ok());
+        assert_eq!(budget.remaining(), (10 << 20) - 3 * charged);
+    }
 }

crates/cos/src/object.rs

@@ -283,32 +283,44 @@ impl Stream {
     /// Like [`Stream::decoded`], for a stream whose decoded size has a tighter bound than
     /// [`MAX_DECODED`]: decoding past `max` bytes is an error.
     pub fn decoded_within(&self, max: usize) -> Result<Vec<u8>, CosError> {
-        self.decoded_within_typed(max).map_err(|e| CosError::Filter(e.to_string()))
+        self.decoded_within_counted(max).map(|(v, _)| v).map_err(|(e, _)| CosError::Filter(e.to_string()))
     }
 
     /// [`Stream::decoded_within`] with the filter error kept typed, so a caller can tell a
-    /// spent limit from corrupt data.
-    pub(crate) fn decoded_within_typed(&self, max: usize) -> Result<Vec<u8>, pdfcraft_filters::FilterError> {
+    /// spent limit from corrupt data, and the bytes every filter stage decoded (the work done,
+    /// on success and on failure).
+    pub(crate) fn decoded_within_counted(&self, max: usize) -> Result<(Vec<u8>, usize), (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(|v| {
+                    let n = v.len();
+                    (v, n)
+                })
+                .map_err(|e| (e, 0));
         }
-        pdfcraft_filters::decode_tolerant(&chain, &self.raw, max).map(|(v, _)| v)
+        pdfcraft_filters::decode_tolerant_counted(&chain, &self.raw, max).map(|(v, _, spent)| (v, spent))
     }
 
     /// 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_counted(MAX_DECODED).map(|(v, _)| v).map_err(|(e, _)| CosError::Filter(e.to_string()))
     }
 
-    /// [`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)> {
+    /// [`Stream::decoded_strict`] bounded by `max`, counted like [`Stream::decoded_within_counted`].
+    pub(crate) fn decoded_strict_within_counted(&self, max: usize) -> Result<(Vec<u8>, usize), (pdfcraft_filters::FilterError, usize)> {
         let max = max.min(MAX_DECODED);
         let chain = self.readable_filters().map_err(|e| (e, 0))?;
         if chain.is_empty() {
-            return self.raw_within_typed(max).map_err(|e| (e, 0));
+            return self
+                .raw_within_typed(max)
+                .map(|v| {
+                    let n = v.len();
+                    (v, n)
+                })
+                .map_err(|e| (e, 0));
         }
         pdfcraft_filters::decode_counted(&chain, &self.raw, max)
     }

crates/filters/README.md

@@ -17,8 +17,9 @@ pub struct Params { predictor, colors, bits_per_component, columns, early_change
 pub enum FilterError { Corrupt { filter, detail }, LimitExceeded(usize), Unsupported(String) }
 
 decode(chain, data, max_output) -> Result<Vec<u8>, FilterError>
-decode_counted(chain, data, max_output) -> Result<Vec<u8>, (FilterError, usize /*decoded before the error*/)>
+decode_counted(chain, data, max_output) -> Result<(Vec<u8>, usize /*decoded*/), (FilterError, usize /*decoded*/)>
 decode_tolerant(chain, data, max_output) -> Result<(Vec<u8>, bool /*partial*/), FilterError>
+decode_tolerant_counted(chain, data, max_output) -> Result<(Vec<u8>, bool /*partial*/, usize /*decoded*/), (FilterError, usize /*decoded*/)>
 encode(filter, params, data) -> Result<Vec<u8>, FilterError>
 encode_flate(data) -> Vec<u8>   // zlib, level 6
 ```
@@ -26,7 +27,7 @@ encode_flate(data) -> Vec<u8>   // zlib, level 6
 - `chain` is applied first to last, as in a `/Filter` array.
 - Image codecs (DCT, JPX, JBIG2, CCITT), `Crypt` and unknown filters anywhere in a chain give `Unsupported`: callers keep such data encoded. So does a chain longer than `MAX_CHAIN` (8) filters, before any stage runs. Decryption belongs to `crypt`.
 - `max_output` bounds every intermediate buffer (including the pre-predictor buffer) and gives `LimitExceeded(max_output)`. This also holds when Flate retries raw deflate after a misleading zlib header, in both strict and tolerant decoding.
-- `decode_counted` is `decode` that also gives, on failure, the bytes the stages decoded before it, so a caller charging a budget (`pdfcraft_cos::DecodeBudget`) pays for a failed decode too.
+- `decode_counted` and `decode_tolerant_counted` also give the bytes every stage that ran decoded, on success and on failure, so a caller charging a budget (`pdfcraft_cos::DecodeBudget`) pays for the work: a stage's output counts even when a later stage shrinks it, and a failed decode pays for what it decoded first.
 - `decode_tolerant` turns `Corrupt` into "what was decoded so far" plus `partial = true`, then runs the rest of the chain on it (predictors are applied to partial output too). `LimitExceeded` and `Unsupported` are still errors.
 
 ## Behaviour

crates/filters/src/lib.rs

@@ -140,28 +140,36 @@ 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)
+    decode_counted(chain, data, max_output).map(|(v, _)| v).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)
+/// Like [`decode`], also giving the bytes every stage that ran decoded, on success and on
+/// failure: the work a caller charging an allowance spent. A stage's output counts even when a
+/// later stage shrinks it.
+pub fn decode_counted(chain: &[(Filter, Params)], data: &[u8], max_output: usize) -> Result<(Vec<u8>, usize), (FilterError, usize)> {
+    run(chain, data, max_output, true).map(|(v, _, spent)| (v, spent))
 }
 
 /// Like [`decode`], but on corrupt data returns what was decoded before the error and
 /// 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).map_err(|(e, _)| e)
+    decode_tolerant_counted(chain, data, max_output).map(|(v, partial, _)| (v, partial)).map_err(|(e, _)| e)
+}
+
+/// [`decode_tolerant`], also giving the bytes every stage that ran decoded, as
+/// [`decode_counted`] does.
+pub fn decode_tolerant_counted(chain: &[(Filter, Params)], data: &[u8], max_output: usize) -> Result<(Vec<u8>, bool, usize), (FilterError, usize)> {
+    run(chain, data, max_output, false)
 }
 
 /// 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;
 
-/// 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)> {
+/// The decoded data, whether it is partial and the bytes all stages decoded, or the error and
+/// the bytes decoded before it.
+fn run(chain: &[(Filter, Params)], data: &[u8], max: usize, strict: bool) -> Result<(Vec<u8>, bool, usize), (FilterError, usize)> {
     if chain.len() > MAX_CHAIN {
         return Err((FilterError::Unsupported(format!("a chain of {} filters", chain.len())), 0));
     }
@@ -184,7 +192,7 @@ fn run(chain: &[(Filter, Params)], data: &[u8], max: usize, strict: bool) -> Res
         });
         spent = spent.saturating_add(cur.len());
     }
-    Ok((cur.into_owned(), partial))
+    Ok((cur.into_owned(), partial, spent))
 }
 
 fn decode_one(filter: &Filter, params: &Params, data: &[u8], max: usize) -> Step {