fix: Make WASI polyfill more robust in image classification (#850)

This replaces `ic_wasi_polyfill::init` with a safer
`ic_wasi_polyfill::init_with_memory` that passes a virtual stable
memory to WASI polyfill to use instead of the global stable memory.

rust/image-classification/README.md

@@ -39,6 +39,12 @@ Download MobileNet v2-7 to `src/backend/assets/mobilenetv2-7.onnx`:
 ./downdload_model.sh
 ```
 
+Install NodeJS dependencies for the frontend:
+
+```
+npm install
+```
+
 # Build
 
 ```

rust/image-classification/src/backend/Cargo.toml

@@ -17,4 +17,5 @@ prost-types = "0.11.0"
 image = { version = "0.24", features = ["png"], default-features = false }
 serde = { version = "1.0", features = ["derive"] }
 tract-onnx = { git = "https://github.com/sonos/tract", version = "=0.21.2-pre" }
+ic-stable-structures = "0.6"
 ic-wasi-polyfill = { git = "https://github.com/wasm-forge/ic-wasi-polyfill", version = "0.3.17" }

rust/image-classification/src/backend/src/lib.rs

@@ -1,7 +1,19 @@
+use std::cell::RefCell;
 use candid::{CandidType, Deserialize};
+use ic_stable_structures::{memory_manager::{MemoryId, MemoryManager}, DefaultMemoryImpl};
 
 mod onnx;
 
+// WASI polyfill requires a virtual stable memory to store the file system.
+// You can replace `0` with any index up to `254`.
+const WASI_MEMORY_ID: MemoryId = MemoryId::new(0);
+
+thread_local! {
+    // The memory manager is used for simulating multiple memories.
+    static MEMORY_MANAGER: RefCell<MemoryManager<DefaultMemoryImpl>> =
+        RefCell::new(MemoryManager::init(DefaultMemoryImpl::default()));
+}
+
 #[derive(CandidType, Deserialize)]
 struct Classification {
     label: String,
@@ -32,11 +44,14 @@ fn classify(image: Vec<u8>) -> ClassificationResult {
 
 #[ic_cdk::init]
 fn init() {
-    ic_wasi_polyfill::init(&[0u8; 32], &[]);
+    let wasi_memory = MEMORY_MANAGER.with(|m| m.borrow().get(WASI_MEMORY_ID));
+    ic_wasi_polyfill::init_with_memory(&[0u8; 32], &[], wasi_memory);
     onnx::setup().unwrap();
 }
 
 #[ic_cdk::post_upgrade]
 fn post_upgrade() {
+    let wasi_memory = MEMORY_MANAGER.with(|m| m.borrow().get(WASI_MEMORY_ID));
+    ic_wasi_polyfill::init_with_memory(&[0u8; 32], &[], wasi_memory);
     onnx::setup().unwrap();
 }