From 093d9d082fbb01a0738ddebd81f7600e0f2f7912 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     |  3 +++
 crates/gateway/src/errors.rs  |  3 +++
 crates/gateway/src/gateway.rs | 12 ++++++++++++
 5 files changed, 23 insertions(+)

diff --git a/Cargo.lock b/Cargo.lock
index 0491c040..4881ce20 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -5491,9 +5491,12 @@ dependencies = [
  "assert_matches",
  "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",
  "pretty_assertions",
  "reqwest",
diff --git a/Cargo.toml b/Cargo.toml
index 63e2a21d..c1101bac 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -40,6 +40,7 @@ blockifier = { git = "https://github.com/starkware-libs/blockifier.git", branch
     "testing",
 ] }
 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"
@@ -54,6 +55,7 @@ indexmap = "2.1.0"
 itertools = "0.13.0"
 lazy_static = "1.4.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 = "050e470f" }
diff --git a/crates/gateway/Cargo.toml b/crates/gateway/Cargo.toml
index c9514273..0cb38606 100644
--- a/crates/gateway/Cargo.toml
+++ b/crates/gateway/Cargo.toml
@@ -15,9 +15,12 @@ testing = []
 axum.workspace = true
 assert_matches.workspace = true
 blockifier.workspace = true
+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
 reqwest.workspace = true
 serde.workspace = true
diff --git a/crates/gateway/src/errors.rs b/crates/gateway/src/errors.rs
index d3effbbc..d720e008 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;
@@ -31,6 +32,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)]
diff --git a/crates/gateway/src/gateway.rs b/crates/gateway/src/gateway.rs
index b8df34a4..2e57b61e 100644
--- a/crates/gateway/src/gateway.rs
+++ b/crates/gateway/src/gateway.rs
@@ -8,9 +8,13 @@ use axum::routing::{get, post};
 use axum::{Json, Router};
 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, RPCTransaction};
 use starknet_api::transaction::TransactionHash;
@@ -216,5 +220,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(())
 }