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

half done?

Browse Source
This commit is contained in:
lda
2026-04-12 19:05:55 +07:00 Unverified
parent a57535c6a2
commit 8abad25380
5 changed files with 186 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
REVIEW_CONTROL_PLANE_2026-04-12.md
+14
View File
@@ -0,0 +1,14 @@
# Control Plane Review
## Findings
1. High: the admin/control API is effectively unauthenticated, so anyone who can reach the server can issue enroll tokens, list/revoke tokens, inspect audit/alerts, and send live commands to agents. In [`wakey-control-plane/src/runtime/mod.rs:61`](c:/Users/Admin/Documents/realshit/wakey/wakey-control-plane/src/runtime/mod.rs#L61) through [`wakey-control-plane/src/runtime/mod.rs:109`](c:/Users/Admin/Documents/realshit/wakey/wakey-control-plane/src/runtime/mod.rs#L109), `control_api_routes()` is merged straight into the app with no auth middleware despite the comment saying these routes are “admin-only.” The UI also calls those endpoints directly with plain `fetch` and no auth material in [`ui/src/api.ts:60`](c:/Users/Admin/Documents/realshit/wakey/ui/src/api.ts#L60) through [`ui/src/api.ts:123`](c:/Users/Admin/Documents/realshit/wakey/ui/src/api.ts#L123). This is a full remote-takeover issue for the control plane, not just a missing polish item.
2. High: a second websocket connection for the same `agent_id` silently replaces the current session, but the old authenticated socket is left alive and can still submit `result`/`error` frames against pending requests. In [`wakey-control-plane/src/ws.rs:195`](c:/Users/Admin/Documents/realshit/wakey/wakey-control-plane/src/ws.rs#L195) through [`wakey-control-plane/src/ws.rs:200`](c:/Users/Admin/Documents/realshit/wakey/wakey-control-plane/src/ws.rs#L200), a successful auth simply overwrites `sessions[agent_id] = tx.clone()`. The previous socket is not closed or demoted. Later, any authenticated socket can satisfy pending requests in [`wakey-control-plane/src/ws.rs:229`](c:/Users/Admin/Documents/realshit/wakey/wakey-control-plane/src/ws.rs#L229) through [`wakey-control-plane/src/ws.rs:247`](c:/Users/Admin/Documents/realshit/wakey/wakey-control-plane/src/ws.rs#L247), while requests are correlated only by `request_id` created in [`wakey-control-plane/src/api/commands.rs:91`](c:/Users/Admin/Documents/realshit/wakey/wakey-control-plane/src/api/commands.rs#L91) through [`wakey-control-plane/src/api/commands.rs:97`](c:/Users/Admin/Documents/realshit/wakey/wakey-control-plane/src/api/commands.rs#L97). That creates a split-brain/race condition where a stale or malicious prior session for the same agent can inject or win replies.
3. Medium: configured seed enroll tokens are reinserted into the database on every startup, so “one-time” tokens become reusable after a restart if they remain in config. In [`wakey-control-plane/src/state/store.rs:52`](c:/Users/Admin/Documents/realshit/wakey/wakey-control-plane/src/state/store.rs#L52) through [`wakey-control-plane/src/state/store.rs:68`](c:/Users/Admin/Documents/realshit/wakey/wakey-control-plane/src/state/store.rs#L68), `load_or_init()` blindly seeds `daemon.enroll_tokens` into sled every time the process starts. Enrollment consumes tokens in [`wakey-control-plane/src/state/store.rs:91`](c:/Users/Admin/Documents/realshit/wakey/wakey-control-plane/src/state/store.rs#L91) through [`wakey-control-plane/src/state/store.rs:109`](c:/Users/Admin/Documents/realshit/wakey/wakey-control-plane/src/state/store.rs#L109), but that consumption is undone on the next boot if the token still exists in config. That breaks the “short-lived/one-time” assumption and makes operational mistakes much more likely.
## Residual Risks
- The control-plane/UI surface is growing quickly and currently assumes a trusted environment in multiple places. Even after adding admin auth, I would expect more authz/session-boundary issues to surface.
- The command relay path is conceptually good, but it needs stronger session ownership rules before it is trustworthy under reconnect races or duplicated agents.
wakey-control-plane/src/api/commands.rs
+1 -1
View File
@@ -77,7 +77,7 @@ pub async fn run_command(
let tx = {
let sessions = state.sessions.read().await;
sessions.get(&agent_id).cloned()
sessions.get(&agent_id).map(|session| session.tx.clone())
}
.ok_or_else(|| {
warn!("command rejected: agent not connected");
wakey-control-plane/src/runtime/mod.rs
+12 -3
View File
@@ -9,10 +9,13 @@ use axum::routing::{get, post};
use axum::routing::get_service;
use tokio::net::TcpListener;
use tokio::sync::{Mutex, RwLock, mpsc, oneshot};
#[cfg(unix)]
use tokio::time::MissedTickBehavior;
use tower_http::services::ServeDir;
use tower_http::services::ServeFile;
use tracing::{info, warn};
use tracing::info;
#[cfg(unix)]
use tracing::warn;
use wakey_agent::protocol::{ErrorPayload, ServerMessage};
use crate::api;
@@ -32,13 +35,19 @@ use process::{remove_pid_file, write_pid_file};
#[derive(Clone)]
pub struct AppState {
pub store: Arc<state::Store>,
pub sessions: Arc<RwLock<HashMap<String, mpsc::UnboundedSender<ServerMessage>>>>,
pub sessions: Arc<RwLock<HashMap<String, AgentSession>>>,
pub pending: Arc<Mutex<HashMap<String, oneshot::Sender<AgentReply>>>>,
pub public_url: String,
pub command_timeout: Duration,
pub enroll_token_ttl: Duration,
}
#[derive(Clone)]
pub struct AgentSession {
pub connection_id: String,
pub tx: mpsc::UnboundedSender<ServerMessage>,
}
pub enum AgentReply {
Result(serde_json::Value),
Error(ErrorPayload),
@@ -110,7 +119,7 @@ pub async fn serve(daemon: config::DaemonConfig) -> Result<()> {
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 {
let server = tokio::spawn(async move {
axum::serve(listener, app)
.await
.context("control-plane server exited unexpectedly")
wakey-control-plane/src/state/store.rs
+82 -11
View File
@@ -21,6 +21,7 @@ pub struct Store {
}
const SCHEMA_VERSION_KEY: &[u8] = b"schema_version";
const SEEDED_ENROLL_TOKEN_PREFIX: &[u8] = b"seeded_enroll_token:";
const SCHEMA_VERSION: u32 = 1;
impl Store {
@@ -60,17 +61,7 @@ impl Store {
store.ensure_schema_version()?;
for token in enroll_tokens {
let token = token.trim();
if token.is_empty() {
continue;
}
let expires_at = now_unix().saturating_add(seed_ttl.as_secs().max(1));
store
.enroll_tokens
.insert(token.as_bytes(), &expires_at.to_le_bytes())
.with_context(|| format!("failed to seed enroll token into {}", store.db_path.display()))?;
}
store.seed_bootstrap_enroll_tokens(&enroll_tokens, seed_ttl)?;
store.gc_expired_enroll_tokens_inner()?;
@@ -206,10 +197,12 @@ impl Store {
})
}
#[cfg_attr(not(unix), allow(dead_code))]
pub async fn gc_expired_enroll_tokens(&self) -> Result<u64> {
self.gc_expired_enroll_tokens_inner()
}
#[cfg_attr(not(unix), allow(dead_code))]
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");
@@ -391,6 +384,9 @@ impl Store {
}
fn flush(&self) -> Result<()> {
self.meta
.flush()
.context("failed to flush meta tree")?;
self.enroll_tokens
.flush()
.context("failed to flush enroll token tree")?;
@@ -410,6 +406,33 @@ impl Store {
Ok(())
}
fn seed_bootstrap_enroll_tokens(&self, enroll_tokens: &[String], seed_ttl: Duration) -> Result<()> {
for token in enroll_tokens {
let token = token.trim();
if token.is_empty() {
continue;
}
let marker_key = seeded_enroll_token_key(token);
if self
.meta
.contains_key(&marker_key)
.with_context(|| format!("failed reading bootstrap marker in {}", self.db_path.display()))?
{
continue;
}
let expires_at = now_unix().saturating_add(seed_ttl.as_secs().max(1));
self.enroll_tokens
.insert(token.as_bytes(), &expires_at.to_le_bytes())
.with_context(|| format!("failed to seed enroll token into {}", self.db_path.display()))?;
self.meta
.insert(marker_key, &expires_at.to_le_bytes())
.with_context(|| format!("failed to persist bootstrap marker into {}", self.db_path.display()))?;
}
Ok(())
}
fn gc_expired_enroll_tokens_inner(&self) -> Result<u64> {
let now = now_unix();
let mut removed = 0u64;
@@ -494,6 +517,13 @@ fn decode_schema(raw: &[u8]) -> Result<u32> {
Ok(u32::from_le_bytes(arr))
}
fn seeded_enroll_token_key(token: &str) -> Vec<u8> {
let mut key = Vec::with_capacity(SEEDED_ENROLL_TOKEN_PREFIX.len() + token.len());
key.extend_from_slice(SEEDED_ENROLL_TOKEN_PREFIX);
key.extend_from_slice(token.as_bytes());
key
}
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)
@@ -720,4 +750,45 @@ mod tests {
assert!(history.len() >= 2);
cleanup_dir(&dir);
}
#[tokio::test]
async fn bootstrap_seed_tokens_are_not_reseeded_after_consumption() {
let dir =
std::env::temp_dir().join(format!("wakey-cp-store-test-{}", uuid::Uuid::new_v4()));
let db_path = dir.join("state.db");
let first = Store::load_or_init(
&db_path,
vec!["enr-bootstrap-once".to_string()],
Duration::from_secs(60),
)
.await
.expect("initial store should initialize");
let issued = first
.enroll("enr-bootstrap-once")
.await
.expect("bootstrap token should enroll once");
assert!(!issued.agent_id.is_empty());
drop(first);
let second = Store::load_or_init(
&db_path,
vec!["enr-bootstrap-once".to_string()],
Duration::from_secs(60),
)
.await
.expect("reloaded store should initialize");
let err = second
.enroll("enr-bootstrap-once")
.await
.expect_err("bootstrap token should not resurrect after restart");
assert!(err
.to_string()
.contains("invalid or already-used enroll token"));
cleanup_dir(&dir);
}
}
wakey-control-plane/src/ws.rs
+77 -9
View File
@@ -10,7 +10,7 @@ use tracing::{debug, info, info_span, warn};
use uuid::Uuid;
use wakey_agent::protocol::{ErrorPayload, RequestId, ServerMessage};
use crate::runtime::{AgentReply, AppState};
use crate::runtime::{AgentReply, AgentSession, AppState};
use crate::state::AuditEventInput;
#[derive(Debug, Deserialize)]
@@ -89,6 +89,7 @@ async fn handle_agent_socket(state: AppState, socket: WebSocket) {
if let Err(err) = process_agent_text(
&state,
&tx,
&connection_id,
&mut authed_agent_id,
&mut hello_at,
connected_at,
@@ -114,7 +115,14 @@ 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);
let mut sessions = state.sessions.write().await;
let should_remove = sessions
.get(&agent_id)
.map(|session| session.connection_id == connection_id)
.unwrap_or(false);
if should_remove {
sessions.remove(&agent_id);
}
if let Err(err) = state
.store
.append_audit_event(AuditEventInput {
@@ -141,6 +149,7 @@ async fn handle_agent_socket(state: AppState, socket: WebSocket) {
async fn process_agent_text(
state: &AppState,
tx: &mpsc::UnboundedSender<ServerMessage>,
connection_id: &str,
authed_agent_id: &mut Option<String>,
hello_at: &mut Option<Instant>,
connected_at: Instant,
@@ -196,7 +205,13 @@ async fn process_agent_text(
.sessions
.write()
.await
.insert(agent_id.clone(), tx.clone());
.insert(
agent_id.clone(),
AgentSession {
connection_id: connection_id.to_string(),
tx: tx.clone(),
},
);
*authed_agent_id = Some(agent_id.clone());
info!(agent_id = %agent_id, connect_to_auth_ms, hello_to_auth_ms = hello_to_auth_ms.unwrap_or(0), "agent authenticated");
if let Err(err) = state
@@ -224,12 +239,14 @@ async fn process_agent_text(
if authed_agent_id.as_deref() != Some(agent_id.as_str()) {
anyhow::bail!("heartbeat for unauthenticated or mismatched agent");
}
ensure_current_session(state, &agent_id, connection_id).await?;
debug!(agent_id = %agent_id, "heartbeat received");
}
IncomingClientMessage::Result { request_id, result } => {
if authed_agent_id.is_none() {
anyhow::bail!("result before auth");
}
let agent_id = authed_agent_id
.as_deref()
.ok_or_else(|| anyhow::anyhow!("result before auth"))?;
ensure_current_session(state, agent_id, connection_id).await?;
let key = request_id.as_str().to_string();
if let Some(waiter) = state.pending.lock().await.remove(&key) {
let _ = waiter.send(AgentReply::Result(result));
@@ -238,9 +255,10 @@ async fn process_agent_text(
}
}
IncomingClientMessage::Error { request_id, error } => {
if authed_agent_id.is_none() {
anyhow::bail!("error before auth");
}
let agent_id = authed_agent_id
.as_deref()
.ok_or_else(|| anyhow::anyhow!("error before auth"))?;
ensure_current_session(state, agent_id, connection_id).await?;
let key = request_id.as_str().to_string();
if let Some(waiter) = state.pending.lock().await.remove(&key) {
let _ = waiter.send(AgentReply::Error(error));
@@ -256,3 +274,53 @@ async fn process_agent_text(
fn now_duration_ms(duration: std::time::Duration) -> u64 {
duration.as_millis() as u64
}
async fn ensure_current_session(state: &AppState, agent_id: &str, connection_id: &str) -> Result<()> {
let sessions = state.sessions.read().await;
if is_current_session(&sessions, agent_id, connection_id) {
Ok(())
} else if sessions.contains_key(agent_id) {
anyhow::bail!("stale agent session")
} else {
anyhow::bail!("agent session not registered")
}
}
fn is_current_session(
sessions: &std::collections::HashMap<String, AgentSession>,
agent_id: &str,
connection_id: &str,
) -> bool {
sessions
.get(agent_id)
.map(|session| session.connection_id == connection_id)
.unwrap_or(false)
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use tokio::sync::mpsc;
use crate::runtime::AgentSession;
use super::is_current_session;
#[test]
fn current_session_check_rejects_stale_connection_ids() {
let (tx, _rx) = mpsc::unbounded_channel();
let mut sessions = HashMap::new();
sessions.insert(
"agent-a".to_string(),
AgentSession {
connection_id: "conn-new".to_string(),
tx,
},
);
assert!(is_current_session(&sessions, "agent-a", "conn-new"));
assert!(!is_current_session(&sessions, "agent-a", "conn-old"));
assert!(!is_current_session(&sessions, "agent-b", "conn-new"));
}
}