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

add remote fleet terminal sessions

Browse Source
This commit is contained in:
lda
2026-07-15 07:33:37 +07:00 Verified
parent 4ecef8da38
commit 4d793136ba
24 changed files with 2386 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wakey-control-plane/src/api/commands.rs
+6
View File
@@ -6,6 +6,7 @@ use serde::{Deserialize, Serialize};
use std::time::Instant;
use tracing::{info, info_span, warn};
use uuid::Uuid;
use wakey_agent::protocol::AgentCapability;
use wakey_agent::protocol::{AgentCommand, ErrorPayload, RequestId, ServerMessage};
use crate::api::ApiError;
@@ -17,6 +18,7 @@ pub struct AgentStatus {
pub agent_id: String,
pub connected: bool,
pub nickname: Option<String>,
pub capabilities: Vec<AgentCapability>,
}
#[derive(Debug, Deserialize)]
@@ -43,6 +45,10 @@ pub async fn list_agents(State(state): State<AppState>) -> Result<impl IntoRespo
.into_iter()
.map(|(agent_id, nickname)| AgentStatus {
connected: sessions.contains_key(&agent_id),
capabilities: sessions
.get(&agent_id)
.map(|session| session.capabilities.clone())
.unwrap_or_default(),
agent_id,
nickname,
})
wakey-control-plane/src/api/mod.rs
+5
View File
@@ -5,6 +5,7 @@ mod alerts;
mod audit;
mod commands;
mod control;
mod terminals;
pub use alerts::{active_alerts, alert_history, alerts_stream};
pub use audit::list_audit_events;
@@ -16,6 +17,10 @@ pub use control::{
list_fleet_devices, list_known_devices, merge_known_device, refresh_fleet_devices,
revoke_agent, revoke_enroll_token, set_agent_nickname, state_stats, wake_fleet_device,
};
pub use terminals::{
agent_terminal_ws, attach_terminal, close_terminal, create_terminal, get_terminal,
operator_terminal_ws,
};
use axum::response::{IntoResponse, Response};
use serde::{Deserialize, Serialize};
wakey-control-plane/src/api/terminals.rs
+617
View File
@@ -0,0 +1,617 @@
use axum::Json;
use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use axum::extract::{Path, State};
use axum::http::StatusCode;
use axum::response::Response;
use futures_util::{SinkExt, StreamExt};
use serde::{Deserialize, Serialize};
use tracing::{info, warn};
use wakey_agent::protocol::{
AgentCapability, ServerMessage, TerminalAgentHandshake, TerminalControl, TerminalId,
TerminalOperatorHandshake,
};
use crate::api::ApiError;
use crate::runtime::terminals::{
TERMINAL_ABSOLUTE_TIMEOUT, TERMINAL_ATTACH_TIMEOUT, TERMINAL_DISCONNECT_GRACE,
TERMINAL_MAX_FRAME_BYTES, TerminalRelayFrame,
};
use crate::runtime::{AppState, SessionEvent};
use crate::state::AuditEventInput;
#[derive(Debug, Deserialize)]
pub struct CreateTerminalRequest {
pub agent_id: String,
#[serde(default = "default_rows")]
pub rows: u16,
#[serde(default = "default_cols")]
pub cols: u16,
}
#[derive(Debug, Serialize)]
pub struct TerminalSessionResponse {
pub terminal_id: String,
pub agent_id: String,
pub created_at_unix: u64,
pub agent_attached: bool,
pub operator_attached: bool,
pub websocket_url: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub attachment_token: Option<String>,
}
pub async fn create_terminal(
State(state): State<AppState>,
Json(request): Json<CreateTerminalRequest>,
) -> Result<(StatusCode, Json<TerminalSessionResponse>), ApiError> {
validate_size(request.rows, request.cols)?;
let agent_tx = {
let sessions = state.sessions.read().await;
let session = sessions.get(&request.agent_id).ok_or_else(|| {
ApiError::new(
StatusCode::NOT_FOUND,
"agent_not_connected",
"agent is not connected",
)
})?;
if !session.capabilities.contains(&AgentCapability::Terminal) {
return Err(ApiError::new(
StatusCode::CONFLICT,
"terminal_not_supported",
"agent has not advertised terminal capability",
));
}
session.tx.clone()
};
let created = state
.terminals
.create(request.agent_id.clone())
.await
.map_err(registry_error)?;
let terminal_id = TerminalId::new(created.terminal_id.clone()).map_err(|message| {
ApiError::new(
StatusCode::INTERNAL_SERVER_ERROR,
"terminal_id_invalid",
message,
)
})?;
if agent_tx
.send(SessionEvent::Message(ServerMessage::OpenTerminal {
terminal_id: terminal_id.clone(),
relay_token: created.relay_token,
rows: request.rows,
cols: request.cols,
}))
.is_err()
{
state.terminals.remove(terminal_id.as_str()).await;
return Err(ApiError::new(
StatusCode::BAD_GATEWAY,
"agent_send_failed",
"failed to send terminal request to agent",
));
}
spawn_absolute_timeout(state.clone(), terminal_id.clone(), request.agent_id.clone());
append_terminal_audit(
&state,
&request.agent_id,
terminal_id.as_str(),
TerminalAudit {
actor_type: "admin_api",
event_type: "terminal_request",
outcome: "sent",
message: "terminal session requested",
metadata: serde_json::json!({ "rows": request.rows, "cols": request.cols }),
},
)
.await;
info!(terminal_id = %terminal_id, agent_id = %request.agent_id, "terminal session requested");
Ok((
StatusCode::CREATED,
Json(TerminalSessionResponse {
websocket_url: operator_ws_path(terminal_id.as_str()),
terminal_id: terminal_id.to_string(),
agent_id: request.agent_id,
created_at_unix: created.created_at_unix,
agent_attached: false,
operator_attached: false,
attachment_token: Some(created.attachment_token),
}),
))
}
pub async fn get_terminal(
State(state): State<AppState>,
Path(terminal_id): Path<String>,
) -> Result<Json<TerminalSessionResponse>, ApiError> {
let (agent_id, created_at_unix, agent_attached, operator_attached) = state
.terminals
.summary(&terminal_id)
.await
.ok_or_else(|| terminal_not_found(&terminal_id))?;
Ok(Json(TerminalSessionResponse {
websocket_url: operator_ws_path(&terminal_id),
terminal_id,
agent_id,
created_at_unix,
agent_attached,
operator_attached,
attachment_token: None,
}))
}
pub async fn attach_terminal(
State(state): State<AppState>,
Path(terminal_id): Path<String>,
) -> Result<Json<TerminalSessionResponse>, ApiError> {
let attachment_token = state
.terminals
.issue_attachment_token(&terminal_id)
.await
.map_err(registry_error)?;
let (agent_id, created_at_unix, agent_attached, operator_attached) = state
.terminals
.summary(&terminal_id)
.await
.ok_or_else(|| terminal_not_found(&terminal_id))?;
Ok(Json(TerminalSessionResponse {
websocket_url: operator_ws_path(&terminal_id),
terminal_id,
agent_id,
created_at_unix,
agent_attached,
operator_attached,
attachment_token: Some(attachment_token),
}))
}
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 {
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));
}
info!(terminal_id, "terminal session was already closed");
}
Ok(StatusCode::NO_CONTENT)
}
pub async fn agent_terminal_ws(
ws: WebSocketUpgrade,
State(state): State<AppState>,
Path(terminal_id): Path<String>,
) -> Response {
ws.max_message_size(TERMINAL_MAX_FRAME_BYTES)
.on_upgrade(move |socket| handle_agent_terminal_socket(state, terminal_id, socket))
}
pub async fn operator_terminal_ws(
ws: WebSocketUpgrade,
State(state): State<AppState>,
Path(terminal_id): Path<String>,
) -> Response {
ws.max_message_size(TERMINAL_MAX_FRAME_BYTES)
.on_upgrade(move |socket| handle_operator_terminal_socket(state, terminal_id, socket))
}
async fn handle_agent_terminal_socket(state: AppState, terminal_id: String, mut socket: WebSocket) {
let auth = match receive_text_handshake(&mut socket).await.and_then(|text| {
serde_json::from_str::<TerminalAgentHandshake>(&text).map_err(|_| "invalid handshake")
}) {
Ok(TerminalAgentHandshake::Auth {
agent_id,
relay_token,
}) => (agent_id, relay_token),
Err(code) => {
close_socket(&mut socket, code).await;
return;
}
};
let (mut outbound, pending) = match state
.terminals
.attach_agent(&terminal_id, &auth.0, &auth.1)
.await
{
Ok(outbound) => outbound,
Err(code) => {
close_socket(&mut socket, code).await;
return;
}
};
info!(terminal_id, agent_id = %auth.0, "agent terminal socket attached");
append_terminal_audit(
&state,
&auth.0,
&terminal_id,
TerminalAudit {
actor_type: "agent",
event_type: "terminal_agent_attach",
outcome: "ok",
message: "agent terminal transport attached",
metadata: serde_json::json!({}),
},
)
.await;
let (mut write, mut read) = socket.split();
for frame in pending {
let closes = matches!(frame, TerminalRelayFrame::Close);
if send_relay_frame(&mut write, frame).await.is_err() || closes {
state.terminals.remove(&terminal_id).await;
return;
}
}
loop {
tokio::select! {
outbound_frame = outbound.recv() => {
let Some(frame) = outbound_frame else { break; };
if send_relay_frame(&mut write, frame).await.is_err() { break; }
}
incoming = read.next() => {
let Some(Ok(message)) = incoming else { break; };
match agent_relay_frame(message) {
Ok(Some(frame)) => {
let closes = matches!(frame, TerminalRelayFrame::Close);
audit_agent_control_frame(&state, &auth.0, &terminal_id, &frame).await;
if state.terminals.relay_from_agent(&terminal_id, frame).await.is_err() { break; }
if closes { break; }
}
Ok(None) => {}
Err(code) => {
warn!(terminal_id, code, "invalid agent terminal frame");
break;
}
}
}
}
}
state.terminals.remove(&terminal_id).await;
info!(
terminal_id,
"agent terminal socket detached; session closed"
);
}
async fn handle_operator_terminal_socket(
state: AppState,
terminal_id: String,
mut socket: WebSocket,
) {
let attachment_token = match receive_text_handshake(&mut socket).await.and_then(|text| {
serde_json::from_str::<TerminalOperatorHandshake>(&text).map_err(|_| "invalid handshake")
}) {
Ok(TerminalOperatorHandshake::Attach { attachment_token }) => attachment_token,
Err(code) => {
close_socket(&mut socket, code).await;
return;
}
};
let (mut outbound, replay) = match state
.terminals
.attach_operator(&terminal_id, &attachment_token)
.await
{
Ok(attached) => attached,
Err(code) => {
close_socket(&mut socket, code).await;
return;
}
};
info!(terminal_id, "operator terminal socket attached");
if let Some((agent_id, _, _, _)) = state.terminals.summary(&terminal_id).await {
append_terminal_audit(
&state,
&agent_id,
&terminal_id,
TerminalAudit {
actor_type: "admin_api",
event_type: "terminal_operator_attach",
outcome: "ok",
message: "operator terminal transport attached",
metadata: serde_json::json!({}),
},
)
.await;
}
let (mut write, mut read) = socket.split();
for frame in replay {
if send_relay_frame(&mut write, frame).await.is_err() {
return;
}
}
let mut explicit_close = false;
loop {
tokio::select! {
outbound_frame = outbound.recv() => {
let Some(frame) = outbound_frame else { break; };
if send_relay_frame(&mut write, frame).await.is_err() { break; }
}
incoming = read.next() => {
let Some(Ok(message)) = incoming else { break; };
match operator_relay_frame(message) {
Ok(Some(frame)) => {
explicit_close = matches!(frame, TerminalRelayFrame::Close);
if state.terminals.relay_to_agent(&terminal_id, frame).await.is_err() { break; }
if explicit_close { break; }
}
Ok(None) => {}
Err(code) => {
warn!(terminal_id, code, "invalid operator terminal frame");
break;
}
}
}
}
}
if explicit_close {
close_registered_terminal(&state, &terminal_id).await;
} else if let Some(detached_at) = state.terminals.detach_operator(&terminal_id).await {
let terminals = state.terminals.clone();
let terminal_id_for_grace = terminal_id.clone();
tokio::spawn(async move {
tokio::time::sleep(TERMINAL_DISCONNECT_GRACE).await;
terminals
.remove_if_still_detached(&terminal_id_for_grace, detached_at)
.await;
});
}
info!(
terminal_id,
explicit_close, "operator terminal socket detached"
);
}
fn agent_relay_frame(message: Message) -> Result<Option<TerminalRelayFrame>, &'static str> {
match message {
Message::Binary(bytes) => Ok(Some(TerminalRelayFrame::Binary(bytes.to_vec()))),
Message::Text(text) => {
let control: TerminalControl =
serde_json::from_str(&text).map_err(|_| "terminal_control_invalid")?;
if !matches!(
control,
TerminalControl::Ready
| TerminalControl::Exited { .. }
| TerminalControl::Error { .. }
| TerminalControl::Close
) {
return Err("terminal_control_direction_invalid");
}
Ok(Some(TerminalRelayFrame::Text(
serde_json::to_string(&control).map_err(|_| "terminal_control_invalid")?,
)))
}
Message::Ping(_) | Message::Pong(_) => Ok(None),
Message::Close(_) => Ok(Some(TerminalRelayFrame::Close)),
}
}
fn operator_relay_frame(message: Message) -> Result<Option<TerminalRelayFrame>, &'static str> {
match message {
Message::Binary(bytes) => Ok(Some(TerminalRelayFrame::Binary(bytes.to_vec()))),
Message::Text(text) => {
let control: TerminalControl =
serde_json::from_str(&text).map_err(|_| "terminal_control_invalid")?;
if !matches!(
control,
TerminalControl::Resize { .. } | TerminalControl::Close
) {
return Err("terminal_control_direction_invalid");
}
if let TerminalControl::Resize { rows, cols } = control {
validate_size(rows, cols).map_err(|_| "terminal_size_invalid")?;
}
Ok(Some(TerminalRelayFrame::Text(
serde_json::to_string(&control).map_err(|_| "terminal_control_invalid")?,
)))
}
Message::Ping(_) | Message::Pong(_) => Ok(None),
Message::Close(_) => Ok(None),
}
}
async fn receive_text_handshake(socket: &mut WebSocket) -> Result<String, &'static str> {
match tokio::time::timeout(TERMINAL_ATTACH_TIMEOUT, socket.recv()).await {
Ok(Some(Ok(Message::Text(text)))) => Ok(text.to_string()),
Ok(_) => Err("terminal_handshake_required"),
Err(_) => Err("terminal_handshake_timeout"),
}
}
async fn close_socket(socket: &mut WebSocket, code: &str) {
let control = TerminalControl::Error {
code: code.into(),
message: code.replace('_', " "),
};
if let Ok(json) = serde_json::to_string(&control) {
let _ = socket.send(Message::Text(json.into())).await;
}
let _ = socket.send(Message::Close(None)).await;
}
async fn send_relay_frame<S>(write: &mut S, frame: TerminalRelayFrame) -> Result<(), axum::Error>
where
S: futures_util::Sink<Message, Error = axum::Error> + Unpin,
{
let message = match frame {
TerminalRelayFrame::Binary(bytes) => Message::Binary(bytes.into()),
TerminalRelayFrame::Text(text) => Message::Text(text.into()),
TerminalRelayFrame::Close => Message::Close(None),
};
write.send(message).await
}
async fn close_registered_terminal(state: &AppState, terminal_id: &str) -> 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()?;
let _ = session
.tx
.send(SessionEvent::Message(ServerMessage::CloseTerminal {
terminal_id,
}));
}
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()
{
info!(terminal_id = %terminal_id, agent_id, "terminal absolute timeout reached");
}
});
}
struct TerminalAudit<'a> {
actor_type: &'a str,
event_type: &'a str,
outcome: &'a str,
message: &'a str,
metadata: serde_json::Value,
}
async fn audit_agent_control_frame(
state: &AppState,
agent_id: &str,
terminal_id: &str,
frame: &TerminalRelayFrame,
) {
let TerminalRelayFrame::Text(text) = frame else {
return;
};
let Ok(control) = serde_json::from_str::<TerminalControl>(text) else {
return;
};
let audit = match control {
TerminalControl::Ready => TerminalAudit {
actor_type: "agent",
event_type: "terminal_ready",
outcome: "ok",
message: "terminal PTY ready",
metadata: serde_json::json!({}),
},
TerminalControl::Exited { exit_code } => TerminalAudit {
actor_type: "agent",
event_type: "terminal_exit",
outcome: "exited",
message: "terminal process exited",
metadata: serde_json::json!({ "exit_code": exit_code }),
},
TerminalControl::Error { code, .. } => TerminalAudit {
actor_type: "agent",
event_type: "terminal_error",
outcome: "error",
message: "terminal worker reported an error",
metadata: serde_json::json!({ "code": code }),
},
_ => return,
};
append_terminal_audit(state, agent_id, terminal_id, audit).await;
}
async fn append_terminal_audit(
state: &AppState,
agent_id: &str,
terminal_id: &str,
audit: TerminalAudit<'_>,
) {
if let Err(err) = state
.store
.append_audit_event(AuditEventInput {
actor_type: audit.actor_type.into(),
actor_id: None,
agent_id: Some(agent_id.to_string()),
request_id: Some(terminal_id.to_string()),
event_type: audit.event_type.into(),
outcome: audit.outcome.into(),
latency_ms: None,
message: audit.message.into(),
metadata: audit.metadata,
})
.await
{
warn!(terminal_id, event_type = audit.event_type, error = %err, "failed to append terminal audit event");
}
}
fn validate_size(rows: u16, cols: u16) -> Result<(), ApiError> {
if !(1..=300).contains(&rows) || !(1..=500).contains(&cols) {
return Err(ApiError::new(
StatusCode::BAD_REQUEST,
"terminal_size_invalid",
"terminal rows must be 1..=300 and columns must be 1..=500",
));
}
Ok(())
}
fn registry_error(code: &'static str) -> ApiError {
let status = match code {
"terminal_not_found" => StatusCode::NOT_FOUND,
"terminal_relay_token_invalid" | "terminal_attachment_token_invalid" => {
StatusCode::UNAUTHORIZED
}
_ => StatusCode::CONFLICT,
};
ApiError::new(status, code, code.replace('_', " "))
}
fn terminal_not_found(terminal_id: &str) -> ApiError {
ApiError::new(
StatusCode::NOT_FOUND,
"terminal_not_found",
format!("terminal session {terminal_id} was not found"),
)
}
fn operator_ws_path(terminal_id: &str) -> String {
format!("/api/v1/control/terminals/{terminal_id}/ws")
}
const fn default_rows() -> u16 {
24
}
const fn default_cols() -> u16 {
80
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn terminal_control_direction_is_enforced() {
let resize = serde_json::to_string(&TerminalControl::Resize { rows: 24, cols: 80 })
.expect("resize json");
assert!(operator_relay_frame(Message::Text(resize.into())).is_ok());
let ready = serde_json::to_string(&TerminalControl::Ready).expect("ready json");
assert!(operator_relay_frame(Message::Text(ready.into())).is_err());
}
}
wakey-control-plane/src/runtime/mod.rs
+22 -1
View File
@@ -16,7 +16,7 @@ use tower_http::services::ServeFile;
use tracing::info;
#[cfg(unix)]
use tracing::warn;
use wakey_agent::protocol::{ErrorPayload, ServerMessage};
use wakey_agent::protocol::{AgentCapability, ErrorPayload, ServerMessage};
use crate::api;
use crate::config;
@@ -25,6 +25,7 @@ use crate::ws;
mod admin;
mod process;
pub mod terminals;
pub use admin::revoke_agent;
pub use admin::{
issue_enroll_token, list_enroll_tokens, migrate_sqlite_state, revoke_enroll_token, state_stats,
@@ -41,12 +42,14 @@ pub struct AppState {
pub public_url: String,
pub command_timeout: Duration,
pub enroll_token_ttl: Duration,
pub terminals: terminals::TerminalRegistry,
}
#[derive(Clone)]
pub struct AgentSession {
pub connection_id: String,
pub tx: mpsc::UnboundedSender<SessionEvent>,
pub capabilities: Vec<AgentCapability>,
}
#[derive(Clone)]
pub enum SessionEvent {
@@ -72,6 +75,10 @@ fn public_api_routes(ui_dist_dir: std::path::PathBuf) -> Router<AppState> {
.route("/healthz", get(api::healthz))
.route("/api/v1/agents/enroll", post(api::enroll))
.route("/api/v1/agent/ws", get(ws::agent_ws))
.route(
"/api/v1/agent/terminals/{terminal_id}/ws",
get(api::agent_terminal_ws),
)
}
fn control_api_routes() -> Router<AppState> {
@@ -135,6 +142,19 @@ fn control_api_routes() -> Router<AppState> {
"/api/v1/control/agents/{agent_id}/command",
post(api::run_command),
)
.route("/api/v1/control/terminals", post(api::create_terminal))
.route(
"/api/v1/control/terminals/{terminal_id}",
get(api::get_terminal).delete(api::close_terminal),
)
.route(
"/api/v1/control/terminals/{terminal_id}/attach",
post(api::attach_terminal),
)
.route(
"/api/v1/control/terminals/{terminal_id}/ws",
get(api::operator_terminal_ws),
)
}
/// Starts the control-plane HTTP and websocket surfaces and manages daemon lifecycle hooks.
@@ -157,6 +177,7 @@ pub async fn serve(daemon: config::DaemonConfig) -> Result<()> {
public_url: daemon.public_url.clone(),
command_timeout: daemon.command_timeout,
enroll_token_ttl: daemon.enroll_token_ttl,
terminals: terminals::TerminalRegistry::new(),
};
// Keep route classes explicit so edge policy can map directly:
wakey-control-plane/src/runtime/terminals.rs
+481
View File
@@ -0,0 +1,481 @@
use std::collections::{HashMap, VecDeque};
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use tokio::sync::{Mutex, mpsc};
use uuid::Uuid;
pub const TERMINAL_RELAY_QUEUE: usize = 32;
pub const TERMINAL_MAX_FRAME_BYTES: usize = 64 * 1024;
pub const TERMINAL_REPLAY_BYTES: usize = 256 * 1024;
pub const TERMINAL_MAX_SESSIONS_PER_AGENT: usize = 2;
pub const TERMINAL_ATTACH_TIMEOUT: Duration = Duration::from_secs(10);
pub const TERMINAL_DISCONNECT_GRACE: Duration = Duration::from_secs(15);
pub const TERMINAL_ABSOLUTE_TIMEOUT: Duration = Duration::from_secs(30 * 60);
const TERMINAL_TOMBSTONE_TTL: Duration = Duration::from_secs(5 * 60);
const TERMINAL_MAX_TOMBSTONES: usize = 1024;
#[derive(Clone, Debug)]
pub enum TerminalRelayFrame {
Binary(Vec<u8>),
Text(String),
Close,
}
#[derive(Clone)]
pub struct TerminalRegistry {
inner: Arc<Mutex<HashMap<String, TerminalSession>>>,
closed: Arc<Mutex<HashMap<String, Instant>>>,
}
struct TerminalSession {
agent_id: String,
created_at_unix: u64,
expires_at: Instant,
relay_token: Option<String>,
attachment_token: Option<String>,
agent_tx: Option<mpsc::Sender<TerminalRelayFrame>>,
pending_agent: VecDeque<TerminalRelayFrame>,
pending_agent_bytes: usize,
operator_tx: Option<mpsc::Sender<TerminalRelayFrame>>,
operator_detached_at: Option<Instant>,
replay: VecDeque<TerminalRelayFrame>,
replay_bytes: usize,
}
pub struct CreatedTerminal {
pub terminal_id: String,
pub relay_token: String,
pub attachment_token: String,
pub created_at_unix: u64,
}
impl Default for TerminalRegistry {
fn default() -> Self {
Self::new()
}
}
impl TerminalRegistry {
pub fn new() -> Self {
Self {
inner: Arc::new(Mutex::new(HashMap::new())),
closed: Arc::new(Mutex::new(HashMap::new())),
}
}
pub async fn create(&self, agent_id: String) -> Result<CreatedTerminal, &'static str> {
let mut sessions = self.inner.lock().await;
if sessions
.values()
.filter(|session| session.agent_id == agent_id)
.count()
>= TERMINAL_MAX_SESSIONS_PER_AGENT
{
return Err("agent_terminal_limit_reached");
}
let terminal_id = Uuid::new_v4().to_string();
let relay_token = new_token();
let attachment_token = new_token();
let created_at_unix = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
sessions.insert(
terminal_id.clone(),
TerminalSession {
agent_id,
created_at_unix,
expires_at: Instant::now() + TERMINAL_ABSOLUTE_TIMEOUT,
relay_token: Some(relay_token.clone()),
attachment_token: Some(attachment_token.clone()),
agent_tx: None,
pending_agent: VecDeque::new(),
pending_agent_bytes: 0,
operator_tx: None,
operator_detached_at: None,
replay: VecDeque::new(),
replay_bytes: 0,
},
);
Ok(CreatedTerminal {
terminal_id,
relay_token,
attachment_token,
created_at_unix,
})
}
pub async fn remove(&self, terminal_id: &str) -> Option<String> {
let session = self.inner.lock().await.remove(terminal_id)?;
self.remember_closed(terminal_id).await;
if let Some(tx) = session.agent_tx {
let _ =
tokio::time::timeout(Duration::from_secs(1), tx.send(TerminalRelayFrame::Close))
.await;
}
if let Some(tx) = session.operator_tx {
let _ =
tokio::time::timeout(Duration::from_secs(1), tx.send(TerminalRelayFrame::Close))
.await;
}
Some(session.agent_id)
}
/// Reports whether a session ID was recently removed. Tombstones make
/// idempotent DELETE distinguishable from a completely unknown ID.
pub async fn was_closed(&self, terminal_id: &str) -> bool {
let mut closed = self.closed.lock().await;
prune_tombstones(&mut closed);
closed.contains_key(terminal_id)
}
async fn remember_closed(&self, terminal_id: &str) {
let mut closed = self.closed.lock().await;
prune_tombstones(&mut closed);
closed.insert(terminal_id.to_string(), Instant::now());
if closed.len() > TERMINAL_MAX_TOMBSTONES
&& let Some(oldest) = closed
.iter()
.min_by_key(|(_, closed_at)| **closed_at)
.map(|(terminal_id, _)| terminal_id.clone())
{
closed.remove(&oldest);
}
}
pub async fn remove_agent(&self, agent_id: &str) {
let terminal_ids = {
let sessions = self.inner.lock().await;
sessions
.iter()
.filter(|(_, session)| session.agent_id == agent_id)
.map(|(terminal_id, _)| terminal_id.clone())
.collect::<Vec<_>>()
};
for terminal_id in terminal_ids {
self.remove(&terminal_id).await;
}
}
pub async fn issue_attachment_token(&self, terminal_id: &str) -> Result<String, &'static str> {
let mut sessions = self.inner.lock().await;
let session = active_session(&mut sessions, terminal_id)?;
if session.operator_tx.is_some() {
return Err("terminal_operator_already_attached");
}
let token = new_token();
session.attachment_token = Some(token.clone());
Ok(token)
}
pub async fn attach_agent(
&self,
terminal_id: &str,
agent_id: &str,
relay_token: &str,
) -> Result<(mpsc::Receiver<TerminalRelayFrame>, Vec<TerminalRelayFrame>), &'static str> {
let mut sessions = self.inner.lock().await;
let session = active_session(&mut sessions, terminal_id)?;
if session.agent_id != agent_id {
return Err("terminal_agent_mismatch");
}
if session.agent_tx.is_some() {
return Err("terminal_agent_already_attached");
}
if session.relay_token.as_deref() != Some(relay_token) {
return Err("terminal_relay_token_invalid");
}
session.relay_token = None;
let (tx, rx) = mpsc::channel(TERMINAL_RELAY_QUEUE);
session.agent_tx = Some(tx);
let pending = session.pending_agent.drain(..).collect();
session.pending_agent_bytes = 0;
Ok((rx, pending))
}
pub async fn attach_operator(
&self,
terminal_id: &str,
attachment_token: &str,
) -> Result<(mpsc::Receiver<TerminalRelayFrame>, Vec<TerminalRelayFrame>), &'static str> {
let mut sessions = self.inner.lock().await;
let session = active_session(&mut sessions, terminal_id)?;
if session.operator_tx.is_some() {
return Err("terminal_operator_already_attached");
}
if session.attachment_token.as_deref() != Some(attachment_token) {
return Err("terminal_attachment_token_invalid");
}
session.attachment_token = None;
session.operator_detached_at = None;
let replay = session.replay.drain(..).collect();
session.replay_bytes = 0;
let (tx, rx) = mpsc::channel(TERMINAL_RELAY_QUEUE);
session.operator_tx = Some(tx);
Ok((rx, replay))
}
pub async fn relay_to_agent(
&self,
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_REPLAY_BYTES {
if let Some(dropped) = session.pending_agent.pop_front() {
session.pending_agent_bytes -= relay_frame_size(&dropped);
} else {
break;
}
}
Ok(())
}
pub async fn relay_from_agent(
&self,
terminal_id: &str,
frame: TerminalRelayFrame,
) -> Result<(), &'static str> {
let operator_tx = self
.inner
.lock()
.await
.get(terminal_id)
.and_then(|session| session.operator_tx.clone());
if let Some(tx) = operator_tx {
match tx.send(frame).await {
Ok(()) => return Ok(()),
Err(err) => {
// The browser task may not have marked itself detached yet.
// Preserve this frame so that race does not kill the PTY.
let mut sessions = self.inner.lock().await;
let session = active_session(&mut sessions, terminal_id)?;
session.operator_tx = None;
session
.operator_detached_at
.get_or_insert_with(Instant::now);
push_replay(session, err.0);
return Ok(());
}
}
}
let mut sessions = self.inner.lock().await;
let session = active_session(&mut sessions, terminal_id)?;
push_replay(session, frame);
Ok(())
}
pub async fn reject(
&self,
terminal_id: &str,
agent_id: &str,
error_json: String,
) -> Result<(), &'static str> {
let matches_agent = self
.inner
.lock()
.await
.get(terminal_id)
.is_some_and(|session| session.agent_id == agent_id);
if !matches_agent {
return Err("terminal_agent_mismatch");
}
self.relay_from_agent(terminal_id, TerminalRelayFrame::Text(error_json))
.await?;
self.relay_from_agent(terminal_id, TerminalRelayFrame::Close)
.await
}
pub async fn detach_operator(&self, terminal_id: &str) -> Option<Instant> {
let mut sessions = self.inner.lock().await;
let session = sessions.get_mut(terminal_id)?;
session.operator_tx = None;
let detached_at = Instant::now();
session.operator_detached_at = Some(detached_at);
Some(detached_at)
}
pub async fn remove_if_still_detached(&self, terminal_id: &str, detached_at: Instant) -> bool {
let should_remove = self
.inner
.lock()
.await
.get(terminal_id)
.is_some_and(|session| session.operator_detached_at == Some(detached_at));
if should_remove {
self.remove(terminal_id).await;
}
should_remove
}
pub async fn summary(&self, terminal_id: &str) -> Option<(String, u64, bool, bool)> {
self.inner.lock().await.get(terminal_id).map(|session| {
(
session.agent_id.clone(),
session.created_at_unix,
session.agent_tx.is_some(),
session.operator_tx.is_some(),
)
})
}
}
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() {
return Err("terminal_expired");
}
Ok(session)
}
fn new_token() -> String {
format!("{}{}", Uuid::new_v4().simple(), Uuid::new_v4().simple())
}
fn prune_tombstones(closed: &mut HashMap<String, Instant>) {
closed.retain(|_, closed_at| closed_at.elapsed() < TERMINAL_TOMBSTONE_TTL);
}
fn relay_frame_size(frame: &TerminalRelayFrame) -> usize {
match frame {
TerminalRelayFrame::Binary(bytes) => bytes.len(),
TerminalRelayFrame::Text(text) => text.len(),
TerminalRelayFrame::Close => 0,
}
}
fn push_replay(session: &mut TerminalSession, frame: TerminalRelayFrame) {
session.replay_bytes += relay_frame_size(&frame);
session.replay.push_back(frame);
while session.replay_bytes > TERMINAL_REPLAY_BYTES {
if let Some(dropped) = session.replay.pop_front() {
session.replay_bytes -= relay_frame_size(&dropped);
} else {
break;
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn credentials_are_scoped_and_single_use() {
let registry = TerminalRegistry::new();
let created = registry.create("router".into()).await.expect("create");
assert!(
registry
.attach_agent(&created.terminal_id, "other", &created.relay_token)
.await
.is_err()
);
let _ = registry
.attach_agent(&created.terminal_id, "router", &created.relay_token)
.await
.expect("attach agent");
assert_eq!(
registry
.attach_agent(&created.terminal_id, "router", &created.relay_token)
.await
.expect_err("relay token is single use"),
"terminal_agent_already_attached"
);
registry
.attach_operator(&created.terminal_id, &created.attachment_token)
.await
.expect("attach operator");
assert!(
registry
.attach_operator(&created.terminal_id, &created.attachment_token)
.await
.is_err()
);
}
#[tokio::test]
async fn detached_output_replay_is_bounded() {
let registry = TerminalRegistry::new();
let created = registry.create("router".into()).await.expect("create");
for _ in 0..10 {
registry
.relay_from_agent(
&created.terminal_id,
TerminalRelayFrame::Binary(vec![0; TERMINAL_REPLAY_BYTES / 4]),
)
.await
.expect("buffer output");
}
let (_, replay) = registry
.attach_operator(&created.terminal_id, &created.attachment_token)
.await
.expect("attach operator");
assert!(replay.iter().map(relay_frame_size).sum::<usize>() <= TERMINAL_REPLAY_BYTES);
}
#[tokio::test]
async fn operator_input_waits_for_agent_attachment() {
let registry = TerminalRegistry::new();
let created = registry.create("router".into()).await.expect("create");
let resize = TerminalRelayFrame::Text(r#"{"type":"resize","rows":30,"cols":120}"#.into());
registry
.relay_to_agent(&created.terminal_id, resize)
.await
.expect("queue resize before agent attachment");
let (_, pending) = registry
.attach_agent(&created.terminal_id, "router", &created.relay_token)
.await
.expect("attach agent");
assert_eq!(pending.len(), 1);
assert!(matches!(pending[0], TerminalRelayFrame::Text(_)));
}
#[tokio::test]
async fn agent_can_hold_two_terminal_sessions() {
let registry = TerminalRegistry::new();
registry.create("router".into()).await.expect("first");
registry.create("router".into()).await.expect("second");
let third = registry.create("router".into()).await;
assert_eq!(third.err(), Some("agent_terminal_limit_reached"));
}
#[tokio::test]
async fn removed_session_is_distinct_from_unknown_session() {
let registry = TerminalRegistry::new();
let created = registry.create("router".into()).await.expect("create");
assert!(!registry.was_closed(&created.terminal_id).await);
assert!(!registry.was_closed("never-existed").await);
registry.remove(&created.terminal_id).await.expect("remove");
assert!(registry.was_closed(&created.terminal_id).await);
assert!(!registry.was_closed("never-existed").await);
}
}
wakey-control-plane/src/ws.rs
+61 -17
View File
@@ -8,7 +8,9 @@ use std::time::Instant;
use tokio::sync::mpsc;
use tracing::{debug, info, info_span, warn};
use uuid::Uuid;
use wakey_agent::protocol::{ErrorPayload, RequestId, ServerMessage};
use wakey_agent::protocol::{
AgentCapability, ErrorPayload, RequestId, ServerMessage, TerminalControl, TerminalId,
};
use wakey_core::Device;
use crate::runtime::{AgentReply, AgentSession, AppState, SessionEvent};
@@ -19,6 +21,8 @@ use crate::state::AuditEventInput;
enum IncomingClientMessage {
Hello {
agent_id: String,
#[serde(default)]
capabilities: Vec<AgentCapability>,
},
Auth {
agent_id: String,
@@ -39,6 +43,17 @@ enum IncomingClientMessage {
request_id: RequestId,
error: ErrorPayload,
},
TerminalRejected {
terminal_id: TerminalId,
error: ErrorPayload,
},
}
#[derive(Default)]
struct AgentConnectionState {
authed_agent_id: Option<String>,
hello_at: Option<Instant>,
capabilities: Vec<AgentCapability>,
}
pub async fn agent_ws(ws: WebSocketUpgrade, State(state): State<AppState>) -> impl IntoResponse {
@@ -80,8 +95,7 @@ async fn handle_agent_socket(state: AppState, socket: WebSocket) {
debug!("websocket writer loop ended");
});
let mut authed_agent_id: Option<String> = None;
let mut hello_at: Option<Instant> = None;
let mut connection = AgentConnectionState::default();
loop {
let frame = read.next().await;
@@ -103,8 +117,7 @@ async fn handle_agent_socket(state: AppState, socket: WebSocket) {
&state,
&tx,
&connection_id,
&mut authed_agent_id,
&mut hello_at,
&mut connection,
connected_at,
&text,
)
@@ -126,7 +139,7 @@ async fn handle_agent_socket(state: AppState, socket: WebSocket) {
}
}
if let Some(agent_id) = authed_agent_id {
if let Some(agent_id) = connection.authed_agent_id {
info!(agent_id = %agent_id, "agent disconnected");
let mut sessions = state.sessions.write().await;
let should_remove = sessions
@@ -136,6 +149,8 @@ async fn handle_agent_socket(state: AppState, socket: WebSocket) {
if should_remove {
sessions.remove(&agent_id);
}
drop(sessions);
state.terminals.remove_agent(&agent_id).await;
if let Err(err) = state
.store
.append_audit_event(AuditEventInput {
@@ -163,8 +178,7 @@ async fn process_agent_text(
state: &AppState,
tx: &mpsc::UnboundedSender<SessionEvent>,
connection_id: &str,
authed_agent_id: &mut Option<String>,
hello_at: &mut Option<Instant>,
connection: &mut AgentConnectionState,
connected_at: Instant,
text: &str,
) -> Result<()> {
@@ -172,20 +186,26 @@ async fn process_agent_text(
serde_json::from_str(text).context("invalid client websocket payload")?;
match message {
IncomingClientMessage::Hello { agent_id } => {
IncomingClientMessage::Hello {
agent_id,
capabilities,
} => {
let now = Instant::now();
if hello_at.is_none() {
*hello_at = Some(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.capabilities = capabilities;
}
IncomingClientMessage::Auth {
agent_id,
agent_token,
} => {
let connect_to_auth_ms = connected_at.elapsed().as_millis() as u64;
let hello_to_auth_ms = hello_at.map(|t| now_duration_ms(t.elapsed()));
let hello_to_auth_ms = connection
.hello_at
.map(|time| now_duration_ms(time.elapsed()));
if !state
.store
.verify_agent_token(&agent_id, &agent_token)
@@ -219,9 +239,10 @@ async fn process_agent_text(
AgentSession {
connection_id: connection_id.to_string(),
tx: tx.clone(),
capabilities: connection.capabilities.clone(),
},
);
*authed_agent_id = Some(agent_id.clone());
connection.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
.store
@@ -246,14 +267,14 @@ async fn process_agent_text(
let _ = tx.send(SessionEvent::Message(ServerMessage::SyncDeviceSnapshot));
}
IncomingClientMessage::Heartbeat { agent_id } => {
if authed_agent_id.as_deref() != Some(agent_id.as_str()) {
if connection.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::DeviceSnapshot { agent_id, devices } => {
if authed_agent_id.as_deref() != Some(agent_id.as_str()) {
if connection.authed_agent_id.as_deref() != Some(agent_id.as_str()) {
anyhow::bail!("device_snapshot for unauthenticated or mismatched agent");
}
ensure_current_session(state, &agent_id, connection_id).await?;
@@ -272,7 +293,8 @@ async fn process_agent_text(
}
}
IncomingClientMessage::Result { request_id, result } => {
let agent_id = authed_agent_id
let agent_id = connection
.authed_agent_id
.as_deref()
.ok_or_else(|| anyhow::anyhow!("result before auth"))?;
ensure_current_session(state, agent_id, connection_id).await?;
@@ -284,7 +306,8 @@ async fn process_agent_text(
}
}
IncomingClientMessage::Error { request_id, error } => {
let agent_id = authed_agent_id
let agent_id = connection
.authed_agent_id
.as_deref()
.ok_or_else(|| anyhow::anyhow!("error before auth"))?;
ensure_current_session(state, agent_id, connection_id).await?;
@@ -295,6 +318,26 @@ async fn process_agent_text(
debug!(request_id = %key, "dropping unsolicited error from agent");
}
}
IncomingClientMessage::TerminalRejected { terminal_id, error } => {
let agent_id = connection
.authed_agent_id
.as_deref()
.ok_or_else(|| anyhow::anyhow!("terminal rejection before auth"))?;
ensure_current_session(state, agent_id, connection_id).await?;
let error_json = serde_json::to_string(&TerminalControl::Error {
code: error.code,
message: error.message,
})?;
if let Err(code) = state
.terminals
.reject(terminal_id.as_str(), agent_id, error_json)
.await
{
debug!(terminal_id = %terminal_id, agent_id, code, "dropping rejection for inactive terminal");
return Ok(());
}
warn!(terminal_id = %terminal_id, agent_id, "agent rejected terminal request");
}
}
Ok(())
@@ -349,6 +392,7 @@ mod tests {
AgentSession {
connection_id: "conn-new".to_string(),
tx,
capabilities: Vec::new(),
},
);