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

admin cmds

Browse Source
This commit is contained in:
lda
2026-04-11 18:57:32 +07:00 Unverified
parent 5e6e2d6d2e
commit 3dfc2b4832
7 changed files with 403 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wakey-control-plane/src/api.rs
+67 -1
View File
@@ -32,6 +32,24 @@ pub struct IssueEnrollTokenQuery {
pub ttl_seconds: Option<u64>,
}
#[derive(Debug, Deserialize)]
pub struct ListEnrollTokenQuery {
pub include_expired: Option<bool>,
}
#[derive(Debug, Serialize)]
pub struct EnrollTokenStatus {
pub enroll_token: String,
pub expires_at_unix: u64,
pub expired: bool,
}
#[derive(Debug, Serialize)]
pub struct RevokeEnrollTokenResponse {
pub token: String,
pub revoked: bool,
}
#[derive(Debug, Serialize)]
pub struct AgentStatus {
pub agent_id: String,
@@ -135,6 +153,54 @@ pub async fn list_agents(
Ok((StatusCode::OK, Json(agents)))
}
pub async fn list_enroll_tokens(
State(state): State<AppState>,
Query(query): Query<ListEnrollTokenQuery>,
) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
let include_expired = query.include_expired.unwrap_or(false);
match state.store.list_enroll_tokens(include_expired).await {
Ok(tokens) => {
let body = tokens
.into_iter()
.map(|t| EnrollTokenStatus {
enroll_token: t.enroll_token,
expires_at_unix: t.expires_at_unix,
expired: t.expired,
})
.collect::<Vec<_>>();
Ok((StatusCode::OK, Json(body)))
}
Err(err) => {
warn!(error = %err, "failed to list enroll tokens");
Err(json_error(
StatusCode::INTERNAL_SERVER_ERROR,
"list_enroll_tokens_failed",
&err.to_string(),
))
}
}
}
pub async fn revoke_enroll_token(
State(state): State<AppState>,
AxumPath(token): AxumPath<String>,
) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
match state.store.revoke_enroll_token(&token).await {
Ok(revoked) => Ok((
StatusCode::OK,
Json(RevokeEnrollTokenResponse { token, revoked }),
)),
Err(err) => {
warn!(error = %err, "failed to revoke enroll token");
Err(json_error(
StatusCode::INTERNAL_SERVER_ERROR,
"revoke_enroll_token_failed",
&err.to_string(),
))
}
}
}
pub async fn run_command(
State(state): State<AppState>,
AxumPath(agent_id): AxumPath<String>,
@@ -265,5 +331,5 @@ pub fn json_error(
AgentCommand::Devs(_) => "devs",
AgentCommand::Inventory(_) => "inventory",
AgentCommand::Wake(_) => "wake",
}
}
}
wakey-control-plane/src/cli.rs
+54
View File
@@ -26,6 +26,12 @@ pub enum Command {
InitConfig(InitConfigArgs),
/// Create a new enroll token for provisioning a router.
IssueEnrollToken(IssueEnrollTokenArgs),
/// List current enroll tokens and their expiration status.
ListEnrollTokens(ListEnrollTokensArgs),
/// Revoke a specific enroll token.
RevokeEnrollToken(RevokeEnrollTokenArgs),
/// Print state backend stats.
StateStats(StateStatsArgs),
/// Send SIGHUP to an already-running daemon.
Reload(ReloadArgs),
}
@@ -123,6 +129,54 @@ pub struct IssueEnrollTokenArgs {
pub ttl_seconds: Option<u64>,
}
#[derive(Args)]
pub struct ListEnrollTokensArgs {
#[arg(long, default_value = DEFAULT_CONFIG_FILE)]
pub config_file: PathBuf,
#[arg(long)]
pub state_file: Option<PathBuf>,
#[arg(long)]
pub data_dir: Option<PathBuf>,
#[arg(long)]
pub include_expired: bool,
#[arg(long)]
pub json: bool,
}
#[derive(Args)]
pub struct RevokeEnrollTokenArgs {
#[arg(long, default_value = DEFAULT_CONFIG_FILE)]
pub config_file: PathBuf,
#[arg(long)]
pub state_file: Option<PathBuf>,
#[arg(long)]
pub data_dir: Option<PathBuf>,
#[arg(long)]
pub token: String,
}
#[derive(Args)]
pub struct StateStatsArgs {
#[arg(long, default_value = DEFAULT_CONFIG_FILE)]
pub config_file: PathBuf,
#[arg(long)]
pub state_file: Option<PathBuf>,
#[arg(long)]
pub data_dir: Option<PathBuf>,
#[arg(long)]
pub json: bool,
}
#[derive(Args)]
pub struct ReloadArgs {
#[arg(long, default_value = DEFAULT_PID_FILE)]
wakey-control-plane/src/config.rs
+48 -17
View File
@@ -6,7 +6,10 @@ use std::time::Duration;
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use crate::cli::{InitConfigArgs, IssueEnrollTokenArgs, ServeArgs};
use crate::cli::{
InitConfigArgs, IssueEnrollTokenArgs, ListEnrollTokensArgs, RevokeEnrollTokenArgs,
ServeArgs, StateStatsArgs,
};
#[derive(Debug, Clone)]
pub struct DaemonConfig {
@@ -279,22 +282,14 @@ pub struct IssueTokenSettings {
pub ttl: Duration,
}
pub struct StateAccessSettings {
pub data_dir: PathBuf,
pub state_file: PathBuf,
}
pub fn resolve_issue_token_settings(args: &IssueEnrollTokenArgs) -> Result<IssueTokenSettings> {
let file = load_file_config(&args.config_file)?;
let data_dir = args
.data_dir
.clone()
.or(file.data_dir)
.unwrap_or_else(|| PathBuf::from(crate::cli::DEFAULT_DATA_DIR));
let state_file = resolve_path(
&data_dir,
args.state_file
.clone()
.or(file.state_file)
.unwrap_or_else(|| PathBuf::from("state.db")),
);
let state = resolve_state_access(&args.config_file, args.data_dir.clone(), args.state_file.clone())?;
let ttl = Duration::from_secs(
args.ttl_seconds
@@ -304,8 +299,44 @@ pub fn resolve_issue_token_settings(args: &IssueEnrollTokenArgs) -> Result<Issue
);
Ok(IssueTokenSettings {
data_dir,
state_file,
data_dir: state.data_dir,
state_file: state.state_file,
ttl,
})
}
pub fn resolve_list_enroll_token_settings(args: &ListEnrollTokensArgs) -> Result<StateAccessSettings> {
resolve_state_access(&args.config_file, args.data_dir.clone(), args.state_file.clone())
}
pub fn resolve_revoke_enroll_token_settings(args: &RevokeEnrollTokenArgs) -> Result<StateAccessSettings> {
resolve_state_access(&args.config_file, args.data_dir.clone(), args.state_file.clone())
}
pub fn resolve_state_stats_settings(args: &StateStatsArgs) -> Result<StateAccessSettings> {
resolve_state_access(&args.config_file, args.data_dir.clone(), args.state_file.clone())
}
fn resolve_state_access(
config_file: &Path,
cli_data_dir: Option<PathBuf>,
cli_state_file: Option<PathBuf>,
) -> Result<StateAccessSettings> {
let file = load_file_config(config_file)?;
let data_dir = cli_data_dir
.or(file.data_dir)
.unwrap_or_else(|| PathBuf::from(crate::cli::DEFAULT_DATA_DIR));
let state_file = resolve_path(
&data_dir,
cli_state_file
.or(file.state_file)
.unwrap_or_else(|| PathBuf::from("state.db")),
);
Ok(StateAccessSettings {
data_dir,
state_file,
})
}
wakey-control-plane/src/main.rs
+12
View File
@@ -34,6 +34,18 @@ async fn main() -> Result<()> {
tracing::init(cli.verbose, &config::TelemetryConfig::default())?;
runtime::issue_enroll_token(args).await
}
Command::ListEnrollTokens(args) => {
tracing::init(cli.verbose, &config::TelemetryConfig::default())?;
runtime::list_enroll_tokens(args).await
}
Command::RevokeEnrollToken(args) => {
tracing::init(cli.verbose, &config::TelemetryConfig::default())?;
runtime::revoke_enroll_token(args).await
}
Command::StateStats(args) => {
tracing::init(cli.verbose, &config::TelemetryConfig::default())?;
runtime::state_stats(args).await
}
Command::Reload(args) => {
tracing::init(cli.verbose, &config::TelemetryConfig::default())?;
runtime::reload_daemon(&args.pid_file)
wakey-control-plane/src/runtime.rs
+72 -1
View File
@@ -8,11 +8,12 @@ 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;
use crate::cli::{IssueEnrollTokenArgs, ListEnrollTokensArgs, RevokeEnrollTokenArgs, StateStatsArgs};
use crate::config;
use crate::state;
use crate::ws;
@@ -57,6 +58,11 @@ pub async fn serve(daemon: config::DaemonConfig) -> Result<()> {
.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/agent/ws", get(ws::agent_ws))
.route("/api/v1/control/agents", get(api::list_agents))
.route(
@@ -77,6 +83,8 @@ pub async fn serve(daemon: config::DaemonConfig) -> Result<()> {
{
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! {
@@ -91,6 +99,16 @@ pub async fn serve(daemon: config::DaemonConfig) -> Result<()> {
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")?;
@@ -166,6 +184,59 @@ pub async fn issue_enroll_token(args: IssueEnrollTokenArgs) -> Result<()> {
Ok(())
}
pub async fn list_enroll_tokens(args: ListEnrollTokensArgs) -> Result<()> {
let settings = config::resolve_list_enroll_token_settings(&args)?;
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!(
"{}",
serde_json::to_string_pretty(&tokens).context("failed to render json")?
);
return Ok(());
}
for token in tokens {
println!(
"token={} expires_at_unix={} expired={}",
token.enroll_token, token.expires_at_unix, token.expired
);
}
Ok(())
}
pub async fn revoke_enroll_token(args: RevokeEnrollTokenArgs) -> Result<()> {
let settings = config::resolve_revoke_enroll_token_settings(&args)?;
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(())
}
pub async fn state_stats(args: StateStatsArgs) -> Result<()> {
let settings = config::resolve_state_stats_settings(&args)?;
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!(
"{}",
serde_json::to_string_pretty(&stats).context("failed to render json")?
);
return Ok(());
}
println!("db_path={}", stats.db_path.display());
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);
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");
wakey-control-plane/src/state.rs
+121 -49
View File
@@ -1,4 +1,3 @@
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
@@ -20,20 +19,34 @@ pub struct IssuedEnrollToken {
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct LegacyPersistedState {
enroll_tokens: std::collections::HashSet<String>,
agents: HashMap<String, String>,
pub struct EnrollTokenInfo {
pub enroll_token: String,
pub expires_at_unix: u64,
pub expired: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StateStats {
pub db_path: PathBuf,
pub schema_version: u32,
pub agent_count: usize,
pub enroll_token_count: usize,
pub expired_enroll_token_count: usize,
}
pub struct Store {
db_path: PathBuf,
meta: sled::Tree,
enroll_tokens: sled::Tree,
agents: sled::Tree,
}
const SCHEMA_VERSION_KEY: &[u8] = b"schema_version";
const SCHEMA_VERSION: u32 = 1;
impl Store {
pub async fn load_or_init(path: &Path, enroll_tokens: Vec<String>, seed_ttl: Duration) -> Result<Self> {
let db_path = canonical_db_path(path);
let db_path = path.to_path_buf();
if let Some(parent) = db_path.parent() {
std::fs::create_dir_all(parent)
.with_context(|| format!("failed to create state dir {}", parent.display()))?;
@@ -41,6 +54,7 @@ impl Store {
let db = sled::open(&db_path)
.with_context(|| format!("failed to open state db {}", db_path.display()))?;
let meta_tree = db.open_tree("meta").context("failed to open meta tree")?;
let enroll_tree = db
.open_tree("enroll_tokens")
.context("failed to open enroll_tokens tree")?;
@@ -48,11 +62,12 @@ impl Store {
let store = Self {
db_path,
meta: meta_tree,
enroll_tokens: enroll_tree,
agents: agents_tree,
};
store.maybe_migrate_legacy_json(path)?;
store.ensure_schema_version()?;
for token in enroll_tokens {
let token = token.trim();
@@ -66,7 +81,7 @@ impl Store {
.with_context(|| format!("failed to seed enroll token into {}", store.db_path.display()))?;
}
store.gc_expired_enroll_tokens()?;
store.gc_expired_enroll_tokens_inner()?;
store
.flush()
@@ -137,6 +152,65 @@ impl Store {
})
}
pub async fn list_enroll_tokens(&self, include_expired: bool) -> Result<Vec<EnrollTokenInfo>> {
let now = now_unix();
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 expired = expires_at_unix <= now;
if !include_expired && expired {
continue;
}
let enroll_token = String::from_utf8(token.to_vec()).context("invalid utf-8 enroll token in db")?;
out.push(EnrollTokenInfo {
enroll_token,
expires_at_unix,
expired,
});
}
out.sort_by(|a, b| a.expires_at_unix.cmp(&b.expires_at_unix).then(a.enroll_token.cmp(&b.enroll_token)));
Ok(out)
}
pub async fn revoke_enroll_token(&self, token: &str) -> Result<bool> {
let removed = self
.enroll_tokens
.remove(token.as_bytes())
.context("failed removing enroll token")?
.is_some();
if removed {
self.flush().context("failed flushing db after enroll token revoke")?;
}
Ok(removed)
}
pub async fn stats(&self) -> Result<StateStats> {
let now = now_unix();
let mut enroll_token_count = 0usize;
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")?;
enroll_token_count = enroll_token_count.saturating_add(1);
if expires_at <= now {
expired_enroll_token_count = expired_enroll_token_count.saturating_add(1);
}
}
Ok(StateStats {
db_path: self.db_path.clone(),
schema_version: self.schema_version()?,
agent_count: self.agents.iter().count(),
enroll_token_count,
expired_enroll_token_count,
})
}
pub async fn gc_expired_enroll_tokens(&self) -> Result<u64> {
self.gc_expired_enroll_tokens_inner()
}
pub async fn reload_from_disk(&self) -> Result<()> {
// sled is durable and read-through; explicit reload is a no-op.
info!(path = %self.db_path.display(), "reload requested; sled backend does not require in-memory reload");
@@ -176,43 +250,7 @@ impl Store {
Ok(())
}
fn maybe_migrate_legacy_json(&self, configured_path: &Path) -> Result<()> {
let legacy_path = match configured_path.extension().and_then(|x| x.to_str()) {
Some("json") => configured_path.to_path_buf(),
_ => configured_path.with_extension("json"),
};
if self.enroll_tokens.iter().next().is_some() || self.agents.iter().next().is_some() {
return Ok(());
}
if !legacy_path.exists() {
return Ok(());
}
let raw = std::fs::read_to_string(&legacy_path)
.with_context(|| format!("failed to read legacy state {}", legacy_path.display()))?;
let legacy: LegacyPersistedState = serde_json::from_str(&raw)
.with_context(|| format!("failed to parse legacy state {}", legacy_path.display()))?;
let expires_at_unix = now_unix().saturating_add(86_400);
for token in legacy.enroll_tokens {
self.enroll_tokens
.insert(token.as_bytes(), &expires_at_unix.to_le_bytes())
.context("failed to migrate enroll token")?;
}
for (agent_id, token) in legacy.agents {
self.agents
.insert(agent_id.as_bytes(), token.as_bytes())
.context("failed to migrate agent token")?;
}
self.flush().context("failed flushing migrated legacy state")?;
std::fs::remove_file(&legacy_path)
.with_context(|| format!("failed to delete legacy state {}", legacy_path.display()))?;
info!(legacy = %legacy_path.display(), db = %self.db_path.display(), "migrated legacy json state into sled db and deleted legacy file");
Ok(())
}
fn gc_expired_enroll_tokens(&self) -> Result<()> {
fn gc_expired_enroll_tokens_inner(&self) -> Result<u64> {
let now = now_unix();
let mut removed = 0u64;
for item in self.enroll_tokens.iter() {
@@ -229,14 +267,39 @@ impl Store {
self.flush().context("failed flushing db after gc")?;
info!(removed, "garbage-collected expired enroll tokens");
}
Ok(removed)
}
fn ensure_schema_version(&self) -> Result<()> {
match self.meta.get(SCHEMA_VERSION_KEY).context("failed reading schema version")? {
Some(raw) => {
let schema = decode_schema(raw.as_ref()).context("failed decoding schema version")?;
if schema != SCHEMA_VERSION {
anyhow::bail!(
"unsupported db schema version {}; expected {}",
schema,
SCHEMA_VERSION
);
}
}
None => {
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");
}
}
Ok(())
}
}
fn canonical_db_path(configured: &Path) -> PathBuf {
match configured.extension().and_then(|x| x.to_str()) {
Some("json") => configured.with_extension("db"),
_ => configured.to_path_buf(),
fn schema_version(&self) -> Result<u32> {
let raw = self
.meta
.get(SCHEMA_VERSION_KEY)
.context("failed reading schema version")?
.ok_or_else(|| anyhow::anyhow!("missing schema version in state db"))?;
decode_schema(raw.as_ref()).context("failed decoding schema version")
}
}
@@ -255,3 +318,12 @@ fn decode_expiry(raw: &[u8]) -> Result<u64> {
arr.copy_from_slice(raw);
Ok(u64::from_le_bytes(arr))
}
fn decode_schema(raw: &[u8]) -> Result<u32> {
if raw.len() != 4 {
anyhow::bail!("invalid schema version length {}", raw.len());
}
let mut arr = [0u8; 4];
arr.copy_from_slice(raw);
Ok(u32::from_le_bytes(arr))
}