跳至主要內容
蒼時弦也
蒼時弦也
資深軟體工程師
發表於

Kobako:資源限制

這篇文章是 Kobako 系列的一部分。

Kobako:WebAssembly 的記憶體交換中解決了 Ruby 和 WebAssembly 互動的問題,也讓 Ruby 和 mruby 的數值能夠有互相交換的能力,基本上已經把「沙盒(Sandbox)」的問題處理乾淨。

然而 Kobako 的開發仍在非常初期,因為還有許多實際使用才會遇到的問題需要解決,其中一個就是「決定怎麼停下來」這樣的問題。

理想狀況

正常狀況下,我們會預期使用者很清楚自己在做什麼,因此會覺得寫一段可以正常運作的程式碼是很基本的,我們會很自然預期看到這樣的實作。

1sandbox = Kobako::Sandbox.new
2sandbox.eval("1 + 2").value  # => 3

然而,現實並非如此,假設一個使用者不小心寫了一個無限迴圈呢?

1sandbox = Kobako::Sandbox.new
2sandbox.eval("loop { puts 1 }").value  # => ?

不管在 Ruby 或者 mruby 都是合理的,在 mruby 的主戰場嵌入式系統中,也很常會需要用這種方式執行某個功能或者服務,因此不能禁止使用者用這種方式執行,但如果不阻止使用者,這個沙盒就會不斷地存在,直到宿主(Host)資源耗盡為止。

執行限制

以 Kobako 使用的 Wasmtime 內建提供了 Fuel 機制,可以很明確的限制使用的資源,但是在實際測試中,Fuel 會讓 Kobako 跑得更慢,因此改為使用 Epoch 來對執行時間進行限制。

Fuel 機制比 Epoch 慢的原因在於他需要對每一條 WebAssembly 指令都做確認,但是 Epoch 只是單純的指定時間,因此整體上會快上不少,在 Kobako 中大致上是這樣做的。

首先,建立 Wasmtime 時,我們把 Epoch 啟用,並且建立一個輕量的 Thread 以 10ms 的間隔把 Epoch 遞增。

 1let mut config = Config::new();
 2config.epoch_interruption(true);
 3let engine = Engine::new(&config)?;
 4{
 5    let engine = engine.clone();
 6    thread::spawn(move || loop {
 7        thread::sleep(Duration::from_millis(10));
 8        engine.increment_epoch();
 9    });
10}

每當我們使用 Sandbox#eval 或者 Sandbox#run 時,都會啟動一個全新的 mruby 沙盒,我們會在啟動的瞬間根據 timeout 設定,該停止的時機點。

 1let deadline = Instant::now() + Duration::from_secs(1);
 2store.set_epoch_deadline(1);
 3
 4store.epoch_deadline_callback(move |_ctx| {
 5    if Instant::now() >= deadline {
 6        Err(anyhow!("wall-clock deadline exceeded"))
 7    } else {
 8        Ok(UpdateDeadline::Continue(1))
 9    }
10});

每當 Epoch 更新時,我們就會檢查是否超出可執行的時間點,如果超過的話就直接丟出錯誤來中斷目前執行的沙盒,這樣就可以確保不會有無限執行的狀態發生(除非刻意設定 timeout: nil 會放行)

這樣一來就可以在相對低的成本下,達成控制執行時間的目的,就不用擔心寫錯的實作,或者 LLM(大型語言模型)生成錯誤的實作,造成佔用資源的狀況。

記憶體限制

既然執行時間會卡住,那麼記憶體也一樣會有相同的問題,像是我們錯誤的載入大量的資料,或者透過某種方式製造問題。

1sandbox = Kobako::Sandbox.new
2sandbox.eval("'S' * 2**1024").value  # => ?

在 Wasmtime 本身就有 ResourceLimiter 這個特性可用,每次需要增加記憶體配置時,就會先進行詢問,被允許後才能使用,因此 Kobako 加入了這樣的機制。

 1struct MemoryLimiter { limit: usize, baseline: usize, active: bool }
 2
 3impl ResourceLimiter for MemoryLimiter {
 4    fn memory_growing(&mut self, _cur: usize, desired: usize, _max: Option<usize>)
 5        -> Result<bool>
 6    {
 7        if !self.active { return Ok(true); }
 8        if desired - self.baseline > self.limit {
 9            bail!("memory usage exceeded memory_limit");
10        }
11        Ok(true)
12    }
13    fn table_growing(&mut self, ..) -> Result<bool> { Ok(true) }
14}
15
16
17store.limiter(|state| &mut state.limiter);
18
19state.limiter.baseline = memory.data_size(&store);
20state.limiter.active = true;
21let result = export.call(&mut store, params);
22state.limiter.active = false;

因為 mruby 本身也會佔用記憶體,假設 mruby 執行本身需要 1 MB 的記憶體,直接用「增加多少記憶體」就會讓上限 1 MB 瞬間被吃滿,因此會有 baselineactive 的設計存在。

當我們初始化完畢後,會先確認目前 mruby 使用掉多少記憶體,把它記錄成 baseline 然後再把 active 標記成啟用,然後讓他繼續往下執行,此時記錄到的 目標記憶體 - 基準 就會是執行腳本實際佔用的記憶體,那麼我們就可以確保呼叫時能有完整的記憶體使用。

雖然執行時間跟記憶體都不是「精確計算」,而是用粗略估計的,但光是這樣就可以阻止錯誤或者惡意的使用,確保宿主不會因為沙盒中異常的實作受到影響,一定程度可以減緩執行不安全的 Ruby 腳本的傷害。

當然,這也一定程度能對 ReDoS(正規表示式導致的阻斷攻擊)這類攻擊有防範,但攻擊的手法仍有很多種,對資源進行限制只是其中一種手段,為了確保能更安全的執行,我們還有更多不同類型的取捨,會陸續地被加入到 Kobako 裡面。