Merge pull request #2384 from AleoHQ/limit_program_size

Limit program size and add unit test

console/network/src/lib.rs

@@ -144,10 +144,19 @@ pub trait Network:
     /// The maximum number of entries in a record.
     const MAX_RECORD_ENTRIES: usize = Self::MIN_RECORD_ENTRIES.saturating_add(Self::MAX_DATA_ENTRIES);
 
+    /// The maximum program size by number of characters.
+    const MAX_PROGRAM_SIZE: usize = 100_000; // 100 KB
+
     /// The maximum number of mappings in a program.
     const MAX_MAPPINGS: usize = 31;
     /// The maximum number of functions in a program.
     const MAX_FUNCTIONS: usize = 31;
+    /// The maximum number of structs in a program.
+    const MAX_STRUCTS: usize = 10 * Self::MAX_FUNCTIONS;
+    /// The maximum number of records in a program.
+    const MAX_RECORDS: usize = 10 * Self::MAX_FUNCTIONS;
+    /// The maximum number of closures in a program.
+    const MAX_CLOSURES: usize = 2 * Self::MAX_FUNCTIONS;
     /// The maximum number of operands in an instruction.
     const MAX_OPERANDS: usize = Self::MAX_INPUTS;
     /// The maximum number of instructions in a closure or function.

synthesizer/process/src/tests/test_execute.rs

@@ -2616,10 +2616,14 @@ fn test_max_imports() {
 #[test]
 fn test_program_exceeding_transaction_spend_limit() {
     // Construct a finalize body whose finalize cost is excessively large.
-    let finalize_body = (0..<CurrentNetwork as Network>::MAX_COMMANDS)
-        .map(|i| format!("hash.bhp256 0field into r{i} as field;"))
-        .collect::<Vec<_>>()
-        .join("\n");
+    let mut finalize_body = r"
+    cast  0u8 0u8 0u8 0u8 0u8 0u8 0u8 0u8 0u8 0u8 0u8 0u8 0u8 0u8 0u8 0u8 into r0 as [u8; 16u32];
+    cast  r0 r0 r0 r0 r0 r0 r0 r0 r0 r0 r0 r0 r0 r0 r0 r0 into r1 as [[u8; 16u32]; 16u32];
+    cast  r1 r1 r1 r1 r1 r1 r1 r1 r1 r1 r1 r1 r1 r1 r1 r1 into r2 as [[[u8; 16u32]; 16u32]; 16u32];"
+        .to_string();
+    (3..500).for_each(|i| {
+        finalize_body.push_str(&format!("hash.bhp256 r2 into r{i} as field;\n"));
+    });
     // Construct the program.
     let program = Program::from_str(&format!(
         r"program test_max_spend_limit.aleo;

synthesizer/program/src/lib.rs

@@ -362,6 +362,9 @@ impl<N: Network, Instruction: InstructionTrait<N>, Command: CommandTrait<N>> Pro
         // Retrieve the struct name.
         let struct_name = *struct_.name();
 
+        // Ensure the program has not exceeded the maximum number of structs.
+        ensure!(self.structs.len() < N::MAX_STRUCTS, "Program exceeds the maximum number of structs.");
+
         // Ensure the struct name is new.
         ensure!(self.is_unique_name(&struct_name), "'{struct_name}' is already in use.");
         // Ensure the struct name is not a reserved opcode.
@@ -420,6 +423,9 @@ impl<N: Network, Instruction: InstructionTrait<N>, Command: CommandTrait<N>> Pro
         // Retrieve the record name.
         let record_name = *record.name();
 
+        // Ensure the program has not exceeded the maximum number of records.
+        ensure!(self.records.len() < N::MAX_RECORDS, "Program exceeds the maximum number of records.");
+
         // Ensure the record name is new.
         ensure!(self.is_unique_name(&record_name), "'{record_name}' is already in use.");
         // Ensure the record name is not a reserved opcode.
@@ -480,6 +486,9 @@ impl<N: Network, Instruction: InstructionTrait<N>, Command: CommandTrait<N>> Pro
         // Retrieve the closure name.
         let closure_name = *closure.name();
 
+        // Ensure the program has not exceeded the maximum number of closures.
+        ensure!(self.closures.len() < N::MAX_CLOSURES, "Program exceeds the maximum number of closures.");
+
         // Ensure the closure name is new.
         ensure!(self.is_unique_name(&closure_name), "'{closure_name}' is already in use.");
         // Ensure the closure name is not a reserved opcode.

synthesizer/program/src/parse.rs

@@ -106,6 +106,9 @@ impl<N: Network, Instruction: InstructionTrait<N>, Command: CommandTrait<N>> Fro
 
     /// Returns a program from a string literal.
     fn from_str(string: &str) -> Result<Self> {
+        // Ensure the raw program string is less than MAX_PROGRAM_SIZE.
+        ensure!(string.len() <= N::MAX_PROGRAM_SIZE, "Program length exceeds N::MAX_PROGRAM_SIZE.");
+
         match Self::parse(string) {
             Ok((remainder, object)) => {
                 // Ensure the remainder is empty.
@@ -255,4 +258,127 @@ function compute:
 
         Ok(())
     }
+
+    #[test]
+    fn test_program_size() {
+        // Define variable name for easy experimentation with program sizes.
+        let var_name = "a";
+
+        // Helper function to generate imports.
+        let gen_import_string = |n: usize| -> String {
+            let mut s = String::new();
+            for i in 0..n {
+                s.push_str(&format!("import foo{i}.aleo;\n"));
+            }
+            s
+        };
+
+        // Helper function to generate large structs.
+        let gen_struct_string = |n: usize| -> String {
+            let mut s = String::with_capacity(CurrentNetwork::MAX_PROGRAM_SIZE);
+            for i in 0..n {
+                s.push_str(&format!("struct m{}:\n", i));
+                for j in 0..10 {
+                    s.push_str(&format!("    {}{} as u128;\n", var_name, j));
+                }
+            }
+            s
+        };
+
+        // Helper function to generate large records.
+        let gen_record_string = |n: usize| -> String {
+            let mut s = String::with_capacity(CurrentNetwork::MAX_PROGRAM_SIZE);
+            for i in 0..n {
+                s.push_str(&format!("record r{}:\n    owner as address.private;\n", i));
+                for j in 0..10 {
+                    s.push_str(&format!("    {}{} as u128.private;\n", var_name, j));
+                }
+            }
+            s
+        };
+
+        // Helper function to generate large mappings.
+        let gen_mapping_string = |n: usize| -> String {
+            let mut s = String::with_capacity(CurrentNetwork::MAX_PROGRAM_SIZE);
+            for i in 0..n {
+                s.push_str(&format!(
+                    "mapping {}{}:\n    key as field.public;\n    value as field.public;\n",
+                    var_name, i
+                ));
+            }
+            s
+        };
+
+        // Helper function to generate large closures.
+        let gen_closure_string = |n: usize| -> String {
+            let mut s = String::with_capacity(CurrentNetwork::MAX_PROGRAM_SIZE);
+            for i in 0..n {
+                s.push_str(&format!("closure c{}:\n    input r0 as u128;\n", i));
+                for j in 0..10 {
+                    s.push_str(&format!("    add r0 r0 into r{};\n", j));
+                }
+                s.push_str(&format!("    output r{} as u128;\n", 4000));
+            }
+            s
+        };
+
+        // Helper function to generate large functions.
+        let gen_function_string = |n: usize| -> String {
+            let mut s = String::with_capacity(CurrentNetwork::MAX_PROGRAM_SIZE);
+            for i in 0..n {
+                s.push_str(&format!("function f{}:\n    add 1u128 1u128 into r0;\n", i));
+                for j in 0..10 {
+                    s.push_str(&format!("    add r0 r0 into r{j};\n"));
+                }
+            }
+            s
+        };
+
+        // Helper function to generate and parse a program.
+        let test_parse = |imports: &str, body: &str, should_succeed: bool| {
+            let program = format!("{imports}\nprogram to_parse.aleo;\n\n{body}");
+            let result = Program::<CurrentNetwork>::from_str(&program);
+            if result.is_ok() != should_succeed {
+                println!("Program failed to parse: {program}");
+            }
+            assert_eq!(result.is_ok(), should_succeed);
+        };
+
+        // A program with MAX_IMPORTS should succeed.
+        test_parse(&gen_import_string(CurrentNetwork::MAX_IMPORTS), &gen_struct_string(1), true);
+        // A program with more than MAX_IMPORTS should fail.
+        test_parse(&gen_import_string(CurrentNetwork::MAX_IMPORTS + 1), &gen_struct_string(1), false);
+        // A program with MAX_STRUCTS should succeed.
+        test_parse("", &gen_struct_string(CurrentNetwork::MAX_STRUCTS), true);
+        // A program with more than MAX_STRUCTS should fail.
+        test_parse("", &gen_struct_string(CurrentNetwork::MAX_STRUCTS + 1), false);
+        // A program with MAX_RECORDS should succeed.
+        test_parse("", &gen_record_string(CurrentNetwork::MAX_RECORDS), true);
+        // A program with more than MAX_RECORDS should fail.
+        test_parse("", &gen_record_string(CurrentNetwork::MAX_RECORDS + 1), false);
+        // A program with MAX_MAPPINGS should succeed.
+        test_parse("", &gen_mapping_string(CurrentNetwork::MAX_MAPPINGS), true);
+        // A program with more than MAX_MAPPINGS should fail.
+        test_parse("", &gen_mapping_string(CurrentNetwork::MAX_MAPPINGS + 1), false);
+        // A program with MAX_CLOSURES should succeed.
+        test_parse("", &gen_closure_string(CurrentNetwork::MAX_CLOSURES), true);
+        // A program with more than MAX_CLOSURES should fail.
+        test_parse("", &gen_closure_string(CurrentNetwork::MAX_CLOSURES + 1), false);
+        // A program with MAX_FUNCTIONS should succeed.
+        test_parse("", &gen_function_string(CurrentNetwork::MAX_FUNCTIONS), true);
+        // A program with more than MAX_FUNCTIONS should fail.
+        test_parse("", &gen_function_string(CurrentNetwork::MAX_FUNCTIONS + 1), false);
+
+        // Initialize a program which is too big.
+        let program_too_big = format!(
+            "{} {} {} {} {}",
+            gen_struct_string(CurrentNetwork::MAX_STRUCTS),
+            gen_record_string(CurrentNetwork::MAX_RECORDS),
+            gen_mapping_string(CurrentNetwork::MAX_MAPPINGS),
+            gen_closure_string(CurrentNetwork::MAX_CLOSURES),
+            gen_function_string(CurrentNetwork::MAX_FUNCTIONS)
+        );
+        // A program which is too big should fail.
+        test_parse("", &program_too_big, false);
+    }
 }