use crate::type_schema::{TypeSchema, TypeSchemaRegistry};
use shape_value::KindedSlot;
use std::collections::HashMap;
use std::ffi::c_void;
use std::sync::Arc;
#[derive(Clone, Copy)]
pub struct RawCallableInvoker {
pub ctx: *mut c_void,
pub invoke: unsafe fn(*mut c_void, &u64, &[u64]) -> Result<u64, String>,
}
impl RawCallableInvoker {
pub unsafe fn call(&self, callable: &u64, args: &[u64]) -> Result<u64, String> {
unsafe { (self.invoke)(self.ctx, callable, args) }
}
}
#[derive(Debug, Clone)]
pub struct FrameInfo {
pub function_id: Option<u16>,
pub function_name: String,
pub blob_hash: Option<[u8; 32]>,
pub local_ip: usize,
pub locals: Vec<KindedSlot>,
pub upvalues: Option<Vec<KindedSlot>>,
pub args: Vec<KindedSlot>,
}
pub trait VmStateAccessor: Send + Sync {
fn current_frame(&self) -> Option<FrameInfo>;
fn all_frames(&self) -> Vec<FrameInfo>;
fn caller_frame(&self) -> Option<FrameInfo>;
fn current_args(&self) -> Vec<KindedSlot>;
fn current_locals(&self) -> Vec<(String, KindedSlot)>;
fn module_bindings(&self) -> Vec<(String, KindedSlot)>;
fn instruction_count(&self) -> usize {
0
}
}
pub struct ModuleContext<'a> {
pub schemas: &'a TypeSchemaRegistry,
pub invoke_callable:
Option<&'a dyn Fn(&KindedSlot, &[KindedSlot]) -> Result<KindedSlot, String>>,
pub raw_invoker: Option<RawCallableInvoker>,
pub function_hashes: Option<&'a [Option<[u8; 32]>]>,
pub vm_state: Option<&'a dyn VmStateAccessor>,
pub granted_permissions: Option<shape_abi_v1::PermissionSet>,
pub scope_constraints: Option<shape_abi_v1::ScopeConstraints>,
pub set_pending_resume: Option<&'a dyn Fn(KindedSlot)>,
pub set_pending_frame_resume: Option<&'a dyn Fn(usize, Vec<KindedSlot>)>,
}
pub fn check_permission(
ctx: &ModuleContext,
permission: shape_abi_v1::Permission,
) -> Result<(), String> {
if let Some(ref granted) = ctx.granted_permissions {
if !granted.contains(&permission) {
return Err(format!(
"Permission denied: {} ({})",
permission.description(),
permission.name()
));
}
}
Ok(())
}
pub fn check_fs_permission(
ctx: &ModuleContext,
permission: shape_abi_v1::Permission,
path: &str,
) -> Result<(), String> {
check_permission(ctx, permission)?;
if let Some(ref constraints) = ctx.scope_constraints {
if !constraints.allowed_paths.is_empty() {
let target = std::path::Path::new(path);
let allowed = constraints.allowed_paths.iter().any(|pattern| {
let pattern = pattern.trim_end_matches("**").trim_end_matches('*');
let prefix = std::path::Path::new(pattern.trim_end_matches('/'));
target.starts_with(prefix)
});
if !allowed {
return Err(format!(
"Scope constraint denied: path '{}' is not in allowed paths",
path
));
}
}
}
Ok(())
}
pub fn check_net_permission(
ctx: &ModuleContext,
permission: shape_abi_v1::Permission,
address: &str,
) -> Result<(), String> {
check_permission(ctx, permission)?;
if let Some(ref constraints) = ctx.scope_constraints {
if !constraints.allowed_hosts.is_empty() {
let target_host = address.split(':').next().unwrap_or(address);
let allowed = constraints.allowed_hosts.iter().any(|pattern| {
let pattern_host = pattern.split(':').next().unwrap_or(pattern);
if let Some(suffix) = pattern_host.strip_prefix("*.") {
target_host.ends_with(suffix) && target_host.len() > suffix.len()
} else {
target_host == pattern_host
}
});
if !allowed {
return Err(format!(
"Scope constraint denied: address '{}' is not in allowed hosts",
address
));
}
}
}
Ok(())
}
pub type ModuleFn = Arc<
dyn for<'ctx> Fn(&[KindedSlot], &ModuleContext<'ctx>) -> Result<KindedSlot, String>
+ Send
+ Sync,
>;
#[derive(Clone)]
pub enum ModuleFnEntry {
Typed(crate::typed_module_exports::TypedModuleFunction),
TypedAsync(crate::typed_module_exports::TypedModuleAsyncFunction),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ModuleExportVisibility {
Public,
ComptimeOnly,
Internal,
}
impl Default for ModuleExportVisibility {
fn default() -> Self {
Self::Public
}
}
#[derive(Debug, Clone)]
pub struct ModuleParam {
pub name: String,
pub type_name: String,
pub required: bool,
pub description: String,
pub default_snippet: Option<String>,
pub allowed_values: Option<Vec<String>>,
pub nested_params: Option<Vec<ModuleParam>>,
}
impl Default for ModuleParam {
fn default() -> Self {
Self {
name: String::new(),
type_name: "any".to_string(),
required: false,
description: String::new(),
default_snippet: None,
allowed_values: None,
nested_params: None,
}
}
}
#[derive(Debug, Clone)]
pub struct ModuleFunction {
pub description: String,
pub params: Vec<ModuleParam>,
pub return_type: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ModuleArtifact {
pub module_path: String,
pub source: Option<String>,
pub compiled: Option<Vec<u8>>,
}
#[derive(Clone)]
pub struct ModuleExports {
pub name: String,
pub description: String,
pub schemas: HashMap<String, ModuleFunction>,
pub export_visibility: HashMap<String, ModuleExportVisibility>,
pub shape_sources: Vec<(String, String)>,
pub module_artifacts: Vec<ModuleArtifact>,
pub method_intrinsics: HashMap<String, HashMap<String, ModuleFn>>,
pub type_schemas: Vec<TypeSchema>,
pub typed_exports: crate::typed_module_exports::TypedModuleExports,
}
impl ModuleExports {
pub fn new(name: impl Into<String>) -> Self {
Self {
name: name.into(),
description: String::new(),
schemas: HashMap::new(),
export_visibility: HashMap::new(),
shape_sources: Vec::new(),
module_artifacts: Vec::new(),
method_intrinsics: HashMap::new(),
type_schemas: Vec::new(),
typed_exports: crate::typed_module_exports::TypedModuleExports::new(),
}
}
pub fn typed_exports_mut(
&mut self,
) -> &mut crate::typed_module_exports::TypedModuleExports {
&mut self.typed_exports
}
pub fn typed_exports(&self) -> &crate::typed_module_exports::TypedModuleExports {
&self.typed_exports
}
pub fn add_schema_only(
&mut self,
name: impl Into<String>,
schema: ModuleFunction,
) -> &mut Self {
let name = name.into();
self.schemas.insert(name.clone(), schema);
self.export_visibility.entry(name).or_default();
self
}
pub fn set_export_visibility(
&mut self,
name: impl Into<String>,
visibility: ModuleExportVisibility,
) -> &mut Self {
self.export_visibility.insert(name.into(), visibility);
self
}
pub fn export_visibility(&self, name: &str) -> ModuleExportVisibility {
self.export_visibility
.get(name)
.copied()
.unwrap_or_default()
}
pub fn is_export_available(&self, name: &str, comptime_mode: bool) -> bool {
match self.export_visibility(name) {
ModuleExportVisibility::Public => true,
ModuleExportVisibility::ComptimeOnly => comptime_mode,
ModuleExportVisibility::Internal => true,
}
}
pub fn is_export_public_surface(&self, name: &str, comptime_mode: bool) -> bool {
match self.export_visibility(name) {
ModuleExportVisibility::Public => true,
ModuleExportVisibility::ComptimeOnly => comptime_mode,
ModuleExportVisibility::Internal => false,
}
}
pub fn export_names_available(&self, comptime_mode: bool) -> Vec<&str> {
self.export_names()
.into_iter()
.filter(|name| self.is_export_available(name, comptime_mode))
.collect()
}
pub fn export_names_public_surface(&self, comptime_mode: bool) -> Vec<&str> {
self.export_names()
.into_iter()
.filter(|name| self.is_export_public_surface(name, comptime_mode))
.collect()
}
pub fn add_shape_source(&mut self, filename: &str, source: &str) -> &mut Self {
self.module_artifacts.push(ModuleArtifact {
module_path: filename.to_string(),
source: Some(source.to_string()),
compiled: None,
});
self.shape_sources
.push((filename.to_string(), source.to_string()));
self
}
pub fn add_shape_artifact(
&mut self,
module_path: impl Into<String>,
source: Option<String>,
compiled: Option<Vec<u8>>,
) -> &mut Self {
self.module_artifacts.push(ModuleArtifact {
module_path: module_path.into(),
source,
compiled,
});
self
}
pub fn add_intrinsic<F>(&mut self, type_name: &str, method_name: &str, f: F) -> &mut Self
where
F: for<'ctx> Fn(&[KindedSlot], &ModuleContext<'ctx>) -> Result<KindedSlot, String>
+ Send
+ Sync
+ 'static,
{
self.method_intrinsics
.entry(type_name.to_string())
.or_default()
.insert(method_name.to_string(), Arc::new(f));
self
}
pub fn add_type_schema(&mut self, schema: TypeSchema) -> crate::type_schema::SchemaId {
let id = schema.id;
self.type_schemas.push(schema);
id
}
pub fn has_export(&self, name: &str) -> bool {
self.typed_exports.functions.contains_key(name)
|| self.typed_exports.async_functions.contains_key(name)
}
pub fn is_async(&self, name: &str) -> bool {
self.typed_exports.async_functions.contains_key(name)
}
pub fn get_schema(&self, name: &str) -> Option<&ModuleFunction> {
self.schemas.get(name)
}
pub fn export_names(&self) -> Vec<&str> {
let mut names: Vec<&str> = self
.typed_exports
.functions
.keys()
.chain(self.typed_exports.async_functions.keys())
.map(|s| s.as_str())
.collect();
names.sort_unstable();
names.dedup();
names
}
pub fn to_parsed_schema(&self) -> crate::extensions::ParsedModuleSchema {
let functions = self
.schemas
.iter()
.filter(|(name, _)| self.is_export_public_surface(name, false))
.map(|(name, schema)| crate::extensions::ParsedModuleFunction {
name: name.clone(),
description: schema.description.clone(),
params: schema.params.iter().map(|p| p.type_name.clone()).collect(),
return_type: schema.return_type.clone(),
})
.collect();
crate::extensions::ParsedModuleSchema {
module_name: self.name.clone(),
functions,
artifacts: Vec::new(),
}
}
pub fn stdlib_module_schemas() -> Vec<crate::extensions::ParsedModuleSchema> {
crate::stdlib::all_stdlib_modules()
.into_iter()
.map(|m| m.to_parsed_schema())
.collect()
}
}
impl std::fmt::Debug for ModuleExports {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ModuleExports")
.field("name", &self.name)
.field("description", &self.description)
.field(
"typed_exports",
&self
.typed_exports
.functions
.keys()
.chain(self.typed_exports.async_functions.keys())
.collect::<Vec<_>>(),
)
.field("schemas", &self.schemas.keys().collect::<Vec<_>>())
.field(
"shape_sources",
&self
.shape_sources
.iter()
.map(|(f, _)| f)
.collect::<Vec<_>>(),
)
.field(
"method_intrinsics",
&self.method_intrinsics.keys().collect::<Vec<_>>(),
)
.finish()
}
}
#[derive(Default)]
pub struct ModuleExportRegistry {
modules: HashMap<String, ModuleExports>,
}
impl ModuleExportRegistry {
pub fn new() -> Self {
Self {
modules: HashMap::new(),
}
}
pub fn register(&mut self, module: ModuleExports) {
let canonical = module.name.clone();
self.modules.insert(canonical, module);
}
pub fn get(&self, name: &str) -> Option<&ModuleExports> {
self.modules.get(name)
}
pub fn has(&self, name: &str) -> bool {
self.get(name).is_some()
}
pub fn module_names(&self) -> Vec<&str> {
self.modules.keys().map(|s| s.as_str()).collect()
}
pub fn modules(&self) -> &HashMap<String, ModuleExports> {
&self.modules
}
}
impl std::fmt::Debug for ModuleExportRegistry {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ModuleExportRegistry")
.field("modules", &self.modules.keys().collect::<Vec<_>>())
.finish()
}
}
#[cfg(test)]
#[path = "module_exports_tests.rs"]
mod tests;