advertise and enforce per-agent terminal session limits
This commit is contained in:
@@ -4,6 +4,8 @@ use std::fmt;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::LazyLock;
|
||||
|
||||
use crate::protocol::DEFAULT_TERMINAL_MAX_SESSIONS;
|
||||
|
||||
pub const DEFAULT_CONFIG_PATH: &str = "/etc/wakey-agent/config.toml";
|
||||
pub const DEFAULT_PID_FILE: &str = "/var/run/wakey-agent.pid";
|
||||
const WAKEY_DHCP_LEASES_ENV: &str = "WAKEY_DHCP_LEASES";
|
||||
@@ -40,6 +42,7 @@ pub struct AgentConfig {
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct TerminalConfig {
|
||||
#[serde(default)]
|
||||
pub enabled: bool,
|
||||
@@ -142,7 +145,7 @@ fn default_terminal_shell() -> PathBuf {
|
||||
}
|
||||
|
||||
const fn default_terminal_max_sessions() -> usize {
|
||||
2
|
||||
DEFAULT_TERMINAL_MAX_SESSIONS
|
||||
}
|
||||
|
||||
impl AgentConfig {
|
||||
@@ -312,4 +315,22 @@ max_sessions = 2
|
||||
assert!(config.terminal.args.is_empty());
|
||||
assert!(config.terminal.current_dir.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn terminal_config_rejects_misspelled_fields() {
|
||||
let error = toml::from_str::<AgentConfig>(
|
||||
r#"
|
||||
server_url = "https://example.com"
|
||||
agent_id = "agent-1"
|
||||
agent_token = "secret"
|
||||
|
||||
[terminal]
|
||||
enabled = true
|
||||
max_session = 67
|
||||
"#,
|
||||
)
|
||||
.expect_err("unknown terminal fields must not silently use defaults");
|
||||
|
||||
assert!(error.to_string().contains("max_session"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -65,6 +65,30 @@ pub enum AgentCapability {
|
||||
Terminal,
|
||||
}
|
||||
|
||||
pub const DEFAULT_TERMINAL_MAX_SESSIONS: usize = 2;
|
||||
|
||||
/// Optional parameters attached to advertised agent capabilities.
|
||||
///
|
||||
/// Keep this separate from `AgentCapability`: the capability list remains a
|
||||
/// compact, backward-compatible feature check, while this object can grow as
|
||||
/// individual capabilities gain configurable limits or modes.
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct AgentCapabilityOptions {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub terminal: Option<TerminalCapabilityOptions>,
|
||||
}
|
||||
|
||||
impl AgentCapabilityOptions {
|
||||
fn is_empty(&self) -> bool {
|
||||
self.terminal.is_none()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct TerminalCapabilityOptions {
|
||||
pub max_sessions: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct AgentTerminalSession {
|
||||
pub terminal_id: TerminalId,
|
||||
@@ -238,6 +262,8 @@ pub enum ClientMessage {
|
||||
agent_id: String,
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
capabilities: Vec<AgentCapability>,
|
||||
#[serde(default, skip_serializing_if = "AgentCapabilityOptions::is_empty")]
|
||||
capability_options: AgentCapabilityOptions,
|
||||
},
|
||||
Auth {
|
||||
agent_id: String,
|
||||
@@ -397,4 +423,18 @@ mod tests {
|
||||
let json = serde_json::to_string(&resume).expect("serialize terminal resume");
|
||||
assert!(json.contains("\"type\":\"resume_terminal\""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hello_serializes_typed_capability_options() {
|
||||
let message = ClientMessage::Hello {
|
||||
agent_id: "router".into(),
|
||||
capabilities: vec![AgentCapability::Terminal],
|
||||
capability_options: AgentCapabilityOptions {
|
||||
terminal: Some(TerminalCapabilityOptions { max_sessions: 3 }),
|
||||
},
|
||||
};
|
||||
|
||||
let value = serde_json::to_value(message).expect("serialize hello");
|
||||
assert_eq!(value["capability_options"]["terminal"]["max_sessions"], 3);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -8,7 +8,10 @@ use tracing::{debug, error, info, info_span, warn};
|
||||
|
||||
use crate::config::AgentConfig;
|
||||
use crate::dispatch::{dispatch_command, inventory_for_config};
|
||||
use crate::protocol::{AgentCapability, AgentCommand, ClientMessage, ErrorPayload, ServerMessage};
|
||||
use crate::protocol::{
|
||||
AgentCapability, AgentCapabilityOptions, AgentCommand, ClientMessage, ErrorPayload,
|
||||
ServerMessage, TerminalCapabilityOptions,
|
||||
};
|
||||
use crate::terminal::TerminalManager;
|
||||
|
||||
pub async fn run(config: AgentConfig) -> Result<()> {
|
||||
@@ -82,6 +85,14 @@ async fn run_once(
|
||||
&ClientMessage::Hello {
|
||||
agent_id: config.agent_id.clone(),
|
||||
capabilities: agent_capabilities(config),
|
||||
capability_options: AgentCapabilityOptions {
|
||||
terminal: config
|
||||
.terminal
|
||||
.enabled
|
||||
.then_some(TerminalCapabilityOptions {
|
||||
max_sessions: config.terminal.max_sessions.max(1),
|
||||
}),
|
||||
},
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
|
||||
Reference in New Issue
Block a user