use std::collections::BTreeMap;
use std::future::Future;
use std::sync::Arc;
use std::time::Duration;
use rquickjs::function::{Func, Opt};
use rquickjs::promise::{MaybePromise, Promised};
use rquickjs::{Ctx, IntoJs, JsLifetime, Module, Object, Value, class::Class, class::Trace};
use super::bdd::{tool_dispatch, tool_names};
use super::http_client::HttpClientJs;
use crate::bindings::convert::{json_to_js, serde_from_js};
use crate::command_spec::CommandSpec;
use crate::engine::SessionProcsUd;
use crate::error::ScriptError;
use crate::session_procs::{self, SessionProcs};
#[derive(Debug, Clone)]
pub struct PluginBinding {
pub bytecode: Arc<[u8]>,
}
#[derive(JsLifetime, Trace)]
#[rquickjs::class(rename = "PluginCommands")]
pub struct PluginCommandsJs {
#[qjs(skip_trace)]
allowed: Arc<BTreeMap<String, CommandSpec>>,
#[qjs(skip_trace)]
procs: Option<Arc<SessionProcs>>,
}
impl PluginCommandsJs {
fn cmd_err(verb: &'static str, msg: impl std::fmt::Display) -> rquickjs::Error {
rquickjs::Error::new_from_js_message(verb, "Error", msg.to_string())
}
fn spec(&self, verb: &'static str, name: &str) -> rquickjs::Result<CommandSpec> {
self.allowed.get(name).cloned().ok_or_else(|| {
Self::cmd_err(
verb,
format!("\"{name}\" is not in the commands allow-list for this tool"),
)
})
}
fn vars_of<'js>(ctx: &Ctx<'js>, vars: Opt<Value<'js>>) -> rquickjs::Result<BTreeMap<String, serde_json::Value>> {
match vars.0 {
Some(v) if !v.is_undefined() && !v.is_null() => serde_from_js(ctx, v),
_ => Ok(BTreeMap::new()),
}
}
fn registry(&self, verb: &'static str) -> rquickjs::Result<&Arc<SessionProcs>> {
self
.procs
.as_ref()
.ok_or_else(|| Self::cmd_err(verb, "persistent commands are unavailable in this context"))
}
}
#[rquickjs::methods]
impl PluginCommandsJs {
#[qjs(rename = "run")]
pub async fn run<'js>(&self, ctx: Ctx<'js>, name: String, vars: Opt<Value<'js>>) -> rquickjs::Result<Value<'js>> {
let spec = self.spec("commands.run", &name)?;
let vars_map = Self::vars_of(&ctx, vars)?;
let resolved = spec
.resolve(&vars_map)
.map_err(|m| Self::cmd_err("commands.run", format!("{name}: {m}")))?;
let value = Box::pin(session_procs::run_oneshot(&resolved))
.await
.map_err(|m| Self::cmd_err("commands.run", format!("{name}: {m}")))?;
json_to_js(&ctx, &value)
}
#[qjs(rename = "start")]
pub fn start<'js>(&self, ctx: Ctx<'js>, name: String, vars: Opt<Value<'js>>) -> rquickjs::Result<Value<'js>> {
let spec = self.spec("commands.start", &name)?;
let vars_map = Self::vars_of(&ctx, vars)?;
let resolved = spec
.resolve(&vars_map)
.map_err(|m| Self::cmd_err("commands.start", format!("{name}: {m}")))?;
let pid = self
.registry("commands.start")?
.start(&name, &resolved)
.map_err(|m| Self::cmd_err("commands.start", format!("{name}: {m}")))?;
json_to_js(&ctx, &serde_json::json!({ "name": name, "pid": pid }))
}
#[qjs(rename = "status")]
pub fn status<'js>(&self, ctx: Ctx<'js>, name: String) -> rquickjs::Result<Value<'js>> {
let value = self
.registry("commands.status")?
.status(&name)
.map_err(|m| Self::cmd_err("commands.status", m))?;
json_to_js(&ctx, &value)
}
#[qjs(rename = "stop")]
pub fn stop(&self, name: String) -> rquickjs::Result<()> {
self
.registry("commands.stop")?
.stop(&name)
.map_err(|m| Self::cmd_err("commands.stop", m))
}
}
fn rq(e: &ScriptError) -> rquickjs::Error {
rquickjs::Error::new_from_js_message("plugins", "Error", e.message.clone())
}
pub fn install_plugins(ctx: &Ctx<'_>, files: &[PluginBinding]) -> rquickjs::Result<()> {
for file in files {
#[allow(unsafe_code)]
let module = unsafe { Module::load(ctx.clone(), &file.bytecode) }?;
let (_evaluated, _promise) = module.eval()?;
}
let names = tool_names(ctx).map_err(|e| rq(&e))?;
let plugins_obj = Object::new(ctx.clone())?;
for (idx, name) in names.into_iter().enumerate() {
let f = Func::from(move |ctx, call_args| dispatch_tool(ctx, idx, call_args));
plugins_obj.set(name.as_str(), f)?;
}
ctx.globals().set("plugins", plugins_obj)?;
Ok(())
}
fn dispatch_tool<'js>(
ctx: Ctx<'js>,
idx: usize,
call_args: Opt<Value<'js>>,
) -> Promised<impl std::future::Future<Output = rquickjs::Result<Value<'js>>> + 'js> {
Promised::from(async move {
let d = tool_dispatch(&ctx, idx).map_err(|e| rq(&e))?;
let arg = Object::new(ctx.clone())?;
let undef = Value::new_undefined(ctx.clone());
arg.set("args", call_args.0.unwrap_or_else(|| undef.clone()))?;
let g = ctx.globals();
arg.set("page", g.get::<_, Value<'js>>("page").unwrap_or_else(|_| undef.clone()))?;
arg.set(
"context",
g.get::<_, Value<'js>>("context").unwrap_or_else(|_| undef.clone()),
)?;
let net_policy: Option<Arc<[String]>> = if d.allowed_net.is_empty() {
None
} else {
Some(Arc::from(d.allowed_net.as_slice()))
};
let req_val: Value<'js> = g.get("request").unwrap_or_else(|_| undef.clone());
let request_out: Value<'js> = match net_policy.clone() {
Some(net) => match Class::<HttpClientJs>::from_value(&req_val) {
Ok(cls) => {
let inner = cls.borrow().inner_arc();
let guarded = Class::instance(ctx.clone(), HttpClientJs::with_net(inner, net))?;
guarded.into_js(&ctx)?
},
Err(_) => req_val,
},
None => req_val,
};
arg.set("request", request_out)?;
let procs = ctx.userdata::<SessionProcsUd>().map(|u| u.0.clone());
let commands = Class::instance(
ctx.clone(),
PluginCommandsJs {
allowed: Arc::new(d.allowed_commands),
procs,
},
)?;
arg.set("commands", commands)?;
let policy_cell = ctx
.userdata::<crate::bindings::fetch::NetPolicyUd>()
.map(|u| u.0.clone());
let handler = d.handler;
let timeout_ms = d.timeout_ms;
let inner = async move {
let mp: MaybePromise<'js> = handler.call((arg,))?;
let fut = mp.into_future::<Value<'js>>();
match timeout_ms {
Some(t) => match tokio::time::timeout(Duration::from_millis(t), fut).await {
Ok(r) => r,
Err(_) => Err(rquickjs::Error::new_from_js_message(
"plugins",
"Error",
format!("tool timed out after {t}ms"),
)),
},
None => fut.await,
}
};
match policy_cell {
None => inner.await,
Some(cell) => {
let mut inner = std::pin::pin!(inner);
std::future::poll_fn(move |cx2| {
let prev = cell.swap(net_policy.clone());
let r = inner.as_mut().poll(cx2);
cell.swap(prev);
r
})
.await
},
}
})
}