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

fix(terminal): harden relay lifecycle and session cleanup

Browse Source 
- bound relay input and detach saturated output transports
- drain PTY output after child exit
- make terminal close lifecycle race-safe
- prune expired CC sessions and audit closure reasons
- require matching agent hello/auth identities
- validate terminal IDs and PTY dimensions
This commit is contained in:
lda
2026-07-16 08:51:55 +07:00 Verified
parent 536b50220f
commit 3476892cd1
10 changed files with 408 additions and 115 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wakey-control-plane/src/api/terminals.rs
+49 -19
View File
@@ -191,21 +191,10 @@ pub async fn close_terminal(
State(state): State<AppState>,
Path(terminal_id): Path<String>,
) -> Result<StatusCode, ApiError> {
if let Some(agent_id) = close_registered_terminal(&state, &terminal_id).await {
if let Some(agent_id) =
close_registered_terminal(&state, &terminal_id, TerminalCloseReason::HttpDelete).await
{
info!(terminal_id, agent_id, "terminal session closed by operator");
append_terminal_audit(
&state,
&agent_id,
&terminal_id,
TerminalAudit {
actor_type: "admin_api",
event_type: "terminal_close",
outcome: "ok",
message: "terminal session closed by operator",
metadata: serde_json::json!({}),
},
)
.await;
} else {
if !state.terminals.was_closed(&terminal_id).await {
return Err(terminal_not_found(&terminal_id));
@@ -409,7 +398,7 @@ async fn handle_operator_terminal_socket(
}
if explicit_close {
close_registered_terminal(&state, &terminal_id).await;
close_registered_terminal(&state, &terminal_id, TerminalCloseReason::BrowserClose).await;
} else {
state.terminals.detach_operator(&terminal_id, &auth.0).await;
}
@@ -510,7 +499,31 @@ where
write.send(message).await
}
async fn close_registered_terminal(state: &AppState, terminal_id: &str) -> Option<String> {
#[derive(Clone, Copy)]
enum TerminalCloseReason {
HttpDelete,
BrowserClose,
AbsoluteTimeout,
}
impl TerminalCloseReason {
fn as_str(self) -> &'static str {
match self {
Self::HttpDelete => "http_delete",
Self::BrowserClose => "browser_close",
Self::AbsoluteTimeout => "absolute_timeout",
}
}
}
/// Removes and audits a terminal as one logical operation. Only the caller
/// that actually removes the session records the close, keeping retries and
/// simultaneous timeout/operator closes from producing duplicate events.
async fn close_registered_terminal(
state: &AppState,
terminal_id: &str,
reason: TerminalCloseReason,
) -> Option<String> {
let agent_id = state.terminals.remove(terminal_id).await?;
if let Some(session) = state.sessions.read().await.get(&agent_id) {
let terminal_id = TerminalId::new(terminal_id.to_string()).ok()?;
@@ -520,15 +533,32 @@ async fn close_registered_terminal(state: &AppState, terminal_id: &str) -> Optio
terminal_id,
}));
}
append_terminal_audit(
state,
&agent_id,
terminal_id,
TerminalAudit {
actor_type: "admin_api",
event_type: "terminal_close",
outcome: "ok",
message: "terminal session closed",
metadata: serde_json::json!({ "reason": reason.as_str() }),
},
)
.await;
Some(agent_id)
}
fn spawn_absolute_timeout(state: AppState, terminal_id: TerminalId, agent_id: String) {
tokio::spawn(async move {
tokio::time::sleep(TERMINAL_ABSOLUTE_TIMEOUT).await;
if close_registered_terminal(&state, terminal_id.as_str())
.await
.is_some()
if close_registered_terminal(
&state,
terminal_id.as_str(),
TerminalCloseReason::AbsoluteTimeout,
)
.await
.is_some()
{
info!(terminal_id = %terminal_id, agent_id, "terminal absolute timeout reached");
}
wakey-control-plane/src/runtime/terminals.rs
+103 -32
View File
@@ -77,6 +77,7 @@ impl TerminalRegistry {
pub async fn create(&self, agent_id: String) -> Result<CreatedTerminal, &'static str> {
let mut sessions = self.inner.lock().await;
prune_expired_sessions(&mut sessions);
if sessions
.values()
.filter(|session| session.agent_id == agent_id)
@@ -168,10 +169,12 @@ impl TerminalRegistry {
) -> Vec<(TerminalId, String)> {
let reported_ids = reported
.iter()
.filter(|session| expires_at_for_created_at(session.created_at_unix).is_some())
.map(|session| session.terminal_id.as_str())
.collect::<std::collections::HashSet<_>>();
let stale_ids = {
let sessions = self.inner.lock().await;
let mut sessions = self.inner.lock().await;
prune_expired_sessions(&mut sessions);
sessions
.iter()
.filter(|(terminal_id, session)| {
@@ -189,13 +192,17 @@ impl TerminalRegistry {
let mut credentials = Vec::new();
let mut sessions = self.inner.lock().await;
for reported_session in reported {
let Some(expires_at) = expires_at_for_created_at(reported_session.created_at_unix)
else {
continue;
};
let terminal_id = reported_session.terminal_id.as_str().to_string();
let session = sessions
.entry(terminal_id)
.or_insert_with(|| TerminalSession {
agent_id: agent_id.to_string(),
created_at_unix: reported_session.created_at_unix,
expires_at: Instant::now() + TERMINAL_ABSOLUTE_TIMEOUT,
expires_at,
agent_confirmed: true,
relay_token: None,
attachment_token: None,
@@ -319,31 +326,30 @@ impl TerminalRegistry {
terminal_id: &str,
frame: TerminalRelayFrame,
) -> Result<(), &'static str> {
let tx = self
.inner
.lock()
.await
.get(terminal_id)
.and_then(|session| session.agent_tx.clone());
if let Some(tx) = tx {
return tx
.send(frame)
.await
.map_err(|_| "terminal_agent_disconnected");
}
let mut sessions = self.inner.lock().await;
let session = active_session(&mut sessions, terminal_id)?;
session.pending_agent_bytes += relay_frame_size(&frame);
session.pending_agent.push_back(frame);
while session.pending_agent_bytes > TERMINAL_PENDING_AGENT_BYTES {
if let Some(dropped) = session.pending_agent.pop_front() {
session.pending_agent_bytes -= relay_frame_size(&dropped);
// Select the attached transport or enqueue the frame while holding one
// lock. Otherwise an agent can attach between those decisions and
// leave input stranded in the pre-attachment queue.
let tx = {
let mut sessions = self.inner.lock().await;
let session = active_session(&mut sessions, terminal_id)?;
if let Some(tx) = session.agent_tx.clone() {
tx
} else {
break;
session.pending_agent_bytes += relay_frame_size(&frame);
session.pending_agent.push_back(frame);
while session.pending_agent_bytes > TERMINAL_PENDING_AGENT_BYTES {
if let Some(dropped) = session.pending_agent.pop_front() {
session.pending_agent_bytes -= relay_frame_size(&dropped);
} else {
break;
}
}
return Ok(());
}
}
Ok(())
};
tx.send(frame)
.await
.map_err(|_| "terminal_agent_disconnected")
}
pub async fn relay_from_agent(
@@ -417,7 +423,9 @@ impl TerminalRegistry {
}
pub async fn summary(&self, terminal_id: &str) -> Option<(String, u64, bool, bool)> {
self.inner.lock().await.get(terminal_id).map(|session| {
let mut sessions = self.inner.lock().await;
prune_expired_sessions(&mut sessions);
sessions.get(terminal_id).map(|session| {
(
session.agent_id.clone(),
session.created_at_unix,
@@ -428,7 +436,8 @@ impl TerminalRegistry {
}
pub async fn summaries(&self) -> Vec<TerminalSummary> {
let sessions = self.inner.lock().await;
let mut sessions = self.inner.lock().await;
prune_expired_sessions(&mut sessions);
let mut summaries = sessions
.iter()
.map(|(terminal_id, session)| TerminalSummary {
@@ -448,11 +457,14 @@ fn active_session<'a>(
sessions: &'a mut HashMap<String, TerminalSession>,
terminal_id: &str,
) -> Result<&'a mut TerminalSession, &'static str> {
let session = sessions.get_mut(terminal_id).ok_or("terminal_not_found")?;
if session.expires_at <= Instant::now() {
if sessions
.get(terminal_id)
.is_some_and(|session| session.expires_at <= Instant::now())
{
sessions.remove(terminal_id);
return Err("terminal_expired");
}
Ok(session)
sessions.get_mut(terminal_id).ok_or("terminal_not_found")
}
fn validate_operator_id(operator_id: &str) -> Result<(), &'static str> {
@@ -485,6 +497,26 @@ fn prune_tombstones(closed: &mut HashMap<String, Instant>) {
closed.retain(|_, closed_at| closed_at.elapsed() < TERMINAL_TOMBSTONE_TTL);
}
fn prune_expired_sessions(sessions: &mut HashMap<String, TerminalSession>) {
let now = Instant::now();
sessions.retain(|_, session| session.expires_at > now);
}
/// Converts the agent's durable creation timestamp into CC's monotonic
/// deadline. Reconciliation must not grant an old PTY a fresh twelve hours.
fn expires_at_for_created_at(created_at_unix: u64) -> Option<Instant> {
let now_unix = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
let age = Duration::from_secs(now_unix.saturating_sub(created_at_unix));
if age >= TERMINAL_ABSOLUTE_TIMEOUT {
return None;
}
let remaining = TERMINAL_ABSOLUTE_TIMEOUT - age;
Some(Instant::now() + remaining)
}
fn relay_frame_size(frame: &TerminalRelayFrame) -> usize {
match frame {
TerminalRelayFrame::Binary(bytes) => bytes.len(),
@@ -627,9 +659,10 @@ mod tests {
async fn agent_inventory_adopts_and_reconnects_live_session() {
let registry = TerminalRegistry::new();
let terminal_id = TerminalId::new("survived-cc").expect("terminal id");
let created_at_unix = u64::MAX;
let reported = AgentTerminalSession {
terminal_id: terminal_id.clone(),
created_at_unix: 42,
created_at_unix,
};
let credentials = registry
@@ -647,10 +680,48 @@ mod tests {
.await
.expect("adopted summary");
assert_eq!(summary.0, "router");
assert_eq!(summary.1, 42);
assert_eq!(summary.1, created_at_unix);
assert!(summary.2);
}
#[tokio::test]
async fn expired_sessions_release_agent_quota() {
let registry = TerminalRegistry::new();
let first = registry.create("router".into()).await.expect("first");
registry.create("router".into()).await.expect("second");
registry
.inner
.lock()
.await
.get_mut(&first.terminal_id)
.expect("first session")
.expires_at = Instant::now();
registry
.create("router".into())
.await
.expect("expired session no longer consumes quota");
assert!(registry.summary(&first.terminal_id).await.is_none());
}
#[tokio::test]
async fn reconciliation_does_not_revive_expired_agent_session() {
let registry = TerminalRegistry::new();
let terminal_id = TerminalId::new("expired-agent-session").expect("terminal id");
let reported = AgentTerminalSession {
terminal_id: terminal_id.clone(),
created_at_unix: 0,
};
assert!(
registry
.reconcile_agent_sessions("router", &[reported])
.await
.is_empty()
);
assert!(registry.summary(terminal_id.as_str()).await.is_none());
}
#[tokio::test]
async fn operator_detach_does_not_remove_session() {
let registry = TerminalRegistry::new();
wakey-control-plane/src/ws.rs
+32 -1
View File
@@ -56,6 +56,7 @@ enum IncomingClientMessage {
#[derive(Default)]
struct AgentConnectionState {
authed_agent_id: Option<String>,
hello_agent_id: Option<String>,
hello_at: Option<Instant>,
capabilities: Vec<AgentCapability>,
}
@@ -193,18 +194,27 @@ async fn process_agent_text(
agent_id,
capabilities,
} => {
if connection
.hello_agent_id
.as_deref()
.is_some_and(|hello_agent_id| hello_agent_id != agent_id)
{
anyhow::bail!("hello changed agent identity");
}
let now = Instant::now();
if connection.hello_at.is_none() {
connection.hello_at = Some(now);
}
let connect_to_hello_ms = connected_at.elapsed().as_millis() as u64;
info!(agent_id = %agent_id, connect_to_hello_ms, "agent hello received");
connection.hello_agent_id = Some(agent_id);
connection.capabilities = capabilities;
}
IncomingClientMessage::Auth {
agent_id,
agent_token,
} => {
validate_auth_identity(connection, &agent_id)?;
let connect_to_auth_ms = connected_at.elapsed().as_millis() as u64;
let hello_to_auth_ms = connection
.hello_at
@@ -372,6 +382,17 @@ fn now_duration_ms(duration: std::time::Duration) -> u64 {
duration.as_millis() as u64
}
fn validate_auth_identity(connection: &AgentConnectionState, agent_id: &str) -> Result<()> {
let hello_agent_id = connection
.hello_agent_id
.as_deref()
.ok_or_else(|| anyhow::anyhow!("auth before hello"))?;
if hello_agent_id != agent_id {
anyhow::bail!("auth agent does not match hello agent");
}
Ok(())
}
async fn ensure_current_session(
state: &AppState,
agent_id: &str,
@@ -406,7 +427,7 @@ mod tests {
use crate::runtime::AgentSession;
use super::is_current_session;
use super::{AgentConnectionState, is_current_session, validate_auth_identity};
#[test]
fn current_session_check_rejects_stale_connection_ids() {
@@ -425,4 +446,14 @@ mod tests {
assert!(!is_current_session(&sessions, "agent-a", "conn-old"));
assert!(!is_current_session(&sessions, "agent-b", "conn-new"));
}
#[test]
fn auth_requires_a_matching_hello_identity() {
let mut connection = AgentConnectionState::default();
assert!(validate_auth_identity(&connection, "agent-a").is_err());
connection.hello_agent_id = Some("agent-a".into());
assert!(validate_auth_identity(&connection, "agent-a").is_ok());
assert!(validate_auth_identity(&connection, "agent-b").is_err());
}
}