From b9ab62c62716b309c46b0cd33a96cca5ce2bc82e Mon Sep 17 00:00:00 2001
From: Esteve Soler Arderiu <esteve.soler@lambdaclass.com>
Date: Thu, 19 Dec 2024 12:38:31 -0300
Subject: [PATCH] Add version check for cached AOT contracts.

---
 runtime/src/lib.rs       | 15 +++++++++++++++
 src/executor/aot.rs      |  3 +--
 src/executor/contract.rs | 34 +++++++++++++++++++---------------
 3 files changed, 35 insertions(+), 17 deletions(-)

diff --git a/runtime/src/lib.rs b/runtime/src/lib.rs
index e292cfba6..c370dbbd5 100644
--- a/runtime/src/lib.rs
+++ b/runtime/src/lib.rs
@@ -23,6 +23,7 @@ use std::{
     mem::ManuallyDrop,
     os::fd::FromRawFd,
     ptr::{self, null, null_mut},
+    slice,
 };
 use std::{ops::Mul, vec::IntoIter};
 
@@ -35,6 +36,20 @@ lazy_static! {
         (DICT_SQUASH_UNIQUE_KEY_COST.cost() - DICT_SQUASH_REPEATED_ACCESS_COST.cost()) as u64;
 }
 
+#[no_mangle]
+#[allow(clippy::missing_safety_doc)]
+pub unsafe extern "C" fn cairo_native__get_version(target: *mut u8, length: usize) -> usize {
+    let target = slice::from_raw_parts_mut(target, length);
+
+    let version = env!("CARGO_PKG_VERSION");
+    assert!(length > version.len(), "version buffer not big enough");
+
+    target.copy_from_slice(version.as_bytes());
+    target[version.len()] = b'\0';
+
+    version.len()
+}
+
 /// Based on `cairo-lang-runner`'s implementation.
 ///
 /// Source: <https://github.com/starkware-libs/cairo/blob/main/crates/cairo-lang-runner/src/casm_run/mod.rs#L1946-L1948>
diff --git a/src/executor/aot.rs b/src/executor/aot.rs
index 84528862a..14a53280f 100644
--- a/src/executor/aot.rs
+++ b/src/executor/aot.rs
@@ -1,5 +1,3 @@
-use std::io;
-
 use crate::{
     error::Error,
     execution_result::{ContractExecutionResult, ExecutionResult},
@@ -20,6 +18,7 @@ use educe::Educe;
 use libc::c_void;
 use libloading::Library;
 use starknet_types_core::felt::Felt;
+use std::io;
 use tempfile::NamedTempFile;
 
 #[derive(Educe)]
diff --git a/src/executor/contract.rs b/src/executor/contract.rs
index 20f2e3424..11c770ae1 100644
--- a/src/executor/contract.rs
+++ b/src/executor/contract.rs
@@ -119,20 +119,7 @@ pub enum BuiltinType {
 
 impl BuiltinType {
     pub const fn size_in_bytes(&self) -> usize {
-        match self {
-            BuiltinType::Bitwise => 8,
-            BuiltinType::EcOp => 8,
-            BuiltinType::RangeCheck => 8,
-            BuiltinType::SegmentArena => 8,
-            BuiltinType::Poseidon => 8,
-            BuiltinType::Pedersen => 8,
-            BuiltinType::RangeCheck96 => 8,
-            BuiltinType::CircuitAdd => 8,
-            BuiltinType::CircuitMul => 8,
-            BuiltinType::Gas => 16,
-            BuiltinType::System => 8,
-            BuiltinType::BuiltinCosts => 8,
-        }
+        size_of::<u64>()
     }
 }
 
@@ -279,8 +266,25 @@ impl AotContractExecutor {
     pub fn load(library_path: &Path) -> Result<Self> {
         let info_str = std::fs::read_to_string(library_path.with_extension("json"))?;
         let contract_info: NativeContractInfo = serde_json::from_str(&info_str)?;
+
+        let library = Arc::new(unsafe { Library::new(library_path)? });
+        unsafe {
+            let get_version = library
+                .get::<extern "C" fn(*mut u8, usize) -> usize>(b"cairo_native__get_version")?;
+
+            let mut version_buffer = [0u8; 16];
+            let version_len = get_version(version_buffer.as_mut_ptr(), version_buffer.len());
+
+            let target_version = env!("CARGO_PKG_VERSION");
+            assert_eq!(
+                &version_buffer[..version_len],
+                target_version.as_bytes(),
+                "aot-compiled contract version mismatch"
+            );
+        };
+
         Ok(Self {
-            library: Arc::new(unsafe { Library::new(library_path)? }),
+            library,
             path: library_path.to_path_buf(),
             is_temp_path: false,
             contract_info,