From 58787d3ec1c53cbe97be405993b6e83d6c0fac48 Mon Sep 17 00:00:00 2001
From: nhtyy <n@push0.dev>
Date: Thu, 21 Nov 2024 12:15:23 -0800
Subject: [PATCH] fix: fp2 invariant checks

---
 .../executor/src/syscalls/precompiles/fptower/fp.rs  |  9 ++++-----
 .../src/syscalls/precompiles/fptower/fp2_addsub.rs   | 12 +++++++-----
 .../src/syscalls/precompiles/fptower/fp2_mul.rs      | 12 +++++++-----
 3 files changed, 18 insertions(+), 15 deletions(-)

diff --git a/crates/core/executor/src/syscalls/precompiles/fptower/fp.rs b/crates/core/executor/src/syscalls/precompiles/fptower/fp.rs
index 2087379076..a6336dc697 100644
--- a/crates/core/executor/src/syscalls/precompiles/fptower/fp.rs
+++ b/crates/core/executor/src/syscalls/precompiles/fptower/fp.rs
@@ -32,14 +32,13 @@ impl<P: FpOpField> Syscall for FpOpSyscall<P> {
     ) -> Option<u32> {
         let clk = rt.clk;
         let x_ptr = arg1;
-        if x_ptr % 4 != 0 {
-            panic!();
-        }
         let y_ptr = arg2;
-        if y_ptr % 4 != 0 {
-            panic!();
+        // Need to check alignment
+        if x_ptr % 4 > 0 || y_ptr % 4 > 0 {
+            return rt.invariant_violated();
         }
 
+
         let num_words = <P as NumWords>::WordsFieldElement::USIZE;
 
         let x = rt.slice_unsafe(x_ptr, num_words);
diff --git a/crates/core/executor/src/syscalls/precompiles/fptower/fp2_addsub.rs b/crates/core/executor/src/syscalls/precompiles/fptower/fp2_addsub.rs
index f583432310..c1f7f27015 100644
--- a/crates/core/executor/src/syscalls/precompiles/fptower/fp2_addsub.rs
+++ b/crates/core/executor/src/syscalls/precompiles/fptower/fp2_addsub.rs
@@ -32,12 +32,10 @@ impl<P: FpOpField> Syscall for Fp2AddSubSyscall<P> {
     ) -> Option<u32> {
         let clk = rt.clk;
         let x_ptr = arg1;
-        if x_ptr % 4 != 0 {
-            panic!();
-        }
         let y_ptr = arg2;
-        if y_ptr % 4 != 0 {
-            panic!();
+        // Need to check alignment
+        if x_ptr % 4 > 0 || y_ptr % 4 > 0 {
+            return rt.invariant_violated();
         }
 
         let num_words = <P as NumWords>::WordsCurvePoint::USIZE;
@@ -55,6 +53,10 @@ impl<P: FpOpField> Syscall for Fp2AddSubSyscall<P> {
         let bc1 = &BigUint::from_slice(bc1);
         let modulus = &BigUint::from_bytes_le(P::MODULUS);
 
+        if ac0 >= modulus || ac1 >= modulus || bc0 >= modulus || bc1 >= modulus {
+            return rt.invariant_violated();
+        }
+
         let (c0, c1) = match self.op {
             FieldOperation::Add => ((ac0 + bc0) % modulus, (ac1 + bc1) % modulus),
             FieldOperation::Sub => {
diff --git a/crates/core/executor/src/syscalls/precompiles/fptower/fp2_mul.rs b/crates/core/executor/src/syscalls/precompiles/fptower/fp2_mul.rs
index 5089c00726..bed5c241d1 100644
--- a/crates/core/executor/src/syscalls/precompiles/fptower/fp2_mul.rs
+++ b/crates/core/executor/src/syscalls/precompiles/fptower/fp2_mul.rs
@@ -32,12 +32,10 @@ impl<P: FpOpField> Syscall for Fp2MulSyscall<P> {
     ) -> Option<u32> {
         let clk = rt.clk;
         let x_ptr = arg1;
-        if x_ptr % 4 != 0 {
-            panic!();
-        }
         let y_ptr = arg2;
-        if y_ptr % 4 != 0 {
-            panic!();
+        // Need to check alignment
+        if x_ptr % 4 > 0 || y_ptr % 4 > 0 {
+            return rt.invariant_violated();
         }
 
         let num_words = <P as NumWords>::WordsCurvePoint::USIZE;
@@ -55,6 +53,10 @@ impl<P: FpOpField> Syscall for Fp2MulSyscall<P> {
         let bc1 = &BigUint::from_slice(bc1);
         let modulus = &BigUint::from_bytes_le(P::MODULUS);
 
+        if ac0 >= modulus || ac1 >= modulus || bc0 >= modulus || bc1 >= modulus {
+            return rt.invariant_violated();
+        }
+
         #[allow(clippy::match_bool)]
         let c0 = match (ac0 * bc0) % modulus < (ac1 * bc1) % modulus {
             true => ((modulus + (ac0 * bc0) % modulus) - (ac1 * bc1) % modulus) % modulus,