modular
This commit is contained in:
@@ -0,0 +1,165 @@
|
||||
use axum::Json;
|
||||
use axum::extract::{Path as AxumPath, State};
|
||||
use axum::http::StatusCode;
|
||||
use axum::response::IntoResponse;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tracing::{info, info_span, warn};
|
||||
use uuid::Uuid;
|
||||
use wakey_agent::protocol::{AgentCommand, ErrorPayload, RequestId, ServerMessage};
|
||||
|
||||
use crate::api::json_error;
|
||||
use crate::runtime::{AgentReply, AppState};
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct AgentStatus {
|
||||
pub agent_id: String,
|
||||
pub connected: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct RelayCommandRequest {
|
||||
pub command: AgentCommand,
|
||||
pub timeout_ms: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct RelayCommandResponse {
|
||||
pub request_id: String,
|
||||
pub status: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub result: Option<serde_json::Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub error: Option<ErrorPayload>,
|
||||
}
|
||||
|
||||
pub async fn list_agents(
|
||||
State(state): State<AppState>,
|
||||
) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
|
||||
let enrolled = state.store.list_agents().await;
|
||||
let sessions = state.sessions.read().await;
|
||||
|
||||
let agents = enrolled
|
||||
.into_iter()
|
||||
.map(|agent_id| AgentStatus {
|
||||
connected: sessions.contains_key(&agent_id),
|
||||
agent_id,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
Ok((StatusCode::OK, Json(agents)))
|
||||
}
|
||||
|
||||
pub async fn run_command(
|
||||
State(state): State<AppState>,
|
||||
AxumPath(agent_id): AxumPath<String>,
|
||||
Json(req): Json<RelayCommandRequest>,
|
||||
) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
|
||||
let request_id_string = format!("req-{}", Uuid::new_v4());
|
||||
let command = command_kind(&req.command);
|
||||
let span = info_span!(
|
||||
"relay_command",
|
||||
agent_id = %agent_id,
|
||||
request_id = %request_id_string,
|
||||
command = %command,
|
||||
);
|
||||
let _span_guard = span.enter();
|
||||
|
||||
let request_id = RequestId::try_from(request_id_string.clone()).map_err(|err| {
|
||||
json_error(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
"invalid_request_id",
|
||||
&err,
|
||||
)
|
||||
})?;
|
||||
|
||||
let tx = {
|
||||
let sessions = state.sessions.read().await;
|
||||
sessions.get(&agent_id).cloned()
|
||||
}
|
||||
.ok_or_else(|| {
|
||||
warn!("command rejected: agent not connected");
|
||||
json_error(
|
||||
StatusCode::NOT_FOUND,
|
||||
"agent_not_connected",
|
||||
"agent is not connected",
|
||||
)
|
||||
})?;
|
||||
|
||||
let (pending_tx, pending_rx) = tokio::sync::oneshot::channel();
|
||||
state
|
||||
.pending
|
||||
.lock()
|
||||
.await
|
||||
.insert(request_id_string.clone(), pending_tx);
|
||||
|
||||
info!("dispatching command to agent");
|
||||
|
||||
if let Err(err) = tx.send(ServerMessage::Command {
|
||||
request_id,
|
||||
command: req.command,
|
||||
}) {
|
||||
state.pending.lock().await.remove(&request_id_string);
|
||||
warn!(error = %err, "failed sending command to agent session");
|
||||
return Err(json_error(
|
||||
StatusCode::BAD_GATEWAY,
|
||||
"agent_send_failed",
|
||||
&format!("failed to send command to agent: {err}"),
|
||||
));
|
||||
}
|
||||
|
||||
let timeout = std::time::Duration::from_millis(
|
||||
req.timeout_ms
|
||||
.unwrap_or(state.command_timeout.as_millis() as u64)
|
||||
.max(1),
|
||||
);
|
||||
let outcome = tokio::time::timeout(timeout, pending_rx).await;
|
||||
let response = match outcome {
|
||||
Ok(Ok(AgentReply::Result(result))) => {
|
||||
info!("agent command completed");
|
||||
RelayCommandResponse {
|
||||
request_id: request_id_string,
|
||||
status: "ok".into(),
|
||||
result: Some(result),
|
||||
error: None,
|
||||
}
|
||||
}
|
||||
Ok(Ok(AgentReply::Error(error))) => {
|
||||
warn!(code = %error.code, "agent command returned error");
|
||||
RelayCommandResponse {
|
||||
request_id: request_id_string,
|
||||
status: "error".into(),
|
||||
result: None,
|
||||
error: Some(error),
|
||||
}
|
||||
}
|
||||
Ok(Err(_)) => {
|
||||
warn!("agent response channel dropped");
|
||||
return Err(json_error(
|
||||
StatusCode::BAD_GATEWAY,
|
||||
"agent_response_dropped",
|
||||
"agent response channel dropped",
|
||||
));
|
||||
}
|
||||
Err(_) => {
|
||||
state.pending.lock().await.remove(&request_id_string);
|
||||
warn!(timeout_ms = timeout.as_millis() as u64, "agent command timed out");
|
||||
return Err(json_error(
|
||||
StatusCode::GATEWAY_TIMEOUT,
|
||||
"agent_timeout",
|
||||
"agent did not answer before timeout",
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
Ok((StatusCode::OK, Json(response)))
|
||||
}
|
||||
|
||||
fn command_kind(command: &AgentCommand) -> &'static str {
|
||||
match command {
|
||||
AgentCommand::Status(_) => "status",
|
||||
AgentCommand::Leases(_) => "leases",
|
||||
AgentCommand::Devs(_) => "devs",
|
||||
AgentCommand::Inventory(_) => "inventory",
|
||||
AgentCommand::Wake(_) => "wake",
|
||||
}
|
||||
}
|
||||
@@ -3,11 +3,10 @@ use axum::extract::{Path as AxumPath, Query, State};
|
||||
use axum::http::StatusCode;
|
||||
use axum::response::IntoResponse;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tracing::{info, info_span, warn};
|
||||
use uuid::Uuid;
|
||||
use wakey_agent::protocol::{AgentCommand, ErrorPayload, RequestId, ServerMessage};
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::runtime::{AgentReply, AppState};
|
||||
use crate::api::json_error;
|
||||
use crate::runtime::AppState;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct EnrollRequest {
|
||||
@@ -59,28 +58,6 @@ pub struct StateStatsResponse {
|
||||
pub expired_enroll_token_count: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct AgentStatus {
|
||||
pub agent_id: String,
|
||||
pub connected: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct RelayCommandRequest {
|
||||
pub command: AgentCommand,
|
||||
pub timeout_ms: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct RelayCommandResponse {
|
||||
pub request_id: String,
|
||||
pub status: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub result: Option<serde_json::Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub error: Option<ErrorPayload>,
|
||||
}
|
||||
|
||||
pub async fn healthz() -> &'static str {
|
||||
"ok"
|
||||
}
|
||||
@@ -125,7 +102,10 @@ pub async fn issue_enroll_token(
|
||||
|
||||
match state.store.issue_enroll_token(ttl).await {
|
||||
Ok(issued) => {
|
||||
info!(expires_at_unix = issued.expires_at_unix, "issued enroll token");
|
||||
info!(
|
||||
expires_at_unix = issued.expires_at_unix,
|
||||
"issued enroll token"
|
||||
);
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
Json(IssueEnrollTokenResponse {
|
||||
@@ -145,23 +125,6 @@ pub async fn issue_enroll_token(
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn list_agents(
|
||||
State(state): State<AppState>,
|
||||
) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
|
||||
let enrolled = state.store.list_agents().await;
|
||||
let sessions = state.sessions.read().await;
|
||||
|
||||
let agents = enrolled
|
||||
.into_iter()
|
||||
.map(|agent_id| AgentStatus {
|
||||
connected: sessions.contains_key(&agent_id),
|
||||
agent_id,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
Ok((StatusCode::OK, Json(agents)))
|
||||
}
|
||||
|
||||
pub async fn list_enroll_tokens(
|
||||
State(state): State<AppState>,
|
||||
Query(query): Query<ListEnrollTokenQuery>,
|
||||
@@ -234,136 +197,3 @@ pub async fn state_stats(
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn run_command(
|
||||
State(state): State<AppState>,
|
||||
AxumPath(agent_id): AxumPath<String>,
|
||||
Json(req): Json<RelayCommandRequest>,
|
||||
) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
|
||||
let request_id_string = format!("req-{}", Uuid::new_v4());
|
||||
let command = command_kind(&req.command);
|
||||
let span = info_span!(
|
||||
"relay_command",
|
||||
agent_id = %agent_id,
|
||||
request_id = %request_id_string,
|
||||
command = %command,
|
||||
);
|
||||
let _span_guard = span.enter();
|
||||
|
||||
let request_id = RequestId::try_from(request_id_string.clone()).map_err(|err| {
|
||||
json_error(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
"invalid_request_id",
|
||||
&err,
|
||||
)
|
||||
})?;
|
||||
|
||||
let tx = {
|
||||
let sessions = state.sessions.read().await;
|
||||
sessions.get(&agent_id).cloned()
|
||||
}
|
||||
.ok_or_else(|| {
|
||||
warn!("command rejected: agent not connected");
|
||||
json_error(
|
||||
StatusCode::NOT_FOUND,
|
||||
"agent_not_connected",
|
||||
"agent is not connected",
|
||||
)
|
||||
})?;
|
||||
|
||||
let (pending_tx, pending_rx) = tokio::sync::oneshot::channel();
|
||||
state
|
||||
.pending
|
||||
.lock()
|
||||
.await
|
||||
.insert(request_id_string.clone(), pending_tx);
|
||||
|
||||
info!(
|
||||
"dispatching command to agent"
|
||||
);
|
||||
|
||||
if let Err(err) = tx.send(ServerMessage::Command {
|
||||
request_id,
|
||||
command: req.command,
|
||||
}) {
|
||||
state.pending.lock().await.remove(&request_id_string);
|
||||
warn!(error = %err, "failed sending command to agent session");
|
||||
return Err(json_error(
|
||||
StatusCode::BAD_GATEWAY,
|
||||
"agent_send_failed",
|
||||
&format!("failed to send command to agent: {err}"),
|
||||
));
|
||||
}
|
||||
|
||||
let timeout = std::time::Duration::from_millis(
|
||||
req.timeout_ms
|
||||
.unwrap_or(state.command_timeout.as_millis() as u64)
|
||||
.max(1),
|
||||
);
|
||||
let outcome = tokio::time::timeout(timeout, pending_rx).await;
|
||||
let response = match outcome {
|
||||
Ok(Ok(AgentReply::Result(result))) => {
|
||||
info!("agent command completed");
|
||||
RelayCommandResponse {
|
||||
request_id: request_id_string,
|
||||
status: "ok".into(),
|
||||
result: Some(result),
|
||||
error: None,
|
||||
}
|
||||
}
|
||||
Ok(Ok(AgentReply::Error(error))) => {
|
||||
warn!(code = %error.code, "agent command returned error");
|
||||
RelayCommandResponse {
|
||||
request_id: request_id_string,
|
||||
status: "error".into(),
|
||||
result: None,
|
||||
error: Some(error),
|
||||
}
|
||||
}
|
||||
Ok(Err(_)) => {
|
||||
warn!("agent response channel dropped");
|
||||
return Err(json_error(
|
||||
StatusCode::BAD_GATEWAY,
|
||||
"agent_response_dropped",
|
||||
"agent response channel dropped",
|
||||
));
|
||||
}
|
||||
Err(_) => {
|
||||
state.pending.lock().await.remove(&request_id_string);
|
||||
warn!(timeout_ms = timeout.as_millis() as u64, "agent command timed out");
|
||||
return Err(json_error(
|
||||
StatusCode::GATEWAY_TIMEOUT,
|
||||
"agent_timeout",
|
||||
"agent did not answer before timeout",
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
Ok((StatusCode::OK, Json(response)))
|
||||
}
|
||||
|
||||
pub fn json_error(
|
||||
status: StatusCode,
|
||||
code: &str,
|
||||
message: &str,
|
||||
) -> (StatusCode, Json<serde_json::Value>) {
|
||||
(
|
||||
status,
|
||||
Json(serde_json::json!({
|
||||
"error": {
|
||||
"code": code,
|
||||
"message": message,
|
||||
}
|
||||
})),
|
||||
)
|
||||
}
|
||||
|
||||
fn command_kind(command: &AgentCommand) -> &'static str {
|
||||
match command {
|
||||
AgentCommand::Status(_) => "status",
|
||||
AgentCommand::Leases(_) => "leases",
|
||||
AgentCommand::Devs(_) => "devs",
|
||||
AgentCommand::Inventory(_) => "inventory",
|
||||
AgentCommand::Wake(_) => "wake",
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
use axum::Json;
|
||||
use axum::http::StatusCode;
|
||||
|
||||
mod commands;
|
||||
mod control;
|
||||
|
||||
pub use commands::{list_agents, run_command};
|
||||
pub use control::{
|
||||
EnrollTokenStatus, IssueEnrollTokenResponse, RevokeEnrollTokenResponse, StateStatsResponse,
|
||||
enroll, healthz, issue_enroll_token, list_enroll_tokens, revoke_enroll_token, state_stats,
|
||||
};
|
||||
|
||||
pub fn json_error(
|
||||
status: StatusCode,
|
||||
code: &str,
|
||||
message: &str,
|
||||
) -> (StatusCode, Json<serde_json::Value>) {
|
||||
(
|
||||
status,
|
||||
Json(serde_json::json!({
|
||||
"error": {
|
||||
"code": code,
|
||||
"message": message,
|
||||
}
|
||||
})),
|
||||
)
|
||||
}
|
||||
@@ -1,134 +1,11 @@
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use axum::Router;
|
||||
use axum::routing::{get, post};
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::sync::{Mutex, RwLock, mpsc, oneshot};
|
||||
use tokio::time::MissedTickBehavior;
|
||||
use tracing::{info, warn};
|
||||
use wakey_agent::protocol::{ErrorPayload, ServerMessage};
|
||||
|
||||
use crate::api;
|
||||
use crate::cli::{IssueEnrollTokenArgs, ListEnrollTokensArgs, RevokeEnrollTokenArgs, StateStatsArgs};
|
||||
use crate::config;
|
||||
use crate::state;
|
||||
use crate::ws;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AppState {
|
||||
pub store: Arc<state::Store>,
|
||||
pub sessions: Arc<RwLock<HashMap<String, mpsc::UnboundedSender<ServerMessage>>>>,
|
||||
pub pending: Arc<Mutex<HashMap<String, oneshot::Sender<AgentReply>>>>,
|
||||
pub public_url: String,
|
||||
pub command_timeout: Duration,
|
||||
pub enroll_token_ttl: Duration,
|
||||
}
|
||||
|
||||
pub enum AgentReply {
|
||||
Result(serde_json::Value),
|
||||
Error(ErrorPayload),
|
||||
}
|
||||
|
||||
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.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,
|
||||
};
|
||||
|
||||
let app = Router::new()
|
||||
.route("/healthz", get(api::healthz))
|
||||
.route("/api/v1/agents/enroll", post(api::enroll))
|
||||
.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/agent/ws", get(ws::agent_ws))
|
||||
.route("/api/v1/control/agents", get(api::list_agents))
|
||||
.route(
|
||||
"/api/v1/control/agents/{agent_id}/command",
|
||||
post(api::run_command),
|
||||
)
|
||||
.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(), "starting control-plane server");
|
||||
let listener = TcpListener::bind(daemon.bind).await?;
|
||||
let mut 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 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);
|
||||
|
||||
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"),
|
||||
}
|
||||
}
|
||||
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(())
|
||||
}
|
||||
|
||||
pub async fn issue_enroll_token(args: IssueEnrollTokenArgs) -> Result<()> {
|
||||
let settings = config::resolve_issue_token_settings(&args)?;
|
||||
@@ -137,7 +14,7 @@ pub async fn issue_enroll_token(args: IssueEnrollTokenArgs) -> Result<()> {
|
||||
let base = config::normalize_public_url(&url);
|
||||
let ttl_seconds = settings.ttl.as_secs().max(1);
|
||||
let endpoint = format!("{}?ttl_seconds={ttl_seconds}", config::issue_token_endpoint(&base));
|
||||
info!(endpoint = %endpoint, "requesting live enroll token from running control-plane daemon");
|
||||
tracing::info!(endpoint = %endpoint, "requesting live enroll token from running control-plane daemon");
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
let response = client
|
||||
@@ -159,7 +36,7 @@ pub async fn issue_enroll_token(args: IssueEnrollTokenArgs) -> Result<()> {
|
||||
.json()
|
||||
.await
|
||||
.context("failed to decode live issuance response")?;
|
||||
info!("received live enroll token response");
|
||||
tracing::info!("received live enroll token response");
|
||||
|
||||
println!("enroll_token={}", payload.enroll_token);
|
||||
println!("expires_at_unix={}", payload.expires_at_unix);
|
||||
@@ -170,8 +47,7 @@ pub async fn issue_enroll_token(args: IssueEnrollTokenArgs) -> Result<()> {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Fallback for offline tooling: writes to state file, requires daemon reload to pick up.
|
||||
info!(data_dir = %settings.data_dir.display(), state_file = %settings.state_file.display(), ttl_seconds = settings.ttl.as_secs(), "issuing enroll token via offline state file fallback");
|
||||
tracing::info!(data_dir = %settings.data_dir.display(), state_file = %settings.state_file.display(), ttl_seconds = settings.ttl.as_secs(), "issuing enroll token via offline state file fallback");
|
||||
let store = state::Store::load_or_init(&settings.state_file, args.enroll_tokens, settings.ttl)
|
||||
.await
|
||||
.with_context(|| format!("failed to initialize store {}", settings.state_file.display()))?;
|
||||
@@ -330,59 +206,3 @@ pub async fn state_stats(args: StateStatsArgs) -> Result<()> {
|
||||
println!("expired_enroll_token_count={}", stats.expired_enroll_token_count);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn reload_daemon(pid_file: &Path) -> Result<()> {
|
||||
let pid = read_pid(pid_file)?;
|
||||
info!(pid, pid_file = %pid_file.display(), "sending control-plane reload signal");
|
||||
send_hup(pid)
|
||||
}
|
||||
|
||||
fn write_pid_file(path: &Path) -> Result<()> {
|
||||
if let Some(parent) = path.parent() {
|
||||
std::fs::create_dir_all(parent)
|
||||
.with_context(|| format!("failed to create pid dir {}", parent.display()))?;
|
||||
}
|
||||
std::fs::write(path, format!("{}\n", std::process::id()))
|
||||
.with_context(|| format!("failed to write pid file {}", path.display()))
|
||||
}
|
||||
|
||||
fn remove_pid_file(path: &Path) -> Result<()> {
|
||||
match std::fs::remove_file(path) {
|
||||
Ok(()) => Ok(()),
|
||||
Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(()),
|
||||
Err(err) => {
|
||||
Err(err).with_context(|| format!("failed to remove pid file {}", path.display()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn read_pid(path: &Path) -> Result<i32> {
|
||||
let raw = std::fs::read_to_string(path)
|
||||
.with_context(|| format!("failed to read pid file {}", path.display()))?;
|
||||
let pid = raw
|
||||
.trim()
|
||||
.parse::<i32>()
|
||||
.with_context(|| format!("invalid pid in {}", path.display()))?;
|
||||
if pid <= 0 {
|
||||
anyhow::bail!("invalid non-positive pid {pid}");
|
||||
}
|
||||
Ok(pid)
|
||||
}
|
||||
|
||||
fn send_hup(pid: i32) -> Result<()> {
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use nix::sys::signal::{Signal, kill};
|
||||
use nix::unistd::Pid;
|
||||
|
||||
kill(Pid::from_raw(pid), Signal::SIGHUP)
|
||||
.with_context(|| format!("failed to send SIGHUP to pid {pid}"))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
{
|
||||
let _ = pid;
|
||||
anyhow::bail!("reload is only supported on Unix (SIGHUP unavailable on this platform)")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,191 @@
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use axum::Router;
|
||||
use axum::routing::{get, post};
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::sync::{Mutex, RwLock, mpsc, oneshot};
|
||||
use tokio::time::MissedTickBehavior;
|
||||
use tracing::{info, warn};
|
||||
use wakey_agent::protocol::{ErrorPayload, ServerMessage};
|
||||
|
||||
use crate::api;
|
||||
use crate::config;
|
||||
use crate::state;
|
||||
use crate::ws;
|
||||
|
||||
mod admin;
|
||||
pub use admin::{
|
||||
issue_enroll_token, list_enroll_tokens, revoke_enroll_token, state_stats,
|
||||
};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AppState {
|
||||
pub store: Arc<state::Store>,
|
||||
pub sessions: Arc<RwLock<HashMap<String, mpsc::UnboundedSender<ServerMessage>>>>,
|
||||
pub pending: Arc<Mutex<HashMap<String, oneshot::Sender<AgentReply>>>>,
|
||||
pub public_url: String,
|
||||
pub command_timeout: Duration,
|
||||
pub enroll_token_ttl: Duration,
|
||||
}
|
||||
|
||||
pub enum AgentReply {
|
||||
Result(serde_json::Value),
|
||||
Error(ErrorPayload),
|
||||
}
|
||||
|
||||
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.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,
|
||||
};
|
||||
|
||||
let app = Router::new()
|
||||
.route("/healthz", get(api::healthz))
|
||||
.route("/api/v1/agents/enroll", post(api::enroll))
|
||||
.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/agent/ws", get(ws::agent_ws))
|
||||
.route("/api/v1/control/agents", get(api::list_agents))
|
||||
.route(
|
||||
"/api/v1/control/agents/{agent_id}/command",
|
||||
post(api::run_command),
|
||||
)
|
||||
.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(), "starting control-plane server");
|
||||
let listener = TcpListener::bind(daemon.bind).await?;
|
||||
let mut 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 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);
|
||||
|
||||
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"),
|
||||
}
|
||||
}
|
||||
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(())
|
||||
}
|
||||
|
||||
pub fn reload_daemon(pid_file: &Path) -> Result<()> {
|
||||
let pid = read_pid(pid_file)?;
|
||||
info!(pid, pid_file = %pid_file.display(), "sending control-plane reload signal");
|
||||
send_hup(pid)
|
||||
}
|
||||
|
||||
fn write_pid_file(path: &Path) -> Result<()> {
|
||||
if let Some(parent) = path.parent() {
|
||||
std::fs::create_dir_all(parent)
|
||||
.with_context(|| format!("failed to create pid dir {}", parent.display()))?;
|
||||
}
|
||||
std::fs::write(path, format!("{}\n", std::process::id()))
|
||||
.with_context(|| format!("failed to write pid file {}", path.display()))
|
||||
}
|
||||
|
||||
fn remove_pid_file(path: &Path) -> Result<()> {
|
||||
match std::fs::remove_file(path) {
|
||||
Ok(()) => Ok(()),
|
||||
Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(()),
|
||||
Err(err) => {
|
||||
Err(err).with_context(|| format!("failed to remove pid file {}", path.display()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn read_pid(path: &Path) -> Result<i32> {
|
||||
let raw = std::fs::read_to_string(path)
|
||||
.with_context(|| format!("failed to read pid file {}", path.display()))?;
|
||||
let pid = raw
|
||||
.trim()
|
||||
.parse::<i32>()
|
||||
.with_context(|| format!("invalid pid in {}", path.display()))?;
|
||||
if pid <= 0 {
|
||||
anyhow::bail!("invalid non-positive pid {pid}");
|
||||
}
|
||||
Ok(pid)
|
||||
}
|
||||
|
||||
fn send_hup(pid: i32) -> Result<()> {
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use nix::sys::signal::{Signal, kill};
|
||||
use nix::unistd::Pid;
|
||||
|
||||
kill(Pid::from_raw(pid), Signal::SIGHUP)
|
||||
.with_context(|| format!("failed to send SIGHUP to pid {pid}"))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
{
|
||||
let _ = pid;
|
||||
anyhow::bail!("reload is only supported on Unix (SIGHUP unavailable on this platform)")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user