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

device page, fmt;

Browse Source
This commit is contained in:
lda
2026-04-12 21:47:23 +07:00 Unverified
parent bb813170d1
commit b576ba9a01
19 changed files with 356 additions and 162 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wakey-control-plane/src/api/commands.rs
+4 -1
View File
@@ -230,7 +230,10 @@ 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");
warn!(
timeout_ms = timeout.as_millis() as u64,
"agent command timed out"
);
if let Err(err) = state
.store
.append_audit_event(AuditEventInput {
wakey-control-plane/src/api/control.rs
+5 -1
View File
@@ -242,7 +242,11 @@ pub async fn revoke_enroll_token(
agent_id: None,
request_id: None,
event_type: "enroll_token_revoke".into(),
outcome: if revoked { "ok".into() } else { "not_found".into() },
outcome: if revoked {
"ok".into()
} else {
"not_found".into()
},
latency_ms: None,
message: if revoked {
"revoked enroll token".into()
wakey-control-plane/src/api/mod.rs
+3 -3
View File
@@ -1,14 +1,14 @@
use axum::Json;
use axum::http::StatusCode;
mod alerts;
mod audit;
mod commands;
mod control;
mod audit;
mod alerts;
pub use commands::{list_agents, run_command};
pub use alerts::{active_alerts, alert_history, alerts_stream};
pub use audit::list_audit_events;
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,
wakey-control-plane/src/runtime/admin.rs
+49 -16
View File
@@ -3,7 +3,9 @@ use std::time::Duration;
use anyhow::{Context, Result};
use crate::api;
use crate::cli::{IssueEnrollTokenArgs, ListEnrollTokensArgs, RevokeEnrollTokenArgs, StateStatsArgs};
use crate::cli::{
IssueEnrollTokenArgs, ListEnrollTokensArgs, RevokeEnrollTokenArgs, StateStatsArgs,
};
use crate::config;
use crate::state;
@@ -13,7 +15,10 @@ pub async fn issue_enroll_token(args: IssueEnrollTokenArgs) -> Result<()> {
if let Some(url) = args.public_url {
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));
let endpoint = format!(
"{}?ttl_seconds={ttl_seconds}",
config::issue_token_endpoint(&base)
);
tracing::info!(endpoint = %endpoint, "requesting live enroll token from running control-plane daemon");
let client = reqwest::Client::new();
@@ -50,7 +55,12 @@ pub async fn issue_enroll_token(args: IssueEnrollTokenArgs) -> Result<()> {
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()))?;
.with_context(|| {
format!(
"failed to initialize store {}",
settings.state_file.display()
)
})?;
let issued = store.issue_enroll_token(settings.ttl).await?;
println!("enroll_token={}", issued.enroll_token);
println!("expires_at_unix={}", issued.expires_at_unix);
@@ -66,8 +76,7 @@ pub async fn list_enroll_tokens(args: ListEnrollTokensArgs) -> Result<()> {
if let Some(base) = settings.public_url.as_deref() {
let url = format!(
"{}/api/v1/control/enroll-tokens?include_expired={}",
base,
args.include_expired
base, args.include_expired
);
let response = reqwest::get(&url)
.await
@@ -100,9 +109,15 @@ pub async fn list_enroll_tokens(args: ListEnrollTokensArgs) -> Result<()> {
return Ok(());
}
let store = state::Store::load_or_init(&settings.state_file, Vec::new(), Duration::from_secs(1))
.await
.with_context(|| format!("failed to initialize store {}", settings.state_file.display()))?;
let store =
state::Store::load_or_init(&settings.state_file, Vec::new(), Duration::from_secs(1))
.await
.with_context(|| {
format!(
"failed to initialize store {}",
settings.state_file.display()
)
})?;
let tokens = store.list_enroll_tokens(args.include_expired).await?;
if args.json {
println!(
@@ -146,9 +161,15 @@ pub async fn revoke_enroll_token(args: RevokeEnrollTokenArgs) -> Result<()> {
return Ok(());
}
let store = state::Store::load_or_init(&settings.state_file, Vec::new(), Duration::from_secs(1))
.await
.with_context(|| format!("failed to initialize store {}", settings.state_file.display()))?;
let store =
state::Store::load_or_init(&settings.state_file, Vec::new(), Duration::from_secs(1))
.await
.with_context(|| {
format!(
"failed to initialize store {}",
settings.state_file.display()
)
})?;
let removed = store.revoke_enroll_token(&args.token).await?;
println!("token={} revoked={}", args.token, removed);
Ok(())
@@ -184,13 +205,22 @@ pub async fn state_stats(args: StateStatsArgs) -> Result<()> {
println!("schema_version={}", body.schema_version);
println!("agent_count={}", body.agent_count);
println!("enroll_token_count={}", body.enroll_token_count);
println!("expired_enroll_token_count={}", body.expired_enroll_token_count);
println!(
"expired_enroll_token_count={}",
body.expired_enroll_token_count
);
return Ok(());
}
let store = state::Store::load_or_init(&settings.state_file, Vec::new(), Duration::from_secs(1))
.await
.with_context(|| format!("failed to initialize store {}", settings.state_file.display()))?;
let store =
state::Store::load_or_init(&settings.state_file, Vec::new(), Duration::from_secs(1))
.await
.with_context(|| {
format!(
"failed to initialize store {}",
settings.state_file.display()
)
})?;
let stats = store.stats().await?;
if args.json {
println!(
@@ -203,6 +233,9 @@ pub async fn state_stats(args: StateStatsArgs) -> Result<()> {
println!("schema_version={}", stats.schema_version);
println!("agent_count={}", stats.agent_count);
println!("enroll_token_count={}", stats.enroll_token_count);
println!("expired_enroll_token_count={}", stats.expired_enroll_token_count);
println!(
"expired_enroll_token_count={}",
stats.expired_enroll_token_count
);
Ok(())
}
wakey-control-plane/src/runtime/mod.rs
+12 -7
View File
@@ -5,8 +5,8 @@ use std::time::Duration;
use anyhow::{Context, Result};
use axum::Router;
use axum::response::Redirect;
use axum::routing::{get, post};
use axum::routing::get_service;
use axum::routing::{get, post};
use tokio::net::TcpListener;
use tokio::sync::{Mutex, RwLock, mpsc, oneshot};
#[cfg(unix)]
@@ -25,9 +25,7 @@ use crate::ws;
mod admin;
mod process;
pub use admin::{
issue_enroll_token, list_enroll_tokens, revoke_enroll_token, state_stats,
};
pub use admin::{issue_enroll_token, list_enroll_tokens, revoke_enroll_token, state_stats};
pub use process::reload_daemon;
use process::{remove_pid_file, write_pid_file};
@@ -55,7 +53,7 @@ pub enum AgentReply {
fn public_api_routes() -> Router<AppState> {
Router::new()
.route("/ui", get(|| async { Redirect::temporary("/ui/") }))
.route("/ui", get(|| async { Redirect::temporary("/ui/") }))
.nest_service(
"/ui/",
get_service(
@@ -69,8 +67,14 @@ fn public_api_routes() -> Router<AppState> {
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-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),
@@ -128,6 +132,7 @@ pub async fn serve(daemon: config::DaemonConfig) -> Result<()> {
#[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);
wakey-control-plane/src/state/store.rs
+75 -35
View File
@@ -25,7 +25,11 @@ const SEEDED_ENROLL_TOKEN_PREFIX: &[u8] = b"seeded_enroll_token:";
const SCHEMA_VERSION: u32 = 1;
impl Store {
pub async fn load_or_init(path: &Path, enroll_tokens: Vec<String>, seed_ttl: Duration) -> Result<Self> {
pub async fn load_or_init(
path: &Path,
enroll_tokens: Vec<String>,
seed_ttl: Duration,
) -> Result<Self> {
let db_path = path.to_path_buf();
if let Some(parent) = db_path.parent() {
std::fs::create_dir_all(parent)
@@ -38,7 +42,9 @@ impl Store {
let enroll_tree = db
.open_tree("enroll_tokens")
.context("failed to open enroll_tokens tree")?;
let agents_tree = db.open_tree("agents").context("failed to open agents tree")?;
let agents_tree = db
.open_tree("agents")
.context("failed to open agents tree")?;
let audit_events_tree = db
.open_tree("audit_events")
.context("failed to open audit_events tree")?;
@@ -92,13 +98,17 @@ impl Store {
anyhow::bail!("invalid or already-used enroll token");
};
let expires_at_unix = decode_expiry(raw_expiry.as_ref())
.context("failed decoding enroll token expiry")?;
let expires_at_unix =
decode_expiry(raw_expiry.as_ref()).context("failed decoding enroll token expiry")?;
let now = now_unix();
if expires_at_unix <= now {
let _ = self.enroll_tokens.remove(enroll_token.as_bytes());
self.flush().ok();
warn!(expires_at_unix, now_unix = now, "rejecting expired enroll token");
warn!(
expires_at_unix,
now_unix = now,
"rejecting expired enroll token"
);
anyhow::bail!("enroll token has expired");
}
@@ -112,7 +122,8 @@ impl Store {
self.agents
.insert(agent_id.as_bytes(), agent_token.as_bytes())
.context("failed persisting agent credentials")?;
self.flush().context("failed flushing state db after enroll")?;
self.flush()
.context("failed flushing state db after enroll")?;
info!(agent_id = %agent_id, "issued persistent agent credentials");
Ok(IssuedAgent {
@@ -141,7 +152,8 @@ impl Store {
let mut out = Vec::new();
for item in self.enroll_tokens.iter() {
let (token, value) = item.context("failed iterating enroll token tree")?;
let expires_at_unix = decode_expiry(value.as_ref()).context("failed decoding token expiry")?;
let expires_at_unix =
decode_expiry(value.as_ref()).context("failed decoding token expiry")?;
let expired = expires_at_unix <= now;
if !include_expired && expired {
continue;
@@ -169,7 +181,8 @@ impl Store {
.context("failed removing enroll token")?
.is_some();
if removed {
self.flush().context("failed flushing db after enroll token revoke")?;
self.flush()
.context("failed flushing db after enroll token revoke")?;
}
Ok(removed)
}
@@ -180,8 +193,8 @@ impl Store {
let mut expired_enroll_token_count = 0usize;
for item in self.enroll_tokens.iter() {
let (_, value) = item.context("failed iterating enroll token tree")?;
let expires_at =
decode_expiry(value.as_ref()).context("failed decoding token expiry during stats")?;
let expires_at = decode_expiry(value.as_ref())
.context("failed decoding token expiry during stats")?;
enroll_token_count = enroll_token_count.saturating_add(1);
if expires_at <= now {
expired_enroll_token_count = expired_enroll_token_count.saturating_add(1);
@@ -372,7 +385,9 @@ impl Store {
let (_, raw) = item.context("failed iterating alert transition tree")?;
let transition: AlertTransition =
serde_json::from_slice(raw.as_ref()).context("failed decoding alert transition")?;
if let Some(since) = since_unix && transition.ts_unix < since {
if let Some(since) = since_unix
&& transition.ts_unix < since
{
continue;
}
out.push(transition);
@@ -384,15 +399,11 @@ impl Store {
}
fn flush(&self) -> Result<()> {
self.meta
.flush()
.context("failed to flush meta tree")?;
self.meta.flush().context("failed to flush meta tree")?;
self.enroll_tokens
.flush()
.context("failed to flush enroll token tree")?;
self.agents
.flush()
.context("failed to flush agents tree")?;
self.agents.flush().context("failed to flush agents tree")?;
self.audit_events
.flush()
.context("failed to flush audit event tree")?;
@@ -406,7 +417,11 @@ impl Store {
Ok(())
}
fn seed_bootstrap_enroll_tokens(&self, enroll_tokens: &[String], seed_ttl: Duration) -> Result<()> {
fn seed_bootstrap_enroll_tokens(
&self,
enroll_tokens: &[String],
seed_ttl: Duration,
) -> Result<()> {
for token in enroll_tokens {
let token = token.trim();
if token.is_empty() {
@@ -414,21 +429,32 @@ impl Store {
}
let marker_key = seeded_enroll_token_key(token);
if self
.meta
.contains_key(&marker_key)
.with_context(|| format!("failed reading bootstrap marker in {}", self.db_path.display()))?
{
if self.meta.contains_key(&marker_key).with_context(|| {
format!(
"failed reading bootstrap marker in {}",
self.db_path.display()
)
})? {
continue;
}
let expires_at = now_unix().saturating_add(seed_ttl.as_secs().max(1));
self.enroll_tokens
.insert(token.as_bytes(), &expires_at.to_le_bytes())
.with_context(|| format!("failed to seed enroll token into {}", self.db_path.display()))?;
.with_context(|| {
format!(
"failed to seed enroll token into {}",
self.db_path.display()
)
})?;
self.meta
.insert(marker_key, &expires_at.to_le_bytes())
.with_context(|| format!("failed to persist bootstrap marker into {}", self.db_path.display()))?;
.with_context(|| {
format!(
"failed to persist bootstrap marker into {}",
self.db_path.display()
)
})?;
}
Ok(())
}
@@ -475,8 +501,12 @@ impl Store {
self.meta
.insert(SCHEMA_VERSION_KEY, &SCHEMA_VERSION.to_le_bytes())
.context("failed writing schema version")?;
self.flush().context("failed flushing db after schema init")?;
info!(schema_version = SCHEMA_VERSION, "initialized state schema version");
self.flush()
.context("failed flushing db after schema init")?;
info!(
schema_version = SCHEMA_VERSION,
"initialized state schema version"
);
}
}
Ok(())
@@ -535,16 +565,24 @@ fn matches_audit_filter(event: &AuditEvent, filter: &AuditEventFilter) -> bool {
{
return false;
}
if let Some(event_type) = filter.event_type.as_deref() && event.event_type != event_type {
if let Some(event_type) = filter.event_type.as_deref()
&& event.event_type != event_type
{
return false;
}
if let Some(outcome) = filter.outcome.as_deref() && event.outcome != outcome {
if let Some(outcome) = filter.outcome.as_deref()
&& event.outcome != outcome
{
return false;
}
if let Some(since_unix) = filter.since_unix && event.ts_unix < since_unix {
if let Some(since_unix) = filter.since_unix
&& event.ts_unix < since_unix
{
return false;
}
if let Some(until_unix) = filter.until_unix && event.ts_unix > until_unix {
if let Some(until_unix) = filter.until_unix
&& event.ts_unix > until_unix
{
return false;
}
@@ -559,7 +597,8 @@ mod tests {
use super::Store;
async fn make_store() -> (Store, std::path::PathBuf) {
let dir = std::env::temp_dir().join(format!("wakey-cp-store-test-{}", uuid::Uuid::new_v4()));
let dir =
std::env::temp_dir().join(format!("wakey-cp-store-test-{}", uuid::Uuid::new_v4()));
let db_path = dir.join("state.db");
let store = Store::load_or_init(&db_path, Vec::new(), Duration::from_secs(60))
.await
@@ -785,9 +824,10 @@ mod tests {
.enroll("enr-bootstrap-once")
.await
.expect_err("bootstrap token should not resurrect after restart");
assert!(err
.to_string()
.contains("invalid or already-used enroll token"));
assert!(
err.to_string()
.contains("invalid or already-used enroll token")
);
cleanup_dir(&dir);
}
wakey-control-plane/src/tracing.rs
+11 -11
View File
@@ -1,6 +1,6 @@
use anyhow::{Context, Result};
use opentelemetry::trace::TracerProvider as _;
use opentelemetry::global;
use opentelemetry::trace::TracerProvider as _;
use opentelemetry_otlp::WithExportConfig;
use opentelemetry_sdk::Resource;
use opentelemetry_sdk::trace::{SdkTracerProvider, Tracer};
@@ -28,7 +28,10 @@ pub fn init(verbose: u8, telemetry: &TelemetryConfig) -> Result<()> {
.with(filter)
.with(fmt::layer().json())
.init();
tracing::info!(json_logs = telemetry.json_logs, "tracing initialized without otlp exporter");
tracing::info!(
json_logs = telemetry.json_logs,
"tracing initialized without otlp exporter"
);
}
} else if let Some(otel_layer) = otel {
tracing_subscriber::registry()
@@ -42,7 +45,10 @@ pub fn init(verbose: u8, telemetry: &TelemetryConfig) -> Result<()> {
.with(filter)
.with(fmt::layer())
.init();
tracing::info!(json_logs = telemetry.json_logs, "tracing initialized without otlp exporter");
tracing::info!(
json_logs = telemetry.json_logs,
"tracing initialized without otlp exporter"
);
}
Ok(())
@@ -50,14 +56,8 @@ pub fn init(verbose: u8, telemetry: &TelemetryConfig) -> Result<()> {
fn build_otel_layer(
telemetry: &TelemetryConfig,
) -> Result<
Option<
tracing_opentelemetry::OpenTelemetryLayer<
tracing_subscriber::Registry,
Tracer,
>,
>,
> {
) -> Result<Option<tracing_opentelemetry::OpenTelemetryLayer<tracing_subscriber::Registry, Tracer>>>
{
let Some(endpoint) = telemetry.otlp_endpoint.as_deref() else {
return Ok(None);
};
wakey-control-plane/src/ws.rs
+12 -12
View File
@@ -201,17 +201,13 @@ async fn process_agent_text(
}
anyhow::bail!("agent auth rejected");
}
state
.sessions
.write()
.await
.insert(
agent_id.clone(),
AgentSession {
connection_id: connection_id.to_string(),
tx: tx.clone(),
},
);
state.sessions.write().await.insert(
agent_id.clone(),
AgentSession {
connection_id: connection_id.to_string(),
tx: tx.clone(),
},
);
*authed_agent_id = Some(agent_id.clone());
info!(agent_id = %agent_id, connect_to_auth_ms, hello_to_auth_ms = hello_to_auth_ms.unwrap_or(0), "agent authenticated");
if let Err(err) = state
@@ -275,7 +271,11 @@ fn now_duration_ms(duration: std::time::Duration) -> u64 {
duration.as_millis() as u64
}
async fn ensure_current_session(state: &AppState, agent_id: &str, connection_id: &str) -> Result<()> {
async fn ensure_current_session(
state: &AppState,
agent_id: &str,
connection_id: &str,
) -> Result<()> {
let sessions = state.sessions.read().await;
if is_current_session(&sessions, agent_id, connection_id) {
Ok(())