Migrate test plugins to Rust and async-ify syscall modules
- Replace Python test plugins (echo, ai_test, claudecode) with Rust implementations driven by the plugin SDK - Convert fs/network/process/timer syscall modules to async/await - Change Network.OpenConn to return stream ConnEvent (conn_id handed to client in first "open" frame for Send/Close routing) - Bind timer callbacks to plugin identity; stream_call subscriptions are identity-checked - Update docs and regenerate Python SDK protobuf stubs Verified: cargo build/test/clippy clean; echo + ai_test full syscall suites pass; hermes & claudecode register and heartbeat; cross-plugin bridge.call routing and auth denial verified end-to-end. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -23,3 +23,5 @@ uuid = { workspace = true }
|
||||
notify = { workspace = true }
|
||||
reqwest = { workspace = true }
|
||||
libc = { workspace = true }
|
||||
cron = "0.12"
|
||||
chrono = "0.4"
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Mutex;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
/// A fired timer event delivered to the plugin that created it.
|
||||
@@ -11,7 +12,8 @@ pub struct TimerEvent {
|
||||
|
||||
/// Manages per-plugin event channels for timer callbacks.
|
||||
pub struct CallbackBus {
|
||||
senders: Mutex<HashMap<String, mpsc::Sender<TimerEvent>>>,
|
||||
senders: Mutex<HashMap<String, (mpsc::Sender<TimerEvent>, u64)>>,
|
||||
generation: AtomicU64,
|
||||
}
|
||||
|
||||
impl Default for CallbackBus {
|
||||
@@ -24,23 +26,27 @@ impl CallbackBus {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
senders: Mutex::new(HashMap::new()),
|
||||
generation: AtomicU64::new(0),
|
||||
}
|
||||
}
|
||||
|
||||
/// Register a plugin and get a receiver for its timer events.
|
||||
pub fn subscribe(&self, plugin_id: &str) -> mpsc::Receiver<TimerEvent> {
|
||||
/// Returns the receiver and a generation counter to be used with
|
||||
/// [`remove_if_generation`] for safe cleanup.
|
||||
pub fn subscribe(&self, plugin_id: &str) -> (mpsc::Receiver<TimerEvent>, u64) {
|
||||
let (tx, rx) = mpsc::channel(64);
|
||||
let gen = self.generation.fetch_add(1, Ordering::Relaxed);
|
||||
self.senders
|
||||
.lock()
|
||||
.unwrap()
|
||||
.insert(plugin_id.to_string(), tx);
|
||||
rx
|
||||
.insert(plugin_id.to_string(), (tx, gen));
|
||||
(rx, gen)
|
||||
}
|
||||
|
||||
/// Deliver a timer event to the target plugin.
|
||||
pub fn fire(&self, plugin_id: &str, event: TimerEvent) {
|
||||
let senders = self.senders.lock().unwrap();
|
||||
if let Some(tx) = senders.get(plugin_id) {
|
||||
if let Some((tx, _gen)) = senders.get(plugin_id) {
|
||||
match tx.try_send(event) {
|
||||
Ok(()) => {}
|
||||
Err(mpsc::error::TrySendError::Full(_)) => {
|
||||
@@ -57,7 +63,20 @@ impl CallbackBus {
|
||||
}
|
||||
|
||||
/// Remove a plugin's channel (on unregister/disconnect).
|
||||
#[allow(dead_code)]
|
||||
pub fn remove(&self, plugin_id: &str) {
|
||||
self.senders.lock().unwrap().remove(plugin_id);
|
||||
}
|
||||
|
||||
/// Remove a plugin's channel only if the generation counter matches.
|
||||
/// This prevents a stale unsubscribe from accidentally removing a newer
|
||||
/// subscription that replaced it.
|
||||
pub fn remove_if_generation(&self, plugin_id: &str, gen: u64) {
|
||||
let mut senders = self.senders.lock().unwrap();
|
||||
if let Some((_, current_gen)) = senders.get(plugin_id) {
|
||||
if *current_gen == gen {
|
||||
senders.remove(plugin_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -46,8 +46,12 @@ pub fn check_permission(
|
||||
syscall_name: &str,
|
||||
) -> Result<String, Status> {
|
||||
let Some(plugin_id) = plugin_id_from_metadata(req_metadata) else {
|
||||
// Direct/internal/manual testing call.
|
||||
return Ok(String::new());
|
||||
// The server listens on TCP: an absent plugin id means an unidentified
|
||||
// local process, not a trusted internal call. Reject instead of allow-all.
|
||||
return Err(Status::unauthenticated(format!(
|
||||
"syscall '{}' requires x-plugin-id metadata from a registered plugin",
|
||||
syscall_name
|
||||
)));
|
||||
};
|
||||
|
||||
let record = require_running_plugin(plugins, &plugin_id)?;
|
||||
|
||||
@@ -120,10 +120,18 @@ impl PluginManager {
|
||||
}
|
||||
|
||||
fn save_to_disk(&self) {
|
||||
let plugins = self.plugins.lock().unwrap();
|
||||
let records: Vec<&PluginRecord> = plugins.values().collect();
|
||||
let records: Vec<PluginRecord> = {
|
||||
let plugins = self.plugins.lock().unwrap();
|
||||
plugins.values().cloned().collect()
|
||||
};
|
||||
// Lock released here — serialize and write outside the lock
|
||||
if let Ok(json) = serde_json::to_string_pretty(&records) {
|
||||
let _ = std::fs::write(Self::state_path(), json);
|
||||
// Write to temp file first, then rename for atomicity
|
||||
let path = Self::state_path();
|
||||
let tmp = path.with_extension("json.tmp");
|
||||
if std::fs::write(&tmp, &json).is_ok() {
|
||||
let _ = std::fs::rename(&tmp, &path);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -9,9 +9,11 @@ use agentsd_proto::agentsd::{
|
||||
plugin_bridge_server::PluginBridge as PluginBridgeService,
|
||||
process_server::Process as ProcessService, timer_server::Timer as TimerService,
|
||||
};
|
||||
use std::sync::Arc;
|
||||
use std::sync::{Arc, OnceLock};
|
||||
use tonic::{Request, Response, Status};
|
||||
|
||||
static REGISTER_TOKEN: OnceLock<Option<String>> = OnceLock::new();
|
||||
|
||||
pub struct AgentsdServer {
|
||||
kernel: Arc<Kernel>,
|
||||
plugins: Arc<PluginManager>,
|
||||
@@ -231,10 +233,9 @@ impl AudioService for AudioSvc {
|
||||
req: Request<RecordStartRequest>,
|
||||
) -> Result<Response<Self::RecordStreamStream>, Status> {
|
||||
auth(&self.plugins, &req, "audio.record_stream")?;
|
||||
let (_tx, rx) = tokio::sync::mpsc::channel(16);
|
||||
Ok(Response::new(tokio_stream::wrappers::ReceiverStream::new(
|
||||
rx,
|
||||
)))
|
||||
Err(Status::unimplemented(
|
||||
"audio.record_stream not yet implemented (requires audio driver plugin)",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -366,7 +367,24 @@ impl FsService for FsSvc {
|
||||
req: Request<FsPathRequest>,
|
||||
) -> Result<Response<Self::WatchStream>, Status> {
|
||||
auth(&self.plugins, &req, "fs.watch")?;
|
||||
let (_tx, rx) = tokio::sync::mpsc::channel(16);
|
||||
let r = req.into_inner();
|
||||
let mut events = self
|
||||
.kernel
|
||||
.fs
|
||||
.watch(&r.path)
|
||||
.map_err(Status::invalid_argument)?;
|
||||
let (tx, rx) = tokio::sync::mpsc::channel(64);
|
||||
tokio::spawn(async move {
|
||||
while let Some(ev) = events.recv().await {
|
||||
let event = FsWatchEvent {
|
||||
path: ev.path,
|
||||
kind: ev.kind,
|
||||
};
|
||||
if tx.send(Ok(event)).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
Ok(Response::new(tokio_stream::wrappers::ReceiverStream::new(
|
||||
rx,
|
||||
)))
|
||||
@@ -518,14 +536,49 @@ impl NetworkService for NetworkSvc {
|
||||
}
|
||||
}
|
||||
|
||||
type OpenConnStream = tokio_stream::wrappers::ReceiverStream<Result<DataChunk, Status>>;
|
||||
type OpenConnStream = tokio_stream::wrappers::ReceiverStream<Result<ConnEvent, Status>>;
|
||||
|
||||
async fn open_conn(
|
||||
&self,
|
||||
req: Request<ConnectRequest>,
|
||||
) -> Result<Response<Self::OpenConnStream>, Status> {
|
||||
auth(&self.plugins, &req, "network.open_conn")?;
|
||||
let (_tx, rx) = tokio::sync::mpsc::channel(16);
|
||||
let r = req.into_inner();
|
||||
let (conn_id, mut inbound) = self
|
||||
.kernel
|
||||
.network
|
||||
.open_conn(&r.addr, &r.protocol)
|
||||
.await
|
||||
.map_err(Status::unavailable)?;
|
||||
let (tx, rx) = tokio::sync::mpsc::channel(64);
|
||||
tokio::spawn(async move {
|
||||
// First frame hands the conn_id to the client for Send/Close routing.
|
||||
let open_event = ConnEvent {
|
||||
conn_id: conn_id.clone(),
|
||||
data: vec![],
|
||||
kind: "open".to_string(),
|
||||
};
|
||||
if tx.send(Ok(open_event)).await.is_err() {
|
||||
return;
|
||||
}
|
||||
while let Some(data) = inbound.recv().await {
|
||||
let event = ConnEvent {
|
||||
conn_id: conn_id.clone(),
|
||||
data,
|
||||
kind: "data".to_string(),
|
||||
};
|
||||
if tx.send(Ok(event)).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
let _ = tx
|
||||
.send(Ok(ConnEvent {
|
||||
conn_id: conn_id.clone(),
|
||||
data: vec![],
|
||||
kind: "closed".to_string(),
|
||||
}))
|
||||
.await;
|
||||
});
|
||||
Ok(Response::new(tokio_stream::wrappers::ReceiverStream::new(
|
||||
rx,
|
||||
)))
|
||||
@@ -538,7 +591,51 @@ impl NetworkService for NetworkSvc {
|
||||
req: Request<ListenRequest>,
|
||||
) -> Result<Response<Self::ListenStream>, Status> {
|
||||
auth(&self.plugins, &req, "network.listen")?;
|
||||
let (_tx, rx) = tokio::sync::mpsc::channel(16);
|
||||
let r = req.into_inner();
|
||||
let mut accepted = self
|
||||
.kernel
|
||||
.network
|
||||
.listen(r.port, &r.protocol)
|
||||
.await
|
||||
.map_err(Status::unavailable)?;
|
||||
let (tx, rx) = tokio::sync::mpsc::channel(64);
|
||||
let network = self.kernel.network.clone();
|
||||
tokio::spawn(async move {
|
||||
while let Some(conn) = accepted.recv().await {
|
||||
let accept_event = ConnEvent {
|
||||
conn_id: conn.conn_id.clone(),
|
||||
data: conn.peer.clone().into_bytes(),
|
||||
kind: "accept".to_string(),
|
||||
};
|
||||
if tx.send(Ok(accept_event)).await.is_err() {
|
||||
break;
|
||||
}
|
||||
// Forward this connection's inbound data onto the same event stream.
|
||||
if let Some(mut inbound) = network.take_inbound(&conn.conn_id).await {
|
||||
let tx_data = tx.clone();
|
||||
let cid = conn.conn_id.clone();
|
||||
tokio::spawn(async move {
|
||||
while let Some(data) = inbound.recv().await {
|
||||
let event = ConnEvent {
|
||||
conn_id: cid.clone(),
|
||||
data,
|
||||
kind: "data".to_string(),
|
||||
};
|
||||
if tx_data.send(Ok(event)).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
let _ = tx_data
|
||||
.send(Ok(ConnEvent {
|
||||
conn_id: cid.clone(),
|
||||
data: vec![],
|
||||
kind: "closed".to_string(),
|
||||
}))
|
||||
.await;
|
||||
});
|
||||
}
|
||||
}
|
||||
});
|
||||
Ok(Response::new(tokio_stream::wrappers::ReceiverStream::new(
|
||||
rx,
|
||||
)))
|
||||
@@ -546,18 +643,26 @@ impl NetworkService for NetworkSvc {
|
||||
|
||||
async fn send(&self, req: Request<SendRequest>) -> Result<Response<NetResponse>, Status> {
|
||||
auth(&self.plugins, &req, "network.send")?;
|
||||
Ok(Response::new(NetResponse {
|
||||
ok: false,
|
||||
error: "send: not yet implemented".to_string(),
|
||||
}))
|
||||
let r = req.into_inner();
|
||||
match self.kernel.network.send(&r.conn_id, &r.data).await {
|
||||
Ok(()) => Ok(Response::new(NetResponse {
|
||||
ok: true,
|
||||
error: String::new(),
|
||||
})),
|
||||
Err(e) => Ok(Response::new(NetResponse { ok: false, error: e })),
|
||||
}
|
||||
}
|
||||
|
||||
async fn close(&self, req: Request<CloseRequest>) -> Result<Response<NetResponse>, Status> {
|
||||
auth(&self.plugins, &req, "network.close")?;
|
||||
Ok(Response::new(NetResponse {
|
||||
ok: false,
|
||||
error: "close: not yet implemented".to_string(),
|
||||
}))
|
||||
let r = req.into_inner();
|
||||
match self.kernel.network.close(&r.conn_id).await {
|
||||
Ok(()) => Ok(Response::new(NetResponse {
|
||||
ok: true,
|
||||
error: String::new(),
|
||||
})),
|
||||
Err(e) => Ok(Response::new(NetResponse { ok: false, error: e })),
|
||||
}
|
||||
}
|
||||
|
||||
async fn available(
|
||||
@@ -617,10 +722,7 @@ impl InputService for InputSvc {
|
||||
|
||||
async fn events(&self, req: Request<Empty>) -> Result<Response<Self::EventsStream>, Status> {
|
||||
auth(&self.plugins, &req, "input.events")?;
|
||||
let (_tx, rx) = tokio::sync::mpsc::channel(16);
|
||||
Ok(Response::new(tokio_stream::wrappers::ReceiverStream::new(
|
||||
rx,
|
||||
)))
|
||||
Err(Status::unimplemented("input.events not yet implemented"))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -792,12 +894,20 @@ impl TimerService for TimerSvc {
|
||||
&self,
|
||||
req: Request<TimerCronRequest>,
|
||||
) -> Result<Response<TimerResponse>, Status> {
|
||||
auth(&self.plugins, &req, "timer.cron")?;
|
||||
Ok(Response::new(TimerResponse {
|
||||
timer_id: String::new(),
|
||||
ok: false,
|
||||
error: "cron: not yet implemented".to_string(),
|
||||
}))
|
||||
let plugin_id = auth(&self.plugins, &req, "timer.cron")?;
|
||||
let r = req.into_inner();
|
||||
match self.kernel.timer.cron(&r.expr, r.callback_id, plugin_id) {
|
||||
Ok(id) => Ok(Response::new(TimerResponse {
|
||||
timer_id: id,
|
||||
ok: true,
|
||||
error: String::new(),
|
||||
})),
|
||||
Err(e) => Ok(Response::new(TimerResponse {
|
||||
timer_id: String::new(),
|
||||
ok: false,
|
||||
error: e,
|
||||
})),
|
||||
}
|
||||
}
|
||||
|
||||
async fn cancel(
|
||||
@@ -921,7 +1031,35 @@ impl PluginBridgeService for BridgeSvc {
|
||||
&self,
|
||||
req: Request<PluginInfo>,
|
||||
) -> Result<Response<RegisterResponse>, Status> {
|
||||
// Require register token when AGENTSD_REGISTER_TOKEN env is set.
|
||||
let env_token =
|
||||
REGISTER_TOKEN.get_or_init(|| std::env::var("AGENTSD_REGISTER_TOKEN").ok());
|
||||
if let Some(expected) = env_token {
|
||||
let provided = req
|
||||
.metadata()
|
||||
.get("x-register-token")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.unwrap_or_default();
|
||||
if provided != expected.as_str() {
|
||||
return Ok(Response::new(RegisterResponse {
|
||||
ok: false,
|
||||
error: "invalid or missing register token".into(),
|
||||
session_id: String::new(),
|
||||
}));
|
||||
}
|
||||
}
|
||||
|
||||
let r = req.into_inner();
|
||||
|
||||
// External registrations must never claim plugin_type="system".
|
||||
if r.plugin_type == "system" {
|
||||
return Ok(Response::new(RegisterResponse {
|
||||
ok: false,
|
||||
error: "external registration of system plugins is not allowed".into(),
|
||||
session_id: String::new(),
|
||||
}));
|
||||
}
|
||||
|
||||
let record = PluginRecord {
|
||||
id: r.id,
|
||||
name: r.name,
|
||||
@@ -956,10 +1094,21 @@ impl PluginBridgeService for BridgeSvc {
|
||||
}
|
||||
|
||||
async fn call(&self, req: Request<CallRequest>) -> Result<Response<CallResponse>, Status> {
|
||||
// bridge.call requires a verified, running plugin identity.
|
||||
let plugin_id = permission::plugin_id_from_metadata(req.metadata())
|
||||
.ok_or_else(|| Status::unauthenticated("bridge.call requires x-plugin-id"))?;
|
||||
permission::require_running_plugin(&self.plugins, &plugin_id)?;
|
||||
let r = req.into_inner();
|
||||
let fn_name = &r.r#fn;
|
||||
match self.plugins.find_export_record(&r.target, fn_name) {
|
||||
Some(record) => {
|
||||
if record.status != PluginStatus::Running {
|
||||
return Ok(Response::new(CallResponse {
|
||||
result: vec![],
|
||||
ok: false,
|
||||
error: format!("plugin '{}' is not running", record.id),
|
||||
}));
|
||||
}
|
||||
let endpoint = record.endpoint.as_deref().unwrap_or_default();
|
||||
if endpoint.is_empty() {
|
||||
return Ok(Response::new(CallResponse {
|
||||
@@ -1002,28 +1151,25 @@ impl PluginBridgeService for BridgeSvc {
|
||||
&self,
|
||||
req: Request<CallRequest>,
|
||||
) -> Result<Response<Self::StreamCallStream>, Status> {
|
||||
let metadata_plugin_id = permission::plugin_id_from_metadata(req.metadata());
|
||||
// Subscriptions are identity-bound: the caller may only subscribe to
|
||||
// its own event channel, so a verified x-plugin-id is mandatory.
|
||||
let Some(plugin_id) = permission::plugin_id_from_metadata(req.metadata()) else {
|
||||
return Err(Status::unauthenticated(
|
||||
"stream_call requires x-plugin-id metadata",
|
||||
));
|
||||
};
|
||||
let r = req.into_inner();
|
||||
let target = r.target;
|
||||
|
||||
if let Some(ref metadata_id) = metadata_plugin_id {
|
||||
if !target.is_empty() && metadata_id != &target {
|
||||
return Err(Status::permission_denied(format!(
|
||||
"plugin '{}' cannot subscribe to '{}' events",
|
||||
metadata_id, target
|
||||
)));
|
||||
}
|
||||
}
|
||||
|
||||
let plugin_id = metadata_plugin_id.unwrap_or(target);
|
||||
if plugin_id.is_empty() {
|
||||
return Err(Status::invalid_argument(
|
||||
"stream_call requires target or x-plugin-id",
|
||||
));
|
||||
if !target.is_empty() && plugin_id != target {
|
||||
return Err(Status::permission_denied(format!(
|
||||
"plugin '{}' cannot subscribe to '{}' events",
|
||||
plugin_id, target
|
||||
)));
|
||||
}
|
||||
permission::require_running_plugin(&self.plugins, &plugin_id)?;
|
||||
|
||||
let mut events = self.kernel.callback_bus.subscribe(&plugin_id);
|
||||
let (mut events, generation) = self.kernel.callback_bus.subscribe(&plugin_id);
|
||||
let (tx, rx) = tokio::sync::mpsc::channel(16);
|
||||
let bus = self.kernel.callback_bus.clone();
|
||||
let subscribed_plugin_id = plugin_id.clone();
|
||||
@@ -1043,7 +1189,7 @@ impl PluginBridgeService for BridgeSvc {
|
||||
break;
|
||||
}
|
||||
}
|
||||
bus.remove(&subscribed_plugin_id);
|
||||
bus.remove_if_generation(&subscribed_plugin_id, generation);
|
||||
});
|
||||
Ok(Response::new(tokio_stream::wrappers::ReceiverStream::new(
|
||||
rx,
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
use std::path::Path;
|
||||
use tokio::fs as async_fs;
|
||||
use tokio::io::{AsyncReadExt, AsyncSeekExt};
|
||||
use notify::{Watcher, RecursiveMode, EventKind};
|
||||
|
||||
const MAX_READ_SIZE: usize = 64 * 1024 * 1024; // 64 MB
|
||||
|
||||
@@ -106,6 +107,57 @@ impl FsSyscall {
|
||||
}
|
||||
Ok(entries)
|
||||
}
|
||||
|
||||
pub fn watch(&self, path: &str) -> Result<tokio::sync::mpsc::Receiver<FsWatchEventInfo>, String> {
|
||||
let (std_tx, std_rx) = std::sync::mpsc::channel::<notify::Event>();
|
||||
|
||||
let mut watcher = notify::recommended_watcher(move |res: Result<notify::Event, notify::Error>| {
|
||||
match res {
|
||||
Ok(event) => {
|
||||
let _ = std_tx.send(event);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("fs watch notify error: {}", e);
|
||||
}
|
||||
}
|
||||
}).map_err(|e| format!("failed to create watcher: {}", e))?;
|
||||
|
||||
watcher
|
||||
.watch(std::path::Path::new(path), RecursiveMode::Recursive)
|
||||
.map_err(|e| format!("failed to watch path '{}': {}", path, e))?;
|
||||
|
||||
let (tx, rx) = tokio::sync::mpsc::channel::<FsWatchEventInfo>(64);
|
||||
|
||||
tokio::task::spawn_blocking(move || {
|
||||
// Move watcher into this task to keep it alive
|
||||
let _watcher = watcher;
|
||||
|
||||
while let Ok(event) = std_rx.recv() {
|
||||
let kind = match event.kind {
|
||||
EventKind::Create(_) => "create",
|
||||
EventKind::Modify(_) => "modify",
|
||||
EventKind::Remove(_) => "remove",
|
||||
EventKind::Access(_) => continue,
|
||||
_ => "other",
|
||||
};
|
||||
let path = event
|
||||
.paths
|
||||
.first()
|
||||
.map(|p| p.to_string_lossy().to_string())
|
||||
.unwrap_or_default();
|
||||
let info = FsWatchEventInfo {
|
||||
path,
|
||||
kind: kind.to_string(),
|
||||
};
|
||||
if tx.blocking_send(info).is_err() {
|
||||
// Receiver dropped (gRPC stream closed), stop watching
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok(rx)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct FsStatInfo {
|
||||
@@ -121,3 +173,8 @@ pub struct FsEntryInfo {
|
||||
pub is_dir: bool,
|
||||
pub size: i64,
|
||||
}
|
||||
|
||||
pub struct FsWatchEventInfo {
|
||||
pub path: String,
|
||||
pub kind: String,
|
||||
}
|
||||
|
||||
@@ -1,12 +1,34 @@
|
||||
use std::collections::HashMap;
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::{TcpListener, TcpStream};
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::task::JoinHandle;
|
||||
|
||||
const MAX_RESPONSE_BODY: u64 = 16 * 1024 * 1024; // 16 MB
|
||||
|
||||
/// Accepted connection info delivered via the listen channel.
|
||||
pub struct ConnAccepted {
|
||||
pub conn_id: String,
|
||||
pub peer: String,
|
||||
}
|
||||
|
||||
/// Internal handle for a managed TCP connection.
|
||||
struct ConnHandle {
|
||||
writer: Arc<tokio::sync::Mutex<tokio::net::tcp::OwnedWriteHalf>>,
|
||||
reader_task: JoinHandle<()>,
|
||||
inbound: Option<mpsc::Receiver<Vec<u8>>>,
|
||||
}
|
||||
|
||||
/// Network syscall: remote communication.
|
||||
pub struct NetworkSyscall {
|
||||
client: reqwest::Client,
|
||||
conns: Arc<tokio::sync::Mutex<HashMap<String, ConnHandle>>>,
|
||||
next_conn_id: Arc<AtomicU64>,
|
||||
}
|
||||
|
||||
impl NetworkSyscall {
|
||||
@@ -18,13 +40,23 @@ impl NetworkSyscall {
|
||||
.redirect(reqwest::redirect::Policy::limited(5))
|
||||
.build()
|
||||
.unwrap_or_else(|_| reqwest::Client::new());
|
||||
Self { client }
|
||||
Self {
|
||||
client,
|
||||
conns: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
|
||||
next_conn_id: Arc::new(AtomicU64::new(1)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Generate the next unique connection ID.
|
||||
fn alloc_conn_id(&self) -> String {
|
||||
let id = self.next_conn_id.fetch_add(1, Ordering::Relaxed);
|
||||
format!("conn_{}", id)
|
||||
}
|
||||
|
||||
pub async fn available(&self) -> bool {
|
||||
tokio::time::timeout(
|
||||
Duration::from_secs(2),
|
||||
tokio::net::TcpStream::connect(SocketAddr::from(([8, 8, 8, 8], 53))),
|
||||
TcpStream::connect(SocketAddr::from(([8, 8, 8, 8], 53))),
|
||||
)
|
||||
.await
|
||||
.map(|result| result.is_ok())
|
||||
@@ -91,4 +123,162 @@ impl NetworkSyscall {
|
||||
|
||||
Ok((status, resp_headers, resp_body.to_vec()))
|
||||
}
|
||||
}
|
||||
|
||||
/// Open an outbound TCP connection to `addr` (e.g. "host:port").
|
||||
/// Returns (conn_id, inbound data receiver).
|
||||
pub async fn open_conn(
|
||||
&self,
|
||||
addr: &str,
|
||||
_protocol: &str,
|
||||
) -> Result<(String, mpsc::Receiver<Vec<u8>>), String> {
|
||||
let stream = TcpStream::connect(addr)
|
||||
.await
|
||||
.map_err(|e| format!("connect failed: {e}"))?;
|
||||
let (read_half, write_half) = stream.into_split();
|
||||
|
||||
let conn_id = self.alloc_conn_id();
|
||||
let (tx, rx) = mpsc::channel::<Vec<u8>>(64);
|
||||
|
||||
let writer = Arc::new(tokio::sync::Mutex::new(write_half));
|
||||
let conns = Arc::clone(&self.conns);
|
||||
let cid = conn_id.clone();
|
||||
|
||||
let reader_task = tokio::spawn(async move {
|
||||
let mut read_half = read_half;
|
||||
let mut buf = vec![0u8; 8192];
|
||||
loop {
|
||||
match read_half.read(&mut buf).await {
|
||||
Ok(0) => break,
|
||||
Ok(n) => {
|
||||
if tx.send(buf[..n].to_vec()).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
conns.lock().await.remove(&cid);
|
||||
});
|
||||
|
||||
let handle = ConnHandle {
|
||||
writer,
|
||||
reader_task,
|
||||
inbound: None,
|
||||
};
|
||||
self.conns.lock().await.insert(conn_id.clone(), handle);
|
||||
|
||||
Ok((conn_id, rx))
|
||||
}
|
||||
|
||||
/// Send data on an existing connection.
|
||||
pub async fn send(&self, conn_id: &str, data: &[u8]) -> Result<(), String> {
|
||||
let writer = {
|
||||
let map = self.conns.lock().await;
|
||||
let handle = map
|
||||
.get(conn_id)
|
||||
.ok_or_else(|| format!("connection not found: {conn_id}"))?;
|
||||
Arc::clone(&handle.writer)
|
||||
};
|
||||
let mut w = writer.lock().await;
|
||||
w.write_all(data)
|
||||
.await
|
||||
.map_err(|e| format!("write failed: {e}"))?;
|
||||
w.flush().await.map_err(|e| format!("flush failed: {e}"))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Close and remove a connection.
|
||||
pub async fn close(&self, conn_id: &str) -> Result<(), String> {
|
||||
let handle = self
|
||||
.conns
|
||||
.lock()
|
||||
.await
|
||||
.remove(conn_id)
|
||||
.ok_or_else(|| format!("connection not found: {conn_id}"))?;
|
||||
handle.reader_task.abort();
|
||||
let mut w = handle.writer.lock().await;
|
||||
let _ = w.shutdown().await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Start listening on a TCP port. Returns a receiver that yields each
|
||||
/// accepted connection's id and peer address.
|
||||
pub async fn listen(
|
||||
&self,
|
||||
port: u32,
|
||||
_protocol: &str,
|
||||
) -> Result<mpsc::Receiver<ConnAccepted>, String> {
|
||||
let listener = TcpListener::bind((std::net::Ipv6Addr::LOCALHOST, port as u16))
|
||||
.await
|
||||
.map_err(|e| format!("bind failed: {e}"))?;
|
||||
|
||||
let (accept_tx, accept_rx) = mpsc::channel::<ConnAccepted>(16);
|
||||
let conns = Arc::clone(&self.conns);
|
||||
let id_gen = Arc::clone(&self.next_conn_id);
|
||||
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
let (stream, peer_addr) = match listener.accept().await {
|
||||
Ok(v) => v,
|
||||
Err(_) => break,
|
||||
};
|
||||
|
||||
let (read_half, write_half) = stream.into_split();
|
||||
let conn_id = {
|
||||
let id = id_gen.fetch_add(1, Ordering::Relaxed);
|
||||
format!("conn_{}", id)
|
||||
};
|
||||
|
||||
let (tx, inbound_rx) = mpsc::channel::<Vec<u8>>(64);
|
||||
let writer = Arc::new(tokio::sync::Mutex::new(write_half));
|
||||
let conns_inner = Arc::clone(&conns);
|
||||
let cid = conn_id.clone();
|
||||
|
||||
let reader_task = tokio::spawn(async move {
|
||||
let mut read_half = read_half;
|
||||
let mut buf = vec![0u8; 8192];
|
||||
loop {
|
||||
match read_half.read(&mut buf).await {
|
||||
Ok(0) => break,
|
||||
Ok(n) => {
|
||||
if tx.send(buf[..n].to_vec()).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
conns_inner.lock().await.remove(&cid);
|
||||
});
|
||||
|
||||
let handle = ConnHandle {
|
||||
writer,
|
||||
reader_task,
|
||||
inbound: Some(inbound_rx),
|
||||
};
|
||||
conns.lock().await.insert(conn_id.clone(), handle);
|
||||
|
||||
let accepted = ConnAccepted {
|
||||
conn_id,
|
||||
peer: peer_addr.to_string(),
|
||||
};
|
||||
if accept_tx.send(accepted).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok(accept_rx)
|
||||
}
|
||||
|
||||
/// Take the inbound data receiver for a passively-accepted connection.
|
||||
/// Returns None if the connection doesn't exist or has no stored receiver.
|
||||
pub async fn take_inbound(
|
||||
&self,
|
||||
conn_id: &str,
|
||||
) -> Option<mpsc::Receiver<Vec<u8>>> {
|
||||
let mut map = self.conns.lock().await;
|
||||
let handle = map.get_mut(conn_id)?;
|
||||
handle.inbound.take()
|
||||
}
|
||||
}
|
||||
@@ -2,7 +2,8 @@ use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::process::Command;
|
||||
use tokio::sync::{mpsc, Mutex, RwLock};
|
||||
use tokio::sync::{mpsc, watch, Mutex, RwLock};
|
||||
use tokio::time::Duration;
|
||||
|
||||
/// Per-process managed state, decoupled from the Child handle.
|
||||
struct ManagedProcess {
|
||||
@@ -10,8 +11,8 @@ struct ManagedProcess {
|
||||
status: Arc<Mutex<String>>,
|
||||
stdin_tx: Option<mpsc::Sender<Vec<u8>>>,
|
||||
stdout_rx: Arc<Mutex<mpsc::Receiver<Vec<u8>>>>,
|
||||
/// Notified when the process exits. Clone the receiver to wait.
|
||||
exit_notify: Arc<tokio::sync::Notify>,
|
||||
/// Watch receiver: becomes `true` when process exits. Clone to wait.
|
||||
exit_rx: watch::Receiver<bool>,
|
||||
/// Handle to abort background tasks on kill.
|
||||
_tasks: Vec<tokio::task::JoinHandle<()>>,
|
||||
}
|
||||
@@ -19,13 +20,13 @@ struct ManagedProcess {
|
||||
/// Process syscall: spawn and manage subprocesses.
|
||||
/// Uses the real OS PID as the identifier.
|
||||
pub struct ProcessSyscall {
|
||||
procs: RwLock<HashMap<u64, Arc<ManagedProcess>>>,
|
||||
procs: Arc<RwLock<HashMap<u64, Arc<ManagedProcess>>>>,
|
||||
}
|
||||
|
||||
impl ProcessSyscall {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
procs: RwLock::new(HashMap::new()),
|
||||
procs: Arc::new(RwLock::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -61,7 +62,7 @@ impl ProcessSyscall {
|
||||
let child_stderr = child.stderr.take();
|
||||
|
||||
let status = Arc::new(Mutex::new("running".to_string()));
|
||||
let exit_notify = Arc::new(tokio::sync::Notify::new());
|
||||
let (exit_tx, exit_rx) = watch::channel(false);
|
||||
|
||||
// Stdin writer task
|
||||
let (stdin_tx, mut stdin_rx) = mpsc::channel::<Vec<u8>>(64);
|
||||
@@ -111,14 +112,17 @@ impl ProcessSyscall {
|
||||
}
|
||||
});
|
||||
|
||||
// Reaper task: waits for exit and updates status
|
||||
// Reaper task: waits for exit, updates status, signals watchers, delayed cleanup
|
||||
let status_clone = status.clone();
|
||||
let exit_notify_clone = exit_notify.clone();
|
||||
let procs_clone = Arc::clone(&self.procs);
|
||||
let reaper_task = tokio::spawn(async move {
|
||||
let exit_status = child.wait().await;
|
||||
let code = exit_status.map(|s| s.code().unwrap_or(-1)).unwrap_or(-1);
|
||||
*status_clone.lock().await = format!("exited({})", code);
|
||||
exit_notify_clone.notify_waiters();
|
||||
let _ = exit_tx.send(true);
|
||||
// Delayed cleanup: give stdout readers time to drain buffered data
|
||||
tokio::time::sleep(Duration::from_secs(5)).await;
|
||||
procs_clone.write().await.remove(&pid);
|
||||
});
|
||||
|
||||
let managed = Arc::new(ManagedProcess {
|
||||
@@ -126,7 +130,7 @@ impl ProcessSyscall {
|
||||
status,
|
||||
stdin_tx: Some(stdin_tx),
|
||||
stdout_rx: Arc::new(Mutex::new(stdout_rx)),
|
||||
exit_notify,
|
||||
exit_rx,
|
||||
_tasks: vec![stdin_task, stdout_task, stderr_task, reaper_task],
|
||||
});
|
||||
|
||||
@@ -136,19 +140,25 @@ impl ProcessSyscall {
|
||||
}
|
||||
|
||||
pub async fn kill(&self, pid: u64) -> Result<(), String> {
|
||||
let procs = self.procs.read().await;
|
||||
let proc = procs
|
||||
.get(&pid)
|
||||
.ok_or_else(|| format!("pid {} not found", pid))?
|
||||
.clone();
|
||||
drop(procs);
|
||||
let proc = {
|
||||
let procs = self.procs.read().await;
|
||||
procs
|
||||
.get(&pid)
|
||||
.ok_or_else(|| format!("pid {} not found", pid))?
|
||||
.clone()
|
||||
};
|
||||
|
||||
// Set status to killed first
|
||||
*proc.status.lock().await = "killed".to_string();
|
||||
|
||||
// Abort all background tasks (stdin writer, stdout reader, stderr drain, reaper)
|
||||
// Aborting reaper drops exit_tx, so any wait() blocked on changed() will wake with Err.
|
||||
for task in proc._tasks.iter() {
|
||||
task.abort();
|
||||
}
|
||||
*proc.status.lock().await = "killed".to_string();
|
||||
proc.exit_notify.notify_waiters();
|
||||
|
||||
// Remove from map
|
||||
self.procs.write().await.remove(&pid);
|
||||
|
||||
// Also send OS kill
|
||||
#[cfg(unix)]
|
||||
@@ -161,46 +171,41 @@ impl ProcessSyscall {
|
||||
// On Windows, aborting the reaper task triggers kill_on_drop on the Child
|
||||
}
|
||||
|
||||
self.procs.write().await.remove(&pid);
|
||||
tracing::info!("process killed: pid={}", pid);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn wait(&self, pid: u64) -> Result<i32, String> {
|
||||
let procs = self.procs.read().await;
|
||||
let proc = procs
|
||||
.get(&pid)
|
||||
.ok_or_else(|| format!("pid {} not found", pid))?
|
||||
.clone();
|
||||
drop(procs);
|
||||
let proc = {
|
||||
let procs = self.procs.read().await;
|
||||
procs
|
||||
.get(&pid)
|
||||
.ok_or_else(|| format!("pid {} not found", pid))?
|
||||
.clone()
|
||||
};
|
||||
|
||||
// Check if already exited before waiting
|
||||
{
|
||||
let status = proc.status.lock().await;
|
||||
if status.starts_with("exited") || *status == "killed" {
|
||||
let s = status.clone();
|
||||
drop(status);
|
||||
self.procs.write().await.remove(&pid);
|
||||
return parse_exit_code(&s);
|
||||
// Use watch channel: no race between checking status and waiting.
|
||||
// watch retains the latest value, so even if send(true) happened before
|
||||
// we subscribe, borrow_and_update() will see `true` immediately.
|
||||
let mut rx = proc.exit_rx.clone();
|
||||
while !*rx.borrow_and_update() {
|
||||
if rx.changed().await.is_err() {
|
||||
break; // sender dropped (e.g. kill aborted reaper)
|
||||
}
|
||||
}
|
||||
|
||||
// Wait for exit notification
|
||||
proc.exit_notify.notified().await;
|
||||
let status = proc.status.lock().await.clone();
|
||||
|
||||
// Clean up from map
|
||||
self.procs.write().await.remove(&pid);
|
||||
parse_exit_code(&status)
|
||||
}
|
||||
|
||||
pub async fn write_stdin(&self, pid: u64, data: &[u8]) -> Result<(), String> {
|
||||
let procs = self.procs.read().await;
|
||||
let proc = procs
|
||||
.get(&pid)
|
||||
.ok_or_else(|| format!("pid {} not found", pid))?
|
||||
.clone();
|
||||
drop(procs);
|
||||
let proc = {
|
||||
let procs = self.procs.read().await;
|
||||
procs
|
||||
.get(&pid)
|
||||
.ok_or_else(|| format!("pid {} not found", pid))?
|
||||
.clone()
|
||||
};
|
||||
|
||||
if let Some(ref tx) = proc.stdin_tx {
|
||||
tx.send(data.to_vec())
|
||||
@@ -212,12 +217,13 @@ impl ProcessSyscall {
|
||||
}
|
||||
|
||||
pub async fn read_stdout(&self, pid: u64, buf: &mut [u8]) -> Result<usize, String> {
|
||||
let procs = self.procs.read().await;
|
||||
let proc = procs
|
||||
.get(&pid)
|
||||
.ok_or_else(|| format!("pid {} not found", pid))?
|
||||
.clone();
|
||||
drop(procs);
|
||||
let proc = {
|
||||
let procs = self.procs.read().await;
|
||||
procs
|
||||
.get(&pid)
|
||||
.ok_or_else(|| format!("pid {} not found", pid))?
|
||||
.clone()
|
||||
};
|
||||
|
||||
let mut rx = proc.stdout_rx.lock().await;
|
||||
match rx.recv().await {
|
||||
@@ -268,9 +274,13 @@ impl ProcessSyscall {
|
||||
}
|
||||
|
||||
pub async fn list(&self) -> Vec<(u64, String, String)> {
|
||||
let procs = self.procs.read().await;
|
||||
let mut result = Vec::with_capacity(procs.len());
|
||||
for (pid, mp) in procs.iter() {
|
||||
let snapshot: Vec<(u64, Arc<ManagedProcess>)> = {
|
||||
let procs = self.procs.read().await;
|
||||
procs.iter().map(|(pid, mp)| (*pid, mp.clone())).collect()
|
||||
};
|
||||
// procs read lock is dropped here
|
||||
let mut result = Vec::with_capacity(snapshot.len());
|
||||
for (pid, mp) in &snapshot {
|
||||
let status = mp.status.lock().await.clone();
|
||||
result.push((*pid, mp.cmd.clone(), status));
|
||||
}
|
||||
@@ -288,3 +298,4 @@ fn parse_exit_code(status: &str) -> Result<i32, String> {
|
||||
Ok(-1)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
use std::collections::HashMap;
|
||||
use std::str::FromStr;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use tokio::time::{self, Duration};
|
||||
|
||||
@@ -75,4 +76,53 @@ impl TimerSyscall {
|
||||
.unwrap_or_default()
|
||||
.as_millis() as i64
|
||||
}
|
||||
|
||||
/// Cron 表达式为 cron crate 7 段格式: sec min hour day-of-month month day-of-week year
|
||||
pub fn cron(&self, expr: &str, callback_id: String, plugin_id: String) -> Result<String, String> {
|
||||
let schedule = cron::Schedule::from_str(expr)
|
||||
.map_err(|e| format!("invalid cron expression: {e}"))?;
|
||||
|
||||
let timer_id = {
|
||||
let mut next = self.next_id.lock().unwrap();
|
||||
let id = format!("cron_{}", *next);
|
||||
*next += 1;
|
||||
id
|
||||
};
|
||||
|
||||
let tid = timer_id.clone();
|
||||
let cb = callback_id.clone();
|
||||
let pid = plugin_id.clone();
|
||||
let bus = self.bus.clone();
|
||||
let timers = self.timers.clone();
|
||||
|
||||
let handle = tokio::spawn(async move {
|
||||
loop {
|
||||
let next = schedule.upcoming(chrono::Utc).next();
|
||||
match next {
|
||||
Some(t) => {
|
||||
let dur = (t - chrono::Utc::now())
|
||||
.to_std()
|
||||
.unwrap_or(std::time::Duration::ZERO);
|
||||
tokio::time::sleep(dur).await;
|
||||
bus.fire(
|
||||
&pid,
|
||||
TimerEvent {
|
||||
timer_id: tid.clone(),
|
||||
callback_id: cb.clone(),
|
||||
},
|
||||
);
|
||||
}
|
||||
None => break,
|
||||
}
|
||||
}
|
||||
timers.lock().unwrap().remove(&tid);
|
||||
});
|
||||
|
||||
self.timers
|
||||
.lock()
|
||||
.unwrap()
|
||||
.insert(timer_id.clone(), TimerEntry { handle });
|
||||
|
||||
Ok(timer_id)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user