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

megalog

Browse Source
This commit is contained in:
lda
2026-04-12 00:45:32 +07:00 Unverified
parent b8b8d951f1
commit ac54ec7c74
12 changed files with 634 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wakey-control-plane/src/api/audit.rs
+81
View File
@@ -0,0 +1,81 @@
use axum::Json;
use axum::extract::{Query, State};
use axum::http::StatusCode;
use axum::response::IntoResponse;
use serde::{Deserialize, Serialize};
use tracing::warn;
use crate::api::json_error;
use crate::runtime::AppState;
use crate::state::AuditEventFilter;
#[derive(Debug, Deserialize)]
pub struct ListAuditEventsQuery {
pub agent_id: Option<String>,
pub request_id: Option<String>,
pub event_type: Option<String>,
pub outcome: Option<String>,
pub since_unix: Option<u64>,
pub until_unix: Option<u64>,
pub limit: Option<usize>,
}
#[derive(Debug, Serialize)]
pub struct AuditEventResponse {
pub event_id: String,
pub ts_unix: u64,
pub actor_type: String,
pub actor_id: Option<String>,
pub agent_id: Option<String>,
pub request_id: Option<String>,
pub event_type: String,
pub outcome: String,
pub latency_ms: Option<u64>,
pub message: String,
pub metadata: serde_json::Value,
}
pub async fn list_audit_events(
State(state): State<AppState>,
Query(query): Query<ListAuditEventsQuery>,
) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
let filter = AuditEventFilter {
agent_id: query.agent_id,
request_id: query.request_id,
event_type: query.event_type,
outcome: query.outcome,
since_unix: query.since_unix,
until_unix: query.until_unix,
limit: query.limit.unwrap_or(100),
};
match state.store.list_audit_events(filter).await {
Ok(events) => {
let body = events
.into_iter()
.map(|event| AuditEventResponse {
event_id: event.event_id,
ts_unix: event.ts_unix,
actor_type: event.actor_type,
actor_id: event.actor_id,
agent_id: event.agent_id,
request_id: event.request_id,
event_type: event.event_type,
outcome: event.outcome,
latency_ms: event.latency_ms,
message: event.message,
metadata: event.metadata,
})
.collect::<Vec<_>>();
Ok((StatusCode::OK, Json(body)))
}
Err(err) => {
warn!(error = %err, "failed to list audit events");
Err(json_error(
StatusCode::INTERNAL_SERVER_ERROR,
"list_audit_events_failed",
&err.to_string(),
))
}
}
}
wakey-control-plane/src/api/commands.rs
+105
View File
@@ -3,12 +3,14 @@ use axum::extract::{Path as AxumPath, State};
use axum::http::StatusCode;
use axum::response::IntoResponse;
use serde::{Deserialize, Serialize};
use std::time::Instant;
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};
use crate::state::AuditEventInput;
#[derive(Debug, Serialize)]
pub struct AgentStatus {
@@ -56,6 +58,7 @@ pub async fn run_command(
) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
let request_id_string = format!("req-{}", Uuid::new_v4());
let command = command_kind(&req.command);
let started = Instant::now();
let span = info_span!(
"relay_command",
agent_id = %agent_id,
@@ -93,6 +96,23 @@ pub async fn run_command(
.insert(request_id_string.clone(), pending_tx);
info!("dispatching command to agent");
if let Err(err) = state
.store
.append_audit_event(AuditEventInput {
actor_type: "admin_api".into(),
actor_id: None,
agent_id: Some(agent_id.clone()),
request_id: Some(request_id_string.clone()),
event_type: "command_dispatch".into(),
outcome: "sent".into(),
latency_ms: None,
message: "dispatched command to connected agent".into(),
metadata: serde_json::json!({ "command": command }),
})
.await
{
warn!(error = %err, "failed to append audit event for command dispatch");
}
if let Err(err) = tx.send(ServerMessage::Command {
request_id,
@@ -100,6 +120,23 @@ pub async fn run_command(
}) {
state.pending.lock().await.remove(&request_id_string);
warn!(error = %err, "failed sending command to agent session");
if let Err(audit_err) = state
.store
.append_audit_event(AuditEventInput {
actor_type: "admin_api".into(),
actor_id: None,
agent_id: Some(agent_id.clone()),
request_id: Some(request_id_string.clone()),
event_type: "command_dispatch".into(),
outcome: "send_failed".into(),
latency_ms: Some(started.elapsed().as_millis() as u64),
message: err.to_string(),
metadata: serde_json::json!({ "command": command }),
})
.await
{
warn!(error = %audit_err, "failed to append audit event for command send failure");
}
return Err(json_error(
StatusCode::BAD_GATEWAY,
"agent_send_failed",
@@ -116,6 +153,23 @@ pub async fn run_command(
let response = match outcome {
Ok(Ok(AgentReply::Result(result))) => {
info!("agent command completed");
if let Err(err) = state
.store
.append_audit_event(AuditEventInput {
actor_type: "admin_api".into(),
actor_id: None,
agent_id: Some(agent_id.clone()),
request_id: Some(request_id_string.clone()),
event_type: "command_result".into(),
outcome: "ok".into(),
latency_ms: Some(started.elapsed().as_millis() as u64),
message: "agent command completed".into(),
metadata: serde_json::json!({ "command": command }),
})
.await
{
warn!(error = %err, "failed to append audit event for command success");
}
RelayCommandResponse {
request_id: request_id_string,
status: "ok".into(),
@@ -125,6 +179,23 @@ pub async fn run_command(
}
Ok(Ok(AgentReply::Error(error))) => {
warn!(code = %error.code, "agent command returned error");
if let Err(err) = state
.store
.append_audit_event(AuditEventInput {
actor_type: "admin_api".into(),
actor_id: None,
agent_id: Some(agent_id.clone()),
request_id: Some(request_id_string.clone()),
event_type: "command_result".into(),
outcome: "error".into(),
latency_ms: Some(started.elapsed().as_millis() as u64),
message: error.message.clone(),
metadata: serde_json::json!({ "command": command, "code": error.code }),
})
.await
{
warn!(error = %err, "failed to append audit event for command error result");
}
RelayCommandResponse {
request_id: request_id_string,
status: "error".into(),
@@ -134,6 +205,23 @@ pub async fn run_command(
}
Ok(Err(_)) => {
warn!("agent response channel dropped");
if let Err(err) = state
.store
.append_audit_event(AuditEventInput {
actor_type: "admin_api".into(),
actor_id: None,
agent_id: Some(agent_id.clone()),
request_id: Some(request_id_string.clone()),
event_type: "command_result".into(),
outcome: "response_dropped".into(),
latency_ms: Some(started.elapsed().as_millis() as u64),
message: "agent response channel dropped".into(),
metadata: serde_json::json!({ "command": command }),
})
.await
{
warn!(error = %err, "failed to append audit event for dropped response");
}
return Err(json_error(
StatusCode::BAD_GATEWAY,
"agent_response_dropped",
@@ -143,6 +231,23 @@ 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");
if let Err(err) = state
.store
.append_audit_event(AuditEventInput {
actor_type: "admin_api".into(),
actor_id: None,
agent_id: Some(agent_id.clone()),
request_id: Some(request_id_string.clone()),
event_type: "command_result".into(),
outcome: "timeout".into(),
latency_ms: Some(started.elapsed().as_millis() as u64),
message: "agent command timed out".into(),
metadata: serde_json::json!({ "command": command, "timeout_ms": timeout.as_millis() as u64 }),
})
.await
{
warn!(error = %err, "failed to append audit event for timeout");
}
return Err(json_error(
StatusCode::GATEWAY_TIMEOUT,
"agent_timeout",
wakey-control-plane/src/api/control.rs
+102 -4
View File
@@ -7,6 +7,7 @@ use tracing::{info, warn};
use crate::api::json_error;
use crate::runtime::AppState;
use crate::state::AuditEventInput;
#[derive(Debug, Deserialize)]
pub struct EnrollRequest {
@@ -69,6 +70,23 @@ pub async fn enroll(
match state.store.enroll(&req.enroll_token).await {
Ok(issued) => {
info!(agent_id = %issued.agent_id, "agent enrollment accepted");
if let Err(err) = state
.store
.append_audit_event(AuditEventInput {
actor_type: "agent".into(),
actor_id: Some(issued.agent_id.clone()),
agent_id: Some(issued.agent_id.clone()),
request_id: None,
event_type: "agent_enroll".into(),
outcome: "ok".into(),
latency_ms: None,
message: "agent enrollment accepted".into(),
metadata: serde_json::json!({}),
})
.await
{
warn!(error = %err, "failed to append audit event for enroll success");
}
Ok((
StatusCode::OK,
Json(EnrollResponse {
@@ -80,6 +98,23 @@ pub async fn enroll(
}
Err(err) => {
warn!(error = %err, "agent enrollment rejected");
if let Err(audit_err) = state
.store
.append_audit_event(AuditEventInput {
actor_type: "agent".into(),
actor_id: None,
agent_id: None,
request_id: None,
event_type: "agent_enroll".into(),
outcome: "rejected".into(),
latency_ms: None,
message: err.to_string(),
metadata: serde_json::json!({}),
})
.await
{
warn!(error = %audit_err, "failed to append audit event for enroll rejection");
}
Err(json_error(
StatusCode::UNAUTHORIZED,
"enrollment_rejected",
@@ -106,6 +141,26 @@ pub async fn issue_enroll_token(
expires_at_unix = issued.expires_at_unix,
"issued enroll token"
);
if let Err(err) = state
.store
.append_audit_event(AuditEventInput {
actor_type: "admin_api".into(),
actor_id: None,
agent_id: None,
request_id: None,
event_type: "enroll_token_issue".into(),
outcome: "ok".into(),
latency_ms: None,
message: "issued enroll token".into(),
metadata: serde_json::json!({
"ttl_seconds": ttl.as_secs(),
"expires_at_unix": issued.expires_at_unix,
}),
})
.await
{
warn!(error = %err, "failed to append audit event for token issuance");
}
Ok((
StatusCode::OK,
Json(IssueEnrollTokenResponse {
@@ -132,6 +187,26 @@ pub async fn list_enroll_tokens(
let include_expired = query.include_expired.unwrap_or(false);
match state.store.list_enroll_tokens(include_expired).await {
Ok(tokens) => {
if let Err(err) = state
.store
.append_audit_event(AuditEventInput {
actor_type: "admin_api".into(),
actor_id: None,
agent_id: None,
request_id: None,
event_type: "enroll_token_list".into(),
outcome: "ok".into(),
latency_ms: None,
message: "listed enroll tokens".into(),
metadata: serde_json::json!({
"include_expired": include_expired,
"count": tokens.len(),
}),
})
.await
{
warn!(error = %err, "failed to append audit event for token listing");
}
let body = tokens
.into_iter()
.map(|t| EnrollTokenStatus {
@@ -158,10 +233,33 @@ pub async fn revoke_enroll_token(
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 }),
)),
Ok(revoked) => {
if let Err(err) = state
.store
.append_audit_event(AuditEventInput {
actor_type: "admin_api".into(),
actor_id: None,
agent_id: None,
request_id: None,
event_type: "enroll_token_revoke".into(),
outcome: if revoked { "ok".into() } else { "not_found".into() },
latency_ms: None,
message: if revoked {
"revoked enroll token".into()
} else {
"enroll token not found".into()
},
metadata: serde_json::json!({ "token": token }),
})
.await
{
warn!(error = %err, "failed to append audit event for token revoke");
}
Ok((
StatusCode::OK,
Json(RevokeEnrollTokenResponse { token, revoked }),
))
}
Err(err) => {
warn!(error = %err, "failed to revoke enroll token");
Err(json_error(
wakey-control-plane/src/api/mod.rs
+2
View File
@@ -3,8 +3,10 @@ use axum::http::StatusCode;
mod commands;
mod control;
mod audit;
pub use commands::{list_agents, run_command};
pub use audit::list_audit_events;
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/mod.rs
+1
View File
@@ -72,6 +72,7 @@ pub async fn serve(daemon: config::DaemonConfig) -> Result<()> {
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/agent/ws", get(ws::agent_ws))
.route("/api/v1/control/agents", get(api::list_agents))
.route(
wakey-control-plane/src/state/mod.rs
+1
View File
@@ -2,3 +2,4 @@ mod store;
mod types;
pub use store::Store;
pub use types::{AuditEventFilter, AuditEventInput};
wakey-control-plane/src/state/store.rs
+146 -1
View File
@@ -5,13 +5,17 @@ use anyhow::{Context, Result};
use tracing::{debug, info, warn};
use uuid::Uuid;
use crate::state::types::{EnrollTokenInfo, IssuedAgent, IssuedEnrollToken, StateStats};
use crate::state::types::{
AuditEvent, AuditEventFilter, AuditEventInput, EnrollTokenInfo, IssuedAgent, IssuedEnrollToken,
StateStats,
};
pub struct Store {
db_path: PathBuf,
meta: sled::Tree,
enroll_tokens: sled::Tree,
agents: sled::Tree,
audit_events: sled::Tree,
}
const SCHEMA_VERSION_KEY: &[u8] = b"schema_version";
@@ -32,12 +36,16 @@ impl Store {
.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 audit_events_tree = db
.open_tree("audit_events")
.context("failed to open audit_events tree")?;
let store = Self {
db_path,
meta: meta_tree,
enroll_tokens: enroll_tree,
agents: agents_tree,
audit_events: audit_events_tree,
};
store.ensure_schema_version()?;
@@ -62,10 +70,12 @@ impl Store {
let enroll_tokens = store.enroll_tokens.iter().count();
let agents = store.agents.iter().count();
let audit_events = store.audit_events.iter().count();
info!(
path = %store.db_path.display(),
enroll_tokens,
agents,
audit_events,
"control-plane store ready"
);
Ok(store)
@@ -218,6 +228,53 @@ impl Store {
out
}
pub async fn append_audit_event(&self, input: AuditEventInput) -> Result<AuditEvent> {
let event = AuditEvent {
event_id: format!("evt-{}", Uuid::new_v4()),
ts_unix: now_unix(),
actor_type: input.actor_type,
actor_id: input.actor_id,
agent_id: input.agent_id,
request_id: input.request_id,
event_type: input.event_type,
outcome: input.outcome,
latency_ms: input.latency_ms,
message: input.message,
metadata: input.metadata,
};
let key = format!("{:020}:{}", event.ts_unix, event.event_id);
let value = serde_json::to_vec(&event).context("failed to encode audit event")?;
self.audit_events
.insert(key.as_bytes(), value)
.context("failed persisting audit event")?;
self.flush()
.context("failed flushing state db after audit append")?;
Ok(event)
}
pub async fn list_audit_events(&self, filter: AuditEventFilter) -> Result<Vec<AuditEvent>> {
let limit = filter.limit.clamp(1, 500);
let mut out = Vec::new();
for item in self.audit_events.iter().rev() {
let (_, raw) = item.context("failed iterating audit event tree")?;
let event: AuditEvent =
serde_json::from_slice(raw.as_ref()).context("failed decoding audit event")?;
if !matches_audit_filter(&event, &filter) {
continue;
}
out.push(event);
if out.len() >= limit {
break;
}
}
Ok(out)
}
fn flush(&self) -> Result<()> {
self.enroll_tokens
.flush()
@@ -225,6 +282,9 @@ impl Store {
self.agents
.flush()
.context("failed to flush agents tree")?;
self.audit_events
.flush()
.context("failed to flush audit event tree")?;
debug!(path = %self.db_path.display(), "flushed sled state db");
Ok(())
}
@@ -313,6 +373,33 @@ fn decode_schema(raw: &[u8]) -> Result<u32> {
Ok(u32::from_le_bytes(arr))
}
fn matches_audit_filter(event: &AuditEvent, filter: &AuditEventFilter) -> bool {
if let Some(agent_id) = filter.agent_id.as_deref()
&& event.agent_id.as_deref() != Some(agent_id)
{
return false;
}
if let Some(request_id) = filter.request_id.as_deref()
&& event.request_id.as_deref() != Some(request_id)
{
return false;
}
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 {
return false;
}
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 {
return false;
}
true
}
#[cfg(test)]
mod tests {
use std::fs;
@@ -416,4 +503,62 @@ mod tests {
assert_eq!(stats.expired_enroll_token_count, 1);
cleanup_dir(&dir);
}
#[tokio::test]
async fn audit_events_append_and_filter() {
let (store, dir) = make_store().await;
store
.append_audit_event(crate::state::AuditEventInput {
actor_type: "admin_api".into(),
actor_id: None,
agent_id: Some("agent-1".into()),
request_id: Some("req-1".into()),
event_type: "command_result".into(),
outcome: "ok".into(),
latency_ms: Some(12),
message: "command completed".into(),
metadata: serde_json::json!({"command":"devs"}),
})
.await
.expect("append first event should succeed");
store
.append_audit_event(crate::state::AuditEventInput {
actor_type: "agent".into(),
actor_id: Some("agent-2".into()),
agent_id: Some("agent-2".into()),
request_id: Some("req-2".into()),
event_type: "agent_ws_auth".into(),
outcome: "rejected".into(),
latency_ms: None,
message: "auth rejected".into(),
metadata: serde_json::json!({}),
})
.await
.expect("append second event should succeed");
let all = store
.list_audit_events(crate::state::AuditEventFilter {
limit: 10,
..Default::default()
})
.await
.expect("list all should succeed");
assert_eq!(all.len(), 2);
let filtered = store
.list_audit_events(crate::state::AuditEventFilter {
agent_id: Some("agent-1".into()),
event_type: Some("command_result".into()),
outcome: Some("ok".into()),
limit: 10,
..Default::default()
})
.await
.expect("filtered list should succeed");
assert_eq!(filtered.len(), 1);
assert_eq!(filtered[0].request_id.as_deref(), Some("req-1"));
cleanup_dir(&dir);
}
}
wakey-control-plane/src/state/types.rs
+40
View File
@@ -29,3 +29,43 @@ pub struct StateStats {
pub enroll_token_count: usize,
pub expired_enroll_token_count: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AuditEvent {
pub event_id: String,
pub ts_unix: u64,
pub actor_type: String,
pub actor_id: Option<String>,
pub agent_id: Option<String>,
pub request_id: Option<String>,
pub event_type: String,
pub outcome: String,
pub latency_ms: Option<u64>,
pub message: String,
#[serde(default)]
pub metadata: serde_json::Value,
}
#[derive(Debug, Clone)]
pub struct AuditEventInput {
pub actor_type: String,
pub actor_id: Option<String>,
pub agent_id: Option<String>,
pub request_id: Option<String>,
pub event_type: String,
pub outcome: String,
pub latency_ms: Option<u64>,
pub message: String,
pub metadata: serde_json::Value,
}
#[derive(Debug, Clone, Default)]
pub struct AuditEventFilter {
pub agent_id: Option<String>,
pub request_id: Option<String>,
pub event_type: Option<String>,
pub outcome: Option<String>,
pub since_unix: Option<u64>,
pub until_unix: Option<u64>,
pub limit: usize,
}
wakey-control-plane/src/ws.rs
+52
View File
@@ -10,6 +10,7 @@ use uuid::Uuid;
use wakey_agent::protocol::{ErrorPayload, RequestId, ServerMessage};
use crate::runtime::{AgentReply, AppState};
use crate::state::AuditEventInput;
#[derive(Debug, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
@@ -103,6 +104,23 @@ async fn handle_agent_socket(state: AppState, socket: WebSocket) {
if let Some(agent_id) = authed_agent_id {
info!(agent_id = %agent_id, "agent disconnected");
state.sessions.write().await.remove(&agent_id);
if let Err(err) = state
.store
.append_audit_event(AuditEventInput {
actor_type: "agent".into(),
actor_id: Some(agent_id.clone()),
agent_id: Some(agent_id),
request_id: None,
event_type: "agent_ws_disconnect".into(),
outcome: "ok".into(),
latency_ms: None,
message: "agent websocket disconnected".into(),
metadata: serde_json::json!({}),
})
.await
{
warn!(error = %err, "failed to append audit event for ws disconnect");
}
}
writer.abort();
@@ -132,6 +150,23 @@ async fn process_agent_text(
.await
{
warn!(agent_id = %agent_id, "agent auth rejected");
if let Err(err) = state
.store
.append_audit_event(AuditEventInput {
actor_type: "agent".into(),
actor_id: Some(agent_id.clone()),
agent_id: Some(agent_id),
request_id: None,
event_type: "agent_ws_auth".into(),
outcome: "rejected".into(),
latency_ms: None,
message: "agent auth rejected".into(),
metadata: serde_json::json!({}),
})
.await
{
warn!(error = %err, "failed to append audit event for auth rejection");
}
anyhow::bail!("agent auth rejected");
}
state
@@ -141,6 +176,23 @@ async fn process_agent_text(
.insert(agent_id.clone(), tx.clone());
*authed_agent_id = Some(agent_id.clone());
info!(agent_id = %agent_id, "agent authenticated");
if let Err(err) = state
.store
.append_audit_event(AuditEventInput {
actor_type: "agent".into(),
actor_id: Some(agent_id.clone()),
agent_id: Some(agent_id),
request_id: None,
event_type: "agent_ws_auth".into(),
outcome: "ok".into(),
latency_ms: None,
message: "agent websocket authenticated".into(),
metadata: serde_json::json!({}),
})
.await
{
warn!(error = %err, "failed to append audit event for auth success");
}
}
IncomingClientMessage::Heartbeat { agent_id } => {
if authed_agent_id.as_deref() != Some(agent_id.as_str()) {