274 lines
9.7 KiB
Rust
274 lines
9.7 KiB
Rust
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
|
|
use anyhow::{Context, Result};
|
|
use axum::Router;
|
|
use axum::response::Redirect;
|
|
use axum::routing::get_service;
|
|
use axum::routing::{get, post};
|
|
use tokio::net::TcpListener;
|
|
use tokio::sync::{Mutex, RwLock, mpsc, oneshot};
|
|
#[cfg(unix)]
|
|
use tokio::time::MissedTickBehavior;
|
|
use tower_http::services::ServeDir;
|
|
use tower_http::services::ServeFile;
|
|
use tracing::info;
|
|
#[cfg(unix)]
|
|
use tracing::warn;
|
|
use wakey_agent::protocol::{AgentCapability, AgentCapabilityOptions, ErrorPayload, ServerMessage};
|
|
|
|
use crate::api;
|
|
use crate::config;
|
|
use crate::state;
|
|
use crate::ws;
|
|
|
|
mod admin;
|
|
mod process;
|
|
pub mod terminals;
|
|
pub use admin::revoke_agent;
|
|
pub use admin::{
|
|
issue_enroll_token, list_enroll_tokens, migrate_sqlite_state, revoke_enroll_token, state_stats,
|
|
};
|
|
pub use process::reload_daemon;
|
|
use process::{remove_pid_file, write_pid_file};
|
|
|
|
/// Shared state for HTTP handlers and websocket relay paths.
|
|
#[derive(Clone)]
|
|
pub struct AppState {
|
|
pub store: Arc<state::Store>,
|
|
pub sessions: Arc<RwLock<HashMap<String, AgentSession>>>,
|
|
pub pending: Arc<Mutex<HashMap<String, oneshot::Sender<AgentReply>>>>,
|
|
pub public_url: String,
|
|
pub command_timeout: Duration,
|
|
pub enroll_token_ttl: Duration,
|
|
pub terminals: terminals::TerminalRegistry,
|
|
pub alerts: api::AlertsHub,
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
pub struct AgentSession {
|
|
pub connection_id: String,
|
|
pub tx: mpsc::UnboundedSender<SessionEvent>,
|
|
pub capabilities: Vec<AgentCapability>,
|
|
pub capability_options: AgentCapabilityOptions,
|
|
}
|
|
#[derive(Clone)]
|
|
pub enum SessionEvent {
|
|
Message(ServerMessage),
|
|
Close,
|
|
}
|
|
|
|
pub enum AgentReply {
|
|
Result(serde_json::Value),
|
|
Error(ErrorPayload),
|
|
}
|
|
|
|
fn public_api_routes(ui_dist_dir: std::path::PathBuf) -> Router<AppState> {
|
|
let index_file = ui_dist_dir.join("index.html");
|
|
|
|
Router::new()
|
|
.route("/ui", get(|| async { Redirect::temporary("/ui/") }))
|
|
.route("/", get(|| async { Redirect::temporary("/ui/") })) // same as caddyfile
|
|
.nest_service(
|
|
"/ui/",
|
|
get_service(ServeDir::new(ui_dist_dir).fallback(ServeFile::new(index_file))),
|
|
)
|
|
.route("/healthz", get(api::healthz))
|
|
.route("/api/v1/agents/enroll", post(api::enroll))
|
|
.route("/api/v1/agent/ws", get(ws::agent_ws))
|
|
.route(
|
|
"/api/v1/agent/terminals/{terminal_id}/ws",
|
|
get(api::agent_terminal_ws),
|
|
)
|
|
}
|
|
|
|
fn control_api_routes() -> Router<AppState> {
|
|
Router::new()
|
|
.route(
|
|
"/api/v1/control/enroll-token",
|
|
post(api::issue_enroll_token),
|
|
)
|
|
.route(
|
|
"/api/v1/control/enroll-tokens",
|
|
get(api::list_enroll_tokens),
|
|
)
|
|
.route(
|
|
"/api/v1/control/enroll-tokens/{token}",
|
|
axum::routing::delete(api::revoke_enroll_token),
|
|
)
|
|
.route("/api/v1/control/state-stats", get(api::state_stats))
|
|
.route(
|
|
"/api/v1/control/fleet/devices",
|
|
get(api::list_fleet_devices),
|
|
)
|
|
.route(
|
|
"/api/v1/control/fleet/devices/refresh",
|
|
post(api::refresh_fleet_devices),
|
|
)
|
|
.route("/api/v1/control/fleet/wake", post(api::wake_fleet_device))
|
|
.route(
|
|
"/api/v1/control/devices",
|
|
get(api::list_known_devices).post(api::create_known_device),
|
|
)
|
|
.route(
|
|
"/api/v1/control/devices/{device_id}",
|
|
get(api::get_known_device).delete(api::forget_known_device),
|
|
)
|
|
.route(
|
|
"/api/v1/control/devices/{device_id}/merge",
|
|
post(api::merge_known_device),
|
|
)
|
|
.route(
|
|
"/api/v1/control/devices/{device_id}/identifiers",
|
|
post(api::attach_device_identifier),
|
|
)
|
|
.route(
|
|
"/api/v1/control/devices/{device_id}/identifiers/{identifier_key}",
|
|
axum::routing::delete(api::detach_device_identifier),
|
|
)
|
|
.route("/api/v1/control/audit/events", get(api::list_audit_events))
|
|
.route("/api/v1/control/alerts", get(api::active_alerts))
|
|
.route("/api/v1/control/alerts/history", get(api::alert_history))
|
|
.route("/api/v1/control/alerts/ws", get(api::alerts_stream))
|
|
.route("/api/v1/control/agents", get(api::list_agents))
|
|
.route(
|
|
"/api/v1/control/agents/{agent_id}",
|
|
axum::routing::delete(api::revoke_agent),
|
|
)
|
|
.route(
|
|
"/api/v1/control/agents/{agent_id}/nickname",
|
|
axum::routing::patch(api::set_agent_nickname),
|
|
)
|
|
.route(
|
|
"/api/v1/control/agents/{agent_id}/command",
|
|
post(api::run_command),
|
|
)
|
|
.route(
|
|
"/api/v1/control/terminals",
|
|
get(api::list_terminals).post(api::create_terminal),
|
|
)
|
|
.route(
|
|
"/api/v1/control/terminals/events/ws",
|
|
get(api::terminal_events_ws),
|
|
)
|
|
.route(
|
|
"/api/v1/control/terminals/{terminal_id}",
|
|
get(api::get_terminal).delete(api::close_terminal),
|
|
)
|
|
.route(
|
|
"/api/v1/control/terminals/{terminal_id}/attach",
|
|
post(api::attach_terminal),
|
|
)
|
|
.route(
|
|
"/api/v1/control/terminals/{terminal_id}/ws",
|
|
get(api::operator_terminal_ws),
|
|
)
|
|
}
|
|
|
|
/// Starts the control-plane HTTP and websocket surfaces and manages daemon lifecycle hooks.
|
|
pub async fn serve(daemon: config::DaemonConfig) -> Result<()> {
|
|
write_pid_file(&daemon.pid_file)?;
|
|
info!(pid_file = %daemon.pid_file.display(), "wrote control-plane pid file");
|
|
|
|
let store = state::Store::load_or_init(
|
|
&daemon.state_file,
|
|
daemon.bootstrap_enroll_tokens.clone(),
|
|
daemon.enroll_token_ttl,
|
|
)
|
|
.await
|
|
.with_context(|| format!("failed to initialize store {}", daemon.state_file.display()))?;
|
|
|
|
let app_state = AppState {
|
|
store: Arc::new(store),
|
|
sessions: Arc::new(RwLock::new(HashMap::new())),
|
|
pending: Arc::new(Mutex::new(HashMap::new())),
|
|
public_url: daemon.public_url.clone(),
|
|
command_timeout: daemon.command_timeout,
|
|
enroll_token_ttl: daemon.enroll_token_ttl,
|
|
terminals: terminals::TerminalRegistry::new(),
|
|
alerts: api::AlertsHub::new(),
|
|
};
|
|
|
|
tokio::spawn(api::run_alert_evaluator(app_state.clone()));
|
|
|
|
// Keep route classes explicit so edge policy can map directly:
|
|
// - public_api_routes: intended internet-facing agent endpoints
|
|
// - control_api_routes: intended admin-only endpoints behind Access
|
|
let app = Router::new()
|
|
.merge(public_api_routes(daemon.ui_dist_dir.clone()))
|
|
.merge(control_api_routes())
|
|
.with_state(app_state.clone());
|
|
|
|
info!(bind = %daemon.bind, data_dir = %daemon.data_dir.display(), pid_file = %daemon.pid_file.display(), state_file = %daemon.state_file.display(), ui_dist_dir = %daemon.ui_dist_dir.display(), "starting control-plane server");
|
|
let listener = TcpListener::bind(daemon.bind).await?;
|
|
let server = tokio::spawn(async move {
|
|
axum::serve(listener, app)
|
|
.await
|
|
.context("control-plane server exited unexpectedly")
|
|
});
|
|
|
|
#[cfg(unix)]
|
|
{
|
|
use tokio::signal::unix::{SignalKind, signal};
|
|
let mut server = server;
|
|
let mut hup = signal(SignalKind::hangup()).context("failed to install SIGHUP handler")?;
|
|
let mut gc_tick = tokio::time::interval(Duration::from_secs(300));
|
|
gc_tick.set_missed_tick_behavior(MissedTickBehavior::Skip);
|
|
|
|
// Unix daemon loop: shutdown signal, config reload trigger, and periodic maintenance.
|
|
loop {
|
|
tokio::select! {
|
|
_ = tokio::signal::ctrl_c() => {
|
|
info!("ctrl-c received; shutting down control-plane");
|
|
server.abort();
|
|
break;
|
|
}
|
|
_ = hup.recv() => {
|
|
match app_state.store.reload_from_disk().await {
|
|
Ok(()) => info!("reloaded state from disk"),
|
|
Err(err) => warn!(error = %err, "failed to reload state from disk"),
|
|
}
|
|
}
|
|
_ = gc_tick.tick() => {
|
|
match app_state.store.gc_expired_enroll_tokens().await {
|
|
Ok(removed) => {
|
|
if removed > 0 {
|
|
info!(removed, "periodic gc removed expired enroll tokens");
|
|
}
|
|
}
|
|
Err(err) => warn!(error = %err, "periodic gc failed"),
|
|
}
|
|
match app_state.store.gc_retention(daemon.observation_retention).await {
|
|
Ok(stats) => {
|
|
if stats.audit_events_removed > 0 || stats.alert_transitions_removed > 0 {
|
|
info!(
|
|
audit_events_removed = stats.audit_events_removed,
|
|
alert_transitions_removed = stats.alert_transitions_removed,
|
|
"periodic gc pruned rows older than retention window"
|
|
);
|
|
}
|
|
}
|
|
Err(err) => warn!(error = %err, "periodic retention gc failed"),
|
|
}
|
|
}
|
|
join = &mut server => {
|
|
let _ = remove_pid_file(&daemon.pid_file);
|
|
return join.context("control-plane join failed")?;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(not(unix))]
|
|
{
|
|
tokio::signal::ctrl_c()
|
|
.await
|
|
.context("failed waiting for ctrl-c")?;
|
|
server.abort();
|
|
}
|
|
|
|
let _ = remove_pid_file(&daemon.pid_file);
|
|
Ok(())
|
|
}
|