2 Commits
13 changed files with 900 additions and 788 deletions

No files matched your search

+165
View File
@@ -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",
}
}
+27
View File
@@ -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,
}
})),
)
}
-394
View File
@@ -1,394 +0,0 @@
use std::net::SocketAddr;
use std::path::Path;
use std::path::PathBuf;
use std::time::Duration;
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use crate::cli::{
InitConfigArgs, IssueEnrollTokenArgs, ListEnrollTokensArgs, RevokeEnrollTokenArgs,
ServeArgs, StateStatsArgs,
};
#[derive(Debug, Clone)]
pub struct DaemonConfig {
pub data_dir: PathBuf,
pub bind: SocketAddr,
pub public_url: String,
pub state_file: PathBuf,
pub command_timeout: Duration,
pub enroll_token_ttl: 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 {
data_dir: Option<PathBuf>,
bind: Option<String>,
public_url: Option<String>,
state_file: Option<PathBuf>,
command_timeout_ms: Option<u64>,
enroll_token_ttl_seconds: 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>,
}
#[derive(Debug, Serialize)]
struct WritableConfig {
data_dir: PathBuf,
bind: String,
public_url: String,
state_file: PathBuf,
command_timeout_ms: u64,
enroll_token_ttl_seconds: u64,
pid_file: PathBuf,
#[serde(skip_serializing_if = "Vec::is_empty")]
enroll_tokens: Vec<String>,
telemetry: WritableTelemetry,
}
#[derive(Debug, Serialize)]
struct WritableTelemetry {
#[serde(skip_serializing_if = "Option::is_none")]
otlp_endpoint: Option<String>,
service_name: String,
json_logs: bool,
}
impl DaemonConfig {
pub fn from_serve_args(args: &ServeArgs) -> Result<Self> {
let file = load_file_config(&args.config_file)?;
let data_dir = args
.data_dir
.clone()
.or(file.data_dir.clone())
.unwrap_or_else(|| PathBuf::from(crate::cli::DEFAULT_DATA_DIR));
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_raw = args
.state_file
.clone()
.or(file.state_file)
.unwrap_or_else(|| PathBuf::from("state.db"));
let state_file = resolve_path(&data_dir, state_file_raw);
let command_timeout = Duration::from_millis(
args.command_timeout_ms
.or(file.command_timeout_ms)
.unwrap_or(30_000)
.max(1),
);
let enroll_token_ttl = Duration::from_secs(
args.enroll_token_ttl_seconds
.or(file.enroll_token_ttl_seconds)
.unwrap_or(86_400)
.max(1),
);
let pid_file_raw = args
.pid_file
.clone()
.or(file.pid_file)
.unwrap_or_else(|| PathBuf::from("wakey-control-plane.pid"));
let pid_file = resolve_path(&data_dir, pid_file_raw);
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 {
data_dir,
bind,
public_url,
state_file,
command_timeout,
enroll_token_ttl,
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 {
url.trim_end_matches('/').to_string()
}
pub fn issue_token_endpoint(base_url: &str) -> String {
format!(
"{}/api/v1/control/enroll-token",
normalize_public_url(base_url)
)
}
pub fn write_init_config(args: &InitConfigArgs) -> Result<()> {
if args.config_file.exists() && !args.force {
anyhow::bail!(
"config {} already exists; re-run with --force to overwrite",
args.config_file.display()
);
}
let bind = args
.bind
.unwrap_or_else(|| "0.0.0.0:8080".parse().expect("static default bind should parse"));
let public_url = normalize_public_url(
args.public_url
.as_deref()
.unwrap_or("http://127.0.0.1:8080"),
);
let data_dir = args
.data_dir
.clone()
.unwrap_or_else(|| PathBuf::from(crate::cli::DEFAULT_DATA_DIR));
let state_file_raw = args
.state_file
.clone()
.unwrap_or_else(|| PathBuf::from("state.db"));
let state_file = resolve_path(&data_dir, state_file_raw);
let pid_file_raw = args
.pid_file
.clone()
.unwrap_or_else(|| PathBuf::from("wakey-control-plane.pid"));
let pid_file = resolve_path(&data_dir, pid_file_raw);
let body = WritableConfig {
data_dir,
bind: bind.to_string(),
public_url,
state_file,
command_timeout_ms: args.command_timeout_ms.unwrap_or(30_000).max(1),
enroll_token_ttl_seconds: args.enroll_token_ttl_seconds.unwrap_or(86_400).max(1),
pid_file,
enroll_tokens: args.enroll_tokens.clone(),
telemetry: WritableTelemetry {
otlp_endpoint: args.telemetry_otlp_endpoint.clone(),
service_name: args
.telemetry_service_name
.clone()
.unwrap_or_else(|| "wakey-control-plane".to_string()),
json_logs: args.telemetry_json_logs,
},
};
if let Some(parent) = args.config_file.parent() {
std::fs::create_dir_all(parent)
.with_context(|| format!("failed to create config dir {}", parent.display()))?;
}
let rendered = toml::to_string_pretty(&body).context("failed to render config template")?;
std::fs::write(&args.config_file, rendered)
.with_context(|| format!("failed to write config {}", args.config_file.display()))?;
Ok(())
}
pub fn bootstrap_config_if_missing(args: &ServeArgs) -> Result<bool> {
if args.config_file.exists() {
return Ok(false);
}
let init = InitConfigArgs {
config_file: args.config_file.clone(),
data_dir: args.data_dir.clone(),
bind: args.bind,
public_url: args.public_url.clone(),
state_file: args.state_file.clone(),
pid_file: args.pid_file.clone(),
command_timeout_ms: args.command_timeout_ms,
enroll_token_ttl_seconds: args.enroll_token_ttl_seconds,
enroll_tokens: args.enroll_tokens.clone(),
telemetry_otlp_endpoint: None,
telemetry_service_name: None,
telemetry_json_logs: false,
force: false,
};
write_init_config(&init)?;
Ok(true)
}
pub fn resolve_path(data_dir: &Path, candidate: PathBuf) -> PathBuf {
if candidate.is_absolute() {
candidate
} else {
data_dir.join(candidate)
}
}
pub struct IssueTokenSettings {
pub data_dir: PathBuf,
pub state_file: PathBuf,
pub ttl: Duration,
}
pub struct StateAccessSettings {
pub data_dir: PathBuf,
pub state_file: PathBuf,
pub public_url: Option<String>,
}
pub fn resolve_issue_token_settings(args: &IssueEnrollTokenArgs) -> Result<IssueTokenSettings> {
let file = load_file_config(&args.config_file)?;
let state = resolve_state_access(
&args.config_file,
args.data_dir.clone(),
args.state_file.clone(),
args.public_url.clone(),
)?;
let ttl = Duration::from_secs(
args.ttl_seconds
.or(file.enroll_token_ttl_seconds)
.unwrap_or(86_400)
.max(1),
);
Ok(IssueTokenSettings {
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(),
args.public_url.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(),
args.public_url.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(),
args.public_url.clone(),
)
}
fn resolve_state_access(
config_file: &Path,
cli_data_dir: Option<PathBuf>,
cli_state_file: Option<PathBuf>,
cli_public_url: Option<String>,
) -> 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")),
);
let public_url = cli_public_url
.or(file.public_url)
.map(|url| normalize_public_url(&url));
Ok(StateAccessSettings {
data_dir,
state_file,
public_url,
})
}
+98
View File
@@ -0,0 +1,98 @@
use std::path::PathBuf;
use anyhow::{Context, Result};
use crate::cli::{InitConfigArgs, ServeArgs};
use crate::config::resolve::{normalize_public_url, resolve_path};
use crate::config::types::{WritableConfig, WritableTelemetry};
pub fn write_init_config(args: &InitConfigArgs) -> Result<()> {
if args.config_file.exists() && !args.force {
anyhow::bail!(
"config {} already exists; re-run with --force to overwrite",
args.config_file.display()
);
}
let bind = args.bind.unwrap_or_else(|| {
"0.0.0.0:8080"
.parse()
.expect("static default bind should parse")
});
let public_url = normalize_public_url(
args.public_url
.as_deref()
.unwrap_or("http://127.0.0.1:8080"),
);
let data_dir = args
.data_dir
.clone()
.unwrap_or_else(|| PathBuf::from(crate::cli::DEFAULT_DATA_DIR));
let state_file_raw = args
.state_file
.clone()
.unwrap_or_else(|| PathBuf::from("state.db"));
let state_file = resolve_path(&data_dir, state_file_raw);
let pid_file_raw = args
.pid_file
.clone()
.unwrap_or_else(|| PathBuf::from("wakey-control-plane.pid"));
let pid_file = resolve_path(&data_dir, pid_file_raw);
let body = WritableConfig {
data_dir,
bind: bind.to_string(),
public_url,
state_file,
command_timeout_ms: args.command_timeout_ms.unwrap_or(30_000).max(1),
enroll_token_ttl_seconds: args.enroll_token_ttl_seconds.unwrap_or(86_400).max(1),
pid_file,
enroll_tokens: args.enroll_tokens.clone(),
telemetry: WritableTelemetry {
otlp_endpoint: args.telemetry_otlp_endpoint.clone(),
service_name: args
.telemetry_service_name
.clone()
.unwrap_or_else(|| "wakey-control-plane".to_string()),
json_logs: args.telemetry_json_logs,
},
};
if let Some(parent) = args.config_file.parent() {
std::fs::create_dir_all(parent)
.with_context(|| format!("failed to create config dir {}", parent.display()))?;
}
let rendered = toml::to_string_pretty(&body).context("failed to render config template")?;
std::fs::write(&args.config_file, rendered)
.with_context(|| format!("failed to write config {}", args.config_file.display()))?;
Ok(())
}
pub fn bootstrap_config_if_missing(args: &ServeArgs) -> Result<bool> {
if args.config_file.exists() {
return Ok(false);
}
let init = InitConfigArgs {
config_file: args.config_file.clone(),
data_dir: args.data_dir.clone(),
bind: args.bind,
public_url: args.public_url.clone(),
state_file: args.state_file.clone(),
pid_file: args.pid_file.clone(),
command_timeout_ms: args.command_timeout_ms,
enroll_token_ttl_seconds: args.enroll_token_ttl_seconds,
enroll_tokens: args.enroll_tokens.clone(),
telemetry_otlp_endpoint: None,
telemetry_service_name: None,
telemetry_json_logs: false,
force: false,
};
write_init_config(&init)?;
Ok(true)
}
+11
View File
@@ -0,0 +1,11 @@
mod init;
mod resolve;
mod types;
pub use init::{bootstrap_config_if_missing, write_init_config};
pub use resolve::{
issue_token_endpoint, normalize_public_url, resolve_issue_token_settings,
resolve_list_enroll_token_settings, resolve_revoke_enroll_token_settings,
resolve_state_stats_settings,
};
pub use types::{DaemonConfig, TelemetryConfig};
+255
View File
@@ -0,0 +1,255 @@
use std::net::SocketAddr;
use std::path::{Path, PathBuf};
use std::time::Duration;
use anyhow::{Context, Result};
use crate::cli::{
IssueEnrollTokenArgs, ListEnrollTokensArgs, RevokeEnrollTokenArgs, ServeArgs, StateStatsArgs,
};
use crate::config::types::{
DaemonConfig, FileConfig, FileTelemetryConfig, IssueTokenSettings, StateAccessSettings,
TelemetryConfig,
};
impl DaemonConfig {
pub fn from_serve_args(args: &ServeArgs) -> Result<Self> {
let file = load_file_config(&args.config_file)?;
let data_dir = args
.data_dir
.clone()
.or(file.data_dir.clone())
.unwrap_or_else(|| PathBuf::from(crate::cli::DEFAULT_DATA_DIR));
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_raw = args
.state_file
.clone()
.or(file.state_file)
.unwrap_or_else(|| PathBuf::from("state.db"));
let state_file = resolve_path(&data_dir, state_file_raw);
let command_timeout = Duration::from_millis(
args.command_timeout_ms
.or(file.command_timeout_ms)
.unwrap_or(30_000)
.max(1),
);
let enroll_token_ttl = Duration::from_secs(
args.enroll_token_ttl_seconds
.or(file.enroll_token_ttl_seconds)
.unwrap_or(86_400)
.max(1),
);
let pid_file_raw = args
.pid_file
.clone()
.or(file.pid_file)
.unwrap_or_else(|| PathBuf::from("wakey-control-plane.pid"));
let pid_file = resolve_path(&data_dir, pid_file_raw);
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 {
data_dir,
bind,
public_url,
state_file,
command_timeout,
enroll_token_ttl,
pid_file,
enroll_tokens,
telemetry,
})
}
}
pub fn normalize_public_url(url: &str) -> String {
url.trim_end_matches('/').to_string()
}
pub fn issue_token_endpoint(base_url: &str) -> String {
format!(
"{}/api/v1/control/enroll-token",
normalize_public_url(base_url)
)
}
pub fn resolve_path(data_dir: &Path, candidate: PathBuf) -> PathBuf {
if candidate.is_absolute() {
candidate
} else {
data_dir.join(candidate)
}
}
pub fn resolve_issue_token_settings(args: &IssueEnrollTokenArgs) -> Result<IssueTokenSettings> {
let file = load_file_config(&args.config_file)?;
let state = resolve_state_access(
&args.config_file,
args.data_dir.clone(),
args.state_file.clone(),
args.public_url.clone(),
)?;
let ttl = Duration::from_secs(
args.ttl_seconds
.or(file.enroll_token_ttl_seconds)
.unwrap_or(86_400)
.max(1),
);
Ok(IssueTokenSettings {
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(),
args.public_url.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(),
args.public_url.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(),
args.public_url.clone(),
)
}
pub(crate) 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
}
fn resolve_state_access(
config_file: &Path,
cli_data_dir: Option<PathBuf>,
cli_state_file: Option<PathBuf>,
cli_public_url: Option<String>,
) -> 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")),
);
let public_url = cli_public_url
.or(file.public_url)
.map(|url| normalize_public_url(&url));
Ok(StateAccessSettings {
data_dir,
state_file,
public_url,
})
}
#[cfg(test)]
mod tests {
use super::{normalize_public_url, resolve_path};
use std::path::Path;
use std::path::PathBuf;
#[test]
fn normalize_public_url_trims_trailing_slash() {
assert_eq!(
normalize_public_url("https://cp.example.com/"),
"https://cp.example.com"
);
}
#[test]
fn resolve_path_joins_relative_path() {
let out = resolve_path(
Path::new("/var/lib/wakey-control-plane"),
PathBuf::from("state.db"),
);
assert_eq!(out, PathBuf::from("/var/lib/wakey-control-plane/state.db"));
}
}
+89
View File
@@ -0,0 +1,89 @@
use std::net::SocketAddr;
use std::path::PathBuf;
use std::time::Duration;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone)]
pub struct DaemonConfig {
pub data_dir: PathBuf,
pub bind: SocketAddr,
pub public_url: String,
pub state_file: PathBuf,
pub command_timeout: Duration,
pub enroll_token_ttl: 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)]
pub(crate) struct FileConfig {
pub(crate) data_dir: Option<PathBuf>,
pub(crate) bind: Option<String>,
pub(crate) public_url: Option<String>,
pub(crate) state_file: Option<PathBuf>,
pub(crate) command_timeout_ms: Option<u64>,
pub(crate) enroll_token_ttl_seconds: Option<u64>,
pub(crate) pid_file: Option<PathBuf>,
pub(crate) enroll_tokens: Option<Vec<String>>,
pub(crate) telemetry: Option<FileTelemetryConfig>,
}
#[derive(Debug, Deserialize, Default)]
pub(crate) struct FileTelemetryConfig {
pub(crate) otlp_endpoint: Option<String>,
pub(crate) service_name: Option<String>,
pub(crate) json_logs: Option<bool>,
}
#[derive(Debug, Serialize)]
pub(crate) struct WritableConfig {
pub(crate) data_dir: PathBuf,
pub(crate) bind: String,
pub(crate) public_url: String,
pub(crate) state_file: PathBuf,
pub(crate) command_timeout_ms: u64,
pub(crate) enroll_token_ttl_seconds: u64,
pub(crate) pid_file: PathBuf,
#[serde(skip_serializing_if = "Vec::is_empty")]
pub(crate) enroll_tokens: Vec<String>,
pub(crate) telemetry: WritableTelemetry,
}
#[derive(Debug, Serialize)]
pub(crate) struct WritableTelemetry {
#[serde(skip_serializing_if = "Option::is_none")]
pub(crate) otlp_endpoint: Option<String>,
pub(crate) service_name: String,
pub(crate) json_logs: bool,
}
pub struct IssueTokenSettings {
pub data_dir: PathBuf,
pub state_file: PathBuf,
pub ttl: Duration,
}
pub struct StateAccessSettings {
pub data_dir: PathBuf,
pub state_file: PathBuf,
pub public_url: Option<String>,
}
@@ -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)")
}
}
+191
View File
@@ -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)")
}
}
+4
View File
@@ -0,0 +1,4 @@
mod store;
mod types;
pub use store::Store;
@@ -2,37 +2,10 @@ use std::path::{Path, PathBuf};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use tracing::{debug, info, warn};
use uuid::Uuid;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IssuedAgent {
pub agent_id: String,
pub agent_token: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IssuedEnrollToken {
pub enroll_token: String,
pub expires_at_unix: u64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
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,
}
use crate::state::types::{EnrollTokenInfo, IssuedAgent, IssuedEnrollToken, StateStats};
pub struct Store {
db_path: PathBuf,
@@ -162,14 +135,19 @@ impl Store {
if !include_expired && expired {
continue;
}
let enroll_token = String::from_utf8(token.to_vec()).context("invalid utf-8 enroll token in db")?;
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)));
out.sort_by(|a, b| {
a.expires_at_unix
.cmp(&b.expires_at_unix)
.then(a.enroll_token.cmp(&b.enroll_token))
});
Ok(out)
}
@@ -191,7 +169,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);
@@ -255,7 +234,8 @@ impl Store {
let mut removed = 0u64;
for item in self.enroll_tokens.iter() {
let (token, value) = item.context("failed iterating enroll token tree")?;
let expires_at = decode_expiry(value.as_ref()).context("failed decoding token expiry during gc")?;
let expires_at =
decode_expiry(value.as_ref()).context("failed decoding token expiry during gc")?;
if expires_at <= now {
self.enroll_tokens
.remove(token)
@@ -271,9 +251,14 @@ impl Store {
}
fn ensure_schema_version(&self) -> Result<()> {
match self.meta.get(SCHEMA_VERSION_KEY).context("failed reading schema version")? {
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")?;
let schema =
decode_schema(raw.as_ref()).context("failed decoding schema version")?;
if schema != SCHEMA_VERSION {
anyhow::bail!(
"unsupported db schema version {}; expected {}",
+31
View File
@@ -0,0 +1,31 @@
use std::path::PathBuf;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IssuedAgent {
pub agent_id: String,
pub agent_token: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IssuedEnrollToken {
pub enroll_token: String,
pub expires_at_unix: u64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
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,
}