lda/wakey
1
0
Fork 0
Code Issues 5 Pull Requests Actions Packages Projects Releases 18 Wiki Activity

logging and config sweep

Browse Source
This commit is contained in:
lda
2026-04-11 18:02:12 +07:00 Unverified
parent ce1c0c0558
commit 78b12ad0ee
16 changed files with 839 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wakey-control-plane/src/api.rs
+82 -36
View File
@@ -3,6 +3,7 @@ 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};
@@ -56,19 +57,25 @@ pub async fn enroll(
Json(req): Json<EnrollRequest>,
) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
match state.store.enroll(&req.enroll_token).await {
Ok(issued) => Ok((
StatusCode::OK,
Json(EnrollResponse {
agent_id: issued.agent_id,
agent_token: issued.agent_token,
server_url: state.public_url,
}),
)),
Err(err) => Err(json_error(
StatusCode::UNAUTHORIZED,
"enrollment_rejected",
&err.to_string(),
)),
Ok(issued) => {
info!(agent_id = %issued.agent_id, "agent enrollment accepted");
Ok((
StatusCode::OK,
Json(EnrollResponse {
agent_id: issued.agent_id,
agent_token: issued.agent_token,
server_url: state.public_url,
}),
))
}
Err(err) => {
warn!(error = %err, "agent enrollment rejected");
Err(json_error(
StatusCode::UNAUTHORIZED,
"enrollment_rejected",
&err.to_string(),
))
}
}
}
@@ -76,17 +83,23 @@ pub async fn issue_enroll_token(
State(state): State<AppState>,
) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
match state.store.issue_enroll_token().await {
Ok(token) => Ok((
StatusCode::OK,
Json(IssueEnrollTokenResponse {
enroll_token: token,
}),
)),
Err(err) => Err(json_error(
StatusCode::INTERNAL_SERVER_ERROR,
"issue_enroll_token_failed",
&err.to_string(),
)),
Ok(token) => {
info!("issued enroll token");
Ok((
StatusCode::OK,
Json(IssueEnrollTokenResponse {
enroll_token: token,
}),
))
}
Err(err) => {
warn!(error = %err, "failed to issue enroll token");
Err(json_error(
StatusCode::INTERNAL_SERVER_ERROR,
"issue_enroll_token_failed",
&err.to_string(),
))
}
}
}
@@ -113,6 +126,15 @@ pub async fn run_command(
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,
@@ -126,6 +148,7 @@ pub async fn run_command(
sessions.get(&agent_id).cloned()
}
.ok_or_else(|| {
warn!("command rejected: agent not connected");
json_error(
StatusCode::NOT_FOUND,
"agent_not_connected",
@@ -140,11 +163,16 @@ pub async fn run_command(
.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",
@@ -159,19 +187,26 @@ pub async fn run_command(
);
let outcome = tokio::time::timeout(timeout, pending_rx).await;
let response = match outcome {
Ok(Ok(AgentReply::Result(result))) => RelayCommandResponse {
request_id: request_id_string,
status: "ok".into(),
result: Some(result),
error: None,
},
Ok(Ok(AgentReply::Error(error))) => RelayCommandResponse {
request_id: request_id_string,
status: "error".into(),
result: None,
error: Some(error),
},
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",
@@ -180,6 +215,7 @@ pub async fn run_command(
}
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",
@@ -206,3 +242,13 @@ pub fn json_error(
})),
)
}
fn command_kind(command: &AgentCommand) -> &'static str {
match command {
AgentCommand::Status(_) => "status",
AgentCommand::Leases(_) => "leases",
AgentCommand::Devs(_) => "devs",
AgentCommand::Inventory(_) => "inventory",
AgentCommand::Wake(_) => "wake",
}
}
wakey-control-plane/src/cli.rs
+14 -10
View File
@@ -5,6 +5,7 @@ use clap::{ArgAction, Args, Parser, Subcommand};
pub const DEFAULT_STATE_FILE: &str = "/var/lib/wakey-control-plane/state.json";
pub const DEFAULT_PID_FILE: &str = "/var/run/wakey-control-plane.pid";
pub const DEFAULT_CONFIG_FILE: &str = "/etc/wakey-control-plane/config.toml";
#[derive(Parser)]
#[command(name = "wakey-control-plane")]
@@ -29,23 +30,26 @@ pub enum Command {
#[derive(Args, Clone)]
pub struct ServeArgs {
#[arg(long, default_value = "0.0.0.0:8080")]
pub bind: SocketAddr,
#[arg(long)]
pub bind: Option<SocketAddr>,
#[arg(long, default_value = "http://127.0.0.1:8080")]
pub public_url: String,
#[arg(long)]
pub public_url: Option<String>,
#[arg(long, default_value = DEFAULT_STATE_FILE)]
pub state_file: PathBuf,
#[arg(long)]
pub state_file: Option<PathBuf>,
#[arg(long = "enroll-token")]
pub enroll_tokens: Vec<String>,
#[arg(long, default_value_t = 30_000)]
pub command_timeout_ms: u64,
#[arg(long)]
pub command_timeout_ms: Option<u64>,
#[arg(long, default_value = DEFAULT_PID_FILE)]
pub pid_file: PathBuf,
#[arg(long)]
pub pid_file: Option<PathBuf>,
#[arg(long, default_value = DEFAULT_CONFIG_FILE)]
pub config_file: PathBuf,
}
#[derive(Args)]
wakey-control-plane/src/config.rs
+129 -7
View File
@@ -1,7 +1,11 @@
use std::net::SocketAddr;
use std::path::Path;
use std::path::PathBuf;
use std::time::Duration;
use anyhow::{Context, Result};
use serde::Deserialize;
use crate::cli::ServeArgs;
#[derive(Debug, Clone)]
@@ -11,18 +15,136 @@ pub struct DaemonConfig {
pub state_file: PathBuf,
pub command_timeout: Duration,
pub pid_file: PathBuf,
pub enroll_tokens: Vec<String>,
pub telemetry: TelemetryConfig,
}
#[derive(Debug, Clone)]
pub struct TelemetryConfig {
pub otlp_endpoint: Option<String>,
pub service_name: String,
pub json_logs: bool,
}
impl Default for TelemetryConfig {
fn default() -> Self {
Self {
otlp_endpoint: None,
service_name: "wakey-control-plane".to_string(),
json_logs: false,
}
}
}
#[derive(Debug, Deserialize, Default)]
struct FileConfig {
bind: Option<String>,
public_url: Option<String>,
state_file: Option<PathBuf>,
command_timeout_ms: Option<u64>,
pid_file: Option<PathBuf>,
enroll_tokens: Option<Vec<String>>,
telemetry: Option<FileTelemetryConfig>,
}
#[derive(Debug, Deserialize, Default)]
struct FileTelemetryConfig {
otlp_endpoint: Option<String>,
service_name: Option<String>,
json_logs: Option<bool>,
}
impl DaemonConfig {
pub fn from_serve_args(args: &ServeArgs) -> Self {
Self {
bind: args.bind,
public_url: normalize_public_url(&args.public_url),
state_file: args.state_file.clone(),
command_timeout: Duration::from_millis(args.command_timeout_ms.max(1)),
pid_file: args.pid_file.clone(),
pub fn from_serve_args(args: &ServeArgs) -> Result<Self> {
let file = load_file_config(&args.config_file)?;
let bind = match args.bind {
Some(bind) => bind,
None => match file.bind {
Some(ref bind) => bind
.parse::<SocketAddr>()
.with_context(|| format!("invalid bind address `{bind}` in {}", args.config_file.display()))?,
None => "0.0.0.0:8080".parse().expect("static default bind should parse"),
},
};
let public_url = normalize_public_url(
args.public_url
.as_deref()
.or(file.public_url.as_deref())
.unwrap_or("http://127.0.0.1:8080"),
);
let state_file = args
.state_file
.clone()
.or(file.state_file)
.unwrap_or_else(|| PathBuf::from(crate::cli::DEFAULT_STATE_FILE));
let command_timeout = Duration::from_millis(
args.command_timeout_ms
.or(file.command_timeout_ms)
.unwrap_or(30_000)
.max(1),
);
let pid_file = args
.pid_file
.clone()
.or(file.pid_file)
.unwrap_or_else(|| PathBuf::from(crate::cli::DEFAULT_PID_FILE));
let enroll_tokens = if args.enroll_tokens.is_empty() {
file.enroll_tokens.unwrap_or_default()
} else {
args.enroll_tokens.clone()
};
let telemetry = resolve_telemetry(file.telemetry);
Ok(Self {
bind,
public_url,
state_file,
command_timeout,
pid_file,
enroll_tokens,
telemetry,
})
}
}
fn load_file_config(path: &Path) -> Result<FileConfig> {
if !path.exists() {
return Ok(FileConfig::default());
}
let raw = std::fs::read_to_string(path)
.with_context(|| format!("failed to read config file {}", path.display()))?;
toml::from_str::<FileConfig>(&raw)
.with_context(|| format!("failed to parse config file {}", path.display()))
}
fn resolve_telemetry(file: Option<FileTelemetryConfig>) -> TelemetryConfig {
let mut out = TelemetryConfig::default();
if let Some(file) = file {
if let Some(endpoint) = file.otlp_endpoint {
let trimmed = endpoint.trim().to_string();
if !trimmed.is_empty() {
out.otlp_endpoint = Some(trimmed);
}
}
if let Some(name) = file.service_name {
let trimmed = name.trim();
if !trimmed.is_empty() {
out.service_name = trimmed.to_string();
}
}
if let Some(json_logs) = file.json_logs {
out.json_logs = json_logs;
}
}
out
}
pub fn normalize_public_url(url: &str) -> String {
wakey-control-plane/src/main.rs
+13 -4
View File
@@ -13,11 +13,20 @@ use cli::{Cli, Command};
#[tokio::main]
async fn main() -> Result<()> {
let cli = Cli::parse();
tracing::init(cli.verbose);
match cli.command {
Command::Serve(args) => runtime::serve(args).await,
Command::IssueEnrollToken(args) => runtime::issue_enroll_token(args).await,
Command::Reload(args) => runtime::reload_daemon(&args.pid_file),
Command::Serve(args) => {
let daemon = config::DaemonConfig::from_serve_args(&args)?;
tracing::init(cli.verbose, &daemon.telemetry)?;
runtime::serve(daemon).await
}
Command::IssueEnrollToken(args) => {
tracing::init(cli.verbose, &config::TelemetryConfig::default())?;
runtime::issue_enroll_token(args).await
}
Command::Reload(args) => {
tracing::init(cli.verbose, &config::TelemetryConfig::default())?;
runtime::reload_daemon(&args.pid_file)
}
}
}
wakey-control-plane/src/runtime.rs
+8 -4
View File
@@ -12,7 +12,7 @@ use tracing::{info, warn};
use wakey_agent::protocol::{ErrorPayload, ServerMessage};
use crate::api;
use crate::cli::{IssueEnrollTokenArgs, ServeArgs};
use crate::cli::IssueEnrollTokenArgs;
use crate::config;
use crate::state;
use crate::ws;
@@ -31,11 +31,11 @@ pub enum AgentReply {
Error(ErrorPayload),
}
pub async fn serve(args: ServeArgs) -> Result<()> {
let daemon = config::DaemonConfig::from_serve_args(&args);
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, args.enroll_tokens)
let store = state::Store::load_or_init(&daemon.state_file, daemon.enroll_tokens.clone())
.await
.with_context(|| format!("failed to initialize store {}", daemon.state_file.display()))?;
@@ -109,6 +109,7 @@ pub async fn issue_enroll_token(args: IssueEnrollTokenArgs) -> Result<()> {
if let Some(url) = args.public_url {
let base = config::normalize_public_url(&url);
let endpoint = config::issue_token_endpoint(&base);
info!(endpoint = %endpoint, "requesting live enroll token from running control-plane daemon");
let client = reqwest::Client::new();
let response = client
@@ -130,6 +131,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");
println!("enroll_token={}", payload.enroll_token);
println!(
@@ -140,6 +142,7 @@ pub async fn issue_enroll_token(args: IssueEnrollTokenArgs) -> Result<()> {
}
// Fallback for offline tooling: writes to state file, requires daemon reload to pick up.
info!(state_file = %args.state_file.display(), "issuing enroll token via offline state file fallback");
let store = state::Store::load_or_init(&args.state_file, args.enroll_tokens)
.await
.with_context(|| format!("failed to initialize store {}", args.state_file.display()))?;
@@ -154,6 +157,7 @@ pub async fn issue_enroll_token(args: IssueEnrollTokenArgs) -> Result<()> {
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)
}
wakey-control-plane/src/state.rs
+28
View File
@@ -4,6 +4,7 @@ use std::path::{Path, PathBuf};
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use tokio::sync::RwLock;
use tracing::{debug, info, warn};
use uuid::Uuid;
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -35,9 +36,11 @@ impl Store {
let raw = tokio::fs::read_to_string(path)
.await
.with_context(|| format!("failed to read store {}", path.display()))?;
debug!(path = %path.display(), bytes = raw.len(), "loading persisted control-plane store");
serde_json::from_str::<PersistedState>(&raw)
.with_context(|| format!("failed to decode store {}", path.display()))?
} else {
info!(path = %path.display(), seeded_enroll_tokens = seeded_tokens.len(), "initializing new control-plane store");
PersistedState {
enroll_tokens: seeded_tokens,
agents: HashMap::new(),
@@ -49,12 +52,23 @@ impl Store {
state: RwLock::new(initial),
};
store.save().await?;
let (enroll_tokens, agents) = {
let snapshot = store.state.read().await;
(snapshot.enroll_tokens.len(), snapshot.agents.len())
};
info!(
path = %store.path.display(),
enroll_tokens,
agents,
"control-plane store ready"
);
Ok(store)
}
pub async fn enroll(&self, enroll_token: &str) -> Result<IssuedAgent> {
let mut state = self.state.write().await;
if !state.enroll_tokens.remove(enroll_token) {
warn!("rejecting enroll attempt with invalid or consumed token");
anyhow::bail!("invalid or already-used enroll token");
}
@@ -63,6 +77,7 @@ impl Store {
state.agents.insert(agent_id.clone(), agent_token.clone());
drop(state);
self.save().await?;
info!(agent_id = %agent_id, "issued persistent agent credentials");
Ok(IssuedAgent {
agent_id,
@@ -76,11 +91,13 @@ impl Store {
state.enroll_tokens.insert(token.clone());
drop(state);
self.save().await?;
info!("persisted new enroll token");
Ok(token)
}
pub async fn reload_from_disk(&self) -> Result<()> {
if !self.path.exists() {
warn!(path = %self.path.display(), "reload requested but store file is missing");
return Ok(());
}
@@ -90,7 +107,10 @@ impl Store {
let decoded = serde_json::from_str::<PersistedState>(&raw)
.with_context(|| format!("failed to decode store {}", self.path.display()))?;
let enroll_tokens = decoded.enroll_tokens.len();
let agents = decoded.agents.len();
*self.state.write().await = decoded;
info!(path = %self.path.display(), enroll_tokens, agents, "reloaded control-plane store from disk");
Ok(())
}
@@ -121,6 +141,7 @@ impl Store {
let snapshot = self.state.read().await;
let body =
serde_json::to_string_pretty(&*snapshot).context("failed to serialize store state")?;
let body_len = body.len();
if let Some(parent) = self.path.parent() {
tokio::fs::create_dir_all(parent)
@@ -139,6 +160,13 @@ impl Store {
self.path.display()
)
})?;
debug!(
path = %self.path.display(),
bytes = body_len,
enroll_tokens = snapshot.enroll_tokens.len(),
agents = snapshot.agents.len(),
"saved control-plane store"
);
Ok(())
}
}
wakey-control-plane/src/tracing.rs
+76 -5
View File
@@ -1,14 +1,85 @@
use anyhow::{Context, Result};
use opentelemetry::trace::TracerProvider as _;
use opentelemetry::global;
use opentelemetry_otlp::WithExportConfig;
use opentelemetry_sdk::Resource;
use opentelemetry_sdk::trace::{SdkTracerProvider, Tracer};
use tracing_subscriber::{EnvFilter, fmt, layer::SubscriberExt, util::SubscriberInitExt};
pub fn init(verbose: u8) {
use crate::config::TelemetryConfig;
pub fn init(verbose: u8, telemetry: &TelemetryConfig) -> Result<()> {
let filter = EnvFilter::try_from_default_env()
.or_else(|_| EnvFilter::try_new(default_filter(verbose)))
.expect("static tracing filter should parse");
tracing_subscriber::registry()
.with(filter)
.with(fmt::layer())
.init();
let otel = build_otel_layer(telemetry)?;
if telemetry.json_logs {
if let Some(otel_layer) = otel {
tracing_subscriber::registry()
.with(otel_layer)
.with(filter)
.with(fmt::layer().json())
.init();
tracing::info!(endpoint = ?telemetry.otlp_endpoint, json_logs = telemetry.json_logs, "tracing initialized with otlp exporter");
} else {
tracing_subscriber::registry()
.with(filter)
.with(fmt::layer().json())
.init();
tracing::info!(json_logs = telemetry.json_logs, "tracing initialized without otlp exporter");
}
} else if let Some(otel_layer) = otel {
tracing_subscriber::registry()
.with(otel_layer)
.with(filter)
.with(fmt::layer())
.init();
tracing::info!(endpoint = ?telemetry.otlp_endpoint, json_logs = telemetry.json_logs, "tracing initialized with otlp exporter");
} else {
tracing_subscriber::registry()
.with(filter)
.with(fmt::layer())
.init();
tracing::info!(json_logs = telemetry.json_logs, "tracing initialized without otlp exporter");
}
Ok(())
}
fn build_otel_layer(
telemetry: &TelemetryConfig,
) -> Result<
Option<
tracing_opentelemetry::OpenTelemetryLayer<
tracing_subscriber::Registry,
Tracer,
>,
>,
> {
let Some(endpoint) = telemetry.otlp_endpoint.as_deref() else {
return Ok(None);
};
let exporter = opentelemetry_otlp::SpanExporter::builder()
.with_tonic()
.with_endpoint(endpoint)
.build()
.context("failed to build OTLP span exporter")?;
let provider = SdkTracerProvider::builder()
.with_batch_exporter(exporter)
.with_resource(
Resource::builder_empty()
.with_service_name(telemetry.service_name.clone())
.build(),
)
.build();
let tracer = provider.tracer(telemetry.service_name.clone());
global::set_tracer_provider(provider);
Ok(Some(tracing_opentelemetry::layer().with_tracer(tracer)))
}
fn default_filter(verbose: u8) -> &'static str {
wakey-control-plane/src/ws.rs
+22 -3
View File
@@ -5,7 +5,8 @@ use axum::response::IntoResponse;
use futures_util::{SinkExt, StreamExt};
use serde::Deserialize;
use tokio::sync::mpsc;
use tracing::{debug, info, warn};
use tracing::{debug, info, info_span, warn};
use uuid::Uuid;
use wakey_agent::protocol::{ErrorPayload, RequestId, ServerMessage};
use crate::runtime::{AgentReply, AppState};
@@ -38,6 +39,11 @@ pub async fn agent_ws(ws: WebSocketUpgrade, State(state): State<AppState>) -> im
}
async fn handle_agent_socket(state: AppState, socket: WebSocket) {
let connection_id = Uuid::new_v4().to_string();
let span = info_span!("agent_ws_connection", connection_id = %connection_id);
let _span_guard = span.enter();
info!("agent websocket upgraded");
let (mut write, mut read) = socket.split();
let (tx, mut rx) = mpsc::unbounded_channel::<ServerMessage>();
@@ -55,6 +61,7 @@ async fn handle_agent_socket(state: AppState, socket: WebSocket) {
break;
}
}
debug!("websocket writer loop ended");
});
let mut authed_agent_id: Option<String> = None;
@@ -67,7 +74,10 @@ async fn handle_agent_socket(state: AppState, socket: WebSocket) {
warn!(error = %err, "agent websocket receive error");
break;
}
None => break,
None => {
info!("agent websocket stream ended by peer");
break;
}
};
match msg {
@@ -80,7 +90,10 @@ async fn handle_agent_socket(state: AppState, socket: WebSocket) {
}
Message::Ping(_) => {}
Message::Pong(_) => {}
Message::Close(_) => break,
Message::Close(frame) => {
info!(close = ?frame, "agent websocket close frame received");
break;
}
Message::Binary(_) => {
debug!("ignoring unexpected binary websocket frame");
}
@@ -93,6 +106,7 @@ async fn handle_agent_socket(state: AppState, socket: WebSocket) {
}
writer.abort();
debug!("agent websocket connection cleanup complete");
}
async fn process_agent_text(
@@ -117,6 +131,7 @@ async fn process_agent_text(
.verify_agent_token(&agent_id, &agent_token)
.await
{
warn!(agent_id = %agent_id, "agent auth rejected");
anyhow::bail!("agent auth rejected");
}
state
@@ -140,6 +155,8 @@ async fn process_agent_text(
let key = request_id.as_str().to_string();
if let Some(waiter) = state.pending.lock().await.remove(&key) {
let _ = waiter.send(AgentReply::Result(result));
} else {
debug!(request_id = %key, "dropping unsolicited result from agent");
}
}
IncomingClientMessage::Error { request_id, error } => {
@@ -149,6 +166,8 @@ async fn process_agent_text(
let key = request_id.as_str().to_string();
if let Some(waiter) = state.pending.lock().await.remove(&key) {
let _ = waiter.send(AgentReply::Error(error));
} else {
debug!(request_id = %key, "dropping unsolicited error from agent");
}
}
}