use std::mem::align_of; use std::{borrow::Cow, process::Command}; use chrono::Offset; use wasmtime::{Caller, Memory}; use crate::{ScriptManager, WasmScriptState}; use super::LIBRARY_NAME; pub fn _get_memory<'a>(caller: &'a mut Caller) -> (Memory, &'a mut [u8]) { let memory = caller.get_export("memory").unwrap().into_memory().unwrap(); let mem_ptr = memory.data_mut(caller); (memory, mem_ptr) } fn _get_string<'a>( caller: &'a mut Caller, offset: u32, length: u32, ) -> Option> { let (_, mem_ptr) = _get_memory(caller); if (offset as usize + length as usize) > mem_ptr.len() { return None; } Some(String::from_utf8_lossy( &mem_ptr[offset as usize..length as usize], )) } fn _log(caller: &mut Caller, level: log::Level, message: u32, length: u32) -> u32 { let Some(message) = _get_string(caller, message, length) else { return 4; }; log::log!(level, "{}", message); 0 } pub fn _alloc<'a>( mut caller: &'a mut Caller, size: u32, align: u32, ) -> Option<(u32, &'a mut [u8])> { // find bumper location in memory let bump_offset = caller.data().bump_pointer; let (memory, memory_ptr) = _get_memory(caller); let bump_ptr = &memory_ptr[bump_offset as usize..]; // calculate pointer to the new allocation let offset = bump_ptr.as_ptr().align_offset(align as usize); let new_offset = bump_offset + offset as u32; // update bump pointer let bump_offset = &mut caller.data_mut().bump_pointer; *bump_offset = new_offset + size; // grow memory if necessary let bump_offset = *bump_offset; if memory.data_size(&caller) < bump_offset as usize { let delta = (size + offset as u32) / (64 * 1024) + 1; if memory.grow(&mut caller, delta as u64).is_err() { return None; } } let new_ptr = &mut _get_memory(caller).1[new_offset as usize..size as usize]; Some((new_offset, new_ptr)) } fn allocate(mut caller: Caller, size: u32, align: u32) -> u32 { match _alloc(&mut caller, size, align) { Some((ptr, _)) => ptr, None => 0, } } fn free_sized(mut caller: Caller<'_, WasmScriptState>, ptr: u32, size: u32) -> u32 { let bump_offset = &mut caller.data_mut().bump_pointer; if ptr + size == *bump_offset { *bump_offset = ptr; } 0 } fn free_aligned_sized( caller: Caller<'_, WasmScriptState>, ptr: u32, _alignment: u32, size: u32, ) -> u32 { free_sized(caller, ptr, size) } fn env_var(mut caller: Caller<'_, WasmScriptState>, name: u32, length: u32) -> u32 { let Some(name) = _get_string(&mut caller, name, length) else { return 4; }; let Ok(value) = std::env::var(name.to_string()) else { return 1; }; let Some((offset, ptr)) = _alloc(&mut caller, value.len() as u32, align_of::() as u32) else { return 0; }; ptr.clone_from_slice(value.as_bytes()); offset } fn random_number() -> u64 { rand::random() } fn current_time_utc() -> (i64, u32) { let duration = chrono::Utc::now() - chrono::DateTime::::MIN_UTC; ( duration.num_seconds(), duration.to_std().unwrap().subsec_nanos(), ) } fn local_offset() -> i32 { chrono::Local::now().offset().fix().local_minus_utc() } fn log(mut caller: Caller<'_, WasmScriptState>, message: u32, length: u32) { _log(&mut caller, log::Level::Info, message, length); } fn log_warn(mut caller: Caller<'_, WasmScriptState>, message: u32, length: u32) { _log(&mut caller, log::Level::Warn, message, length); } fn log_error(mut caller: Caller<'_, WasmScriptState>, message: u32, length: u32) { _log(&mut caller, log::Level::Error, message, length); } pub fn library(manager: &mut ScriptManager) -> Option<()> { manager.add_library_function(LIBRARY_NAME, "aligned_alloc", allocate)?; manager.add_library_function(LIBRARY_NAME, "free_sized", free_sized)?; manager.add_library_function(LIBRARY_NAME, "free_aligned_sized", free_aligned_sized)?; manager.add_library_function(LIBRARY_NAME, "env_var", env_var)?; manager.add_library_function(LIBRARY_NAME, "random_number", random_number)?; manager.add_library_function(LIBRARY_NAME, "current_time_utc", current_time_utc)?; manager.add_library_function(LIBRARY_NAME, "local_offset", local_offset)?; manager.add_library_function(LIBRARY_NAME, "log", log)?; manager.add_library_function(LIBRARY_NAME, "log_warn", log_warn)?; manager.add_library_function(LIBRARY_NAME, "log_error", log_error)?; Some(()) }