fix(fmt): budget amplified output before allocation

crates/bashkit/src/builtins/fmt/linebreak.rs

@@ -1,23 +1,24 @@
 // Vendored from uutils/coreutils (MIT, Copyright (c) uutils developers),
 // src/uu/fmt/src/linebreak.rs at rev 6579cc2. Changes: output goes to a
-// `Vec<u8>`, input is any `BufRead`, options come from `super::FmtOptions`.
+// budgeted sink, search steps poll the shared execution budget, and options
+// come from `super::FmtOptions`.
 // Keep in sync with upstream; see knowledge/foundations/builtins.md "fmt".
 
 // spell-checker:ignore (ToDO) INFTY MULT accum breakwords linebreak linebreaking linebreaks linelen maxlength minlength nchars ostream overlen parasplit plass posn powf punct signum slen sstart tabwidth tlen underlen winfo wlen wordlen
 
-use std::io::Write;
 use std::mem;
 
-use super::FmtOptions;
 use super::parasplit::{ParaWords, Paragraph, WordInfo};
+use super::{FmtOptions, FmtOutput};
+use crate::error::Result;
 
 struct BreakArgs<'a> {
     opts: &'a FmtOptions,
     init_len: usize,
     indent: &'a [u8],
     indent_len: usize,
     uniform: bool,
-    ostream: &'a mut Vec<u8>,
+    ostream: &'a mut FmtOutput,
 }
 
 impl BreakArgs<'_> {
@@ -36,11 +37,7 @@ impl BreakArgs<'_> {
     }
 }
 
-pub fn break_lines(
-    para: &Paragraph,
-    opts: &FmtOptions,
-    ostream: &mut Vec<u8>,
-) -> std::io::Result<()> {
+pub fn break_lines(para: &Paragraph, opts: &FmtOptions, ostream: &mut FmtOutput) -> Result<()> {
     // indent
     let p_indent = &para.indent_str;
     let p_indent_len = para.indent_len;
@@ -97,7 +94,7 @@ pub fn break_lines(
 fn break_simple<'a, T: Iterator<Item = &'a WordInfo<'a>>>(
     mut iter: T,
     args: &mut BreakArgs<'a>,
-) -> std::io::Result<()> {
+) -> Result<()> {
     iter.try_fold((args.init_len, false), |(l, prev_punct), winfo| {
         accum_words_simple(args, l, prev_punct, winfo)
     })?;
@@ -109,7 +106,7 @@ fn accum_words_simple<'a>(
     l: usize,
     prev_punct: bool,
     winfo: &'a WordInfo<'a>,
-) -> std::io::Result<(usize, bool)> {
+) -> Result<(usize, bool)> {
     // compute the length of this word, considering how tabs will expand at this position on the line
     let wlen = winfo.word_nchars + args.compute_width(winfo, l, false);
 
@@ -137,12 +134,12 @@ fn accum_words_simple<'a>(
 fn break_knuth_plass<'a, T: Clone + Iterator<Item = &'a WordInfo<'a>>>(
     mut iter: T,
     args: &mut BreakArgs<'a>,
-) -> std::io::Result<()> {
+) -> Result<()> {
     // run the algorithm to get the breakpoints
-    let breakpoints = find_kp_breakpoints(iter.clone(), args);
+    let breakpoints = find_kp_breakpoints(iter.clone(), args)?;
 
     // iterate through the breakpoints (note that breakpoints is in reverse break order, so we .rev() it
-    let result: std::io::Result<(bool, bool)> = breakpoints.iter().rev().try_fold(
+    let result: Result<(bool, bool)> = breakpoints.iter().rev().try_fold(
         (false, false),
         |(mut prev_punct, mut fresh), &(next_break, break_before)| {
             if fresh {
@@ -218,7 +215,7 @@ struct LineBreak<'a> {
 fn find_kp_breakpoints<'a, T: Clone + Iterator<Item = &'a WordInfo<'a>>>(
     iter: T,
     args: &BreakArgs<'a>,
-) -> Vec<(&'a WordInfo<'a>, bool)> {
+) -> Result<Vec<(&'a WordInfo<'a>, bool)>> {
     let mut iter = iter.peekable();
     // set up the initial null linebreak
     let mut linebreaks = vec![LineBreak {
@@ -244,6 +241,7 @@ fn find_kp_breakpoints<'a, T: Clone + Iterator<Item = &'a WordInfo<'a>>>(
     let mut new_linebreaks = vec![];
     let mut is_sentence_start = false;
     while let Some(w) = iter.next() {
+        args.ostream.consume_work(1)?;
         let next_word_sentence_final = is_next_word_sentence_final(iter.clone());
         // if this is the last word, we don't add additional demerits for this break
         let (is_last_word, is_sentence_end) = match iter.peek() {
@@ -268,6 +266,7 @@ fn find_kp_breakpoints<'a, T: Clone + Iterator<Item = &'a WordInfo<'a>>>(
         // break if we are above the minimum required length
         #[allow(clippy::explicit_iter_loop)]
         for &i in active_breaks.iter() {
+            args.ostream.consume_work(1)?;
             let active = &mut linebreaks[i];
             if active.demerits < best_active_demerits {
                 best_active_demerits = active.demerits;
@@ -381,7 +380,7 @@ fn find_kp_breakpoints<'a, T: Clone + Iterator<Item = &'a WordInfo<'a>>>(
     }
 
     // return the best path
-    build_best_path(&linebreaks, &active_breaks)
+    Ok(build_best_path(&linebreaks, &active_breaks))
 }
 
 fn build_best_path<'a>(paths: &[LineBreak<'a>], active: &[usize]) -> Vec<(&'a WordInfo<'a>, bool)> {
@@ -531,13 +530,13 @@ fn slice_if_fresh(
 }
 
 /// Write a newline and add the indent.
-fn write_newline(indent: &[u8], ostream: &mut Vec<u8>) -> std::io::Result<()> {
+fn write_newline(indent: &[u8], ostream: &mut FmtOutput) -> Result<()> {
     ostream.write_all(b"\n")?;
     ostream.write_all(indent)
 }
 
 /// Write the word, along with slen spaces.
-fn write_with_spaces(word: &[u8], slen: usize, ostream: &mut Vec<u8>) -> std::io::Result<()> {
+fn write_with_spaces(word: &[u8], slen: usize, ostream: &mut FmtOutput) -> Result<()> {
     if slen == 2 {
         ostream.write_all(b"  ")?;
     } else if slen == 1 {

crates/bashkit/src/builtins/fmt/mod.rs

@@ -4,9 +4,11 @@
 //! from uutils/coreutils `fmt` (MIT, `linebreak.rs`, `parasplit.rs`), a
 //! Knuth-Plass optimal-fit filler, not a greedy one. Bashkit is MIT, so GNU
 //! `fmt.c` (GPL) must not be ported; uutils already targets GNU-compatible
-//! output. Only option parsing and I/O are bashkit's. Width and goal follow
-//! GNU's documented rules: goal defaults to 93% of the width, `-g` alone
+//! output. Option parsing, I/O and budget enforcement are bashkit's. Width and
+//! goal follow GNU's documented rules: goal defaults to 93% of the width, `-g` alone
 //! sets the width to goal + 10, width is capped at 2500.
+//! Output capacity is leased before allocation; synchronous rendering and
+//! optimal-fit searches charge the shared work/cancellation/deadline budget.
 
 mod linebreak;
 mod parasplit;
@@ -16,6 +18,7 @@ use async_trait::async_trait;
 use super::{Builtin, BuiltinHelper, Context, resolve_path};
 use crate::error::Result;
 use crate::interpreter::ExecResult;
+use crate::limits::{BudgetedBytes, ExecutionBudget};
 
 /// The fmt builtin.
 ///
@@ -226,21 +229,49 @@ fn parse_tab_width(value: &str) -> std::result::Result<usize, ExecResult> {
         .map_err(|_| err(format!("invalid tab width: '{value}'")))
 }
 
+// THREAT[TM-DOS-103]: indentation and prefixes can multiply output size.
+// Keep the capacity lease with the sink, and check work before each copy.
+struct FmtOutput {
+    bytes: BudgetedBytes,
+    budget: Option<ExecutionBudget>,
+}
+
+impl FmtOutput {
+    fn new(budget: Option<ExecutionBudget>) -> Result<Self> {
+        Ok(Self {
+            bytes: BudgetedBytes::new(budget.as_ref())?,
+            budget,
+        })
+    }
+
+    fn consume_work(&self, units: u64) -> Result<()> {
+        if let Some(budget) = &self.budget {
+            budget.consume_work(units)?;
+        }
+        Ok(())
+    }
+
+    fn write_all(&mut self, bytes: &[u8]) -> Result<()> {
+        self.consume_work(1 + bytes.len() as u64)?;
+        self.bytes.try_extend_from_slice(bytes)?;
+        Ok(())
+    }
+}
+
 /// Format one input into `out` with the vendored engine.
-fn format_input(opts: &FmtOptions, input: Vec<u8>, out: &mut Vec<u8>) {
+fn format_input(opts: &FmtOptions, input: Vec<u8>, out: &mut FmtOutput) -> Result<()> {
     let mut reader: FileOrStdReader = std::io::Cursor::new(input);
     for para in parasplit::ParagraphStream::new(opts, &mut reader) {
+        out.consume_work(1)?;
         match para {
             Err(line) => {
-                out.extend_from_slice(&line);
-                out.push(b'\n');
-            }
-            // Writes to a Vec never fail.
-            Ok(para) => {
-                let _ = linebreak::break_lines(&para, opts, out);
+                out.write_all(&line)?;
+                out.write_all(b"\n")?;
             }
+            Ok(para) => linebreak::break_lines(&para, opts, out)?,
         }
     }
+    Ok(())
 }
 
 #[async_trait]
@@ -253,7 +284,15 @@ impl Builtin for Fmt {
             Ok(v) => v,
             Err(e) => return Ok(e),
         };
-        let mut out = Vec::new();
+        let budget = ctx
+            .execution_budget()
+            .map(|budget| {
+                budget
+                    .try_with(Clone::clone)
+                    .map_err(|_| crate::Error::Cancelled)
+            })
+            .transpose()?;
+        let mut out = FmtOutput::new(budget)?;
         let mut stderr = String::new();
         let mut exit_code = 0;
         for file in files {
@@ -274,9 +313,10 @@ impl Builtin for Fmt {
                 }
             };
             ctx.consume_budget_work(1 + (input.len() / 1024) as u64)?;
-            format_input(&opts, input, &mut out);
+            format_input(&opts, input, &mut out)?;
         }
-        let mut result = ExecResult::with_code(crate::StreamData::from(out), exit_code);
+        let (bytes, _output_lease) = out.bytes.into_parts();
+        let mut result = ExecResult::with_code(crate::StreamData::from(bytes), exit_code);
         result.stderr = stderr.into();
         Ok(result)
     }
@@ -291,9 +331,61 @@ mod tests {
         let Ok((opts, _)) = parse_args(&args) else {
             panic!("valid args");
         };
-        let mut out = Vec::new();
-        format_input(&opts, input.as_bytes().to_vec(), &mut out);
-        String::from_utf8(out).unwrap()
+        let mut out = FmtOutput::new(None).unwrap();
+        format_input(&opts, input.as_bytes().to_vec(), &mut out).unwrap();
+        String::from_utf8(out.bytes.into_parts().0).unwrap()
+    }
+
+    #[test]
+    fn output_capacity_boundary_is_checked_before_copy() {
+        use crate::{ExecutionBudgetExceeded, ExecutionLimits, LimitExceeded};
+        let token = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
+        let budget = ExecutionBudget::new(
+            &ExecutionLimits::new().max_live_intermediate_bytes(8),
+            token,
+        );
+        let mut out = FmtOutput::new(Some(budget)).unwrap();
+        out.write_all(b"12345678").unwrap();
+        assert!(matches!(
+            out.write_all(b"9"),
+            Err(crate::Error::ResourceLimit(LimitExceeded::ExecutionBudget(
+                ExecutionBudgetExceeded::LiveBytes { used: 16, limit: 8 }
+            )))
+        ));
+        assert_eq!(&*out.bytes, b"12345678");
+    }
+
+    #[test]
+    fn synchronous_sink_observes_cancellation_and_deadline() {
+        use crate::{ExecutionLimits, LimitExceeded};
+        use std::sync::{
+            Arc,
+            atomic::{AtomicBool, Ordering},
+        };
+        for cancel in [true, false] {
+            let token = Arc::new(AtomicBool::new(false));
+            let budget = ExecutionBudget::new(
+                &ExecutionLimits::new().timeout(std::time::Duration::from_millis(1)),
+                token.clone(),
+            );
+            let mut out = FmtOutput::new(Some(budget)).unwrap();
+            out.write_all(b"a").unwrap();
+            if cancel {
+                token.store(true, Ordering::Relaxed);
+            } else {
+                std::thread::sleep(std::time::Duration::from_millis(120));
+            }
+            let error = out.write_all(b"b").unwrap_err();
+            if cancel {
+                assert!(matches!(error, crate::Error::Cancelled));
+            } else {
+                assert!(matches!(
+                    error,
+                    crate::Error::ResourceLimit(LimitExceeded::Timeout(_))
+                ));
+            }
+            assert_eq!(&*out.bytes, b"a");
+        }
     }
 
     #[test]