use std::sync::{Arc, RwLock};
use std::time::Duration;
use anyhow::Result;
use mcpmesh_local_api::{
API_NAME, API_VERSION, Hello, PeerInfo, ServiceInfo, StatusResult, method_of,
};
use mcpmesh_net::framing::{FrameReader, Inbound, write_frame};
use serde_json::{Value, json};
use tokio::net::{UnixListener, UnixStream};
use tokio::sync::Notify;
use crate::daemon::MeshState;
use crate::ipc::{self, MAX_FRAME_BYTES};
pub struct DaemonState {
pub stack_version: String,
pub services: RwLock<Vec<ServiceInfo>>,
pub peers: RwLock<Vec<PeerInfo>>,
pub(crate) mesh: Option<Arc<MeshState>>,
shutdown: Notify,
}
impl DaemonState {
pub fn new(stack_version: impl Into<String>) -> Self {
Self {
stack_version: stack_version.into(),
services: RwLock::new(Vec::new()),
peers: RwLock::new(Vec::new()),
mesh: None,
shutdown: Notify::new(),
}
}
pub fn with_mesh(
stack_version: impl Into<String>,
mesh: Arc<MeshState>,
services: Vec<ServiceInfo>,
peers: Vec<PeerInfo>,
) -> Self {
Self {
stack_version: stack_version.into(),
services: RwLock::new(services),
peers: RwLock::new(peers),
mesh: Some(mesh),
shutdown: Notify::new(),
}
}
pub(crate) fn mesh(&self) -> Option<&Arc<MeshState>> {
self.mesh.as_ref()
}
}
pub async fn serve_control(listener: UnixListener, state: Arc<DaemonState>) -> Result<()> {
loop {
tokio::select! {
() = state.shutdown.notified() => {
tracing::info!("shutdown requested; control server stopping");
return Ok(());
}
accepted = listener.accept() => {
let (stream, _addr) = match accepted {
Ok(pair) => pair,
Err(e) => {
tracing::warn!(%e, "control accept failed; backing off");
tokio::time::sleep(Duration::from_millis(50)).await;
continue;
}
};
let state = state.clone();
tokio::spawn(async move {
if let Err(e) = handle_conn(stream, state).await {
tracing::debug!(%e, "control connection ended");
}
});
}
}
}
}
async fn handle_conn(stream: UnixStream, state: Arc<DaemonState>) -> Result<()> {
if let Err(e) = ipc::check_peer_uid(&stream) {
tracing::warn!(%e, "refused control connection from a different uid");
return Ok(());
}
let (read_half, mut write_half) = stream.into_split();
let hello = Hello {
api: API_NAME.into(),
api_version: API_VERSION.into(),
stack_version: state.stack_version.clone(),
};
write_frame(&mut write_half, &serde_json::to_value(&hello)?).await?;
let mut reader = FrameReader::new(tokio::io::BufReader::new(read_half), MAX_FRAME_BYTES);
loop {
match reader.next().await? {
None => return Ok(()), Some(Inbound::Violation(v)) => {
let resp = error(Value::Null, -32700, format!("invalid request frame: {v:?}"));
write_frame(&mut write_half, &resp).await?;
}
Some(Inbound::Frame(req)) => {
if method_of(&req) == Some("shutdown") {
state.shutdown.notify_one();
let resp = dispatch(&req, &state);
let _ = write_frame(&mut write_half, &resp).await;
return Ok(());
}
if method_of(&req) == Some("open_session") {
let params = req.get("params").cloned().unwrap_or(Value::Null);
let peer = str_param(¶ms, "peer");
let service = str_param(¶ms, "service");
return crate::daemon::open_session(
&state, &peer, &service, reader, write_half,
)
.await;
}
if method_of(&req) == Some("subscribe") {
return run_subscription(&state, write_half).await;
}
let resp = handle_request(&req, &state).await;
write_frame(&mut write_half, &resp).await?;
}
}
}
}
async fn run_subscription(
state: &Arc<DaemonState>,
mut w: impl tokio::io::AsyncWrite + Unpin,
) -> Result<()> {
use crate::stream::StreamFrame;
let (audit, mesh) = match state.mesh() {
Some(mesh) => (mesh.audit(), Some(mesh)),
None => (crate::audit::AuditSink::disabled(), None),
};
let rx = audit.subscribe();
let snapshot = StreamFrame::Snapshot {
active_sessions: audit.active_sessions(),
reachability: mesh.map(crate::daemon::reachability_of).unwrap_or_default(),
};
write_frame(&mut w, &serde_json::to_value(&snapshot)?).await?;
let Some(mut rx) = rx else {
return Ok(());
};
use tokio::sync::broadcast::error::RecvError;
loop {
let frame = match rx.recv().await {
Ok(record) => StreamFrame::Event {
record: Box::new(record),
},
Err(RecvError::Lagged(n)) => StreamFrame::Lagged { dropped: n },
Err(RecvError::Closed) => return Ok(()),
};
if write_frame(&mut w, &serde_json::to_value(&frame)?)
.await
.is_err()
{
return Ok(()); }
}
}
async fn handle_request(req: &Value, state: &DaemonState) -> Value {
let id = req.get("id").cloned().unwrap_or(Value::Null);
let params = req.get("params").cloned().unwrap_or(Value::Null);
match method_of(req) {
Some("register_service") => respond(
id,
"register_service",
crate::daemon::register_service(state, ¶ms)
.await
.map(unit),
),
Some("peer_add") => respond(
id,
"peer_add",
crate::daemon::add_peer(state, ¶ms).await.map(unit),
),
Some("peer_remove") => respond(
id,
"peer_remove",
crate::daemon::remove_peer(state, ¶ms).await.map(unit),
),
Some("peer_rename") => respond(
id,
"peer_rename",
crate::daemon::rename_peer(state, ¶ms).await.map(unit),
),
Some("invite") => {
let services: Vec<String> = params
.get("services")
.and_then(Value::as_array)
.map(|a| {
a.iter()
.filter_map(|v| v.as_str().map(String::from))
.collect()
})
.unwrap_or_default();
let Some(mesh) = state.mesh() else {
return error(id, -32000, "daemon has no mesh (control-only mode)");
};
respond(
id,
"invite",
crate::daemon::mint_invite(services, mesh).await,
)
}
Some("pair") => {
let invite_line = str_param(¶ms, "invite_line");
respond(id, "pair", crate::daemon::redeem(state, invite_line).await)
}
Some("roster_install") => {
let path = str_param(¶ms, "path");
let org_root_pk = params
.get("org_root_pk")
.and_then(Value::as_str)
.map(String::from);
respond(
id,
"roster_install",
crate::daemon::install_roster(state, path, org_root_pk).await,
)
}
Some("org_join") => respond(
id,
"org_join",
crate::daemon::org_join(
state,
str_param(¶ms, "org_id"),
str_param(¶ms, "org_root_pk"),
str_param(¶ms, "user_id"),
str_param(¶ms, "user_key"),
)
.await,
),
Some("set_roster_url") => respond(
id,
"set_roster_url",
crate::daemon::set_roster_url(state, str_param(¶ms, "url"))
.await
.map(unit),
),
Some("blob_publish") => respond(
id,
"blob_publish",
crate::daemon::blob_publish(
state,
str_param(¶ms, "scope"),
str_param(¶ms, "path"),
)
.await,
),
Some("blob_grant") => respond(
id,
"blob_grant",
crate::daemon::blob_grant(
state,
str_param(¶ms, "scope"),
str_param(¶ms, "principal"),
)
.await
.map(unit),
),
Some("blob_list") => respond(id, "blob_list", crate::daemon::blob_list(state).await),
Some("blob_fetch") => respond(
id,
"blob_fetch",
crate::daemon::blob_fetch(
state,
str_param(¶ms, "ticket"),
str_param(¶ms, "dest_path"),
)
.await,
),
Some("audit_summary") => {
match tokio::task::spawn_blocking(|| {
let dir = mcpmesh_trust::paths::default_audit_dir()?;
crate::audit::read_all_records(&dir)
.map(|recs| crate::audit::summarize_sessions(&recs))
})
.await
{
Ok(r) => respond(id, "audit_summary", r.map_err(anyhow::Error::from)),
Err(e) => error(id, -32000, format!("audit_summary task failed: {e}")),
}
}
_ => dispatch(req, state),
}
}
fn respond<T: serde::Serialize>(id: Value, method: &str, r: anyhow::Result<T>) -> Value {
match r {
Ok(v) => ok(
id,
serde_json::to_value(v).expect("control result serializes"),
),
Err(e) => error(id, -32000, format!("{method} failed: {e}")),
}
}
fn unit((): ()) -> Value {
json!({})
}
fn dispatch(req: &Value, state: &DaemonState) -> Value {
let id = req.get("id").cloned().unwrap_or(Value::Null);
match method_of(req) {
Some("status") => {
let result =
serde_json::to_value(status_result(state)).expect("StatusResult serializes");
ok(id, result)
}
Some("shutdown") => ok(id, json!({})),
Some(other) => error(id, -32601, format!("unknown method: {other}")),
None => error(id, -32600, "request is missing a `method`"),
}
}
pub(crate) fn status_result(state: &DaemonState) -> StatusResult {
let (services, peers, roster) = match state.mesh() {
Some(mesh) => {
let cfg = crate::config::Config::load(&mesh.config_path).ok();
let services = cfg
.as_ref()
.map(crate::daemon::service_infos)
.unwrap_or_else(|| {
state
.services
.read()
.expect("services lock not poisoned")
.clone()
});
let roster = crate::daemon::roster_status(mesh, cfg.as_ref());
(services, crate::daemon::peer_infos(&mesh.store), roster)
}
None => (
state
.services
.read()
.expect("services lock not poisoned")
.clone(),
state.peers.read().expect("peers lock not poisoned").clone(),
None,
),
};
let presence = state
.mesh()
.map(crate::daemon::presence_peers)
.unwrap_or_default();
let self_user_id = state
.mesh()
.and_then(|mesh| mesh.self_binding())
.map(|binding| binding.user_pk);
let recent_pairings = state
.mesh()
.map(|mesh| mesh.recent_pairings())
.unwrap_or_default();
let reachability = state
.mesh()
.map(crate::daemon::reachability_of)
.unwrap_or_default();
StatusResult {
stack_version: state.stack_version.clone(),
services,
peers,
roster,
presence,
self_user_id,
recent_pairings,
reachability,
}
}
fn str_param(params: &Value, key: &str) -> String {
params
.get(key)
.and_then(Value::as_str)
.unwrap_or_default()
.to_string()
}
fn ok(id: Value, result: Value) -> Value {
json!({ "jsonrpc": "2.0", "id": id, "result": result })
}
fn error(id: Value, code: i64, message: impl Into<String>) -> Value {
json!({ "jsonrpc": "2.0", "id": id, "error": { "code": code, "message": message.into() } })
}
#[cfg(test)]
mod tests {
use super::*;
fn control_only() -> Arc<DaemonState> {
Arc::new(DaemonState::new("0.1.0-test"))
}
fn req(method: &str, params: Value) -> Value {
json!({ "jsonrpc": "2.0", "id": 1, "method": method, "params": params })
}
#[test]
fn dispatch_status_answers_from_the_snapshot() {
let st = control_only();
let r = dispatch(&req("status", json!({})), &st);
assert_eq!(r["result"]["stack_version"], "0.1.0-test");
assert!(r["result"]["peers"].as_array().unwrap().is_empty());
assert!(r["result"]["services"].as_array().unwrap().is_empty());
assert!(r["result"]["roster"].is_null());
}
#[test]
fn dispatch_status_tolerates_any_params_shape() {
let st = control_only();
for p in [json!({}), Value::Null, json!({ "junk": true })] {
assert!(dispatch(&req("status", p), &st).get("result").is_some());
}
let omitted = json!({ "jsonrpc": "2.0", "id": 1, "method": "status" });
assert!(dispatch(&omitted, &st).get("result").is_some());
}
#[test]
fn dispatch_shutdown_acks_and_unknown_methods_error() {
let st = control_only();
assert_eq!(
dispatch(&req("shutdown", json!({})), &st)["result"],
json!({})
);
assert_eq!(
dispatch(&req("frobnicate", json!({})), &st)["error"]["code"],
-32601
);
let no_method = json!({ "jsonrpc": "2.0", "id": 1 });
assert_eq!(dispatch(&no_method, &st)["error"]["code"], -32600);
}
#[tokio::test]
async fn mesh_methods_error_gracefully_without_a_mesh() {
let st = control_only();
for method in [
"register_service",
"peer_add",
"peer_remove",
"peer_rename",
"invite",
"pair",
"roster_install",
"org_join",
"set_roster_url",
"blob_publish",
"blob_grant",
"blob_list",
"blob_fetch",
] {
let r = handle_request(&req(method, json!({})), &st).await;
assert_eq!(
r["error"]["code"], -32000,
"method {method} should error in control-only mode, got {r}"
);
}
}
#[tokio::test]
async fn audit_summary_works_in_control_only_mode() {
let st = control_only();
let r = handle_request(&req("audit_summary", json!({})), &st).await;
assert!(
r.get("result").is_some(),
"audit_summary should succeed: {r}"
);
}
#[tokio::test]
async fn handle_request_delegates_status_to_dispatch() {
let st = control_only();
let r = handle_request(&req("status", json!({})), &st).await;
assert_eq!(r["result"]["stack_version"], "0.1.0-test");
}
}