A lexicon-driven AppView for ATProto.
Something went wrong. Try again.
Rust
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449use std::collections::HashMap;use std::sync::Arc;use wasmtime::{Linker, Memory, TypedFunc};
/// State stored in wasmtime's Store during plugin executionpub struct PluginState { pub plugin_id: String, pub scope: String, pub secrets: HashMap<String, String>, pub config: serde_json::Value, pub db: Option<sqlx::AnyPool>, pub db_backend: crate::db::DatabaseBackend, pub http_client: reqwest::Client, pub lexicons: Arc<crate::lexicon::LexiconRegistry>, pub usage: super::ResourceUsage, pub memory: Option<Memory>, pub alloc: Option<TypedFunc<u32, u32>>, pub dealloc: Option<TypedFunc<(u32, u32), ()>>,}
/// Check that a memory access is within boundsfn check_bounds(offset: usize, length: usize, mem_size: usize) -> Result<(usize, usize), ()> { if length == 0 { return Ok((offset, offset)); } let end = offset.checked_add(length).ok_or(())?; if end > mem_size { return Err(()); } Ok((offset, end))}
/// Register all host functions with the linkerpub fn register_host_functions(linker: &mut Linker<PluginState>) -> Result<(), wasmtime::Error> { // Sync functions linker.func_wrap("env", "host_log", host_log)?;
// host_get_secret must be async because it calls alloc on an async store linker.func_wrap_async( "env", "host_get_secret", |mut caller: wasmtime::Caller<'_, PluginState>, (name_ptr, name_len): (i32, i32)| { Box::new(async move { host_get_secret_impl(&mut caller, name_ptr, name_len).await }) }, )?;
// Async functions - HTTP linker.func_wrap_async( "env", "host_http_request", |mut caller: wasmtime::Caller<'_, PluginState>, (req_ptr, req_len): (i32, i32)| { Box::new(async move { host_http_request_impl(&mut caller, req_ptr, req_len).await }) }, )?;
// Async functions - KV linker.func_wrap_async( "env", "host_kv_get", |mut caller: wasmtime::Caller<'_, PluginState>, (key_ptr, key_len): (i32, i32)| { Box::new(async move { host_kv_get_impl(&mut caller, key_ptr, key_len).await }) }, )?;
linker.func_wrap_async( "env", "host_kv_set", |mut caller: wasmtime::Caller<'_, PluginState>, (key_ptr, key_len, val_ptr, val_len, ttl): (i32, i32, i32, i32, i32)| { Box::new(async move { host_kv_set_impl(&mut caller, key_ptr, key_len, val_ptr, val_len, ttl).await }) }, )?;
linker.func_wrap_async( "env", "host_kv_delete", |mut caller: wasmtime::Caller<'_, PluginState>, (key_ptr, key_len): (i32, i32)| { Box::new(async move { host_kv_delete_impl(&mut caller, key_ptr, key_len).await }) }, )?;
// Async functions - Record lookup linker.func_wrap_async( "env", "host_lookup_record", |mut caller: wasmtime::Caller<'_, PluginState>, (req_ptr, req_len): (i32, i32)| { Box::new(async move { host_lookup_record_impl(&mut caller, req_ptr, req_len).await }) }, )?;
Ok(())}
/// Read a string from guest memoryfn read_guest_string( caller: &wasmtime::Caller<'_, PluginState>, ptr: i32, len: i32,) -> Option<String> { let memory = caller.data().memory?; let mem_data = memory.data(caller); let (start, end) = check_bounds(ptr as usize, len as usize, mem_data.len()).ok()?; std::str::from_utf8(&mem_data[start..end]) .ok() .map(String::from)}
/// Read raw bytes from guest memoryfn read_guest_bytes( caller: &wasmtime::Caller<'_, PluginState>, ptr: i32, len: i32,) -> Option<Vec<u8>> { let memory = caller.data().memory?; let mem_data = memory.data(caller); let (start, end) = check_bounds(ptr as usize, len as usize, mem_data.len()).ok()?; Some(mem_data[start..end].to_vec())}
/// Write response data to guest memory, returning packed (ptr << 32) | lenasync fn write_guest_response(caller: &mut wasmtime::Caller<'_, PluginState>, data: &[u8]) -> i64 { let memory = match caller.data().memory { Some(m) => m, None => return 0, }; let alloc = match &caller.data().alloc { Some(a) => a.clone(), None => return 0, };
let len = data.len() as u32; let ptr = match alloc.call_async(&mut *caller, len).await { Ok(p) if p != 0 => p, _ => return 0, };
let mem_data = memory.data_mut(caller); if check_bounds(ptr as usize, len as usize, mem_data.len()).is_err() { return 0; }
mem_data[ptr as usize..(ptr as usize + len as usize)].copy_from_slice(data); ((ptr as i64) << 32) | (len as i64)}
/// Host function: log a message from the pluginfn host_log( caller: wasmtime::Caller<'_, PluginState>, level_ptr: i32, level_len: i32, msg_ptr: i32, msg_len: i32,) { let memory = match caller.data().memory { Some(m) => m, None => return, };
let mem_data = memory.data(&caller); let mem_size = mem_data.len();
let (level_start, level_end) = match check_bounds(level_ptr as usize, level_len as usize, mem_size) { Ok(bounds) => bounds, Err(_) => return, };
let (msg_start, msg_end) = match check_bounds(msg_ptr as usize, msg_len as usize, mem_size) { Ok(bounds) => bounds, Err(_) => return, };
let level = std::str::from_utf8(&mem_data[level_start..level_end]).unwrap_or("info"); let msg = std::str::from_utf8(&mem_data[msg_start..msg_end]).unwrap_or("");
let plugin_id = caller.data().plugin_id.clone(); let db = caller.data().db.clone(); let db_backend = caller.data().db_backend; let log_level: super::LogLevel = level.parse().unwrap_or_default(); super::log(&plugin_id, log_level, msg, db, db_backend);}
/// Host function: get a secret value by name/// Returns a packed i64: (ptr << 32) | len, or 0 on error////// The response is JSON-encoded as `{"ok": "<value>"}` so plugins can/// deserialize it with the same `Response<String>` envelope they use for/// other host calls.async fn host_get_secret_impl( caller: &mut wasmtime::Caller<'_, PluginState>, name_ptr: i32, name_len: i32,) -> i64 { let name = match read_guest_string(caller, name_ptr, name_len) { Some(n) => n, None => return 0, };
let value = match caller.data().secrets.get(&name) { Some(v) => v.clone(), None => return 0, };
let response_bytes = serde_json::to_vec(&serde_json::json!({"ok": value})).unwrap_or_default(); write_guest_response(caller, &response_bytes).await}
// ============================================================================// Async host function implementations// ============================================================================
/// Build a HostContext from PluginState, requires db to be presentfn build_host_context(state: &PluginState) -> Option<super::HostContext> { let db = state.db.clone()?; Some(super::HostContext { plugin_id: state.plugin_id.clone(), scope: state.scope.clone(), secrets: state.secrets.clone(), config: state.config.clone(), db, db_backend: state.db_backend, http_client: state.http_client.clone(), lexicons: state.lexicons.clone(), })}
/// Host function: make an HTTP requestasync fn host_http_request_impl( caller: &mut wasmtime::Caller<'_, PluginState>, req_ptr: i32, req_len: i32,) -> i64 { let req_bytes = match read_guest_bytes(caller, req_ptr, req_len) { Some(b) => b, None => { tracing::error!( "host_http_request: failed to read guest memory (ptr={req_ptr}, len={req_len})" ); return 0; } };
let request: super::HttpRequest = match serde_json::from_slice(&req_bytes) { Ok(r) => r, Err(e) => { tracing::error!("host_http_request: failed to parse request JSON: {e}"); return 0; } };
let url = request.url.clone(); let method = request.method.clone();
let ctx = match build_host_context(caller.data()) { Some(c) => c, None => { tracing::error!("host_http_request: failed to build host context (db missing?)"); return 0; } };
let result = { let usage = &mut caller.data_mut().usage; super::http_request(&ctx, usage, request).await };
let response_bytes = match result { Ok(resp) => serde_json::to_vec(&serde_json::json!({"ok": resp})).unwrap_or_default(), Err(e) => { tracing::warn!("host_http_request: HTTP {method} {url} failed: {e}"); serde_json::to_vec(&serde_json::json!({ "error": {"code": "HTTP_ERROR", "message": e.to_string(), "retryable": false} })) .unwrap_or_default() } };
let packed = write_guest_response(caller, &response_bytes).await; if packed == 0 { tracing::error!( "host_http_request: write_guest_response returned 0 for {} {} (response_len={})", method, url, response_bytes.len() ); } packed}
/// Host function: get a value from KV storeasync fn host_kv_get_impl( caller: &mut wasmtime::Caller<'_, PluginState>, key_ptr: i32, key_len: i32,) -> i64 { let key = match read_guest_string(caller, key_ptr, key_len) { Some(k) => k, None => return 0, };
let ctx = match build_host_context(caller.data()) { Some(c) => c, None => return 0, };
let result = super::kv_get(&ctx, &key).await;
let response_bytes = match result { Ok(Some(value)) => { serde_json::to_vec(&serde_json::json!({"ok": value})).unwrap_or_default() } Ok(None) => return 0, Err(e) => serde_json::to_vec(&serde_json::json!({ "error": {"code": "KV_ERROR", "message": e.to_string(), "retryable": false} })) .unwrap_or_default(), };
write_guest_response(caller, &response_bytes).await}
/// Host function: set a value in KV storeasync fn host_kv_set_impl( caller: &mut wasmtime::Caller<'_, PluginState>, key_ptr: i32, key_len: i32, val_ptr: i32, val_len: i32, ttl: i32,) -> i32 { let key = match read_guest_string(caller, key_ptr, key_len) { Some(k) => k, None => return -1, }; let value = match read_guest_bytes(caller, val_ptr, val_len) { Some(v) => v, None => return -1, };
let ttl_secs = if ttl > 0 { Some(ttl as u32) } else { None };
let ctx = match build_host_context(caller.data()) { Some(c) => c, None => return -1, };
let usage = &mut caller.data_mut().usage; match super::kv_set(&ctx, usage, &key, value, ttl_secs).await { Ok(()) => 0, Err(_) => -1, }}
/// Host function: delete a value from KV storeasync fn host_kv_delete_impl( caller: &mut wasmtime::Caller<'_, PluginState>, key_ptr: i32, key_len: i32,) -> i32 { let key = match read_guest_string(caller, key_ptr, key_len) { Some(k) => k, None => return -1, };
let ctx = match build_host_context(caller.data()) { Some(c) => c, None => return -1, };
match super::kv_delete(&ctx, &key).await { Ok(()) => 0, Err(_) => -1, }}
/// Host function: look up an AT Protocol recordasync fn host_lookup_record_impl( caller: &mut wasmtime::Caller<'_, PluginState>, req_ptr: i32, req_len: i32,) -> i64 { let req_bytes = match read_guest_bytes(caller, req_ptr, req_len) { Some(b) => b, None => return 0, };
let request: super::LookupRequest = match serde_json::from_slice(&req_bytes) { Ok(r) => r, Err(_) => return 0, };
let ctx = match build_host_context(caller.data()) { Some(c) => c, None => return 0, };
let result = super::lookup_record_by_request(&ctx, request).await;
let response_bytes = match result { Ok(record) => serde_json::to_vec(&serde_json::json!({"ok": record})).unwrap_or_default(), Err(e) => serde_json::to_vec(&serde_json::json!({ "error": {"code": "LOOKUP_ERROR", "message": e.to_string(), "retryable": false} })) .unwrap_or_default(), };
write_guest_response(caller, &response_bytes).await}
#[cfg(test)]mod tests { use super::*;
#[test] fn test_plugin_state_fields_exist() { fn _check_fields(state: &PluginState) { let _ = &state.plugin_id; let _ = &state.scope; let _ = &state.secrets; let _ = &state.config; let _ = &state.usage; let _ = &state.memory; let _ = &state.alloc; let _ = &state.dealloc; } }
#[test] fn test_pack_ptr_len() { let ptr: u32 = 0x1000; let len: u32 = 0x0100; let packed: i64 = ((ptr as i64) << 32) | (len as i64); let unpacked_ptr = (packed >> 32) as u32; let unpacked_len = (packed & 0xFFFFFFFF) as u32; assert_eq!(unpacked_ptr, ptr); assert_eq!(unpacked_len, len); }
#[test] fn test_bounds_check_helper() { assert!(check_bounds(0, 10, 100).is_ok()); assert!(check_bounds(90, 10, 100).is_ok()); assert!(check_bounds(91, 10, 100).is_err()); assert!(check_bounds(0, 0, 100).is_ok()); }}