use std::sync::{Arc, OnceLock};
use mercurio_server::{
admin::{Admin, AdminError, AdminExtension, ErrorCode, Paging},
plugin::BoxFuture,
};
use serde::{de::DeserializeOwned, Deserialize};
use serde_json::{json, Value};
use tracing::info;
use crate::bridge::BridgeStatus;
use crate::logging::LogLevel;
use crate::reload::{self, Reload};
pub struct DaemonCommands {
log_level: LogLevel,
reloadable: Arc<[Arc<dyn Reload>]>,
config: Option<Value>,
bridges: Vec<Arc<BridgeStatus>>,
admin: OnceLock<Admin>,
}
#[derive(Deserialize)]
struct LevelArgs {
level: String,
}
#[derive(Deserialize)]
struct PageArgs {
offset: Option<usize>,
count: Option<usize>,
}
impl PageArgs {
fn paging(&self) -> Result<Paging, AdminError> {
let default = Paging::default();
Paging::new(
self.offset.unwrap_or(default.offset),
self.count.unwrap_or(default.count),
)
}
}
impl DaemonCommands {
pub fn new(log_level: LogLevel) -> Self {
DaemonCommands {
log_level,
reloadable: Arc::new([]),
config: None,
bridges: Vec::new(),
admin: OnceLock::new(),
}
}
pub fn with_reload(mut self, files: Arc<[Arc<dyn Reload>]>) -> Self {
self.reloadable = files;
self
}
pub fn with_config(mut self, config: Option<Value>) -> Self {
self.config = config;
self
}
pub fn with_bridges(mut self, bridges: Vec<Arc<BridgeStatus>>) -> Self {
self.bridges = bridges;
self
}
pub fn set_admin(&self, admin: Admin) {
let _ = self.admin.set(admin);
}
async fn list_bridges(&self, args: &Value) -> Result<Value, AdminError> {
let paging = arguments::<PageArgs>(args)?.paging()?;
let mut bridges = Vec::new();
for status in self.bridges.iter().skip(paging.offset).take(paging.count) {
let queued = match self.admin.get() {
Some(admin) => admin
.client(status.client_id())
.await
.map_or(0, |client| client.queued),
None => 0,
};
bridges.push(status.entry(queued));
}
Ok(json!({ "bridges": bridges, "totalCount": self.bridges.len() }))
}
fn reload(&self) -> Value {
info!("Reloading on request");
json!({ "items": reload::reload_all(&self.reloadable) })
}
fn set_log_level(&self, args: &Value) -> Result<Value, AdminError> {
let level = arguments::<LevelArgs>(args)?.level;
self.log_level.set(&level)?;
info!("Log level set to `{}` until restart", level);
Ok(Value::Null)
}
}
impl AdminExtension for DaemonCommands {
fn handle<'a>(
&'a self,
command: &'a str,
args: &'a Value,
) -> BoxFuture<'a, Option<Result<Value, AdminError>>> {
Box::pin(async move {
Some(match command {
"listBridges" => self.list_bridges(args).await,
"getLogLevel" => Ok(json!({ "level": self.log_level.get() })),
"setLogLevel" => self.set_log_level(args),
"reload" => Ok(self.reload()),
"getConfig" => self.config.clone().ok_or_else(|| {
AdminError::new(
ErrorCode::Unavailable,
"the configuration can't be shown as JSON",
)
}),
_ => return None,
})
})
}
}
fn arguments<T: DeserializeOwned>(args: &Value) -> Result<T, AdminError> {
T::deserialize(args).map_err(|e| AdminError::invalid_argument(e.to_string()))
}
#[cfg(test)]
mod tests {
use std::{sync::Arc, time::Duration};
use mercurio_client::{ConnectOptions, Event, MqttClient, QoS};
use mercurio_server::{
admin::{AdminExtension, ErrorCode},
server::{LocalConnector, LocalListener, Server, ServerConfig},
};
use mercurio_storage::memory::MemoryStore;
use serde_json::{json, Value};
use tokio::net::TcpListener;
use super::DaemonCommands;
use crate::bridge::{self, BridgeStatus};
use crate::config::BridgeConfig;
use crate::logging::LogLevel;
use crate::reload::{Reload, Reloadable};
async fn start(commands: Arc<DaemonCommands>, local_listener: LocalListener) {
let mut config = ServerConfig::default();
config.local_listener = Some(local_listener);
let extension: Arc<dyn AdminExtension> = commands.clone();
config.admin_extension = Some(extension);
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let server = Server::new(listener, config, Arc::new(MemoryStore::new()))
.await
.unwrap();
commands.set_admin(server.admin());
tokio::spawn(server.run(std::future::pending::<()>()));
}
async fn request(connector: &LocalConnector, command: Value) -> Value {
let connector = connector.clone();
let options = ConnectOptions::new("localhost", 0)
.client_id("ctl")
.connector(move || connector.connect());
let client = MqttClient::connect(options).await.unwrap();
client
.subscribe(&[("$CONTROL/broker/v1/response", QoS::AtMostOnce)])
.await
.unwrap();
let request = json!({ "commands": [command] }).to_string();
client
.publish("$CONTROL/broker/v1", request, QoS::AtMostOnce)
.await
.unwrap();
let payload = loop {
let event = tokio::time::timeout(Duration::from_secs(5), client.recv())
.await
.expect("no response in time");
match event {
Some(Event::Message { payload, .. }) => break payload,
Some(_) => continue,
None => panic!("disconnected before the response"),
}
};
let _ = client.disconnect().await;
serde_json::from_slice::<Value>(&payload).unwrap()["responses"][0].clone()
}
#[tokio::test(flavor = "multi_thread")]
async fn test_log_level_commands_are_served_through_the_broker() {
let (log_level, _filter) = LogLevel::new("info,mercurio_server=debug", true);
let commands = Arc::new(DaemonCommands::new(log_level));
let (connector, local_listener) = LocalConnector::new();
start(commands, local_listener).await;
let response = request(&connector, json!({ "command": "getLogLevel" })).await;
assert_eq!(response["data"]["level"], "info,mercurio_server=debug");
let response = request(
&connector,
json!({ "command": "setLogLevel", "level": "info,x=loud" }),
)
.await;
assert_eq!(response["code"], "invalidArgument", "{response}");
let response = request(&connector, json!({ "command": "setLogLevel" })).await;
assert_eq!(response["code"], "invalidArgument", "{response}");
let response = request(&connector, json!({ "command": "listBridgez" })).await;
assert_eq!(response["code"], "invalidArgument", "{response}");
}
#[tokio::test]
async fn test_set_log_level_changes_what_get_log_level_returns() {
let (log_level, _filter) = LogLevel::new("info", true);
let commands = DaemonCommands::new(log_level);
let args = json!({ "command": "setLogLevel", "level": "warn" });
assert_eq!(
commands.handle("setLogLevel", &args).await.unwrap(),
Ok(Value::Null)
);
let level = commands.handle("getLogLevel", &json!({})).await.unwrap();
assert_eq!(level.unwrap()["level"], "warn");
assert!(commands.handle("getlogLevel", &json!({})).await.is_none());
}
#[tokio::test(flavor = "multi_thread")]
async fn test_reload_reports_each_file() {
let path =
std::env::temp_dir().join(format!("mercuriod-commands-acl-{}", std::process::id()));
std::fs::write(&path, "topic read public/#\n").unwrap();
let path_str = path.to_str().unwrap().to_string();
let acl: Arc<dyn Reload> =
Arc::new(Reloadable::load(&path_str, crate::load_acl_file).unwrap());
let (log_level, _filter) = LogLevel::new("info", true);
let commands = DaemonCommands::new(log_level).with_reload(vec![acl].into());
let args = json!({ "command": "reload" });
let result = commands.handle("reload", &args).await.unwrap().unwrap();
assert_eq!(
result,
json!({ "items": [{ "item": "acl-file", "path": path_str, "ok": true }] })
);
std::fs::write(&path, "bogus\n").unwrap();
let result = commands.handle("reload", &args).await.unwrap().unwrap();
assert_eq!(result["items"][0]["ok"], false, "{result}");
let error = result["items"][0]["error"].as_str().unwrap();
assert!(error.contains("unknown keyword"), "{error}");
std::fs::remove_file(&path).unwrap();
let (log_level, _filter) = LogLevel::new("info", true);
let none = DaemonCommands::new(log_level);
assert_eq!(
none.handle("reload", &args).await.unwrap(),
Ok(json!({ "items": [] }))
);
}
#[tokio::test]
async fn test_get_config_returns_the_config_given() {
let args = json!({ "command": "getConfig" });
let config = json!({ "server": { "port": 1883 } });
let (log_level, _filter) = LogLevel::new("info", true);
let commands = DaemonCommands::new(log_level).with_config(Some(config.clone()));
assert_eq!(
commands.handle("getConfig", &args).await.unwrap(),
Ok(config)
);
let (log_level, _filter) = LogLevel::new("info", true);
let without = DaemonCommands::new(log_level);
let error = without
.handle("getConfig", &args)
.await
.unwrap()
.unwrap_err();
assert_eq!(error.code, ErrorCode::Unavailable);
}
#[tokio::test]
async fn test_list_bridges_pages_in_configuration_order() {
let statuses = vec![
Arc::new(BridgeStatus::new("a", "mercurio-bridge-a")),
Arc::new(BridgeStatus::new("b", "mercurio-bridge-b")),
];
let (log_level, _filter) = LogLevel::new("info", true);
let commands = DaemonCommands::new(log_level).with_bridges(statuses);
let all = commands
.handle("listBridges", &json!({}))
.await
.unwrap()
.unwrap();
assert_eq!(all["totalCount"], 2);
assert_eq!(all["bridges"][0]["name"], "a");
assert_eq!(all["bridges"][1]["queued"], 0);
let second = commands
.handle("listBridges", &json!({ "offset": 1, "count": 1 }))
.await
.unwrap()
.unwrap();
assert_eq!(second["bridges"].as_array().unwrap().len(), 1);
assert_eq!(second["bridges"][0]["name"], "b");
assert_eq!(second["totalCount"], 2);
let error = commands
.handle("listBridges", &json!({ "count": 1001 }))
.await
.unwrap()
.unwrap_err();
assert_eq!(error.code, ErrorCode::InvalidArgument);
}
async fn bridges_when(commands: &DaemonCommands, done: impl Fn(&Value) -> bool) -> Value {
let deadline = tokio::time::Instant::now() + Duration::from_secs(20);
loop {
let listed = commands
.handle("listBridges", &json!({}))
.await
.unwrap()
.unwrap();
if done(&listed["bridges"][0]) {
return listed;
}
assert!(
tokio::time::Instant::now() < deadline,
"timed out: {listed}"
);
tokio::time::sleep(Duration::from_millis(50)).await;
}
}
#[tokio::test(flavor = "multi_thread")]
async fn test_list_bridges_follows_a_real_bridge() {
let remote = TcpListener::bind("127.0.0.1:0").await.unwrap();
let port = remote.local_addr().unwrap().port();
let (stop_remote, stopped) = tokio::sync::oneshot::channel::<()>();
let remote = tokio::spawn(mercurio_server::server::run(remote, async move {
let _ = stopped.await;
}));
let config: BridgeConfig = toml::from_str(&format!(
"name = \"cloud\"\nhost = \"127.0.0.1\"\nport = {port}\n[[topics]]\npattern = \"a/#\"\n"
))
.unwrap();
let (connector, local_listener) = LocalConnector::new();
let status = bridge::spawn(config, connector.clone()).unwrap();
let (log_level, _filter) = LogLevel::new("info", true);
let commands = Arc::new(DaemonCommands::new(log_level).with_bridges(vec![status]));
start(Arc::clone(&commands), local_listener).await;
let listed = bridges_when(&commands, |b| b["connected"] == true).await;
assert_eq!(listed["totalCount"], 1);
let bridge = &listed["bridges"][0];
assert_eq!(bridge["name"], "cloud");
assert_eq!(bridge["remoteAddress"], format!("127.0.0.1:{port}"));
assert!(bridge["connectedSince"].is_u64(), "{bridge}");
assert_eq!(bridge["queued"], 0);
let response = request(&connector, json!({ "command": "listBridges" })).await;
assert_eq!(
response["data"]["bridges"][0]["connected"], true,
"{response}"
);
stop_remote.send(()).unwrap();
remote.await.unwrap().unwrap();
let listed = bridges_when(&commands, |b| b["connected"] == false).await;
let bridge = &listed["bridges"][0];
assert!(bridge["disconnectedSince"].is_u64(), "{bridge}");
assert_eq!(bridge["remoteAddress"], format!("127.0.0.1:{port}"));
}
}