From 2df1439d351c874196d393032ba2d169c6cd977c Mon Sep 17 00:00:00 2001
From: wborgeaud <williamborgeaud@gmail.com>
Date: Tue, 21 Mar 2023 20:15:46 +0100
Subject: [PATCH] Return error instead of panic in memory operation (#928)

---
 evm/src/witness/errors.rs    | 11 +++++++++++
 evm/src/witness/memory.rs    | 29 +++++++++++++++++++----------
 evm/src/witness/operation.rs |  4 ++--
 3 files changed, 32 insertions(+), 12 deletions(-)

diff --git a/evm/src/witness/errors.rs b/evm/src/witness/errors.rs
index 53263675..e28e7d0f 100644
--- a/evm/src/witness/errors.rs
+++ b/evm/src/witness/errors.rs
@@ -1,3 +1,5 @@
+use ethereum_types::U256;
+
 #[allow(dead_code)]
 #[derive(Debug)]
 pub enum ProgramError {
@@ -8,4 +10,13 @@ pub enum ProgramError {
     InvalidJumpiDestination,
     StackOverflow,
     KernelPanic,
+    MemoryError(MemoryError),
+}
+
+#[allow(clippy::enum_variant_names)]
+#[derive(Debug)]
+pub enum MemoryError {
+    ContextTooLarge { context: U256 },
+    SegmentTooLarge { segment: U256 },
+    VirtTooLarge { virt: U256 },
 }
diff --git a/evm/src/witness/memory.rs b/evm/src/witness/memory.rs
index a2885796..127aac33 100644
--- a/evm/src/witness/memory.rs
+++ b/evm/src/witness/memory.rs
@@ -12,6 +12,9 @@ use MemoryChannel::{Code, GeneralPurpose};
 
 use crate::cpu::kernel::constants::global_metadata::GlobalMetadata;
 use crate::memory::segments::Segment;
+use crate::witness::errors::MemoryError::{ContextTooLarge, SegmentTooLarge, VirtTooLarge};
+use crate::witness::errors::ProgramError;
+use crate::witness::errors::ProgramError::MemoryError;
 
 impl MemoryChannel {
     pub fn index(&self) -> usize {
@@ -41,19 +44,25 @@ impl MemoryAddress {
         }
     }
 
-    pub(crate) fn new_u256s(context: U256, segment: U256, virt: U256) -> Self {
-        assert!(context.bits() <= 32, "context too large: {}", context);
-        assert!(
-            segment < Segment::COUNT.into(),
-            "segment too large: {}",
-            segment
-        );
-        assert!(virt.bits() <= 32, "virt too large: {}", virt);
-        Self {
+    pub(crate) fn new_u256s(
+        context: U256,
+        segment: U256,
+        virt: U256,
+    ) -> Result<Self, ProgramError> {
+        if context.bits() > 32 {
+            return Err(MemoryError(ContextTooLarge { context }));
+        }
+        if segment >= Segment::COUNT.into() {
+            return Err(MemoryError(SegmentTooLarge { segment }));
+        }
+        if virt.bits() > 32 {
+            return Err(MemoryError(VirtTooLarge { virt }));
+        }
+        Ok(Self {
             context: context.as_usize(),
             segment: segment.as_usize(),
             virt: virt.as_usize(),
-        }
+        })
     }
 
     pub(crate) fn increment(&mut self) {
diff --git a/evm/src/witness/operation.rs b/evm/src/witness/operation.rs
index 6d42eb48..43576ff4 100644
--- a/evm/src/witness/operation.rs
+++ b/evm/src/witness/operation.rs
@@ -111,7 +111,7 @@ pub(crate) fn generate_keccak_general<F: Field>(
         stack_pop_with_log_and_fill::<4, _>(state, &mut row)?;
     let len = len.as_usize();
 
-    let base_address = MemoryAddress::new_u256s(context, segment, base_virt);
+    let base_address = MemoryAddress::new_u256s(context, segment, base_virt)?;
     let input = (0..len)
         .map(|i| {
             let address = MemoryAddress {
@@ -608,7 +608,7 @@ pub(crate) fn generate_mload_general<F: Field>(
 
     let (val, log_read) = mem_read_gp_with_log_and_fill(
         3,
-        MemoryAddress::new_u256s(context, segment, virt),
+        MemoryAddress::new_u256s(context, segment, virt)?,
         state,
         &mut row,
     );