half done?

REVIEW_CONTROL_PLANE_2026-04-12.md

@@ -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

@@ -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

@@ -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

@@ -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 {
+        store.seed_bootstrap_enroll_tokens(&enroll_tokens, seed_ttl)?;
-            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.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

@@ -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() {
+            let agent_id = authed_agent_id
-                anyhow::bail!("result before auth");
+                .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() {
+            let agent_id = authed_agent_id
-                anyhow::bail!("error before auth");
+                .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"));
+    }
+}