use super::environment::ScriptingEnvironment;
use super::process_io::{read_line_bounded, ProcessGroup, MAX_STDERR_BYTES, MAX_STDOUT_BYTES};
use super::types::{
parse_script_response, ScriptConfig, ScriptInput, ScriptLanguage, ScriptResponse,
};
use crate::utils::clock::Instant;
use anyhow::{anyhow, Context as AnyhowContext, Result};
use once_cell::sync::Lazy;
use std::collections::HashMap;
use std::hash::{Hash, Hasher};
use std::process::Stdio;
use std::sync::Arc;
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt, BufReader};
use tokio::process::{Child, ChildStdin, ChildStdout, Command};
use tokio::sync::Mutex;
use tracing::{debug, info, warn};
const IDLE_TTL: Duration = Duration::from_secs(300);
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ResidentScope {
Server,
Connection,
}
impl ResidentScope {
pub fn parse(s: Option<&str>) -> Self {
match s.map(|s| s.to_ascii_lowercase()).as_deref() {
Some("connection") | Some("conn") => ResidentScope::Connection,
_ => ResidentScope::Server,
}
}
}
pub fn resident_language_supported(language: ScriptLanguage) -> bool {
matches!(
language,
ScriptLanguage::Python | ScriptLanguage::JavaScript | ScriptLanguage::Perl
)
}
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
struct ScopeKey {
owner_id: u32,
owner_is_client: bool,
connection_id: Option<String>,
language: ScriptLanguage,
code_hash: u64,
private_payloads: bool,
}
impl ScopeKey {
fn new(
scope: ResidentScope,
input: &ScriptInput,
language: ScriptLanguage,
code: &str,
) -> Self {
let connection_id = match scope {
ResidentScope::Connection => input.connection.as_ref().map(|c| c.id.clone()),
ResidentScope::Server => None,
};
let mut hasher = std::collections::hash_map::DefaultHasher::new();
language.as_str().hash(&mut hasher);
code.hash(&mut hasher);
let (owner_id, owner_is_client) = match (&input.server, &input.client) {
(Some(server), _) => (server.id, false),
(None, Some(client)) => (client.id, true),
(None, None) => (0, false),
};
Self {
owner_id,
owner_is_client,
connection_id,
language,
code_hash: hasher.finish(),
private_payloads: false,
}
}
fn describe(&self) -> String {
let owner = if self.owner_is_client {
format!("client #{}", self.owner_id)
} else {
format!("server #{}", self.owner_id)
};
match &self.connection_id {
Some(c) => format!("{} conn {} [{}]", owner, c, self.language.as_str()),
None => format!("{} [{}]", owner, self.language.as_str()),
}
}
}
struct ResidentIo {
group: ProcessGroup,
stderr_task: tokio::task::JoinHandle<()>,
child: Child,
stdin: ChildStdin,
stdout: BufReader<ChildStdout>,
}
impl Drop for ResidentIo {
fn drop(&mut self) {
self.group.kill();
self.stderr_task.abort();
}
}
pub struct ResidentScript {
language: ScriptLanguage,
io: Mutex<Option<ResidentIo>>,
describe: String,
}
impl ResidentScript {
fn spawn(
language: ScriptLanguage,
code: &str,
describe: String,
private_payloads: bool,
) -> Result<Arc<Self>> {
let (command, args, wrapped): (&str, Vec<&str>, String) = match language {
ScriptLanguage::Python => ("python3", vec!["-u", "-c"], python_harness(code)),
ScriptLanguage::JavaScript => ("node", vec!["-e"], javascript_harness(code)),
ScriptLanguage::Perl => ("perl", vec!["-e"], perl_harness(code)),
ScriptLanguage::Go => {
anyhow::bail!("Go is not supported in resident script mode");
}
};
let mut builder = Command::new(command);
builder.args(&args);
builder.arg(&wrapped);
builder
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.kill_on_drop(true);
ProcessGroup::configure(&mut builder);
let mut child = builder
.spawn()
.map_err(|e| spawn_error(language, command, e))?;
let group = ProcessGroup::new(&child)?;
let stdin = child
.stdin
.take()
.context("resident: failed to capture stdin")?;
let stdout = BufReader::new(
child
.stdout
.take()
.context("resident: failed to capture stdout")?,
);
let stderr = child
.stderr
.take()
.context("resident: failed to capture stderr")?;
let describe_for_log = describe.clone();
let stderr_task = tokio::spawn(async move {
let mut stderr = stderr;
let mut chunk = vec![0; 8192];
let mut logged = 0usize;
while let Ok(n) = stderr.read(&mut chunk).await {
if n == 0 {
break;
}
let keep = n.min(MAX_STDERR_BYTES.saturating_sub(logged));
if keep > 0 {
let shown = if private_payloads {
std::borrow::Cow::Borrowed(crate::utils::redact::REDACTED)
} else {
String::from_utf8_lossy(&chunk[..keep])
};
warn!("resident script {}: {}", describe_for_log, shown);
logged += keep;
if logged == MAX_STDERR_BYTES {
warn!("resident script {}: stderr log cap reached; draining further diagnostics without logging", describe_for_log);
}
}
}
});
info!(
"Spawned resident {} script ({})",
language.as_str(),
describe
);
Ok(Arc::new(Self {
language,
io: Mutex::new(Some(ResidentIo {
group,
stderr_task,
child,
stdin,
stdout,
})),
describe,
}))
}
async fn round_trip(&self, input: &ScriptInput, timeout: Duration) -> Result<ScriptResponse> {
let deadline = tokio::time::Instant::now()
.checked_add(timeout)
.ok_or_else(|| {
anyhow!(
"resident script timeout {:?} exceeds the clock range",
timeout
)
})?;
let mut guard = tokio::time::timeout_at(deadline, self.io.lock())
.await
.map_err(|_| {
anyhow!(
"resident script event exceeded {:?} timeout waiting for its turn",
timeout
)
})?;
let mut io = guard
.take()
.ok_or_else(|| anyhow!("resident script process is not alive"))?;
let mut line = serde_json::to_string(input).context("resident: serialize event")?;
line.push('\n');
let result: Result<String> = {
let interaction = async {
io.stdin
.write_all(line.as_bytes())
.await
.context("resident: write event to stdin")?;
io.stdin.flush().await.context("resident: flush stdin")?;
read_line_bounded(&mut io.stdout, MAX_STDOUT_BYTES)
.await
.context("resident: read response line")
};
match tokio::time::timeout_at(deadline, interaction).await {
Ok(inner) => inner,
Err(_) => Err(anyhow!(
"resident script event exceeded {:?} timeout",
timeout
)),
}
};
match result {
Ok(response) => {
*guard = Some(io);
parse_script_response(response.trim())
.with_context(|| format!("resident script {} response", self.describe))
}
Err(e) => {
io.group.kill();
let _ = io.child.start_kill();
let _ = tokio::time::timeout(Duration::from_secs(2), io.child.wait()).await;
warn!("resident script {} failed: {}", self.describe, e);
Err(e)
}
}
}
async fn kill(&self) {
let mut guard = self.io.lock().await;
if let Some(mut io) = guard.take() {
io.group.kill();
let _ = io.child.start_kill();
let _ = tokio::time::timeout(Duration::from_secs(2), io.child.wait()).await;
debug!("Killed resident script {}", self.describe);
}
}
fn is_dead(&self) -> bool {
self.io.try_lock().is_ok_and(|io| io.is_none())
}
pub fn language(&self) -> ScriptLanguage {
self.language
}
}
struct ResidentEntry {
script: Arc<ResidentScript>,
last_used: Instant,
}
static REGISTRY: Lazy<Mutex<HashMap<ScopeKey, ResidentEntry>>> =
Lazy::new(|| Mutex::new(HashMap::new()));
pub struct ResidentScriptManager;
impl ResidentScriptManager {
pub async fn dispatch(
config: &ScriptConfig,
input: &ScriptInput,
scope: ResidentScope,
) -> Result<ScriptResponse> {
Self::dispatch_with_timeout(
config,
input,
scope,
super::executor::DEFAULT_SCRIPT_TIMEOUT,
)
.await
}
pub(crate) async fn dispatch_private(
config: &ScriptConfig,
input: &ScriptInput,
scope: ResidentScope,
) -> Result<ScriptResponse> {
Self::dispatch_with_timeout_and_privacy(
config,
input,
scope,
super::executor::DEFAULT_SCRIPT_TIMEOUT,
true,
)
.await
}
pub async fn dispatch_with_timeout(
config: &ScriptConfig,
input: &ScriptInput,
scope: ResidentScope,
timeout: Duration,
) -> Result<ScriptResponse> {
Self::dispatch_with_timeout_and_privacy(config, input, scope, timeout, false).await
}
async fn dispatch_with_timeout_and_privacy(
config: &ScriptConfig,
input: &ScriptInput,
scope: ResidentScope,
timeout: Duration,
private_payloads: bool,
) -> Result<ScriptResponse> {
if !resident_language_supported(config.language) {
anyhow::bail!(
"language '{}' is not supported in resident mode",
config.language.as_str()
);
}
let deadline = tokio::time::Instant::now()
.checked_add(timeout)
.ok_or_else(|| {
anyhow!(
"resident script timeout {:?} exceeds the clock range",
timeout
)
})?;
let code = tokio::time::timeout_at(deadline, config.source.get_code_async())
.await
.context("resident: timed out loading script code")?
.context("resident: failed to load script code")?;
let mut key = ScopeKey::new(scope, input, config.language, &code);
key.private_payloads = private_payloads;
let script = {
let mut registry = tokio::time::timeout_at(deadline, REGISTRY.lock())
.await
.context("resident: timed out waiting for registry")?;
evict_idle(&mut registry);
if registry
.get(&key)
.is_some_and(|entry| entry.script.is_dead())
{
registry.remove(&key);
}
match registry.get_mut(&key) {
Some(entry) => {
entry.last_used = Instant::now();
entry.script.clone()
}
None => {
let script = ResidentScript::spawn(
config.language,
&code,
key.describe(),
private_payloads,
)?;
registry.insert(
key.clone(),
ResidentEntry {
script: script.clone(),
last_used: Instant::now(),
},
);
script
}
}
};
match script
.round_trip(
input,
deadline.saturating_duration_since(tokio::time::Instant::now()),
)
.await
{
Ok(response) => {
debug!(
"resident script {} handled '{}' ({} actions)",
key.describe(),
input.event_type_id,
response.actions.len()
);
Ok(response)
}
Err(e) => {
if script.is_dead() {
let mut registry = REGISTRY.lock().await;
if let Some(entry) = registry.get(&key) {
if Arc::ptr_eq(&entry.script, &script) {
registry.remove(&key);
}
}
}
Err(e)
}
}
}
pub async fn shutdown_server(server_id: u32) -> usize {
let count = Self::shutdown_owner(server_id, false).await;
if count > 0 {
info!(
"Shut down {} resident script(s) for server #{}",
count, server_id
);
}
count
}
pub async fn shutdown_client(client_id: u32) -> usize {
let count = Self::shutdown_owner(client_id, true).await;
if count > 0 {
info!(
"Shut down {} resident script(s) for client #{}",
count, client_id
);
}
count
}
async fn shutdown_owner(owner_id: u32, owner_is_client: bool) -> usize {
let scripts = {
let mut registry = REGISTRY.lock().await;
let keys: Vec<ScopeKey> = registry
.keys()
.filter(|k| k.owner_id == owner_id && k.owner_is_client == owner_is_client)
.cloned()
.collect();
keys.into_iter()
.filter_map(|k| registry.remove(&k).map(|e| e.script))
.collect::<Vec<_>>()
};
let count = scripts.len();
for script in scripts {
script.kill().await;
}
count
}
pub async fn shutdown_connection(server_id: u32, connection_id: &str) -> usize {
let scripts = {
let mut registry = REGISTRY.lock().await;
let keys: Vec<ScopeKey> = registry
.keys()
.filter(|k| {
k.owner_id == server_id
&& !k.owner_is_client
&& k.connection_id.as_deref() == Some(connection_id)
})
.cloned()
.collect();
keys.into_iter()
.filter_map(|k| registry.remove(&k).map(|e| e.script))
.collect::<Vec<_>>()
};
let count = scripts.len();
for script in scripts {
script.kill().await;
}
count
}
pub async fn shutdown_all() -> usize {
let scripts = {
let mut registry = REGISTRY.lock().await;
registry.drain().map(|(_, e)| e.script).collect::<Vec<_>>()
};
let count = scripts.len();
for script in scripts {
script.kill().await;
}
count
}
pub async fn active_count() -> usize {
REGISTRY.lock().await.len()
}
}
fn evict_idle(registry: &mut HashMap<ScopeKey, ResidentEntry>) {
let now = Instant::now();
let stale: Vec<ScopeKey> = registry
.iter()
.filter(|(_, e)| now.duration_since(e.last_used) > IDLE_TTL)
.map(|(k, _)| k.clone())
.collect();
for key in stale {
debug!("Evicting idle resident script {}", key.describe());
registry.remove(&key);
}
}
pub fn resident_available(env: &ScriptingEnvironment, language: ScriptLanguage) -> bool {
resident_language_supported(language) && env.is_available(language)
}
fn python_harness(user_code: &str) -> String {
format!(
r#"import sys, json, traceback
def _ng_emit(obj):
sys.stdout.write(json.dumps(obj) + "\n")
sys.stdout.flush()
# ===== user code: must define handle(event_type, event, message) =====
{user_code}
# ===== end user code =====
for _ng_line in sys.stdin:
_ng_line = _ng_line.strip()
if not _ng_line:
continue
try:
_ng_msg = json.loads(_ng_line)
except Exception as _ng_err:
sys.stderr.write("resident: could not parse input line: %s\n" % _ng_err)
_ng_emit({{"error": "bad input line"}})
continue
try:
_ng_result = handle(_ng_msg.get("event_type_id"), _ng_msg.get("event"), _ng_msg)
except Exception:
sys.stderr.write("resident handle() raised:\n%s" % traceback.format_exc())
_ng_emit({{"error": "handle() raised"}})
continue
if isinstance(_ng_result, dict) and "actions" in _ng_result:
_ng_emit(_ng_result)
elif isinstance(_ng_result, list):
_ng_emit({{"actions": _ng_result}})
elif _ng_result is None:
_ng_emit({{"actions": []}})
else:
sys.stderr.write("resident handle() returned unexpected type; expected list or dict\n")
_ng_emit({{"error": "handle() returned an unexpected type"}})
"#,
user_code = user_code
)
}
fn javascript_harness(user_code: &str) -> String {
format!(
r#"const readline = require('readline');
function _ngEmit(obj) {{
process.stdout.write(JSON.stringify(obj) + "\n");
}}
// ===== user code: must define handle(event_type, event, message) =====
{user_code}
// ===== end user code =====
const _ngRl = readline.createInterface({{ input: process.stdin }});
_ngRl.on('line', async (line) => {{
line = line.trim();
if (!line) return;
let msg;
try {{
msg = JSON.parse(line);
}} catch (e) {{
process.stderr.write("resident: could not parse input line: " + e + "\n");
_ngEmit({{ error: "bad input line" }});
return;
}}
let result;
try {{
result = await handle(msg.event_type_id, msg.event, msg);
}} catch (e) {{
process.stderr.write("resident handle() raised: " + (e && e.stack ? e.stack : e) + "\n");
_ngEmit({{ error: "handle() raised" }});
return;
}}
if (Array.isArray(result)) {{
_ngEmit({{ actions: result }});
}} else if (result && typeof result === 'object' && 'actions' in result) {{
_ngEmit(result);
}} else if (result === undefined || result === null) {{
_ngEmit({{ actions: [] }});
}} else {{
process.stderr.write("resident handle() returned unexpected type; expected array or object\n");
_ngEmit({{ error: "handle() returned an unexpected type" }});
}}
}});
"#,
user_code = user_code
)
}
fn perl_harness(user_code: &str) -> String {
format!(
r#"use strict;
use warnings;
use JSON::PP;
$| = 1;
my $_ng_json = JSON::PP->new->utf8->canonical;
sub _ng_emit {{
my ($obj) = @_;
print $_ng_json->encode($obj) . "\n";
}}
# ===== user code: must define sub handle($event_type, $event, $message) =====
{user_code}
# ===== end user code =====
while (my $_ng_line = <STDIN>) {{
$_ng_line =~ s/^\s+|\s+$//g;
next if $_ng_line eq '';
my $_ng_msg = eval {{ $_ng_json->decode($_ng_line) }};
if ($@) {{
print STDERR "resident: could not parse input line: $@\n";
_ng_emit({{ error => "bad input line" }});
next;
}}
my $_ng_result = eval {{ handle($_ng_msg->{{event_type_id}}, $_ng_msg->{{event}}, $_ng_msg) }};
if ($@) {{
print STDERR "resident handle() raised: $@\n";
_ng_emit({{ error => "handle() raised" }});
next;
}}
if (ref($_ng_result) eq 'HASH' && exists $_ng_result->{{actions}}) {{
_ng_emit($_ng_result);
}} elsif (ref($_ng_result) eq 'ARRAY') {{
_ng_emit({{ actions => $_ng_result }});
}} elsif (!defined $_ng_result) {{
_ng_emit({{ actions => [] }});
}} else {{
print STDERR "resident handle() returned unexpected type; expected array-ref or hash-ref\n";
_ng_emit({{ error => "handle() returned an unexpected type" }});
}}
}}
"#,
user_code = user_code
)
}
fn spawn_error(language: ScriptLanguage, command: &str, err: std::io::Error) -> anyhow::Error {
if err.kind() == std::io::ErrorKind::NotFound {
anyhow!(
"Interpreter `{command}` for resident '{lang}' script was not found on PATH. \
Install it, or switch this handler to a language whose interpreter is installed.",
command = command,
lang = language.as_str(),
)
} else {
anyhow!(
"Failed to start interpreter `{command}` for resident '{lang}' script: {err}",
command = command,
lang = language.as_str(),
err = err,
)
}
}