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

pih

Browse Source
This commit is contained in:
lda
2026-04-12 01:13:54 +07:00 Unverified
parent f2ac690f17
commit 729fd47327
6 changed files with 339 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wakey-control-plane/src/api/alerts.rs
+123 -8
View File
@@ -2,10 +2,12 @@ use std::collections::{BTreeMap, HashSet};
use std::time::{SystemTime, UNIX_EPOCH};
use axum::Json;
use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use axum::extract::{Query, State};
use axum::http::StatusCode;
use axum::response::IntoResponse;
use serde::Deserialize;
use tokio::time::Duration;
use tracing::warn;
use crate::api::json_error;
@@ -20,14 +22,129 @@ pub struct ActiveAlertsQuery {
pub enroll_rejected_threshold: Option<u64>,
}
#[derive(Debug, Deserialize)]
pub struct AlertHistoryQuery {
pub since_unix: Option<u64>,
pub limit: Option<usize>,
}
#[derive(Debug, Deserialize, Clone)]
pub struct AlertRuleConfig {
pub lookback_seconds: Option<u64>,
pub timeout_threshold: Option<u64>,
pub auth_rejected_threshold: Option<u64>,
pub enroll_rejected_threshold: Option<u64>,
}
pub async fn active_alerts(
State(state): State<AppState>,
Query(query): Query<ActiveAlertsQuery>,
) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
let lookback_seconds = query.lookback_seconds.unwrap_or(900).clamp(60, 86_400);
let timeout_threshold = query.timeout_threshold.unwrap_or(3).max(1);
let auth_rejected_threshold = query.auth_rejected_threshold.unwrap_or(3).max(1);
let enroll_rejected_threshold = query.enroll_rejected_threshold.unwrap_or(5).max(1);
let alerts = evaluate_alerts(
&state,
AlertRuleConfig {
lookback_seconds: query.lookback_seconds,
timeout_threshold: query.timeout_threshold,
auth_rejected_threshold: query.auth_rejected_threshold,
enroll_rejected_threshold: query.enroll_rejected_threshold,
},
)
.await?;
if let Err(err) = state.store.sync_alert_transitions(&alerts).await {
warn!(error = %err, "failed to sync alert transitions");
}
Ok((StatusCode::OK, Json(alerts)))
}
pub async fn alert_history(
State(state): State<AppState>,
Query(query): Query<AlertHistoryQuery>,
) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
let limit = query.limit.unwrap_or(100).clamp(1, 500);
let history = state
.store
.list_alert_transitions(query.since_unix, limit)
.await
.map_err(|err| {
warn!(error = %err, "failed reading alert transition history");
json_error(
StatusCode::INTERNAL_SERVER_ERROR,
"alert_history_failed",
&err.to_string(),
)
})?;
Ok((StatusCode::OK, Json(history)))
}
pub async fn alerts_stream(
ws: WebSocketUpgrade,
State(state): State<AppState>,
Query(query): Query<ActiveAlertsQuery>,
) -> impl IntoResponse {
let config = AlertRuleConfig {
lookback_seconds: query.lookback_seconds,
timeout_threshold: query.timeout_threshold,
auth_rejected_threshold: query.auth_rejected_threshold,
enroll_rejected_threshold: query.enroll_rejected_threshold,
};
ws.on_upgrade(move |socket| alerts_stream_socket(state, socket, config))
}
async fn alerts_stream_socket(state: AppState, mut socket: WebSocket, config: AlertRuleConfig) {
let mut tick = tokio::time::interval(Duration::from_secs(5));
loop {
tick.tick().await;
let alerts = match evaluate_alerts(&state, config.clone()).await {
Ok(alerts) => alerts,
Err(err) => {
warn!(code = %err.0, "failed to evaluate alerts for stream");
continue;
}
};
if let Err(err) = state.store.sync_alert_transitions(&alerts).await {
warn!(error = %err, "failed to sync alert transitions in stream");
}
let history = match state.store.list_alert_transitions(None, 20).await {
Ok(h) => h,
Err(err) => {
warn!(error = %err, "failed to load alert transition history for stream");
Vec::new()
}
};
let payload = serde_json::json!({
"type": "alerts_snapshot",
"ts_unix": now_unix(),
"alerts": alerts,
"recent_transitions": history,
});
let encoded = match serde_json::to_string(&payload) {
Ok(s) => s,
Err(err) => {
warn!(error = %err, "failed to encode alerts stream payload");
continue;
}
};
if socket.send(Message::Text(encoded.into())).await.is_err() {
break;
}
}
}
async fn evaluate_alerts(
state: &AppState,
config: AlertRuleConfig,
) -> Result<Vec<AlertState>, (StatusCode, Json<serde_json::Value>)> {
let lookback_seconds = config.lookback_seconds.unwrap_or(900).clamp(60, 86_400);
let timeout_threshold = config.timeout_threshold.unwrap_or(3).max(1);
let auth_rejected_threshold = config.auth_rejected_threshold.unwrap_or(3).max(1);
let enroll_rejected_threshold = config.enroll_rejected_threshold.unwrap_or(5).max(1);
let now = now_unix();
let since_unix = now.saturating_sub(lookback_seconds);
@@ -98,7 +215,7 @@ pub async fn active_alerts(
)
})?;
let alerts = build_alerts(
Ok(build_alerts(
now,
&enrolled_agents,
&connected_agents,
@@ -108,9 +225,7 @@ pub async fn active_alerts(
timeout_threshold,
auth_rejected_threshold,
enroll_rejected_threshold,
);
Ok((StatusCode::OK, Json(alerts)))
))
}
fn build_alerts(
wakey-control-plane/src/api/mod.rs
+1 -1
View File
@@ -7,7 +7,7 @@ mod audit;
mod alerts;
pub use commands::{list_agents, run_command};
pub use alerts::active_alerts;
pub use alerts::{active_alerts, alert_history, alerts_stream};
pub use audit::list_audit_events;
pub use control::{
EnrollTokenStatus, IssueEnrollTokenResponse, RevokeEnrollTokenResponse, StateStatsResponse,
wakey-control-plane/src/runtime/mod.rs
+32 -17
View File
@@ -40,6 +40,33 @@ pub enum AgentReply {
Error(ErrorPayload),
}
fn public_api_routes() -> Router<AppState> {
Router::new()
.route("/healthz", get(api::healthz))
.route("/api/v1/agents/enroll", post(api::enroll))
.route("/api/v1/agent/ws", get(ws::agent_ws))
}
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-tokens/{token}",
axum::routing::delete(api::revoke_enroll_token),
)
.route("/api/v1/control/state-stats", get(api::state_stats))
.route("/api/v1/control/audit/events", get(api::list_audit_events))
.route("/api/v1/control/alerts", get(api::active_alerts))
.route("/api/v1/control/alerts/history", get(api::alert_history))
.route("/api/v1/control/alerts/ws", get(api::alerts_stream))
.route("/api/v1/control/agents", get(api::list_agents))
.route(
"/api/v1/control/agents/{agent_id}/command",
post(api::run_command),
)
}
/// Starts the control-plane HTTP and websocket surfaces and manages daemon lifecycle hooks.
pub async fn serve(daemon: config::DaemonConfig) -> Result<()> {
write_pid_file(&daemon.pid_file)?;
@@ -62,24 +89,12 @@ pub async fn serve(daemon: config::DaemonConfig) -> Result<()> {
enroll_token_ttl: daemon.enroll_token_ttl,
};
// Keep route classes explicit so edge policy can map directly:
// - public_api_routes: intended internet-facing agent endpoints
// - control_api_routes: intended admin-only endpoints behind Access
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/control/audit/events", get(api::list_audit_events))
.route("/api/v1/control/alerts", get(api::active_alerts))
.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),
)
.merge(public_api_routes())
.merge(control_api_routes())
.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");
wakey-control-plane/src/state/store.rs
+161 -2
View File
@@ -6,8 +6,8 @@ use tracing::{debug, info, warn};
use uuid::Uuid;
use crate::state::types::{
AuditEvent, AuditEventFilter, AuditEventInput, EnrollTokenInfo, IssuedAgent, IssuedEnrollToken,
StateStats,
AlertState, AlertTransition, AuditEvent, AuditEventFilter, AuditEventInput, EnrollTokenInfo,
IssuedAgent, IssuedEnrollToken, StateStats,
};
pub struct Store {
@@ -16,6 +16,8 @@ pub struct Store {
enroll_tokens: sled::Tree,
agents: sled::Tree,
audit_events: sled::Tree,
active_alerts: sled::Tree,
alert_transitions: sled::Tree,
}
const SCHEMA_VERSION_KEY: &[u8] = b"schema_version";
@@ -39,6 +41,12 @@ impl Store {
let audit_events_tree = db
.open_tree("audit_events")
.context("failed to open audit_events tree")?;
let active_alerts_tree = db
.open_tree("active_alerts")
.context("failed to open active_alerts tree")?;
let alert_transitions_tree = db
.open_tree("alert_transitions")
.context("failed to open alert_transitions tree")?;
let store = Self {
db_path,
@@ -46,6 +54,8 @@ impl Store {
enroll_tokens: enroll_tree,
agents: agents_tree,
audit_events: audit_events_tree,
active_alerts: active_alerts_tree,
alert_transitions: alert_transitions_tree,
};
store.ensure_schema_version()?;
@@ -275,6 +285,111 @@ impl Store {
Ok(out)
}
pub async fn sync_alert_transitions(
&self,
current: &[AlertState],
) -> Result<Vec<AlertTransition>> {
let mut previous = std::collections::HashMap::<String, AlertState>::new();
for item in self.active_alerts.iter() {
let (_, raw) = item.context("failed iterating active_alerts tree")?;
let alert: AlertState =
serde_json::from_slice(raw.as_ref()).context("failed decoding active alert")?;
previous.insert(alert.alert_id.clone(), alert);
}
let mut current_map = std::collections::HashMap::<String, AlertState>::new();
for alert in current {
current_map.insert(alert.alert_id.clone(), alert.clone());
}
let now = now_unix();
let mut transitions = Vec::new();
for (alert_id, current_alert) in &current_map {
if previous.contains_key(alert_id) {
continue;
}
transitions.push(AlertTransition {
transition_id: format!("atr-{}", Uuid::new_v4()),
ts_unix: now,
alert_id: current_alert.alert_id.clone(),
kind: current_alert.kind.clone(),
agent_id: current_alert.agent_id.clone(),
from_status: None,
to_status: "active".into(),
message: current_alert.message.clone(),
metadata: current_alert.metadata.clone(),
});
}
for (alert_id, previous_alert) in &previous {
if current_map.contains_key(alert_id) {
continue;
}
transitions.push(AlertTransition {
transition_id: format!("atr-{}", Uuid::new_v4()),
ts_unix: now,
alert_id: previous_alert.alert_id.clone(),
kind: previous_alert.kind.clone(),
agent_id: previous_alert.agent_id.clone(),
from_status: Some("active".into()),
to_status: "resolved".into(),
message: format!("resolved alert {}", previous_alert.alert_id),
metadata: previous_alert.metadata.clone(),
});
}
for item in self.active_alerts.iter() {
let (key, _) = item.context("failed iterating active_alerts keys")?;
self.active_alerts
.remove(key)
.context("failed clearing active alert snapshot")?;
}
for alert in current {
let key = alert.alert_id.as_bytes();
let value = serde_json::to_vec(alert).context("failed encoding active alert")?;
self.active_alerts
.insert(key, value)
.context("failed writing active alert snapshot")?;
}
for transition in &transitions {
let key = format!("{:020}:{}", transition.ts_unix, transition.transition_id);
let value =
serde_json::to_vec(transition).context("failed encoding alert transition")?;
self.alert_transitions
.insert(key.as_bytes(), value)
.context("failed persisting alert transition")?;
}
self.flush()
.context("failed flushing state db after alert sync")?;
Ok(transitions)
}
pub async fn list_alert_transitions(
&self,
since_unix: Option<u64>,
limit: usize,
) -> Result<Vec<AlertTransition>> {
let limit = limit.clamp(1, 500);
let mut out = Vec::new();
for item in self.alert_transitions.iter().rev() {
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 {
continue;
}
out.push(transition);
if out.len() >= limit {
break;
}
}
Ok(out)
}
fn flush(&self) -> Result<()> {
self.enroll_tokens
.flush()
@@ -285,6 +400,12 @@ impl Store {
self.audit_events
.flush()
.context("failed to flush audit event tree")?;
self.active_alerts
.flush()
.context("failed to flush active alerts tree")?;
self.alert_transitions
.flush()
.context("failed to flush alert transitions tree")?;
debug!(path = %self.db_path.display(), "flushed sled state db");
Ok(())
}
@@ -561,4 +682,42 @@ mod tests {
assert_eq!(filtered[0].request_id.as_deref(), Some("req-1"));
cleanup_dir(&dir);
}
#[tokio::test]
async fn alert_transitions_track_open_and_resolve() {
let (store, dir) = make_store().await;
let alert = crate::state::AlertState {
alert_id: "agent_offline:agent-a".into(),
kind: "agent_offline".into(),
severity: "warning".into(),
status: "active".into(),
agent_id: Some("agent-a".into()),
message: "agent agent-a offline".into(),
value: 1,
threshold: 1,
last_seen_unix: 10,
metadata: serde_json::json!({}),
};
let opened = store
.sync_alert_transitions(std::slice::from_ref(&alert))
.await
.expect("open transition should succeed");
assert_eq!(opened.len(), 1);
assert_eq!(opened[0].to_status, "active");
let resolved = store
.sync_alert_transitions(&[])
.await
.expect("resolve transition should succeed");
assert_eq!(resolved.len(), 1);
assert_eq!(resolved[0].to_status, "resolved");
let history = store
.list_alert_transitions(None, 10)
.await
.expect("history should load");
assert!(history.len() >= 2);
cleanup_dir(&dir);
}
}
wakey-control-plane/src/state/types.rs
+13
View File
@@ -83,3 +83,16 @@ pub struct AlertState {
pub last_seen_unix: u64,
pub metadata: serde_json::Value,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AlertTransition {
pub transition_id: String,
pub ts_unix: u64,
pub alert_id: String,
pub kind: String,
pub agent_id: Option<String>,
pub from_status: Option<String>,
pub to_status: String,
pub message: String,
pub metadata: serde_json::Value,
}