From 665036a0c77b4bf3923262d1578dd2c3fe47766b Mon Sep 17 00:00:00 2001
From: Aumetra Weisman <aumetra@cryptolab.net>
Date: Mon, 4 Nov 2024 13:37:11 +0100
Subject: [PATCH] introduce memory limiter

---
 crates/kitsune-wasm-mrf/src/ctx.rs | 55 +++++++++++++++++++++++-------
 1 file changed, 42 insertions(+), 13 deletions(-)

diff --git a/crates/kitsune-wasm-mrf/src/ctx.rs b/crates/kitsune-wasm-mrf/src/ctx.rs
index e018f1cb5..86d59b783 100644
--- a/crates/kitsune-wasm-mrf/src/ctx.rs
+++ b/crates/kitsune-wasm-mrf/src/ctx.rs
@@ -1,12 +1,17 @@
-use crate::{kv_storage, mrf_wit::v1::fep::mrf::keyvalue};
+use crate::{
+    kv_storage,
+    mrf_wit::v1::fep::mrf::{http, keyvalue},
+};
 use slab::Slab;
 use triomphe::Arc;
 use wasmtime::{
     component::{Resource, ResourceTable},
-    Engine, Store,
+    Engine, Store, StoreLimits, StoreLimitsBuilder,
 };
 use wasmtime_wasi::{WasiCtx, WasiCtxBuilder, WasiView};
 
+const TABLE_ELEMENT_SIZE: usize = std::mem::size_of::<usize>();
+
 pub struct KvContext {
     pub module_name: Option<String>,
     pub storage: Arc<kv_storage::BackendDispatch>,
@@ -23,8 +28,22 @@ impl KvContext {
     }
 }
 
+pub struct HttpContext {
+    pub client: kitsune_http_client::Client,
+    pub bodies: Slab<todo!()>,
+}
+
+impl HttpContext {
+    #[inline]
+    pub fn get_body(&self, rep: &Resource<http::ResponseBody>) -> &todo!() {
+        &self.bodies[rep.rep() as usize]
+    }
+}
+
 pub struct Context {
+    pub http_ctx: HttpContext,
     pub kv_ctx: KvContext,
+    pub resource_limiter: StoreLimits,
     pub resource_table: ResourceTable,
     pub wasi_ctx: WasiCtx,
 }
@@ -50,16 +69,26 @@ pub fn construct_store(
         .allow_udp(false)
         .build();
 
-    Store::new(
-        engine,
-        Context {
-            kv_ctx: KvContext {
-                module_name: None,
-                storage,
-                buckets: Slab::new(),
-            },
-            resource_table: ResourceTable::new(),
-            wasi_ctx,
+    let data = Context {
+        http_ctx: HttpContext {
+            client: kitsune_http_client::Client::builder()
+                .content_length_limit(None)
+                .build(),
+            bodies: Slab::new(),
         },
-    )
+        kv_ctx: KvContext {
+            module_name: None,
+            storage,
+            buckets: Slab::new(),
+        },
+        resource_limiter: StoreLimitsBuilder::new()
+            .memory_size(100 * 1024 * 1024)
+            .build(),
+        resource_table: ResourceTable::new(),
+        wasi_ctx,
+    };
+
+    let mut store = Store::new(engine, data);
+    store.limiter(|store| &mut store.resource_limiter);
+    store
 }