feat: Enhance memory limit validation and resource management in WASM components

Signed-off-by: Jiaxiao Zhou <duibao55328@gmail.com>

Author: Mossaka

crates/policy/src/types.rs

@@ -262,9 +262,13 @@ impl CpuLimit {
 }
 
 impl MemoryLimit {
+    // Memory limit bounds: minimum 64KB, maximum 64GB (reasonable for WASM components)
+    const MIN_MEMORY_BYTES: u64 = 64 * 1024; // 64KB
+    const MAX_MEMORY_BYTES: u64 = 64u64 * 1024 * 1024 * 1024; // 64GB
+
     /// Validate and convert memory limit to bytes
     pub fn to_bytes(&self) -> PolicyResult<u64> {
-        match self {
+        let bytes = match self {
             MemoryLimit::String(s) => {
                 if s.is_empty() {
                     bail!("Memory limit string cannot be empty");
@@ -287,16 +291,41 @@ impl MemoryLimit {
                     .parse()
                     .map_err(|_| anyhow::anyhow!("Invalid memory value: {}", s))?;
 
+                if value == 0 {
+                    bail!("Memory limit cannot be zero: {}", s);
+                }
+
                 value
                     .checked_mul(multiplier)
-                    .ok_or_else(|| anyhow::anyhow!("Memory value too large: {}", s))
+                    .ok_or_else(|| anyhow::anyhow!("Memory value too large: {}", s))?
             }
             MemoryLimit::Number(n) => {
+                if *n == 0 {
+                    bail!("Memory limit cannot be zero");
+                }
                 // Assume legacy numeric values are in MB
                 n.checked_mul(1024 * 1024)
-                    .ok_or_else(|| anyhow::anyhow!("Memory value too large: {}", n))
+                    .ok_or_else(|| anyhow::anyhow!("Memory value too large: {}", n))?
             }
+        };
+
+        // Validate bounds
+        if bytes < Self::MIN_MEMORY_BYTES {
+            bail!(
+                "Memory limit {} bytes is below minimum {} KB",
+                bytes,
+                Self::MIN_MEMORY_BYTES / 1024
+            );
         }
+        if bytes > Self::MAX_MEMORY_BYTES {
+            bail!(
+                "Memory limit {} bytes exceeds maximum {} GB",
+                bytes,
+                Self::MAX_MEMORY_BYTES / (1024 * 1024 * 1024)
+            );
+        }
+
+        Ok(bytes)
     }
 }
 
@@ -735,13 +764,16 @@ mod tests {
         let memory_gi = MemoryLimit::String("2Gi".to_string());
         assert_eq!(memory_gi.to_bytes().unwrap(), 2 * 1024 * 1024 * 1024);
 
-        // Test Ti format
-        let memory_ti = MemoryLimit::String("1Ti".to_string());
-        assert_eq!(memory_ti.to_bytes().unwrap(), 1024u64 * 1024 * 1024 * 1024);
+        // Test Gi format (larger value)
+        let memory_gi_large = MemoryLimit::String("32Gi".to_string());
+        assert_eq!(
+            memory_gi_large.to_bytes().unwrap(),
+            32u64 * 1024 * 1024 * 1024
+        );
 
-        // Test plain bytes
-        let memory_bytes = MemoryLimit::String("1024".to_string());
-        assert_eq!(memory_bytes.to_bytes().unwrap(), 1024);
+        // Test plain bytes (above minimum)
+        let memory_bytes = MemoryLimit::String("131072".to_string()); // 128KB
+        assert_eq!(memory_bytes.to_bytes().unwrap(), 131072);
 
         // Test numeric format (legacy, assumes MB)
         let memory_numeric = MemoryLimit::Number(512);
@@ -756,6 +788,21 @@ mod tests {
 
         let invalid_number = MemoryLimit::String("invalidMi".to_string());
         assert!(invalid_number.to_bytes().is_err());
+
+        // Test bounds validation - too small
+        let too_small = MemoryLimit::String("32Ki".to_string()); // 32KB < 64KB minimum
+        assert!(too_small.to_bytes().is_err());
+
+        // Test bounds validation - zero
+        let zero_bytes = MemoryLimit::String("0".to_string());
+        assert!(zero_bytes.to_bytes().is_err());
+
+        let zero_numeric = MemoryLimit::Number(0);
+        assert!(zero_numeric.to_bytes().is_err());
+
+        // Test bounds validation - too large (128GB > 64GB maximum)
+        let too_large = MemoryLimit::String("128Gi".to_string());
+        assert!(too_large.to_bytes().is_err());
     }
 
     #[test]

crates/wassette/src/lib.rs

@@ -490,8 +490,11 @@ impl LifecycleManager {
         if resource_limiter.is_some() {
             store.limiter(|state: &mut WassetteWasiState<WasiState>| {
                 // Extract the resource limiter from the inner state
-                // We know it exists because we checked above
-                state.inner.resource_limiter.as_mut().unwrap()
+                state
+                    .inner
+                    .resource_limiter
+                    .as_mut()
+                    .expect("Resource limiter should be present - checked above")
             });
         }
 

crates/wassette/src/policy_internal.rs

@@ -297,14 +297,18 @@ impl crate::LifecycleManager {
                     .and_then(|v| v.as_str())
                     .ok_or_else(|| anyhow!("Missing 'memory' field for resource permission"))?;
 
-                // Store as a custom rule with the specific structure expected by the policy system
-                let resource_details = serde_json::json!({
-                    "resources": {
-                        "limits": {
-                            "memory": memory
-                        }
-                    }
-                });
+                // Create structured resource limits instead of hardcoded JSON
+                let resource_limits = policy::ResourceLimits {
+                    limits: Some(policy::ResourceLimitValues::new(
+                        None,
+                        Some(policy::MemoryLimit::String(memory.to_string())),
+                    )),
+                    ..Default::default()
+                };
+
+                // Convert to JSON for storage in PermissionRule::Custom
+                let resource_details = serde_json::to_value(resource_limits)
+                    .map_err(|e| anyhow!("Failed to serialize resource limits: {}", e))?;
                 PermissionRule::Custom("resource".to_string(), resource_details)
             }
             other => {

crates/wassette/src/wasistate.rs

@@ -183,11 +183,16 @@ pub fn create_wasi_state_template_from_policy(
     let preopened_dirs = extract_storage_permissions(policy, plugin_dir)?;
     let allowed_hosts = extract_allowed_hosts(policy);
     let memory_limit = extract_memory_limit(policy)?;
-    let store_limits = memory_limit.map(|limit| {
-        wasmtime::StoreLimitsBuilder::new()
-            .memory_size(limit.try_into().unwrap_or(usize::MAX))
-            .build()
-    });
+    let store_limits = memory_limit
+        .map(|limit| -> anyhow::Result<wasmtime::StoreLimits> {
+            let limit_usize = limit.try_into().map_err(|_| {
+                anyhow::anyhow!("Memory limit {} too large for target architecture", limit)
+            })?;
+            Ok(wasmtime::StoreLimitsBuilder::new()
+                .memory_size(limit_usize)
+                .build())
+        })
+        .transpose()?;
 
     Ok(WasiStateTemplate {
         network_perms,