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-agent/src/config.rs
+38
View File
@@ -35,6 +35,28 @@ pub struct AgentConfig {
pub mac_name_cache_path: PathBuf,
#[serde(default = "default_observation_store_path")]
pub observation_store_path: PathBuf,
#[serde(default)]
pub terminal: TerminalConfig,
}
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct TerminalConfig {
#[serde(default)]
pub enabled: bool,
#[serde(default = "default_terminal_shell")]
pub shell: PathBuf,
#[serde(default = "default_terminal_max_sessions")]
pub max_sessions: usize,
}
impl Default for TerminalConfig {
fn default() -> Self {
Self {
enabled: false,
shell: default_terminal_shell(),
max_sessions: default_terminal_max_sessions(),
}
}
}
pub static DEFAULT_CONFIG: LazyLock<AgentConfig> = LazyLock::new(|| AgentConfig {
@@ -49,6 +71,7 @@ pub static DEFAULT_CONFIG: LazyLock<AgentConfig> = LazyLock::new(|| AgentConfig
dhcp_leases_path: default_dhcp_leases_path(),
mac_name_cache_path: default_mac_name_cache_path(),
observation_store_path: default_observation_store_path(),
terminal: TerminalConfig::default(),
});
impl fmt::Debug for AgentConfig {
@@ -71,6 +94,7 @@ impl fmt::Debug for AgentConfig {
.field("dhcp_leases_path", &self.dhcp_leases_path)
.field("mac_name_cache_path", &self.mac_name_cache_path)
.field("observation_store_path", &self.observation_store_path)
.field("terminal", &self.terminal)
.finish()
}
}
@@ -107,6 +131,14 @@ fn default_observation_store_path() -> PathBuf {
DEFAULT_OBSERVATION_STORE_PATH.into()
}
fn default_terminal_shell() -> PathBuf {
"/bin/ash".into()
}
const fn default_terminal_max_sessions() -> usize {
2
}
impl AgentConfig {
pub fn local_path_envs(&self) -> Vec<(&'static str, &Path)> {
vec![
@@ -214,6 +246,11 @@ mod tests {
dhcp_leases_path: "/tmp/test-dhcp.leases".into(),
mac_name_cache_path: "/tmp/test-names.json".into(),
observation_store_path: "/tmp/test-observations.json".into(),
terminal: TerminalConfig {
enabled: true,
shell: "/bin/sh".into(),
max_sessions: 2,
},
};
save_config(&path, &config).expect("save");
@@ -245,5 +282,6 @@ agent_token = "secret"
config.observation_retention_days,
DEFAULT_OBSERVATION_RETENTION_DAYS
);
assert_eq!(config.terminal, TerminalConfig::default());
}
}
wakey-agent/src/enroll.rs
+6
View File
@@ -151,6 +151,11 @@ mod tests {
dhcp_leases_path: "/tmp/custom-dhcp.leases".into(),
mac_name_cache_path: "/tmp/custom-names.json".into(),
observation_store_path: "/tmp/custom-observations.json".into(),
terminal: crate::config::TerminalConfig {
enabled: true,
shell: "/bin/ash".into(),
max_sessions: 2,
},
};
let outcome = enroll(&server_url, "enroll-abc", &path, Some(&base_config))
@@ -164,6 +169,7 @@ mod tests {
assert_eq!(config.observation_retention_days, 11);
assert_eq!(config.pid_file, base_config.pid_file);
assert_eq!(config.dhcp_leases_path, base_config.dhcp_leases_path);
assert_eq!(config.terminal, base_config.terminal);
assert!(outcome.backup_path.is_none());
let persisted = crate::config::load_config(&path).expect("load persisted config");
wakey-agent/src/main.rs
+2 -1
View File
@@ -2,14 +2,15 @@ mod cli;
mod config;
mod dispatch;
mod enroll;
mod protocol;
mod serve;
mod session;
mod terminal;
mod tracing;
use anyhow::Result;
use clap::Parser;
use cli::{Cli, Command, InitConfigArgs, ObserveCommand};
use wakey_agent::protocol;
#[tokio::main]
async fn main() -> Result<()> {
wakey-agent/src/protocol.rs
+90
View File
@@ -17,6 +17,61 @@ impl RequestId {
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(transparent)]
pub struct TerminalId(String);
impl TerminalId {
pub fn new(value: impl Into<String>) -> Result<Self, String> {
let value = value.into();
if value.trim().is_empty() {
return Err("terminal_id must not be empty".into());
}
Ok(Self(value))
}
pub fn as_str(&self) -> &str {
&self.0
}
}
impl fmt::Display for TerminalId {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.0.fmt(f)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AgentCapability {
Terminal,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum TerminalControl {
Resize { rows: u16, cols: u16 },
Ready,
Exited { exit_code: Option<i32> },
Error { code: String, message: String },
Close,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum TerminalAgentHandshake {
Auth {
agent_id: String,
relay_token: String,
},
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum TerminalOperatorHandshake {
Attach { attachment_token: String },
}
impl TryFrom<String> for RequestId {
type Error = String;
@@ -153,6 +208,8 @@ pub enum CommandResult {
pub enum ClientMessage {
Hello {
agent_id: String,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
capabilities: Vec<AgentCapability>,
},
Auth {
agent_id: String,
@@ -173,6 +230,10 @@ pub enum ClientMessage {
request_id: RequestId,
error: ErrorPayload,
},
TerminalRejected {
terminal_id: TerminalId,
error: ErrorPayload,
},
}
#[derive(Clone, Debug, Serialize, Deserialize)]
@@ -183,6 +244,15 @@ pub enum ServerMessage {
command: AgentCommand,
},
SyncDeviceSnapshot,
OpenTerminal {
terminal_id: TerminalId,
relay_token: String,
rows: u16,
cols: u16,
},
CloseTerminal {
terminal_id: TerminalId,
},
}
#[cfg(test)]
@@ -232,4 +302,24 @@ mod tests {
assert!(json.contains("\"type\":\"device_snapshot\""));
assert!(json.contains("\"devices\""));
}
#[test]
fn terminal_messages_serialize_with_explicit_frame_types() {
let message = ServerMessage::OpenTerminal {
terminal_id: TerminalId::new("term-1").expect("terminal id"),
relay_token: "secret".into(),
rows: 30,
cols: 120,
};
let json = serde_json::to_string(&message).expect("serialize open terminal");
assert!(json.contains("\"type\":\"open_terminal\""));
assert!(json.contains("\"terminal_id\":\"term-1\""));
let resize = TerminalControl::Resize {
rows: 40,
cols: 160,
};
let json = serde_json::to_string(&resize).expect("serialize resize");
assert_eq!(json, r#"{"type":"resize","rows":40,"cols":160}"#);
}
}
wakey-agent/src/session.rs
+58 -2
View File
@@ -8,7 +8,8 @@ use tracing::{debug, error, info, info_span, warn};
use crate::config::AgentConfig;
use crate::dispatch::{dispatch_command, inventory_for_config};
use crate::protocol::{AgentCommand, ClientMessage, ErrorPayload, ServerMessage};
use crate::protocol::{AgentCapability, AgentCommand, ClientMessage, ErrorPayload, ServerMessage};
use crate::terminal::TerminalManager;
pub async fn run(config: AgentConfig) -> Result<()> {
let mut backoff = config.reconnect_base_ms.max(100);
@@ -71,6 +72,7 @@ async fn run_once(config: &AgentConfig) -> Result<()> {
&mut sink,
&ClientMessage::Hello {
agent_id: config.agent_id.clone(),
capabilities: agent_capabilities(config),
},
)
.await?;
@@ -91,9 +93,23 @@ async fn run_once(config: &AgentConfig) -> Result<()> {
));
snapshot_sync.set_missed_tick_behavior(MissedTickBehavior::Skip);
snapshot_sync.reset();
let (terminal_manager, mut terminal_events) = TerminalManager::new(config);
loop {
tokio::select! {
Some(event) = terminal_events.recv() => {
send_json(
&mut sink,
&ClientMessage::TerminalRejected {
terminal_id: event.terminal_id,
error: ErrorPayload {
code: "terminal_worker_failed".into(),
message: event.error,
retryable: Some(false),
},
},
).await?;
}
_ = heartbeat.tick() => {
send_json(&mut sink, &ClientMessage::Heartbeat {
agent_id: config.agent_id.clone(),
@@ -118,7 +134,7 @@ async fn run_once(config: &AgentConfig) -> Result<()> {
Message::Text(text) => {
match serde_json::from_str::<ServerMessage>(&text) {
Ok(message) => {
handle_server_message(config, &mut sink, &mut snapshot_sync, message).await?;
handle_server_message(config, &terminal_manager, &mut sink, &mut snapshot_sync, message).await?;
}
Err(err) => {
// Allow the server to introduce extra frame types without
@@ -181,6 +197,7 @@ pub fn next_backoff_ms(current_ms: u64, max_ms: u64) -> u64 {
async fn handle_server_message<S>(
config: &AgentConfig,
terminal_manager: &TerminalManager,
sink: &mut S,
snapshot_sync: &mut tokio::time::Interval,
message: ServerMessage,
@@ -223,6 +240,35 @@ where
send_device_snapshot_ws(sink, config).await?;
snapshot_sync.reset();
}
ServerMessage::OpenTerminal {
terminal_id,
relay_token,
rows,
cols,
} => {
info!(terminal_id = %terminal_id, rows, cols, "received terminal open request");
if let Err(err) =
terminal_manager.open(config, terminal_id.clone(), relay_token, rows, cols)
{
error!(terminal_id = %terminal_id, error = %err, "terminal open request rejected");
send_json(
sink,
&ClientMessage::TerminalRejected {
terminal_id,
error: ErrorPayload {
code: "terminal_open_rejected".into(),
message: err.to_string(),
retryable: Some(false),
},
},
)
.await?;
}
}
ServerMessage::CloseTerminal { terminal_id } => {
info!(terminal_id = %terminal_id, "received terminal close request");
terminal_manager.close(&terminal_id);
}
}
Ok(())
}
@@ -301,9 +347,19 @@ fn client_message_kind(message: &ClientMessage) -> &'static str {
ClientMessage::DeviceSnapshot { .. } => "device_snapshot",
ClientMessage::Result { .. } => "result",
ClientMessage::Error { .. } => "error",
ClientMessage::TerminalRejected { .. } => "terminal_rejected",
}
}
fn agent_capabilities(config: &AgentConfig) -> Vec<AgentCapability> {
#[cfg(unix)]
if config.terminal.enabled {
return vec![AgentCapability::Terminal];
}
Vec::new()
}
pub fn websocket_url(server_url: &str) -> Result<url::Url> {
let base = url::Url::parse(server_url).context("invalid server_url")?;
let scheme = match base.scheme() {
wakey-agent/src/terminal.rs
+345
View File
@@ -0,0 +1,345 @@
use std::collections::HashMap;
use std::path::Path;
use std::sync::{Arc, Mutex, Weak};
use std::time::Duration;
use crate::config::AgentConfig;
use crate::protocol::{TerminalAgentHandshake, TerminalControl, TerminalId};
use anyhow::{Context, Result};
use futures_util::{SinkExt, StreamExt};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::sync::{mpsc, oneshot};
use tokio_tungstenite::tungstenite::Message;
use tracing::{info, warn};
const MAX_TERMINAL_FRAME_BYTES: usize = 64 * 1024;
const PROCESS_SIGNAL_GRACE: Duration = Duration::from_secs(1);
/// Owns cancellation handles for terminal workers started by the control socket.
pub struct TerminalManager {
active: Arc<Mutex<HashMap<String, oneshot::Sender<()>>>>,
max_sessions: usize,
events: mpsc::UnboundedSender<TerminalManagerEvent>,
}
pub struct TerminalManagerEvent {
pub terminal_id: TerminalId,
pub error: String,
}
impl TerminalManager {
pub fn new(config: &AgentConfig) -> (Self, mpsc::UnboundedReceiver<TerminalManagerEvent>) {
let (events, event_rx) = mpsc::unbounded_channel();
(
Self {
active: Arc::new(Mutex::new(HashMap::new())),
max_sessions: config.terminal.max_sessions.max(1),
events,
},
event_rx,
)
}
pub fn open(
&self,
config: &AgentConfig,
terminal_id: TerminalId,
relay_token: String,
rows: u16,
cols: u16,
) -> Result<()> {
if !config.terminal.enabled {
anyhow::bail!("terminal capability is disabled");
}
let terminal_key = terminal_id.to_string();
let (cancel_tx, cancel_rx) = oneshot::channel();
{
let mut active = self.active.lock().expect("terminal manager poisoned");
if active.contains_key(&terminal_key) {
anyhow::bail!("terminal session {terminal_key} is already active");
}
if active.len() >= self.max_sessions {
anyhow::bail!("agent terminal session limit reached");
}
active.insert(terminal_key.clone(), cancel_tx);
}
let config = config.clone();
let active = Arc::downgrade(&self.active);
let events = self.events.clone();
tokio::spawn(async move {
if let Err(err) =
run_terminal(&config, &terminal_id, &relay_token, rows, cols, cancel_rx).await
{
warn!(terminal_id = %terminal_id, error = %err, "terminal worker failed");
let _ = events.send(TerminalManagerEvent {
terminal_id: terminal_id.clone(),
error: err.to_string(),
});
}
remove_completed(&active, terminal_id.as_str());
});
Ok(())
}
pub fn close(&self, terminal_id: &TerminalId) -> bool {
self.active
.lock()
.expect("terminal manager poisoned")
.remove(terminal_id.as_str())
.is_some_and(|cancel| cancel.send(()).is_ok())
}
}
impl Drop for TerminalManager {
fn drop(&mut self) {
if Arc::strong_count(&self.active) == 1
&& let Ok(mut active) = self.active.lock()
{
for (_, cancel) in active.drain() {
let _ = cancel.send(());
}
}
}
}
fn remove_completed(active: &Weak<Mutex<HashMap<String, oneshot::Sender<()>>>>, terminal_id: &str) {
if let Some(active) = active.upgrade()
&& let Ok(mut active) = active.lock()
{
active.remove(terminal_id);
}
}
#[cfg(unix)]
async fn run_terminal(
config: &AgentConfig,
terminal_id: &TerminalId,
relay_token: &str,
rows: u16,
cols: u16,
mut cancel: oneshot::Receiver<()>,
) -> Result<()> {
let ws_url = terminal_websocket_url(&config.server_url, terminal_id)?;
let (stream, _) = tokio::select! {
_ = &mut cancel => return Ok(()),
result = tokio_tungstenite::connect_async(ws_url.as_str()) => {
result.context("failed to connect terminal relay websocket")?
}
};
let (mut sink, mut source) = stream.split();
send_json(
&mut sink,
&TerminalAgentHandshake::Auth {
agent_id: config.agent_id.clone(),
relay_token: relay_token.to_string(),
},
)
.await?;
let terminal = match wakey::wakey_linux::terminal::TerminalPty::spawn(
Path::new(&config.terminal.shell),
rows,
cols,
) {
Ok(terminal) => terminal,
Err(err) => {
let _ = send_json(
&mut sink,
&TerminalControl::Error {
code: "terminal_spawn_failed".into(),
message: err.to_string(),
},
)
.await;
let _ = sink.send(Message::Close(None)).await;
return Err(err);
}
};
let wakey::wakey_linux::terminal::TerminalPty {
mut reader,
mut writer,
mut child,
} = terminal;
let process_group = child.id();
send_json(&mut sink, &TerminalControl::Ready).await?;
info!(terminal_id = %terminal_id, shell = %config.terminal.shell.display(), "terminal PTY ready");
let mut output = [0_u8; 16 * 1024];
let mut requested_close = false;
let mut observed_status = None;
loop {
tokio::select! {
_ = &mut cancel => {
requested_close = true;
break;
}
status = child.wait() => {
observed_status = Some(status.context("failed waiting for terminal child")?);
break;
}
read = reader.read(&mut output) => {
match read {
Ok(0) => break,
Ok(count) => sink
.send(Message::Binary(output[..count].to_vec().into()))
.await
.context("failed to send PTY output")?,
// Linux PTY masters commonly report EIO after the slave closes.
Err(err) if err.raw_os_error() == Some(5) => break,
Err(err) => return Err(err).context("failed to read PTY output"),
}
}
incoming = source.next() => {
let Some(message) = incoming else { break; };
match message.context("terminal relay websocket receive failed")? {
Message::Binary(bytes) => {
if bytes.len() > MAX_TERMINAL_FRAME_BYTES {
anyhow::bail!("terminal input frame exceeds size limit");
}
writer.write_all(&bytes).await.context("failed to write PTY input")?;
}
Message::Text(text) => {
match serde_json::from_str::<TerminalControl>(&text)
.context("invalid terminal control frame")?
{
TerminalControl::Resize { rows, cols } => {
validate_size(rows, cols)?;
wakey::wakey_linux::terminal::resize_terminal(
&writer, rows, cols,
)?;
}
TerminalControl::Close => {
requested_close = true;
break;
}
_ => anyhow::bail!("terminal control frame has invalid direction"),
}
}
Message::Ping(payload) => sink.send(Message::Pong(payload)).await?,
Message::Pong(_) => {}
Message::Close(_) => {
requested_close = true;
break;
}
Message::Frame(_) => {}
}
}
}
}
let status = match observed_status {
Some(status) => status,
None => terminate_process_group(&mut child, process_group).await?,
};
let _ = send_json(
&mut sink,
&TerminalControl::Exited {
exit_code: status.code(),
},
)
.await;
let _ = sink.send(Message::Close(None)).await;
info!(terminal_id = %terminal_id, exit_code = ?status.code(), requested_close, "terminal worker exited");
Ok(())
}
#[cfg(not(unix))]
async fn run_terminal(
_config: &AgentConfig,
_terminal_id: &TerminalId,
_relay_token: &str,
_rows: u16,
_cols: u16,
_cancel: oneshot::Receiver<()>,
) -> Result<()> {
anyhow::bail!("terminal sessions are unsupported on this platform")
}
#[cfg(unix)]
async fn terminate_process_group(
child: &mut tokio::process::Child,
process_group: Option<u32>,
) -> Result<std::process::ExitStatus> {
use nix::sys::signal::{Signal, killpg};
use nix::unistd::Pid;
let Some(process_group) = process_group else {
child
.start_kill()
.context("failed to kill terminal child")?;
return child
.wait()
.await
.context("failed waiting for terminal child");
};
let pid = Pid::from_raw(process_group as i32);
let _ = killpg(pid, Signal::SIGHUP);
if let Ok(status) = tokio::time::timeout(PROCESS_SIGNAL_GRACE, child.wait()).await {
return status.context("failed waiting for terminal child after SIGHUP");
}
let _ = killpg(pid, Signal::SIGTERM);
if let Ok(status) = tokio::time::timeout(PROCESS_SIGNAL_GRACE, child.wait()).await {
return status.context("failed waiting for terminal child after SIGTERM");
}
let _ = killpg(pid, Signal::SIGKILL);
child
.wait()
.await
.context("failed waiting for terminal child after SIGKILL")
}
async fn send_json<S, T>(sink: &mut S, value: &T) -> Result<()>
where
S: futures_util::Sink<Message> + Unpin,
S::Error: std::error::Error + Send + Sync + 'static,
T: serde::Serialize,
{
let json = serde_json::to_string(value).context("failed to encode terminal frame")?;
sink.send(Message::Text(json.into()))
.await
.context("failed to send terminal frame")
}
fn terminal_websocket_url(server_url: &str, terminal_id: &TerminalId) -> Result<url::Url> {
let mut url = url::Url::parse(server_url).context("invalid server_url")?;
let scheme = match url.scheme() {
"http" => "ws",
"https" => "wss",
"ws" => "ws",
"wss" => "wss",
other => anyhow::bail!("unsupported server_url scheme `{other}`"),
};
url.set_scheme(scheme)
.map_err(|_| anyhow::anyhow!("failed to convert server_url scheme"))?;
url.set_path(&format!(
"/api/v1/agent/terminals/{}/ws",
terminal_id.as_str()
));
url.set_query(None);
url.set_fragment(None);
Ok(url)
}
fn validate_size(rows: u16, cols: u16) -> Result<()> {
if !(1..=300).contains(&rows) || !(1..=500).contains(&cols) {
anyhow::bail!("terminal size is outside supported bounds");
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn terminal_url_uses_dedicated_agent_path() {
let id = TerminalId::new("term-1").expect("terminal id");
let url = terminal_websocket_url("https://example.com/base", &id).expect("url");
assert_eq!(
url.as_str(),
"wss://example.com/api/v1/agent/terminals/term-1/ws"
);
}
}