mimobox-wasm 0.1.0-alpha

MimoBox WebAssembly sandbox backend using Wasmtime
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
#![cfg_attr(docsrs, feature(doc_cfg))]
#![warn(missing_docs)]
//! mimobox-wasm: Wasm sandbox backend.
//!
//! Implements a Wasm sandbox on top of the Wasmtime runtime with WASI Preview 1 support.
//! Core design:
//! - Globally shared Engine, owned by `WasmSandbox` and reused across multiple `execute` calls.
//! - Module compilation cache based on SHA256 hashes of file content to avoid repeated compilation.
//! - Independent Store per `execute`, with a fresh WASI context and resource limits.
//! - stdout/stderr captured into in-memory buffers through `MemoryOutputPipe` with a built-in capacity limit.
//! - Dual execution-time limits with fuel and epoch interruption.

use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Instant, UNIX_EPOCH};

use sha2::{Digest, Sha256};
use wasmtime::{
    Config, Engine, Linker, Module, OptLevel, Store, StoreLimits, StoreLimitsBuilder, Trap,
};
use wasmtime_wasi::I32Exit;
use wasmtime_wasi::p1::{self, WasiP1Ctx};
use wasmtime_wasi::p2::pipe::MemoryOutputPipe;
use wasmtime_wasi::{DirPerms, FilePerms, WasiCtxBuilder};

use mimobox_core::{Sandbox, SandboxConfig, SandboxError, SandboxResult};

/// Sandbox Store data combining the WASI context and resource limits.
///
/// [FATAL-01 fix] Embeds `StoreLimits` in the Store data type so that the `store.limiter()`
/// callback correctly returns `&mut dyn ResourceLimiter`, allowing `memory_limit_mb` to be
/// applied to the Wasm runtime.
struct StoreData {
    wasi: WasiP1Ctx,
    limits: StoreLimits,
}

/// Logging macro.
fn wasm_logging_enabled() -> bool {
    std::env::var_os("MIMOBOX_WASM_QUIET").is_none()
}

macro_rules! log_info {
    ($($arg:tt)*) => {
        if wasm_logging_enabled() {
            eprintln!("[mimobox:wasm:info] {}", format!($($arg)*))
        }
    };
}

macro_rules! log_warn {
    ($($arg:tt)*) => {
        if wasm_logging_enabled() {
            eprintln!("[mimobox:wasm:warn] {}", format!($($arg)*))
        }
    };
}

/// Fuel estimation factor: about 15 million Wasm instructions (fuel) per second, including 50% headroom.
const FUEL_PER_SECOND: u64 = 15_000_000;

/// Default fuel limit when no timeout is configured, roughly equivalent to 10 million Wasm instructions.
const DEFAULT_FUEL_LIMIT: u64 = 10_000_000;

/// Maximum stdout/stderr buffer capacity: 1 MB.
/// The `MemoryOutputPipe` capacity parameter also acts as the write limit;
/// writes beyond this capacity return `StreamError::Closed`.
const OUTPUT_MAX_CAPACITY: usize = 1024 * 1024;

/// Maximum returned size for a single output stream: 4 MB; excess data is truncated and logged as a warning.
const MAX_OUTPUT_SIZE: usize = 4 * 1024 * 1024;

/// Maximum Wasm module file size: 100 MB.
const MAX_WASM_FILE_SIZE: u64 = 100 * 1024 * 1024;

/// Default memory limit: 64 MB, used when the config does not specify one.
const DEFAULT_MEMORY_LIMIT_MB: u64 = 64;

/// Epoch tick interval: 10 ms.
const EPOCH_TICK_INTERVAL_MS: u64 = 10;

/// Dynamically calculates the fuel limit from `timeout_secs`.
///
/// [IMPORTANT-01 fix] Roughly maps `timeout_secs` to a fuel quota.
/// Fuel is consumed only while executing pure Wasm instructions; wait time during WASI I/O does
/// not count. Therefore, fuel is an approximate timeout mechanism paired with epoch interruption
/// to enforce wall-clock timeout.
fn fuel_from_timeout(timeout_secs: Option<u64>) -> u64 {
    match timeout_secs {
        Some(secs) => secs.saturating_mul(FUEL_PER_SECOND),
        None => DEFAULT_FUEL_LIMIT,
    }
}

/// Reads output and truncates it to the maximum size.
fn truncate_output(data: Vec<u8>, label: &str) -> Vec<u8> {
    if data.len() > MAX_OUTPUT_SIZE {
        log_warn!(
            "{} output exceeded limit ({} > {} bytes), truncated",
            label,
            data.len(),
            MAX_OUTPUT_SIZE
        );
        data[..MAX_OUTPUT_SIZE].to_vec()
    } else {
        data
    }
}

/// Wasm sandbox backend.
///
/// Holds the globally shared Engine and module cache directory path, and creates an independent
/// Store for each `execute` call.
pub struct WasmSandbox {
    engine: Arc<Engine>,
    config: SandboxConfig,
    cache_dir: PathBuf,
    epoch_running: Arc<AtomicBool>,
    epoch_thread: Option<std::thread::JoinHandle<()>>,
}

/// Calculates the SHA256 hash of file content.
///
/// [IMPORTANT-02 fix] Uses SHA256 instead of `DefaultHasher`, generating cache keys from file
/// content rather than path plus modification time to eliminate TOCTOU race conditions.
fn content_hash(data: &[u8]) -> String {
    let hash = Sha256::digest(data);
    format!("{:x}", hash)
}

/// Gets a lightweight metadata fingerprint for a file: size plus modification time.
///
/// Used to quickly determine whether a file may have changed, avoiding a full file read and
/// SHA256 calculation on every `execute` call.
fn file_fingerprint(path: &Path) -> Option<(u64, u64)> {
    let meta = std::fs::metadata(path).ok()?;
    let size = meta.len();
    let mtime = meta
        .modified()
        .ok()?
        .duration_since(UNIX_EPOCH)
        .ok()?
        .as_nanos() as u64;
    Some((size, mtime))
}

fn compile_module_from_bytes(
    engine: &Engine,
    wasm_path: &Path,
    bytes: &[u8],
) -> Result<Module, SandboxError> {
    // SECURITY: 调用方在读取字节后立刻使用同一份不可变切片编译,
    // 避免“先读取算哈希、再按路径重新打开编译”的 TOCTOU 竞态。
    Module::from_binary(engine, bytes).map_err(|e| {
        SandboxError::ExecutionFailed(format!(
            "Failed to load Wasm module ({:?}): {}",
            wasm_path, e
        ))
    })
}

/// Gets or compiles a module with a disk cache.
///
/// Uses a hybrid cache strategy:
/// 1. First looks up the cache mapping file using the lightweight metadata fingerprint (size plus modification time).
/// 2. When the cache mapping hits, loads directly with the corresponding SHA256 cache key.
/// 3. When the cache mapping misses, calculates the file content SHA256 and updates the cache.
///
/// This avoids reading the whole file to calculate a hash on the cache-hit hot path.
fn get_cached_module(
    engine: &Engine,
    wasm_path: &Path,
    cache_dir: &Path,
) -> Result<Module, SandboxError> {
    let _ = std::fs::create_dir_all(cache_dir);

    let fingerprint = match file_fingerprint(wasm_path) {
        Some(fp) => fp,
        None => {
            // 无法获取元数据时也只读取一次文件,避免在读取与编译之间被路径替换。
            let file_data = std::fs::read(wasm_path).map_err(|e| {
                SandboxError::ExecutionFailed(format!("Failed to read Wasm file: {}", e))
            })?;
            return compile_module_from_bytes(engine, wasm_path, &file_data);
        }
    };

    // 缓存映射文件:记录 "fingerprint -> sha256_hash" 的映射
    // 文件名格式: {size}_{mtime_nanos}.map
    let map_file = cache_dir.join(format!("{}_{}.map", fingerprint.0, fingerprint.1));

    // 尝试通过映射文件找到对应的缓存
    if let Ok(hash) = std::fs::read_to_string(&map_file) {
        let cache_path = cache_dir.join(format!("{}.cwasm", hash.trim()));
        match std::fs::read(&cache_path) {
            Ok(cached) => {
                // SAFETY: 缓存文件由本系统生成,Engine 配置未变。
                // Module::deserialize 要求输入数据来自相同 Engine 配置的 serialize() 输出。
                // 我们在缓存写入时确保了这一点,因此反序列化是安全的。
                match unsafe { Module::deserialize(engine, &cached) } {
                    Ok(module) => {
                        log_info!("Loaded module from cache: {:?}", wasm_path);
                        return Ok(module);
                    }
                    Err(e) => {
                        // 反序列化失败:缓存可能损坏或 Engine 配置变更,静默降级重新编译
                        log_warn!("Cache deserialization failed, recompiling: {}", e);
                        let _ = std::fs::remove_file(&cache_path);
                        let _ = std::fs::remove_file(&map_file);
                    }
                }
            }
            Err(_) => {
                // 缓存文件不存在,映射过期,清理并重新编译
                let _ = std::fs::remove_file(&map_file);
            }
        }
    }

    // 缓存未命中:需要计算文件内容的 SHA256
    let file_data = std::fs::read(wasm_path)
        .map_err(|e| SandboxError::ExecutionFailed(format!("Failed to read Wasm file: {}", e)))?;
    let hash = content_hash(&file_data);
    let cache_path = cache_dir.join(format!("{}.cwasm", hash));

    // 检查是否已有相同内容的缓存(文件内容相同但元数据不同)
    if let Ok(cached) = std::fs::read(&cache_path) {
        // SAFETY: 缓存文件由本系统生成,Engine 配置未变。
        // Module::deserialize 要求输入数据来自相同 Engine 配置的 serialize() 输出。
        // 我们在缓存写入时确保了这一点,因此反序列化是安全的。
        match unsafe { Module::deserialize(engine, &cached) } {
            Ok(module) => {
                // 更新映射文件
                let _ = std::fs::write(&map_file, &hash);
                log_info!("Loaded module from cache (content match): {:?}", wasm_path);
                return Ok(module);
            }
            Err(e) => {
                log_warn!("Cache deserialization failed, recompiling: {}", e);
                let _ = std::fs::remove_file(&cache_path);
            }
        }
    }

    // 编译模块
    let module = compile_module_from_bytes(engine, wasm_path, &file_data)?;

    // 序列化到缓存目录(原子写入:先写临时文件再 rename,避免并发读到不完整数据)
    if let Ok(serialized) = module.serialize() {
        let tmp_path = cache_path.with_extension("cwasm.tmp");
        if std::fs::write(&tmp_path, &serialized).is_ok() {
            // rename 在同一文件系统上是原子的
            if let Err(e) = std::fs::rename(&tmp_path, &cache_path) {
                log_warn!("Failed to rename cache file: {}", e);
                let _ = std::fs::remove_file(&tmp_path);
            }
        }
        // 更新映射文件
        let _ = std::fs::write(&map_file, &hash);
    }

    log_info!("Compiled and cached module: {:?}", wasm_path);
    Ok(module)
}

/// Creates the Wasmtime Engine configuration for sandbox execution.
fn create_engine_config() -> Config {
    let mut config = Config::new();
    config.cranelift_opt_level(OptLevel::Speed);
    config.consume_fuel(true);
    config.epoch_interruption(true);
    config.max_wasm_stack(512 * 1024); // 512KB Wasm 栈
    config.parallel_compilation(true);
    config
}

/// Builds a WASI Preview 1 context.
///
/// Configures filesystem access, environment variables, and related settings from `SandboxConfig`.
/// stdout/stderr are captured into in-memory buffers through `MemoryOutputPipe`.
fn build_wasi_ctx(
    config: &SandboxConfig,
    args: &[String],
    stdout_pipe: MemoryOutputPipe,
    stderr_pipe: MemoryOutputPipe,
) -> WasiP1Ctx {
    let mut builder = WasiCtxBuilder::new();

    // 设置命令行参数
    for arg in args {
        builder.arg(arg);
    }

    // 设置最小必要环境变量
    builder.env("HOME", "/tmp");
    builder.env("PATH", "/usr/bin:/bin");
    builder.env("TERM", "dumb");
    builder.env("SANDBOX", "wasm");

    // 配置 stdout/stderr 捕获
    builder.stdout(Box::new(stdout_pipe));
    builder.stderr(Box::new(stderr_pipe));

    // 文件系统访问:仅允许 config 中配置的路径
    for path in &config.fs_readonly {
        if let Some(path_str) = path.to_str() {
            if path.exists() {
                // 只授予 READ 权限,WASI 的目录创建/删除等写操作会被拒绝。
                if let Err(e) =
                    builder.preopened_dir(path, path_str, DirPerms::READ, FilePerms::READ)
                {
                    log_warn!("Failed to preopen read-only dir {:?}: {}", path, e);
                }
            } else {
                log_warn!("Read-only path does not exist: {:?}", path);
            }
        }
    }
    for path in &config.fs_readwrite {
        if let Some(path_str) = path.to_str() {
            if path.exists() {
                if let Err(e) =
                    builder.preopened_dir(path, path_str, DirPerms::all(), FilePerms::all())
                {
                    log_warn!("Failed to preopen read-write dir {:?}: {}", path, e);
                }
            } else {
                log_warn!("Read-write path does not exist: {:?}", path);
            }
        }
    }

    if config.deny_network {
        log_info!("WASI network denied by SandboxConfig; no sockets are preopened");
    } else {
        log_warn!(
            "SandboxConfig allows network, but WASI backend cannot enable network access; keeping network denied"
        );
    }

    // 时钟能力:WasiCtxBuilder 仅暴露 wall_clock/monotonic_clock 替换点,
    // 没有 WASI Preview 1 clocks 的禁用/白名单 API;执行时限由 fuel、epoch 和 timeout 控制。

    builder.build_p1()
}

impl Sandbox for WasmSandbox {
    fn new(config: SandboxConfig) -> Result<Self, SandboxError> {
        let engine_config = create_engine_config();
        let engine = Arc::new(Engine::new(&engine_config).map_err(|e| {
            SandboxError::ExecutionFailed(format!("Failed to create Wasmtime Engine: {}", e))
        })?);

        let epoch_running = Arc::new(AtomicBool::new(true));
        let epoch_thread_engine = engine.clone();
        let epoch_thread_running = epoch_running.clone();
        let epoch_thread = std::thread::Builder::new()
            .name("mimobox-wasm-epoch-ticker".to_string())
            .spawn(move || {
                let tick_interval = std::time::Duration::from_millis(EPOCH_TICK_INTERVAL_MS);
                while epoch_thread_running.load(Ordering::Relaxed) {
                    std::thread::sleep(tick_interval);
                    epoch_thread_engine.increment_epoch();
                }
            })
            .map_err(|e| {
                SandboxError::ExecutionFailed(format!("Failed to start Wasm epoch ticker: {}", e))
            })?;

        // [IMPORTANT-02 修复] 使用用户专属缓存目录,避免不同用户之间的缓存污染
        // SAFETY: geteuid() 是无副作用的系统调用,始终返回有效的 uid。
        let uid = unsafe { libc::geteuid() };
        let cache_dir = std::env::temp_dir().join(format!("mimobox-cache-{}", uid));

        log_info!(
            "Created Wasm sandbox backend, memory_limit={:?}MB, timeout={:?}s, cache_dir={:?}",
            config.memory_limit_mb,
            config.timeout_secs,
            cache_dir,
        );

        Ok(Self {
            engine,
            config,
            cache_dir,
            epoch_running,
            epoch_thread: Some(epoch_thread),
        })
    }

    fn execute(&mut self, cmd: &[String]) -> Result<SandboxResult, SandboxError> {
        let start = Instant::now();

        if cmd.is_empty() {
            return Err(SandboxError::ExecutionFailed("Command is empty".into()));
        }

        let wasm_path = Path::new(&cmd[0]);
        if !wasm_path.exists() {
            return Err(SandboxError::ExecutionFailed(format!(
                "Wasm file does not exist: {:?}",
                wasm_path
            )));
        }

        // [MINOR-07] 预检查文件大小,防止超大文件导致编译时 OOM
        if let Ok(meta) = std::fs::metadata(wasm_path)
            && meta.len() > MAX_WASM_FILE_SIZE
        {
            return Err(SandboxError::ExecutionFailed(format!(
                "Wasm file too large: {} bytes (limit {} bytes)",
                meta.len(),
                MAX_WASM_FILE_SIZE
            )));
        }

        // 1. 获取或编译模块(带缓存)
        let module = get_cached_module(&self.engine, wasm_path, &self.cache_dir)?;

        // 2. [IMPORTANT-03 说明] stdout/stderr 缓冲区容量限制
        // MemoryOutputPipe 的 capacity 参数是写入上限而非初始容量,
        // 超过此容量后 OutputStream::write() 返回 StreamError::Trap,
        // check_write() 返回 StreamError::Closed。
        let stdout_pipe = MemoryOutputPipe::new(OUTPUT_MAX_CAPACITY);
        let stdout_reader = stdout_pipe.clone(); // 保留读取端
        let stderr_pipe = MemoryOutputPipe::new(OUTPUT_MAX_CAPACITY);
        let stderr_reader = stderr_pipe.clone(); // 保留读取端

        // 3. 构建 WASI 上下文
        let wasi_ctx = build_wasi_ctx(&self.config, cmd, stdout_pipe, stderr_pipe);

        // 4. 创建 Linker 并注册 WASI Preview 1
        // 注意:Linker 的类型参数必须与 Store data type 一致
        let mut linker: Linker<StoreData> = Linker::new(&self.engine);
        p1::add_to_linker_sync(&mut linker, |data| &mut data.wasi).map_err(|e| {
            SandboxError::ExecutionFailed(format!("Failed to register WASI Preview 1: {}", e))
        })?;

        // 5. [FATAL-01 修复] 创建带资源限制的 Store
        // 将 memory_limit_mb 通过 StoreLimitsBuilder 实际应用到 Wasm 运行时
        let memory_limit_bytes: usize = self
            .config
            .memory_limit_mb
            .map(|mb| mb * 1024 * 1024)
            .unwrap_or(DEFAULT_MEMORY_LIMIT_MB * 1024 * 1024)
            .try_into()
            .unwrap_or(usize::MAX);

        let limits = StoreLimitsBuilder::new()
            .memory_size(memory_limit_bytes)
            .memories(1) // 单个线性内存
            .tables(4) // 限制间接调用表数量
            .instances(1) // 单实例
            .trap_on_grow_failure(true) // 内存增长失败时 trap 而非返回 -1
            .build();

        let store_data = StoreData {
            wasi: wasi_ctx,
            limits,
        };
        let mut store = Store::new(&self.engine, store_data);
        store.limiter(|data| &mut data.limits);

        // 6. [IMPORTANT-01 修复] 根据 timeout_secs 动态设置 fuel 上限
        let fuel_limit = fuel_from_timeout(self.config.timeout_secs);
        store
            .set_fuel(fuel_limit)
            .map_err(|e| SandboxError::ExecutionFailed(format!("Failed to set fuel: {}", e)))?;

        // 7. [IMPORTANT-01 补充] 设置 epoch deadline 实现墙钟超时
        // epoch_interruption 可中断包括 WASI I/O 阻塞在内的执行,
        // 弥补 fuel 仅在纯 Wasm 指令执行时消耗的不足。
        store.epoch_deadline_trap();
        let epoch_deadline_ticks = self
            .config
            .timeout_secs
            .map(|s| s.saturating_mul(100)) // 每 10ms 一个 epoch tick
            .unwrap_or(3000); // 默认 30s
        store.set_epoch_deadline(epoch_deadline_ticks);

        // 8. 实例化模块
        let instance = linker.instantiate(&mut store, &module).map_err(|e| {
            SandboxError::ExecutionFailed(format!("Failed to instantiate Wasm module: {}", e))
        })?;

        // 9. 调用 _start 函数(WASI Command 模式)
        // WASI Command 通过 _start 进入,正常退出时调用 proc_exit(code),
        // 这会触发 I32Exit 错误,其中包含退出码。
        let exit_code = match instance.get_typed_func::<(), ()>(&mut store, "_start") {
            Ok(start_func) => {
                match start_func.call(&mut store, ()) {
                    Ok(()) => Some(0),
                    Err(e) => {
                        // 检查是否是 WASI 正常退出(I32Exit)
                        // I32Exit 可能被包装在 error chain 中,需要遍历查找
                        if let Some(exit) = find_exit_code(&e) {
                            Some(exit)
                        } else if is_fuel_exhausted(&store) || is_epoch_interrupt(&e) {
                            log_warn!(
                                "Execution timed out (fuel exhausted or epoch deadline exceeded)"
                            );
                            let elapsed = start.elapsed();
                            let stdout =
                                truncate_output(stdout_reader.contents().to_vec(), "stdout");
                            let stderr =
                                truncate_output(stderr_reader.contents().to_vec(), "stderr");
                            return Ok(SandboxResult {
                                stdout,
                                stderr,
                                exit_code: None,
                                elapsed,
                                timed_out: true,
                            });
                        } else if is_memory_trap(&e) {
                            log_info!("Wasm memory limit exceeded, mapping to exit code 1");
                            Some(1)
                        } else {
                            log_warn!("Wasm execution error: {}", e);
                            None
                        }
                    }
                }
            }
            Err(_) => {
                // 没有 _start 函数,尝试查找 main 函数
                match instance.get_typed_func::<(), i32>(&mut store, "main") {
                    Ok(main_func) => match main_func.call(&mut store, ()) {
                        Ok(code) => Some(code),
                        Err(e) => {
                            if let Some(exit) = find_exit_code(&e) {
                                Some(exit)
                            } else if is_memory_trap(&e) {
                                log_info!("Wasm memory limit exceeded, mapping to exit code 1");
                                Some(1)
                            } else {
                                log_warn!("main function execution failed: {}", e);
                                None
                            }
                        }
                    },
                    Err(_) => {
                        return Err(SandboxError::ExecutionFailed(
                            "Wasm module has no _start or main export function".into(),
                        ));
                    }
                }
            }
        };

        let elapsed = start.elapsed();

        // 10. 从 clone 的 MemoryOutputPipe 中读取捕获的输出(带截断保护)
        let stdout = truncate_output(stdout_reader.contents().to_vec(), "stdout");
        let stderr = truncate_output(stderr_reader.contents().to_vec(), "stderr");

        log_info!(
            "Wasm execution completed, exit_code={:?}, elapsed={:.2}ms",
            exit_code,
            elapsed.as_secs_f64() * 1000.0
        );

        Ok(SandboxResult {
            stdout,
            stderr,
            exit_code,
            elapsed,
            timed_out: false,
        })
    }

    fn destroy(self) -> Result<(), SandboxError> {
        log_info!("Destroying Wasm sandbox backend");
        Ok(())
    }
}

impl Drop for WasmSandbox {
    fn drop(&mut self) {
        self.epoch_running.store(false, Ordering::Relaxed);
        if let Some(epoch_thread) = self.epoch_thread.take()
            && let Err(e) = epoch_thread.join()
        {
            log_warn!("Wasm epoch ticker thread join failed: {:?}", e);
        }
    }
}

/// Checks whether the Store fuel is exhausted.
fn is_fuel_exhausted(store: &Store<StoreData>) -> bool {
    store.get_fuel().is_ok_and(|f| f == 0)
}

/// Checks whether an error is an epoch interruption, indicating wall-clock timeout.
fn is_epoch_interrupt(error: &wasmtime::Error) -> bool {
    if let Some(trap) = error.downcast_ref::<Trap>() {
        matches!(trap, Trap::Interrupt)
    } else {
        false
    }
}

/// Checks whether an error is caused by Wasm memory access or growth limits.
fn is_memory_trap(error: &wasmtime::Error) -> bool {
    if let Some(trap) = error.downcast_ref::<Trap>()
        && matches!(trap, Trap::MemoryOutOfBounds)
    {
        return true;
    }

    let message = format!("{error:#}").to_ascii_lowercase();
    message.contains("memory")
        && (message.contains("out of bounds")
            || message.contains("grow")
            || message.contains("growth"))
}

/// Finds a WASI `I32Exit` exit code in the wasmtime error chain.
///
/// WASI `proc_exit` propagates the exit code through an `I32Exit` error, but it may be wrapped
/// by intermediate layers such as `WasmBacktrace`. This function walks the entire error chain.
fn find_exit_code(error: &wasmtime::Error) -> Option<i32> {
    // 直接 downcast
    if let Some(exit) = error.downcast_ref::<I32Exit>() {
        return Some(exit.0);
    }

    // 检查 root cause
    let root = error.root_cause();
    if let Some(exit) = root.downcast_ref::<I32Exit>() {
        return Some(exit.0);
    }

    None
}

/// Runs the Wasm sandbox cold-start benchmark.
pub fn run_wasm_benchmark(
    wasm_path: &str,
    iterations: usize,
) -> Result<(), Box<dyn std::error::Error>> {
    println!("=== mimobox Wasm Sandbox Benchmark ===\n");

    let mut config = SandboxConfig::default();
    config.deny_network = true;
    config.memory_limit_mb = Some(64);
    config.timeout_secs = Some(30);
    config.fs_readonly = vec![];
    config.fs_readwrite = vec![];
    config.seccomp_profile = mimobox_core::SeccompProfile::Essential;
    config.allow_fork = false;
    config.allowed_http_domains = vec![];

    if !Path::new(wasm_path).exists() {
        return Err(format!("Wasm file does not exist: {}", wasm_path).into());
    }

    let cmd = vec![wasm_path.to_string()];

    // Phase 1: Engine 创建开销(一次性)
    println!("Testing Engine creation overhead...");
    let engine_start = Instant::now();
    let mut sb = WasmSandbox::new(config.clone())?;
    let engine_elapsed = engine_start.elapsed();
    println!(
        "  Engine creation: {:.2}ms",
        engine_elapsed.as_secs_f64() * 1000.0
    );

    // Phase 2: 模块编译开销(首次)
    println!("\nTesting first module compilation...");
    let compile_start = Instant::now();
    let result = sb.execute(&cmd)?;
    let compile_elapsed = compile_start.elapsed();
    println!(
        "  First execution (with compilation): {:.2}ms, exit_code={:?}",
        compile_elapsed.as_secs_f64() * 1000.0,
        result.exit_code
    );

    // Phase 3: 冷启动测试(每次 new + execute)
    println!(
        "\nCold start test ({} iterations, each with new + execute)...",
        iterations
    );
    let mut cold_times = Vec::with_capacity(iterations);
    for i in 0..iterations {
        let start = Instant::now();
        let mut sb = WasmSandbox::new(config.clone())?;
        let result = sb.execute(&cmd)?;
        let elapsed = start.elapsed();
        cold_times.push(elapsed.as_micros() as f64 / 1000.0);

        if result.exit_code != Some(0) {
            eprintln!("Iteration {} failed: exit code {:?}", i, result.exit_code);
        }
    }

    // Phase 4: 热路径测试(复用 Engine,仅 execute)
    println!(
        "\nHot path test ({} iterations, reusing Engine)...",
        iterations
    );
    let mut hot_times = Vec::with_capacity(iterations);
    for _ in 0..iterations {
        let start = Instant::now();
        let result = sb.execute(&cmd)?;
        let elapsed = start.elapsed();
        hot_times.push(elapsed.as_micros() as f64 / 1000.0);

        if result.exit_code != Some(0) {
            eprintln!("Hot path execution failed: {:?}", result.exit_code);
        }
    }

    // 统计输出
    cold_times.sort_by(f64::total_cmp);
    hot_times.sort_by(f64::total_cmp);

    fn print_stats(label: &str, times: &[f64]) {
        let n = times.len();
        if n == 0 {
            println!("{}  no data", label);
            return;
        }
        let p50 = times[n / 2];
        let p95_idx = ((n as f64 * 0.95) as usize).min(n - 1);
        let p99_idx = ((n as f64 * 0.99) as usize).min(n - 1);
        let avg: f64 = times.iter().sum::<f64>() / n as f64;

        println!("\n{} latency:", label);
        println!("  Min:  {:.2}ms", times[0]);
        println!("  P50:  {:.2}ms", p50);
        println!("  P95:  {:.2}ms", times[p95_idx]);
        println!("  P99:  {:.2}ms", times[p99_idx]);
        println!("  Avg:  {:.2}ms", avg);
        println!("  Max:  {:.2}ms", times[n - 1]);
    }

    print_stats("Cold start ", &cold_times);
    print_stats("Hot path ", &hot_times);

    // 目标检查
    let cold_p50 = cold_times[cold_times.len() / 2];
    let hot_p50 = hot_times[hot_times.len() / 2];
    println!("\nTarget check:");
    println!(
        "  Cold start P50: {:.2}ms {}",
        cold_p50,
        if cold_p50 < 5.0 { "[PASS]" } else { "[FAIL]" }
    );
    println!(
        "  Hot path P50: {:.2}ms {}",
        hot_p50,
        if hot_p50 < 1.0 { "[PASS]" } else { "[FAIL]" }
    );

    println!("\n=== Test completed ===");
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use mimobox_core::Sandbox;
    use wasmtime::{Instance, Store};

    fn test_config() -> SandboxConfig {
        let mut config = SandboxConfig::default();
        config.timeout_secs = Some(10);
        config.memory_limit_mb = Some(64);
        config.fs_readonly = vec![];
        config.fs_readwrite = vec![];
        config.deny_network = true;
        config.seccomp_profile = mimobox_core::SeccompProfile::Essential;
        config.allow_fork = false;
        config.allowed_http_domains = vec![];
        config
    }

    #[test]
    fn test_wasm_sandbox_create() {
        let sb = WasmSandbox::new(test_config());
        assert!(sb.is_ok(), "Failed to create Wasm sandbox: {:?}", sb.err());
    }

    #[test]
    fn test_wasm_sandbox_empty_command() {
        let mut sb = WasmSandbox::new(test_config()).expect("Failed to create");
        let result = sb.execute(&[]);
        assert!(result.is_err(), "Empty command should return error");
    }

    #[test]
    fn test_wasm_sandbox_nonexistent_file() {
        let mut sb = WasmSandbox::new(test_config()).expect("Failed to create");
        let result = sb.execute(&["/nonexistent/file.wasm".to_string()]);
        assert!(result.is_err(), "Nonexistent file should return error");
    }

    #[test]
    fn test_wasm_sandbox_destroy() {
        let sb = WasmSandbox::new(test_config()).expect("Failed to create");
        let result = sb.destroy();
        assert!(
            result.is_ok(),
            "Failed to destroy sandbox: {:?}",
            result.err()
        );
    }

    #[test]
    fn test_compile_module_from_bytes_is_not_affected_by_path_swap() {
        let engine = Engine::default();
        let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
        let wasm_path = temp_dir.path().join("swap.wasm");
        let module_a = wat::parse_str(
            r#"
                (module
                  (func (export "main") (result i32)
                    i32.const 1))
            "#,
        )
        .expect("Failed to compile module A WAT");
        let module_b = wat::parse_str(
            r#"
                (module
                  (func (export "main") (result i32)
                    i32.const 2))
            "#,
        )
        .expect("Failed to compile module B WAT");

        std::fs::write(&wasm_path, &module_a).expect("Failed to write initial module");
        let module = compile_module_from_bytes(&engine, &wasm_path, &module_a)
            .expect("Should compile from read bytes");
        std::fs::write(&wasm_path, &module_b).expect("Failed to overwrite module");

        let mut store = Store::new(&engine, ());
        let instance = Instance::new(&mut store, &module, &[])
            .expect("Failed to instantiate module from read bytes");
        let main = instance
            .get_typed_func::<(), i32>(&mut store, "main")
            .expect("Failed to get main export");

        assert_eq!(main.call(&mut store, ()).expect("Failed to call main"), 1);
    }
}