Merge pull request #2352 from AleoHQ/fix/import-depth

[HackerOne-2354265] Restrict number of imports, program depth, and call depth.

console/network/src/lib.rs

@@ -160,6 +160,11 @@ pub trait Network:
     /// The maximum number of outputs per transition.
     const MAX_OUTPUTS: usize = 16;
 
+    /// The maximum program depth.
+    const MAX_PROGRAM_DEPTH: usize = 64;
+    /// The maximum number of imports.
+    const MAX_IMPORTS: usize = 64;
+
     /// The state root type.
     type StateRoot: Bech32ID<Field<Self>>;
     /// The block hash type.

ledger/block/src/transaction/merkle.rs

@@ -16,7 +16,7 @@ use super::*;
 
 impl<N: Network> Transaction<N> {
     /// The maximum number of transitions allowed in a transaction.
-    const MAX_TRANSITIONS: usize = usize::pow(2, TRANSACTION_DEPTH as u32);
+    pub const MAX_TRANSITIONS: usize = usize::pow(2, TRANSACTION_DEPTH as u32);
 
     /// Returns the transaction root, by computing the root for a Merkle tree of the transition IDs.
     pub fn to_root(&self) -> Result<Field<N>> {

synthesizer/process/Cargo.toml

@@ -45,6 +45,11 @@ wasm = [
 ]
 timer = [ "aleo-std/timer" ]
 
+[[bench]]
+name = "stack_operations"
+path = "benches/stack_operations.rs"
+harness = false
+
 [dependencies.console]
 package = "snarkvm-console"
 path = "../../console"
@@ -119,6 +124,9 @@ features = [ "preserve_order" ]
 [dev-dependencies.bincode]
 version = "1.3"
 
+[dev-dependencies.criterion]
+version = "0.5"
+
 [dev-dependencies.ledger-committee]
 package = "snarkvm-ledger-committee"
 path = "../../ledger/committee"

synthesizer/process/benches/stack_operations.rs

@@ -0,0 +1,178 @@
+// Copyright (C) 2019-2023 Aleo Systems Inc.
+// This file is part of the snarkVM library.
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at:
+// http://www.apache.org/licenses/LICENSE-2.0
+
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#[macro_use]
+extern crate criterion;
+
+use console::{
+    network::MainnetV0,
+    program::{Identifier, ProgramID},
+    types::Field,
+};
+use snarkvm_synthesizer_process::{Process, Stack};
+use synthesizer_program::{Program, StackProgram};
+
+use circuit::prelude::bail;
+use console::{network::Network, prelude::SizeInDataBits};
+use criterion::{BatchSize, Criterion};
+use rand::{distributions::Alphanumeric, Rng};
+use std::str::FromStr;
+use utilities::TestRng;
+
+type CurrentNetwork = MainnetV0;
+
+fn bench_stack_new(c: &mut Criterion) {
+    // The depths to benchmark.
+    const DEPTHS: [usize; 6] = [1, 2, 4, 8, 16, 31];
+
+    // Initialize an RNG.
+    let mut rng = TestRng::default();
+
+    // Initialize a process.
+    let mut process = Process::load().unwrap();
+
+    // Benchmark the base case.
+    c.bench_function("Depth 0 | Stack::new", |b| {
+        b.iter_batched_ref(
+            || {
+                // Sample a random identifier.
+                let identifier = sample_identifier_as_string::<CurrentNetwork>(&mut rng).unwrap();
+                // Construct the program.
+                Program::from_str(&format!("program {identifier}.aleo; function foo:")).unwrap()
+            },
+            |program| Stack::<CurrentNetwork>::new(&process, program),
+            BatchSize::PerIteration,
+        )
+    });
+
+    // Add the 0th program to the process.
+    add_program_at_depth(&mut process, 0);
+
+    // Track the depth.
+    let mut depth = 1;
+
+    for i in DEPTHS {
+        // Add programs up to the current depth.
+        while depth < i {
+            // Add the program to the process.
+            add_program_at_depth(&mut process, depth);
+            // Increment the depth.
+            depth += 1;
+        }
+
+        // Benchmark at each depth.
+        c.bench_function(&format!("Depth {i} | Stack::new"), |b| {
+            b.iter_batched_ref(
+                || {
+                    // Sample a random identifier.
+                    let identifier = sample_identifier_as_string::<CurrentNetwork>(&mut rng).unwrap();
+                    // Construct the program.
+                    Program::from_str(&format!(
+                        "program {identifier}.aleo; function foo: call test_{i}.aleo/foo;",
+                        identifier = identifier,
+                        i = i - 1
+                    ))
+                    .unwrap()
+                },
+                |program| Stack::<CurrentNetwork>::new(&process, program),
+                BatchSize::PerIteration,
+            )
+        });
+    }
+}
+
+fn bench_stack_get_number_of_calls(c: &mut Criterion) {
+    // The depths to benchmark.
+    const DEPTHS: [usize; 6] = [1, 2, 4, 8, 16, 30];
+
+    // Initialize a process.
+    let mut process = Process::load().unwrap();
+
+    // Add the 0th program to the process.
+    add_program_at_depth(&mut process, 0);
+
+    // Benchmark the `get_number_of_calls` method for the base case.
+    c.bench_function("Depth 0 | Stack::get_number_of_calls", |b| {
+        b.iter(|| {
+            // Get the `Stack` for the 0th program.
+            let stack = process.get_stack(ProgramID::from_str("test_0.aleo").unwrap()).unwrap();
+            // Benchmark the `get_number_of_calls` method.
+            stack.get_number_of_calls(&Identifier::from_str("foo").unwrap())
+        })
+    });
+
+    // Track the depth.
+    let mut depth = 1;
+
+    for i in DEPTHS {
+        // Add programs up to the current depth.
+        while depth <= i {
+            // Add the program to the process.
+            add_program_at_depth(&mut process, depth);
+            // Increment the depth.
+            depth += 1;
+        }
+
+        // Get the `Stack` for the current test program.
+        let stack = process.get_stack(ProgramID::from_str(&format!("test_{}.aleo", i)).unwrap()).unwrap();
+
+        // Benchmark the `get_number_of_calls` method.
+        c.bench_function(&format!("Depth {i} | Stack::get_number_of_calls"), |b| {
+            b.iter(|| stack.get_number_of_calls(&Identifier::from_str("foo").unwrap()))
+        });
+    }
+}
+
+// Adds a program with a given call depth to the process.
+fn add_program_at_depth(process: &mut Process<CurrentNetwork>, depth: usize) {
+    // Construct the program.
+    let program = if depth == 0 {
+        Program::from_str(r"program test_0.aleo; function foo:").unwrap()
+    } else {
+        Program::from_str(&format!(
+            "import test_{import}.aleo; program test_{current}.aleo; function foo: call test_{import}.aleo/foo;",
+            import = depth - 1,
+            current = depth
+        ))
+        .unwrap()
+    };
+
+    // Add the program to the process.
+    process.add_program(&program).unwrap();
+}
+
+// Samples a random identifier as a string.
+fn sample_identifier_as_string<N: Network>(rng: &mut TestRng) -> console::prelude::Result<String> {
+    // Sample a random fixed-length alphanumeric string, that always starts with an alphabetic character.
+    let string = "a".to_string()
+        + &rng
+            .sample_iter(&Alphanumeric)
+            .take(Field::<N>::size_in_data_bits() / (8 * 2))
+            .map(char::from)
+            .collect::<String>();
+    // Ensure identifier fits within the data capacity of the base field.
+    let max_bytes = Field::<N>::size_in_data_bits() / 8; // Note: This intentionally rounds down.
+    match string.len() <= max_bytes {
+        // Return the identifier.
+        true => Ok(string.to_lowercase()),
+        false => bail!("Identifier exceeds the maximum capacity allowed"),
+    }
+}
+
+criterion_group! {
+    name = stack_operations;
+    config = Criterion::default().sample_size(10);
+    targets = bench_stack_new, bench_stack_get_number_of_calls
+}
+criterion_main!(stack_operations);

synthesizer/process/src/stack/helpers/initialize.rs

@@ -27,9 +27,11 @@ impl<N: Network> Stack<N> {
             universal_srs: process.universal_srs().clone(),
             proving_keys: Default::default(),
             verifying_keys: Default::default(),
+            number_of_calls: Default::default(),
+            program_depth: 0,
         };
 
-        // Add all of the imports into the stack.
+        // Add all the imports into the stack.
         for import in program.imports().keys() {
             // Ensure the program imports all exist in the process already.
             if !process.contains_program(import) {
@@ -39,17 +41,49 @@ impl<N: Network> Stack<N> {
             let external_stack = process.get_stack(import)?;
             // Add the external stack to the stack.
             stack.insert_external_stack(external_stack.clone())?;
+            // Update the program depth, checking that it does not exceed the maximum call depth.
+            stack.program_depth = std::cmp::max(stack.program_depth, external_stack.program_depth() + 1);
+            ensure!(
+                stack.program_depth <= N::MAX_PROGRAM_DEPTH,
+                "Program depth exceeds the maximum allowed call depth"
+            );
         }
         // Add the program closures to the stack.
         for closure in program.closures().values() {
             // Add the closure to the stack.
             stack.insert_closure(closure)?;
         }
+
         // Add the program functions to the stack.
         for function in program.functions().values() {
             // Add the function to the stack.
             stack.insert_function(function)?;
+            // Determine the number of calls for the function.
+            let mut num_calls = 1;
+            for instruction in function.instructions() {
+                if let Instruction::Call(call) = instruction {
+                    // Determine if this is a function call.
+                    if call.is_function_call(&stack)? {
+                        // Increment by the number of calls.
+                        num_calls += match call.operator() {
+                            CallOperator::Locator(locator) => stack
+                                .get_external_stack(locator.program_id())?
+                                .get_number_of_calls(locator.resource())?,
+                            CallOperator::Resource(resource) => stack.get_number_of_calls(resource)?,
+                        };
+                    }
+                }
+            }
+            // Check that the number of calls does not exceed the maximum.
+            // Note that one transition is reserved for the fee.
+            ensure!(
+                num_calls < ledger_block::Transaction::<N>::MAX_TRANSITIONS,
+                "Number of calls exceeds the maximum allowed number of transitions"
+            );
+            // Add the number of calls to the stack.
+            stack.number_of_calls.insert(*function.name(), num_calls);
         }
+
         // Return the stack.
         Ok(stack)
     }
@@ -109,7 +143,6 @@ impl<N: Network> Stack<N> {
             // Add the finalize name and finalize types to the stack.
             self.finalize_types.insert(*name, finalize_types);
         }
-
         // Return success.
         Ok(())
     }

synthesizer/process/src/stack/mod.rs

@@ -185,6 +185,10 @@ pub struct Stack<N: Network> {
     proving_keys: Arc<RwLock<IndexMap<Identifier<N>, ProvingKey<N>>>>,
     /// The mapping of function name to verifying key.
     verifying_keys: Arc<RwLock<IndexMap<Identifier<N>, VerifyingKey<N>>>>,
+    /// The mapping of function names to the number of calls.
+    number_of_calls: IndexMap<Identifier<N>, usize>,
+    /// The program depth.
+    program_depth: usize,
 }
 
 impl<N: Network> Stack<N> {
@@ -226,6 +230,12 @@ impl<N: Network> StackProgram<N> for Stack<N> {
         self.program.id()
     }
 
+    /// Returns the program depth.
+    #[inline]
+    fn program_depth(&self) -> usize {
+        self.program_depth
+    }
+
     /// Returns `true` if the stack contains the external record.
     #[inline]
     fn contains_external_record(&self, locator: &Locator<N>) -> bool {
@@ -279,23 +289,10 @@ impl<N: Network> StackProgram<N> for Stack<N> {
     /// Returns the expected number of calls for the given function name.
     #[inline]
     fn get_number_of_calls(&self, function_name: &Identifier<N>) -> Result<usize> {
-        // Determine the number of calls for this function (including the function itself).
-        let mut num_calls = 1;
-        for instruction in self.get_function(function_name)?.instructions() {
-            if let Instruction::Call(call) = instruction {
-                // Determine if this is a function call.
-                if call.is_function_call(self)? {
-                    // Increment by the number of calls.
-                    num_calls += match call.operator() {
-                        CallOperator::Locator(locator) => {
-                            self.get_external_stack(locator.program_id())?.get_number_of_calls(locator.resource())?
-                        }
-                        CallOperator::Resource(resource) => self.get_number_of_calls(resource)?,
-                    };
-                }
-            }
-        }
-        Ok(num_calls)
+        self.number_of_calls
+            .get(function_name)
+            .copied()
+            .ok_or_else(|| anyhow!("Function '{function_name}' does not exist"))
     }
 
     /// Returns a value for the given value type.