Merge BlockstreamResearch#271: bit machine: fix some overflows with massive types

b2d83e2 test: add unit test which causes addition overflow in type size computation (Andrew Poelstra)
b1f3396 bit machine: check resource bounds on construction (Andrew Poelstra)
d21486a types: use saturating addition in bit width computation during type construction (Andrew Poelstra)

Pull request description:

  There probably are other things we can break with this test program, but these are the ones I encountered just now.

ACKs for top commit:
  uncomputable:
    ACK b2d83e2

Tree-SHA512: 5c2b978d082d47a808d3510e94ed7141b5f84d454cecd5ad204e60aea75e708df7516dfbc43b1c46faaab5b76857cdef8349866f5c26380bdf6b37134915762b

Author: uncomputable

src/bit_machine/limits.rs

@@ -0,0 +1,122 @@
+// SPDX-License-Identifier: CC0-1.0
+
+//! Bit Machine Resource Limits
+//!
+//! Implementation of the Bit Machine, without TCO, as TCO precludes some
+//! frame management optimizations which can be used to great benefit.
+//!
+
+use core::fmt;
+use std::error;
+
+/// The maximum number of cells needed by the bit machine.
+///
+/// This roughly corresponds to the maximum size of any type used by a
+/// program, in bits. In a blockchain context this should be limited by
+/// the transaction's weight budget.
+///
+/// The limit here is an absolute limit enforced by the library to avoid
+/// unbounded allocations.
+// This value must be less than usize::MAX / 2 to avoid panics in this module.
+const MAX_CELLS: usize = 2 * 1024 * 1024 * 1024 - 1;
+
+/// The maximum number of frames needed by the bit machine.
+///
+/// This roughly corresponds to the maximum depth of nested `comp` and
+/// `disconnect` combinators. In a blockchain context this should be
+/// limited by the transaction's weight budget.
+///
+/// The limit here is an absolute limit enforced by the library to avoid
+/// unbounded allocations.
+// This value must be less than usize::MAX / 2 to avoid panics in this module.
+const MAX_FRAMES: usize = 1024 * 1024;
+
+#[non_exhaustive]
+#[derive(Clone, Debug)]
+pub enum LimitError {
+    MaxCellsExceeded {
+        /// The number of cells needed by the program.
+        got: usize,
+        /// The maximum allowed number of cells.
+        max: usize,
+        /// A description of which cell count exceeded the limit.
+        bound: &'static str,
+    },
+    MaxFramesExceeded {
+        /// The number of frames needed by the program.
+        got: usize,
+        /// The maximum allowed number of frames.
+        max: usize,
+        /// A description of which frame count exceeded the limit.
+        bound: &'static str,
+    },
+}
+
+impl fmt::Display for LimitError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        let (limit, got, max, bound) = match self {
+            LimitError::MaxCellsExceeded { got, max, bound } => ("cells", got, max, bound),
+            LimitError::MaxFramesExceeded { got, max, bound } => ("frames", got, max, bound),
+        };
+        write!(
+            f,
+            "maximum number of {} exceeded (needed {}, maximum {}) (bound: {})",
+            limit, got, max, bound,
+        )
+    }
+}
+impl error::Error for LimitError {}
+
+impl LimitError {
+    fn check_max_cells(got: usize, bound: &'static str) -> Result<(), Self> {
+        if got > MAX_CELLS {
+            Err(Self::MaxCellsExceeded {
+                got,
+                max: MAX_CELLS,
+                bound,
+            })
+        } else {
+            Ok(())
+        }
+    }
+
+    fn check_max_frames(got: usize, bound: &'static str) -> Result<(), Self> {
+        if got > MAX_FRAMES {
+            Err(Self::MaxFramesExceeded {
+                got,
+                max: MAX_CELLS,
+                bound,
+            })
+        } else {
+            Ok(())
+        }
+    }
+
+    /// Helper function to check every value and sum for being within bounds.
+    pub(super) fn check_program<J: crate::jet::Jet>(
+        program: &crate::RedeemNode<J>,
+    ) -> Result<(), Self> {
+        let source_ty_width = program.arrow().source.bit_width();
+        let target_ty_width = program.arrow().target.bit_width();
+        let bounds = program.bounds();
+
+        Self::check_max_cells(source_ty_width, "source type width")?;
+        Self::check_max_cells(target_ty_width, "target type width")?;
+        Self::check_max_cells(bounds.extra_cells, "extra cells")?;
+        Self::check_max_cells(
+            source_ty_width + target_ty_width,
+            "source + target type widths",
+        )?;
+        Self::check_max_cells(
+            source_ty_width + target_ty_width + bounds.extra_cells,
+            "source + target type widths + extra cells",
+        )?;
+
+        Self::check_max_frames(bounds.extra_frames, "extra frames")?;
+        Self::check_max_frames(
+            bounds.extra_frames + crate::analysis::IO_EXTRA_FRAMES,
+            "extra frames + fixed overhead",
+        )?;
+        Ok(())
+    }
+}

src/bit_machine/mod.rs

@@ -1,12 +1,13 @@
 // SPDX-License-Identifier: CC0-1.0
 
-//! # Simplicity Execution
+//! Simplicity Execution
 //!
 //! Implementation of the Bit Machine, without TCO, as TCO precludes some
 //! frame management optimizations which can be used to great benefit.
 //!
 
 mod frame;
+mod limits;
 
 use std::collections::HashSet;
 use std::error;
@@ -20,6 +21,8 @@ use crate::{analysis, Imr};
 use crate::{Cmr, FailEntropy, Value};
 use frame::Frame;
 
+pub use self::limits::LimitError;
+
 /// An execution context for a Simplicity program
 pub struct BitMachine {
     /// Space for bytes that read and write frames point to.
@@ -37,16 +40,17 @@ pub struct BitMachine {
 
 impl BitMachine {
     /// Construct a Bit Machine with enough space to execute the given program.
-    pub fn for_program<J: Jet>(program: &RedeemNode<J>) -> Self {
+    pub fn for_program<J: Jet>(program: &RedeemNode<J>) -> Result<Self, LimitError> {
+        LimitError::check_program(program)?;
         let io_width = program.arrow().source.bit_width() + program.arrow().target.bit_width();
 
-        Self {
+        Ok(Self {
             data: vec![0; (io_width + program.bounds().extra_cells + 7) / 8],
             next_frame_start: 0,
             read: Vec::with_capacity(program.bounds().extra_frames + analysis::IO_EXTRA_FRAMES),
             write: Vec::with_capacity(program.bounds().extra_frames + analysis::IO_EXTRA_FRAMES),
             source_ty: program.arrow().source.clone(),
-        }
+        })
     }
 
     #[cfg(test)]
@@ -61,7 +65,7 @@ impl BitMachine {
             .expect("finalizing types")
             .finalize(&mut SimpleFinalizer::new(None.into_iter()))
             .expect("finalizing");
-        let mut mac = BitMachine::for_program(&prog);
+        let mut mac = BitMachine::for_program(&prog).expect("program has reasonable bounds");
         mac.exec(&prog, env)
     }
 
@@ -569,6 +573,8 @@ pub enum ExecutionError {
     ReachedFailNode(FailEntropy),
     /// Reached a pruned branch
     ReachedPrunedBranch(Cmr),
+    /// Exceeded some program limit
+    LimitExceeded(LimitError),
     /// Jet failed during execution
     JetFailed(JetFailed),
 }
@@ -585,12 +591,29 @@ impl fmt::Display for ExecutionError {
             ExecutionError::ReachedPrunedBranch(hash) => {
                 write!(f, "Execution reached a pruned branch: {}", hash)
             }
-            ExecutionError::JetFailed(jet_failed) => fmt::Display::fmt(jet_failed, f),
+            ExecutionError::LimitExceeded(e) => e.fmt(f),
+            ExecutionError::JetFailed(jet_failed) => jet_failed.fmt(f),
         }
     }
 }
 
-impl error::Error for ExecutionError {}
+impl error::Error for ExecutionError {
+    fn source(&self) -> Option<&(dyn error::Error + 'static)> {
+        match self {
+            Self::InputWrongType(..)
+            | Self::ReachedFailNode(..)
+            | Self::ReachedPrunedBranch(..) => None,
+            Self::LimitExceeded(ref e) => Some(e),
+            Self::JetFailed(ref e) => Some(e),
+        }
+    }
+}
+
+impl From<LimitError> for ExecutionError {
+    fn from(e: LimitError) -> Self {
+        ExecutionError::LimitExceeded(e)
+    }
+}
 
 impl From<JetFailed> for ExecutionError {
     fn from(jet_failed: JetFailed) -> Self {
@@ -600,7 +623,6 @@ impl From<JetFailed> for ExecutionError {
 
 #[cfg(test)]
 mod tests {
-    #[cfg(feature = "elements")]
     use super::*;
 
     #[cfg(feature = "elements")]
@@ -656,7 +678,9 @@ mod tests {
 
         // Try to run it on the bit machine and return the result
         let env = ElementsEnv::dummy();
-        BitMachine::for_program(&prog).exec(&prog, &env)
+        BitMachine::for_program(&prog)
+            .expect("program has reasonable bounds")
+            .exec(&prog, &env)
     }
 
     #[test]
@@ -685,4 +709,20 @@ mod tests {
         );
         assert_eq!(res.unwrap(), Value::unit());
     }
+
+    #[test]
+    fn crash_regression2() {
+        use crate::node::{CoreConstructible as _, JetConstructible as _};
+
+        type Node = Arc<crate::ConstructNode<crate::jet::Core>>;
+
+        let mut bomb = Node::jet(
+            &crate::types::Context::new(),
+            crate::jet::Core::Ch8, // arbitrary jet with nonzero output size
+        );
+        for _ in 0..100 {
+            bomb = Node::pair(&bomb, &bomb).unwrap();
+        }
+        let _ = bomb.finalize_pruned(&());
+    }
 }

src/human_encoding/mod.rs

@@ -238,7 +238,7 @@ mod tests {
             .expect("Forest is missing expected root")
             .finalize_pruned(env)
             .expect("Failed to finalize");
-        let mut mac = BitMachine::for_program(&program);
+        let mut mac = BitMachine::for_program(&program).expect("program has reasonable bounds");
         mac.exec(&program, env).expect("Failed to run program");
     }
 
@@ -260,7 +260,7 @@ mod tests {
                 return;
             }
         };
-        let mut mac = BitMachine::for_program(&program);
+        let mut mac = BitMachine::for_program(&program).expect("program has reasonable bounds");
         match mac.exec(&program, env) {
             Ok(_) => panic!("Execution is expected to fail"),
             Err(error) => assert_eq!(&error.to_string(), err_msg),

src/human_encoding/parse/mod.rs

@@ -595,7 +595,8 @@ mod tests {
                     .finalize_unpruned()
                     .expect("finalize");
 
-                let mut mac = BitMachine::for_program(&program);
+                let mut mac =
+                    BitMachine::for_program(&program).expect("program has reasonable bounds");
                 mac.exec(&program, env).expect("execute");
             }
             Err(errs) => {

src/node/commit.rs

@@ -544,7 +544,8 @@ mod tests {
             assert_eq!(counter, 0..98);
 
             // Execute the program to confirm that it worked
-            let mut mac = BitMachine::for_program(&diff1_final);
+            let mut mac =
+                BitMachine::for_program(&diff1_final).expect("program has reasonable bounds");
             mac.exec(&diff1_final, &()).unwrap();
         }
     }

src/node/redeem.rs

@@ -418,7 +418,7 @@ impl<J: Jet> RedeemNode<J> {
 
         // 1) Run the Bit Machine and mark (un)used branches.
         // This is the only fallible step in the pruning process.
-        let mut mac = BitMachine::for_program(self);
+        let mut mac = BitMachine::for_program(self)?;
         let tracker = mac.exec_prune(self, env)?;
 
         // 2) Prune out unused case branches.
@@ -920,7 +920,8 @@ mod tests {
             .finalize_unpruned()
             .expect("expected pruned program should finalize");
 
-        let mut mac = BitMachine::for_program(&unpruned_program);
+        let mut mac = BitMachine::for_program(&unpruned_program)
+            .expect("unpruned program has reasonable bounds");
         let unpruned_output = mac
             .exec(&unpruned_program, env)
             .expect("unpruned program should run without failure");
@@ -934,7 +935,8 @@ mod tests {
             "pruning result differs from expected result"
         );
 
-        let mut mac = BitMachine::for_program(&pruned_program);
+        let mut mac =
+            BitMachine::for_program(&pruned_program).expect("pruned program has reasonable bounds");
         let pruned_output = mac
             .exec(&pruned_program, env)
             .expect("pruned program should run without failure");

src/policy/satisfy.rs

@@ -351,15 +351,15 @@ mod tests {
         program: Arc<RedeemNode<Elements>>,
         env: &ElementsEnv<Arc<elements::Transaction>>,
     ) {
-        let mut mac = BitMachine::for_program(&program);
+        let mut mac = BitMachine::for_program(&program).unwrap();
         assert!(mac.exec(&program, env).is_ok());
     }
 
     fn execute_unsuccessful(
         program: Arc<RedeemNode<Elements>>,
         env: &ElementsEnv<Arc<elements::Transaction>>,
     ) {
-        let mut mac = BitMachine::for_program(&program);
+        let mut mac = BitMachine::for_program(&program).unwrap();
         assert!(mac.exec(&program, env).is_err());
     }
 

src/policy/serialize.rs

@@ -288,7 +288,10 @@ mod tests {
         let finalized = commit
             .finalize(&mut SimpleFinalizer::new(witness.into_iter()))
             .expect("finalize");
-        let mut mac = BitMachine::for_program(&finalized);
+        let mut mac = match BitMachine::for_program(&finalized) {
+            Ok(mac) => mac,
+            Err(_) => return false,
+        };
 
         match mac.exec(&finalized, env) {
             Ok(output) => output == Value::unit(),

src/types/final_data.rs

@@ -194,19 +194,24 @@ impl Final {
 
     /// Create the sum of the given `left` and `right` types.
     pub fn sum(left: Arc<Self>, right: Arc<Self>) -> Arc<Self> {
+        // Use saturating_add for bitwidths. If the user has overflowed usize, even on a 32-bit
+        // system this means that they have a 4-gigabit type and their program should be rejected
+        // by a sanity check somewhere. However, if we panic here, the user cannot finalize their
+        // program and cannot even tell that this resource limit has been hit before panicking.
         Arc::new(Final {
             tmr: Tmr::sum(left.tmr, right.tmr),
-            bit_width: 1 + cmp::max(left.bit_width, right.bit_width),
+            bit_width: cmp::max(left.bit_width, right.bit_width).saturating_add(1),
             has_padding: left.has_padding || right.has_padding || left.bit_width != right.bit_width,
             bound: CompleteBound::Sum(left, right),
         })
     }
 
     /// Create the product of the given `left` and `right` types.
     pub fn product(left: Arc<Self>, right: Arc<Self>) -> Arc<Self> {
+        // See comment in `sum` about use of saturating add.
         Arc::new(Final {
             tmr: Tmr::product(left.tmr, right.tmr),
-            bit_width: left.bit_width + right.bit_width,
+            bit_width: left.bit_width.saturating_add(right.bit_width),
             has_padding: left.has_padding || right.has_padding,
             bound: CompleteBound::Product(left, right),
         })