From 8c26f837615602d2422ab5de30f800a2122b60ff Mon Sep 17 00:00:00 2001
From: venkkatesh-sekar <94061546+venkkatesh-sekar@users.noreply.github.com>
Date: Mon, 23 Sep 2024 09:44:58 +0200
Subject: [PATCH] feat: add data section check when limiting Wasm heap memory
 (#72)

* add data segments check
* validate only against heap memory
* panic instead of logging the error
---
 src/limit_resource.rs | 38 +++++++++++++++++++++++++++++++++++++-
 1 file changed, 37 insertions(+), 1 deletion(-)

diff --git a/src/limit_resource.rs b/src/limit_resource.rs
index af882d1..74183bf 100644
--- a/src/limit_resource.rs
+++ b/src/limit_resource.rs
@@ -175,7 +175,43 @@ fn limit_heap_memory(m: &mut Module, limit: u32) {
         let memory = m.memories.get_mut(memory_id);
         let limit = limit as u64;
         if memory.initial > limit {
-            memory.initial = limit
+            // If memory.initial is greater than the provided limit, it is
+            // possible there is an active data segment with an offset in the
+            // range [limit, memory.initial].
+            //
+            // In that case, we don't restrict the heap memory limit as it could
+            // have undefined behaviour.
+
+            if m.data
+                .iter()
+                .filter_map(|data| {
+                    match data.kind {
+                        DataKind::Passive => None,
+                        DataKind::Active {
+                            memory: data_memory_id,
+                            offset,
+                        } => {
+                            if data_memory_id == memory_id {
+                                match offset {
+                                    ConstExpr::Value(Value::I32(offset)) => Some(offset as u64),
+                                    ConstExpr::Value(Value::I64(offset)) => Some(offset as u64),
+                                    _ => {
+                                        // It wouldn't pass IC wasm validation
+                                        None
+                                    }
+                                }
+                            } else {
+                                None
+                            }
+                        }
+                    }
+                })
+                .all(|offset| offset < limit * 65536)
+            {
+                memory.initial = limit;
+            } else {
+                panic!("Unable to restrict Wasm heap memory to {} pages", limit);
+            }
         }
         memory.maximum = Some(limit);
     }