cred: isolate scripts in WASI and move RSA signing to AWS-LC
This commit is contained in:
@@ -0,0 +1,275 @@
|
||||
//! Credential scripts run in a fresh, memory-limited WASI instance. Only two calls run at once.
|
||||
use super::script::{Error, ErrorKind, Limits, Script};
|
||||
use std::{
|
||||
sync::{Arc, Condvar, Mutex, OnceLock, mpsc},
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
use wasmtime::{
|
||||
Config, Engine, Store, StoreLimits, StoreLimitsBuilder, UpdateDeadline,
|
||||
component::{Component, Linker, ResourceTable},
|
||||
};
|
||||
use wasmtime_wasi::{
|
||||
WasiCtx, WasiCtxBuilder, WasiCtxView, WasiView,
|
||||
clocks::HostWallClock,
|
||||
p2::{
|
||||
bindings::sync::CommandPre,
|
||||
pipe::{MemoryInputPipe, MemoryOutputPipe},
|
||||
},
|
||||
random::Deterministic,
|
||||
};
|
||||
use zeroize::Zeroizing;
|
||||
|
||||
struct State {
|
||||
wasi: WasiCtx,
|
||||
table: ResourceTable,
|
||||
limits: StoreLimits,
|
||||
}
|
||||
impl WasiView for State {
|
||||
fn ctx(&mut self) -> WasiCtxView<'_> {
|
||||
WasiCtxView {
|
||||
ctx: &mut self.wasi,
|
||||
table: &mut self.table,
|
||||
}
|
||||
}
|
||||
}
|
||||
struct FixedClock;
|
||||
impl HostWallClock for FixedClock {
|
||||
fn resolution(&self) -> Duration {
|
||||
Duration::from_millis(1)
|
||||
}
|
||||
fn now(&self) -> Duration {
|
||||
Duration::ZERO
|
||||
}
|
||||
}
|
||||
struct Compiled {
|
||||
engine: Engine,
|
||||
pre: CommandPre<State>,
|
||||
active: Mutex<usize>,
|
||||
available: Condvar,
|
||||
}
|
||||
struct Slot<'a>(&'a Compiled);
|
||||
impl Drop for Slot<'_> {
|
||||
fn drop(&mut self) {
|
||||
*self.0.active.lock().unwrap() -= 1;
|
||||
self.0.available.notify_one();
|
||||
}
|
||||
}
|
||||
impl Compiled {
|
||||
fn new() -> Result<Arc<Self>, Error> {
|
||||
let mut config = Config::new();
|
||||
config.epoch_interruption(true);
|
||||
// The cache also keeps warm test runs from compiling the component again.
|
||||
config.cache(wasmtime::Cache::new(Default::default()).ok());
|
||||
let engine = Engine::new(&config).map_err(|_| failure(ErrorKind::Engine))?;
|
||||
let component = Component::new(&engine, include_bytes!("../../cred-script/script.wasm"))
|
||||
.map_err(|_| failure(ErrorKind::Engine))?;
|
||||
let mut linker = Linker::new(&engine);
|
||||
wasmtime_wasi::p2::add_to_linker_sync(&mut linker)
|
||||
.map_err(|_| failure(ErrorKind::Engine))?;
|
||||
let pre = CommandPre::new(
|
||||
linker
|
||||
.instantiate_pre(&component)
|
||||
.map_err(|_| failure(ErrorKind::Engine))?,
|
||||
)
|
||||
.map_err(|_| failure(ErrorKind::Engine))?;
|
||||
Ok(Arc::new(Self {
|
||||
engine,
|
||||
pre,
|
||||
active: Mutex::new(0),
|
||||
available: Condvar::new(),
|
||||
}))
|
||||
}
|
||||
fn acquire(&self, deadline: Instant) -> Result<Slot<'_>, Error> {
|
||||
let mut active = self.active.lock().map_err(|_| failure(ErrorKind::Engine))?;
|
||||
while *active >= 2 {
|
||||
let remaining = deadline
|
||||
.checked_duration_since(Instant::now())
|
||||
.ok_or_else(|| failure(ErrorKind::Limit))?;
|
||||
(active, _) = self
|
||||
.available
|
||||
.wait_timeout(active, remaining)
|
||||
.map_err(|_| failure(ErrorKind::Engine))?;
|
||||
}
|
||||
if Instant::now() >= deadline {
|
||||
return Err(failure(ErrorKind::Limit));
|
||||
}
|
||||
*active += 1;
|
||||
Ok(Slot(self))
|
||||
}
|
||||
}
|
||||
static COMPILED: OnceLock<Result<Arc<Compiled>, Error>> = OnceLock::new();
|
||||
|
||||
/// Construct once at startup. Compilation is cached; each evaluation gets a fresh Store.
|
||||
#[derive(Clone)]
|
||||
pub struct QuickJs {
|
||||
compiled: Result<Arc<Compiled>, Error>,
|
||||
}
|
||||
impl Default for QuickJs {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
compiled: COMPILED.get_or_init(Compiled::new).clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
fn failure(kind: ErrorKind) -> Error {
|
||||
Error {
|
||||
kind,
|
||||
message: match kind {
|
||||
ErrorKind::Limit => "script time, memory or stack limit exceeded",
|
||||
ErrorKind::Syntax => "invalid script or JSON",
|
||||
ErrorKind::Throw => "script threw or rejected",
|
||||
ErrorKind::NonJson => "result is not JSON",
|
||||
ErrorKind::Engine => "script engine failure",
|
||||
}
|
||||
.into(),
|
||||
}
|
||||
}
|
||||
impl Script for QuickJs {
|
||||
fn evaluate(
|
||||
&self,
|
||||
source: &str,
|
||||
input: &str,
|
||||
export: Option<&str>,
|
||||
limits: Limits,
|
||||
) -> Result<String, Error> {
|
||||
let compiled = self.compiled.as_ref().map_err(Clone::clone)?;
|
||||
let time = limits.max_time_ms.unwrap_or(100).clamp(1, 1000);
|
||||
let heap = limits
|
||||
.max_memory_bytes
|
||||
.unwrap_or(8 * 1024 * 1024)
|
||||
.clamp(256 * 1024, 32 * 1024 * 1024);
|
||||
let deadline = Instant::now() + Duration::from_millis(time);
|
||||
let _slot = compiled.acquire(deadline)?;
|
||||
if source
|
||||
.len()
|
||||
.saturating_add(input.len())
|
||||
.saturating_add(export.map_or(0, str::len))
|
||||
> heap
|
||||
{
|
||||
return Err(failure(ErrorKind::Limit));
|
||||
}
|
||||
// WASI's synchronous adapter must not enter a caller's Tokio runtime.
|
||||
std::thread::scope(|scope| {
|
||||
let worker = std::thread::Builder::new()
|
||||
.name("cred-script".into())
|
||||
.spawn_scoped(scope, move || {
|
||||
#[derive(serde::Serialize)]
|
||||
struct Request<'a> {
|
||||
source: &'a str,
|
||||
input: &'a str,
|
||||
export: Option<&'a str>,
|
||||
max_time_ms: u64,
|
||||
max_memory_bytes: usize,
|
||||
}
|
||||
let wire = Zeroizing::new(
|
||||
serde_json::to_vec(&Request {
|
||||
source,
|
||||
input,
|
||||
export,
|
||||
max_time_ms: time,
|
||||
max_memory_bytes: heap,
|
||||
})
|
||||
.map_err(|_| failure(ErrorKind::Engine))?,
|
||||
);
|
||||
let memory = heap + 8 * 1024 * 1024;
|
||||
let output = MemoryOutputPipe::new(memory);
|
||||
let mut wasi = WasiCtxBuilder::new();
|
||||
wasi.stdin(MemoryInputPipe::new(bytes::Bytes::from_owner(wire)))
|
||||
.stdout(output.clone())
|
||||
.wall_clock(FixedClock)
|
||||
.secure_random(Deterministic::new(vec![0]))
|
||||
.insecure_random(Deterministic::new(vec![0]))
|
||||
.insecure_random_seed(0)
|
||||
.allow_tcp(false)
|
||||
.allow_udp(false)
|
||||
.allow_ip_name_lookup(false);
|
||||
let mut store = Store::new(
|
||||
&compiled.engine,
|
||||
State {
|
||||
wasi: wasi.build(),
|
||||
table: ResourceTable::new(),
|
||||
limits: StoreLimitsBuilder::new()
|
||||
.memory_size(memory)
|
||||
.memories(1)
|
||||
.trap_on_grow_failure(true)
|
||||
.build(),
|
||||
},
|
||||
);
|
||||
store.limiter(|state| &mut state.limits);
|
||||
store.epoch_deadline_callback(move |_| {
|
||||
if Instant::now() >= deadline {
|
||||
Err(wasmtime::Trap::Interrupt.into())
|
||||
} else {
|
||||
Ok(UpdateDeadline::Continue(1))
|
||||
}
|
||||
});
|
||||
store.set_epoch_deadline(1);
|
||||
let (done, wait) = mpsc::channel();
|
||||
let engine = compiled.engine.clone();
|
||||
let watchdog = std::thread::Builder::new()
|
||||
.name("cred-script-deadline".into())
|
||||
.spawn(move || {
|
||||
if wait
|
||||
.recv_timeout(deadline.saturating_duration_since(Instant::now()))
|
||||
.is_err()
|
||||
{
|
||||
engine.increment_epoch();
|
||||
}
|
||||
})
|
||||
.map_err(|_| failure(ErrorKind::Engine))?;
|
||||
let result = compiled
|
||||
.pre
|
||||
.instantiate(&mut store)
|
||||
.and_then(|command| command.wasi_cli_run().call_run(&mut store));
|
||||
let _ = done.send(());
|
||||
let _ = watchdog.join();
|
||||
if Instant::now() >= deadline || !matches!(result, Ok(Ok(()))) {
|
||||
return Err(failure(ErrorKind::Limit));
|
||||
}
|
||||
#[derive(serde::Deserialize)]
|
||||
struct Response {
|
||||
result: Option<String>,
|
||||
kind: Option<String>,
|
||||
}
|
||||
let response: Response = serde_json::from_slice(&output.contents())
|
||||
.map_err(|_| failure(ErrorKind::Engine))?;
|
||||
if Instant::now() >= deadline {
|
||||
return Err(failure(ErrorKind::Limit));
|
||||
}
|
||||
if let Some(result) = response.result {
|
||||
return Ok(result);
|
||||
}
|
||||
Err(failure(match response.kind.as_deref() {
|
||||
Some("Limit") => ErrorKind::Limit,
|
||||
Some("Syntax") => ErrorKind::Syntax,
|
||||
Some("Throw") => ErrorKind::Throw,
|
||||
Some("NonJson") => ErrorKind::NonJson,
|
||||
_ => ErrorKind::Engine,
|
||||
}))
|
||||
})
|
||||
.map_err(|_| failure(ErrorKind::Engine))?;
|
||||
worker.join().map_err(|_| failure(ErrorKind::Engine))?
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn concurrency_is_bounded_and_queue_uses_the_deadline() {
|
||||
let runtime = QuickJs::default();
|
||||
let compiled = runtime.compiled.as_ref().unwrap();
|
||||
let deadline = Instant::now() + Duration::from_secs(1);
|
||||
let first = compiled.acquire(deadline).unwrap();
|
||||
let second = compiled.acquire(deadline).unwrap();
|
||||
assert!(
|
||||
matches!(compiled.acquire(Instant::now()), Err(error) if error.kind == ErrorKind::Limit)
|
||||
);
|
||||
drop(first);
|
||||
let third = compiled.acquire(deadline).unwrap();
|
||||
assert_eq!(*compiled.active.lock().unwrap(), 2);
|
||||
drop((second, third));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user