From fe9d1f4198aea00b9e30ed8a5b1911cc3a26bdfe Mon Sep 17 00:00:00 2001
From: Arni Hod <arnon@starkware.co>
Date: Wed, 26 Jun 2024 16:37:33 +0300
Subject: [PATCH] feat: declare post compilation prime validation

---
 Cargo.lock                        |  3 +++
 Cargo.toml                        |  2 ++
 crates/gateway/Cargo.toml         |  6 +++++-
 crates/gateway/src/compilation.rs | 12 ++++++++++++
 crates/gateway/src/errors.rs      |  3 +++
 5 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/Cargo.lock b/Cargo.lock
index 16f588b39..2edf1b796 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -5231,9 +5231,12 @@ dependencies = [
  "async-trait",
  "axum",
  "blockifier 0.7.0-dev.1 (git+https://github.com/starkware-libs/blockifier.git?branch=main-mempool)",
+ "cairo-felt",
  "cairo-lang-starknet-classes",
  "cairo-vm",
  "hyper",
+ "num-bigint",
+ "num-traits 0.2.19",
  "papyrus_config",
  "papyrus_rpc",
  "pretty_assertions",
diff --git a/Cargo.toml b/Cargo.toml
index 861562759..e73419464 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -38,6 +38,7 @@ axum = "0.6.12"
 # starknet-api version. This should be removed once we have a mono-repo.
 blockifier = { git = "https://github.com/starkware-libs/blockifier.git", branch = "main-mempool" }
 bincode = "1.3.3"
+cairo-felt = "0.9.1"
 cairo-lang-sierra = "2.6.0"
 cairo-lang-starknet-classes = "2.6.0"
 cairo-lang-utils = "2.6.0"
@@ -51,6 +52,7 @@ hyper = { version = "0.14", features = ["client", "http1", "http2"] }
 indexmap = "2.1.0"
 itertools = "0.13.0"
 num-bigint = { version = "0.4.5", default-features = false }
+num-traits = "0.2"
 # TODO(YaelD, 28/5/2024): The special Papyrus version is needed in order to be aligned with the
 # starknet-api version. This should be removed once we have a mono-repo.
 papyrus_common = { git = "https://github.com/starkware-libs/papyrus.git", rev = "56c6fdc2" }
diff --git a/crates/gateway/Cargo.toml b/crates/gateway/Cargo.toml
index 1cb064ca8..d42258ebd 100644
--- a/crates/gateway/Cargo.toml
+++ b/crates/gateway/Cargo.toml
@@ -14,10 +14,14 @@ testing = []
 [dependencies]
 async-trait.workspace = true
 axum.workspace = true
-blockifier= { workspace = true , features = ["testing"] }
+assert_matches.workspace = true
+blockifier = { workspace = true , features = ["testing"] }
+cairo-felt.workspace = true
 cairo-lang-starknet-classes.workspace = true
 cairo-vm.workspace = true
 hyper.workspace = true
+num-bigint.workspace = true
+num-traits.workspace = true
 papyrus_config.workspace = true
 papyrus_rpc.workspace = true
 reqwest.workspace = true
diff --git a/crates/gateway/src/compilation.rs b/crates/gateway/src/compilation.rs
index 7e3f148cf..58103c258 100644
--- a/crates/gateway/src/compilation.rs
+++ b/crates/gateway/src/compilation.rs
@@ -3,9 +3,13 @@ use std::sync::OnceLock;
 
 use blockifier::execution::contract_class::{ClassInfo, ContractClass, ContractClassV1};
 use blockifier::execution::execution_utils::felt_to_stark_felt;
+// TODO(Arni): Consider if you want to import a new crate just for a constant string.
+use cairo_felt::PRIME_STR;
 use cairo_lang_starknet_classes::casm_contract_class::{
     CasmContractClass, CasmContractEntryPoints,
 };
+use num_bigint::BigUint;
+use num_traits::Num;
 use starknet_api::core::CompiledClassHash;
 use starknet_api::rpc_transaction::RPCDeclareTransaction;
 use starknet_api::transaction::Builtin;
@@ -105,5 +109,13 @@ fn validate_casm_class(contract_class: &CasmContractClass) -> Result<(), Gateway
             });
         }
     }
+
+    let prime = contract_class.prime.clone();
+    let expected_prime =
+        BigUint::from_str_radix(&PRIME_STR[2..], 16).expect("Error parsing field prime.");
+    if prime != expected_prime {
+        return Err(GatewayError::InvalidPrime { prime, expected_prime });
+    }
+
     Ok(())
 }
diff --git a/crates/gateway/src/errors.rs b/crates/gateway/src/errors.rs
index 24d5ef74b..f70abdf33 100644
--- a/crates/gateway/src/errors.rs
+++ b/crates/gateway/src/errors.rs
@@ -5,6 +5,7 @@ use blockifier::execution::errors::ContractClassError;
 use blockifier::state::errors::StateError;
 use blockifier::transaction::errors::TransactionExecutionError;
 use cairo_vm::types::errors::program_errors::ProgramError;
+use num_bigint::BigUint;
 use serde_json::{Error as SerdeError, Value};
 use starknet_api::block::{BlockNumber, GasPrice};
 use starknet_api::core::CompiledClassHash;
@@ -32,6 +33,8 @@ pub enum GatewayError {
     DeclaredContractProgramError(#[from] ProgramError),
     #[error("Internal server error: {0}")]
     InternalServerError(#[from] JoinError),
+    #[error("Invalid value for field prime: {prime}. Expected: {expected_prime}")]
+    InvalidPrime { prime: BigUint, expected_prime: BigUint },
     #[error("Error sending message: {0}")]
     MessageSendError(String),
     #[error(transparent)]