From 964054d631e3449f621b332ff81d9f1be19f44c7 Mon Sep 17 00:00:00 2001
From: nhtyy <n@push0.dev>
Date: Thu, 21 Nov 2024 12:22:05 -0800
Subject: [PATCH] fix: uint256 precompile invaraint checks

---
 .../src/syscalls/precompiles/keccak256/permute.rs |  3 ++-
 .../src/syscalls/precompiles/sha256/compress.rs   |  5 ++++-
 .../src/syscalls/precompiles/sha256/extend.rs     |  6 +++++-
 .../executor/src/syscalls/precompiles/uint256.rs  | 15 ++++++++++-----
 4 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/crates/core/executor/src/syscalls/precompiles/keccak256/permute.rs b/crates/core/executor/src/syscalls/precompiles/keccak256/permute.rs
index be3743e1f0..312f651fd5 100644
--- a/crates/core/executor/src/syscalls/precompiles/keccak256/permute.rs
+++ b/crates/core/executor/src/syscalls/precompiles/keccak256/permute.rs
@@ -27,7 +27,8 @@ impl Syscall for Keccak256PermuteSyscall {
         let start_clk = rt.clk;
         let state_ptr = arg1;
         if arg2 != 0 {
-            panic!("Expected arg2 to be 0, got {arg2}");
+            eprintln!("Expected arg2 to be 0, got {arg2}, this violates the Keccak precompile invariant.");
+            return rt.invariant_violated();
         }
 
         let mut state_read_records = Vec::new();
diff --git a/crates/core/executor/src/syscalls/precompiles/sha256/compress.rs b/crates/core/executor/src/syscalls/precompiles/sha256/compress.rs
index db1e9c42ec..46e8b130e7 100644
--- a/crates/core/executor/src/syscalls/precompiles/sha256/compress.rs
+++ b/crates/core/executor/src/syscalls/precompiles/sha256/compress.rs
@@ -32,7 +32,10 @@ impl Syscall for Sha256CompressSyscall {
     ) -> Option<u32> {
         let w_ptr = arg1;
         let h_ptr = arg2;
-        assert_ne!(w_ptr, h_ptr);
+        if w_ptr == h_ptr {
+            eprintln!("w_ptr == h_ptr, violation of the sha256 invariant");
+            return rt.invariant_violated();
+        }
 
         let start_clk = rt.clk;
         let mut h_read_records = Vec::new();
diff --git a/crates/core/executor/src/syscalls/precompiles/sha256/extend.rs b/crates/core/executor/src/syscalls/precompiles/sha256/extend.rs
index cc5e288fb6..4b8499a7af 100644
--- a/crates/core/executor/src/syscalls/precompiles/sha256/extend.rs
+++ b/crates/core/executor/src/syscalls/precompiles/sha256/extend.rs
@@ -19,7 +19,11 @@ impl Syscall for Sha256ExtendSyscall {
     ) -> Option<u32> {
         let clk_init = rt.clk;
         let w_ptr = arg1;
-        assert!(arg2 == 0, "arg2 must be 0");
+        if arg2 != 0 {
+            eprintln!("Warning: sha256_extend syscall arg2 is not zero, this violates the precompile invariants");
+            
+            return rt.invariant_violated();
+        }
 
         let w_ptr_init = w_ptr;
         let mut w_i_minus_15_reads = Vec::with_capacity(48);
diff --git a/crates/core/executor/src/syscalls/precompiles/uint256.rs b/crates/core/executor/src/syscalls/precompiles/uint256.rs
index 769ad0beb8..2731f5cbef 100644
--- a/crates/core/executor/src/syscalls/precompiles/uint256.rs
+++ b/crates/core/executor/src/syscalls/precompiles/uint256.rs
@@ -21,12 +21,11 @@ impl Syscall for Uint256MulSyscall {
         let clk = rt.clk;
 
         let x_ptr = arg1;
-        if x_ptr % 4 != 0 {
-            panic!();
-        }
         let y_ptr = arg2;
-        if y_ptr % 4 != 0 {
-            panic!();
+        
+        // Check alignment.
+        if x_ptr % 4 > 0 || y_ptr % 4 > 0 {
+            return rt.invariant_violated();
         }
 
         // First read the words for the x value. We can read a slice_unsafe here because we write
@@ -45,6 +44,12 @@ impl Syscall for Uint256MulSyscall {
         let uint256_y = BigUint::from_bytes_le(&words_to_bytes_le_vec(&y));
         let uint256_modulus = BigUint::from_bytes_le(&words_to_bytes_le_vec(&modulus));
 
+        if uint256_x >= uint256_modulus || uint256_y >= uint256_modulus {
+            eprintln!("Uint256 Precompile Invariant violated: x or y is greater than or equal to the modulus.");
+
+            return rt.invariant_violated();
+        }
+
         // Perform the multiplication and take the result modulo the modulus.
         let result: BigUint = if uint256_modulus.is_zero() {
             let modulus = BigUint::one() << 256;