configurablility of router side storage

This commit is contained in:
lda
2026-04-27 02:26:57 +07:00 Verified
parent 8a75d82e4a
commit 5b4c5d307f
10 changed files with 359 additions and 31 deletions
+4
View File
@@ -106,6 +106,10 @@ pub struct ReloadArgs {
#[derive(Args)]
pub struct ObserveArgs {
/// Path to the agent config file. If present, local path settings are passed through.
#[arg(long, default_value = config::DEFAULT_CONFIG_PATH)]
pub config: PathBuf,
#[command(subcommand)]
pub command: ObserveCommand,
}
+66
View File
@@ -4,6 +4,12 @@ use std::fmt;
use std::path::{Path, PathBuf};
pub const DEFAULT_CONFIG_PATH: &str = "/etc/wakey-agent/config.toml";
const WAKEY_DHCP_LEASES_ENV: &str = "WAKEY_DHCP_LEASES";
const WAKEY_MAC_NAME_CACHE_ENV: &str = "WAKEY_MAC_NAME_CACHE";
const WAKEY_OBSERVATION_STORE_ENV: &str = "WAKEY_OBSERVATION_STORE";
const DEFAULT_DHCP_LEASES_PATH: &str = "/tmp/dhcp.leases";
const DEFAULT_MAC_NAME_CACHE_PATH: &str = "/tmp/wakey_mac_names.json";
const DEFAULT_OBSERVATION_STORE_PATH: &str = "/tmp/wakey_observations.json";
#[derive(Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct AgentConfig {
@@ -14,6 +20,14 @@ pub struct AgentConfig {
pub reconnect_base_ms: u64,
#[serde(default = "default_reconnect_max_ms")]
pub reconnect_max_ms: u64,
#[serde(default = "default_observation_sync_interval_seconds")]
pub observation_sync_interval_seconds: u64,
#[serde(default = "default_dhcp_leases_path")]
pub dhcp_leases_path: PathBuf,
#[serde(default = "default_mac_name_cache_path")]
pub mac_name_cache_path: PathBuf,
#[serde(default = "default_observation_store_path")]
pub observation_store_path: PathBuf,
}
impl fmt::Debug for AgentConfig {
@@ -24,6 +38,13 @@ impl fmt::Debug for AgentConfig {
.field("agent_token", &"<redacted>")
.field("reconnect_base_ms", &self.reconnect_base_ms)
.field("reconnect_max_ms", &self.reconnect_max_ms)
.field(
"observation_sync_interval_seconds",
&self.observation_sync_interval_seconds,
)
.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)
.finish()
}
}
@@ -36,6 +57,41 @@ const fn default_reconnect_max_ms() -> u64 {
30_000
}
const fn default_observation_sync_interval_seconds() -> u64 {
60
}
fn default_dhcp_leases_path() -> PathBuf {
DEFAULT_DHCP_LEASES_PATH.into()
}
fn default_mac_name_cache_path() -> PathBuf {
DEFAULT_MAC_NAME_CACHE_PATH.into()
}
fn default_observation_store_path() -> PathBuf {
DEFAULT_OBSERVATION_STORE_PATH.into()
}
impl AgentConfig {
pub fn local_path_envs(&self) -> Vec<(&'static str, &Path)> {
vec![
(WAKEY_DHCP_LEASES_ENV, self.dhcp_leases_path.as_path()),
(WAKEY_MAC_NAME_CACHE_ENV, self.mac_name_cache_path.as_path()),
(
WAKEY_OBSERVATION_STORE_ENV,
self.observation_store_path.as_path(),
),
]
}
}
pub fn apply_local_path_env_to_command(cmd: &mut std::process::Command, config: &AgentConfig) {
for (key, path) in config.local_path_envs() {
cmd.env(key, path);
}
}
pub fn load_config(path: &Path) -> Result<AgentConfig> {
let content = std::fs::read_to_string(path)
.with_context(|| format!("failed to read agent config {}", path.display()))?;
@@ -118,12 +174,22 @@ mod tests {
agent_token: "secret".into(),
reconnect_base_ms: 123,
reconnect_max_ms: 456,
observation_sync_interval_seconds: 7,
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(),
};
save_config(&path, &config).expect("save");
let loaded = load_config(&path).expect("load");
assert_eq!(loaded, config);
assert!(format!("{:?}", loaded).contains("<redacted>"));
assert!(
loaded
.local_path_envs()
.iter()
.any(|(key, _)| *key == WAKEY_OBSERVATION_STORE_ENV)
);
let _ = std::fs::remove_file(&path);
let _ = std::fs::remove_dir_all(&dir);
}
+30 -9
View File
@@ -1,27 +1,38 @@
use anyhow::Result;
use tracing::{debug, info, instrument};
use crate::config::AgentConfig;
use crate::protocol::{
AgentCommand, CommandResult, DevsRequest, InventoryRequest, LeasesRequest, WakeRequest,
};
#[instrument(skip_all)]
pub async fn dispatch_command(command: AgentCommand) -> Result<CommandResult> {
pub async fn dispatch_command(
command: AgentCommand,
config: &AgentConfig,
) -> Result<CommandResult> {
let kind = command_kind(&command);
info!(command = %kind, "dispatching command into local wakey services");
match command {
AgentCommand::Leases(req) => dispatch_leases(req).await,
AgentCommand::Leases(req) => dispatch_leases(req, config).await,
AgentCommand::Devs(req) => dispatch_devs(req).await,
AgentCommand::Inventory(req) => dispatch_inventory(req).await,
AgentCommand::Inventory(req) => dispatch_inventory(req, config).await,
AgentCommand::Wake(req) => dispatch_wake(req).await,
}
}
async fn dispatch_leases(req: LeasesRequest) -> Result<CommandResult> {
let leases = wakey::get_leases(wakey_core::LeaseQuery {
include_state: req.include_state,
})
async fn dispatch_leases(req: LeasesRequest, config: &AgentConfig) -> Result<CommandResult> {
let leases = wakey::wakey_linux::dhcp::read_dhcp_leases_with_names_from_paths(
&config.dhcp_leases_path,
&config.observation_store_path,
&config.mac_name_cache_path,
)
.await?;
let leases = if req.include_state {
wakey::wakey_linux::dhcp::enrich_leases_with_nud_state(leases).await
} else {
wakey::leases_without_state(leases)
};
debug!(
rows = leases.len(),
include_state = req.include_state,
@@ -50,8 +61,18 @@ async fn dispatch_devs(req: DevsRequest) -> Result<CommandResult> {
Ok(CommandResult::Devs { rows: devs })
}
async fn dispatch_inventory(req: InventoryRequest) -> Result<CommandResult> {
let inventory = wakey::inventory(req.into_inventory_query()).await?;
async fn dispatch_inventory(req: InventoryRequest, config: &AgentConfig) -> Result<CommandResult> {
let query = req.into_inventory_query();
let neighbors = wakey::wakey_linux::devices::query_neighbors(&query).await?;
let leases = wakey::wakey_linux::dhcp::read_dhcp_leases_with_names_from_paths(
&config.dhcp_leases_path,
&config.observation_store_path,
&config.mac_name_cache_path,
)
.await?;
let inventory = wakey_core::DeviceInventory {
devices: wakey::merge_devices(neighbors, wakey::leases_without_state(leases), &query),
};
debug!(
rows = inventory.devices.len(),
"dispatched inventory command"
+4
View File
@@ -60,6 +60,10 @@ pub async fn enroll(
agent_token: payload.agent_token,
reconnect_base_ms: 1_000,
reconnect_max_ms: 30_000,
observation_sync_interval_seconds: 60,
dhcp_leases_path: "/tmp/dhcp.leases".into(),
mac_name_cache_path: "/tmp/wakey_mac_names.json".into(),
observation_store_path: "/tmp/wakey_observations.json".into(),
};
let backup_path = save_config_with_backup(config_path, &config)?;
info!(agent_id = %config.agent_id, config_path = %config_path.display(), "agent enrollment succeeded and config was written");
+7
View File
@@ -95,6 +95,9 @@ async fn main() -> Result<()> {
fn observe(args: cli::ObserveArgs) -> Result<()> {
let mut cmd = std::process::Command::new(resolve_wakey_binary());
cmd.arg("observe");
if let Ok(config) = config::load_config(&args.config) {
config::apply_local_path_env_to_command(&mut cmd, &config);
}
match args.command {
ObserveCommand::Dhcp(args) => {
@@ -173,6 +176,10 @@ fn init_config(args: InitConfigArgs) -> Result<()> {
.unwrap_or_else(|| "REPLACE_ME_AGENT_TOKEN".to_string()),
reconnect_base_ms: 1_000,
reconnect_max_ms: 30_000,
observation_sync_interval_seconds: 60,
dhcp_leases_path: "/tmp/dhcp.leases".into(),
mac_name_cache_path: "/tmp/wakey_mac_names.json".into(),
observation_store_path: "/tmp/wakey_observations.json".into(),
};
if let Some(path) = &args.config {
+14 -7
View File
@@ -88,7 +88,9 @@ async fn run_once(config: &AgentConfig) -> Result<()> {
let http_client = reqwest::Client::new();
let mut heartbeat = interval(Duration::from_secs(30));
heartbeat.set_missed_tick_behavior(MissedTickBehavior::Skip);
let mut observation_sync = interval(Duration::from_secs(60));
let mut observation_sync = interval(Duration::from_secs(
config.observation_sync_interval_seconds.max(1),
));
observation_sync.set_missed_tick_behavior(MissedTickBehavior::Skip);
loop {
@@ -117,7 +119,7 @@ async fn run_once(config: &AgentConfig) -> Result<()> {
Message::Text(text) => {
match serde_json::from_str::<ServerMessage>(&text) {
Ok(message) => {
handle_server_message(&mut sink, message).await?;
handle_server_message(config, &mut sink, message).await?;
}
Err(err) => {
// Allow the server to introduce extra frame types without
@@ -163,9 +165,10 @@ struct AgentObservationRequest {
}
async fn send_agent_observations(client: &reqwest::Client, config: &AgentConfig) -> Result<()> {
let observations = wakey::list_local_observations()
.await
.context("failed to read local observations")?;
let observations =
wakey::wakey_linux::dhcp::list_local_observations_from_path(&config.observation_store_path)
.await
.context("failed to read local observations")?;
if observations.is_empty() {
return Ok(());
}
@@ -215,7 +218,11 @@ pub fn next_backoff_ms(current_ms: u64, max_ms: u64) -> u64 {
current_ms.saturating_mul(2).min(cap)
}
async fn handle_server_message<S>(sink: &mut S, message: ServerMessage) -> Result<()>
async fn handle_server_message<S>(
config: &AgentConfig,
sink: &mut S,
message: ServerMessage,
) -> Result<()>
where
S: SinkExt<Message> + Unpin,
<S as futures_util::Sink<Message>>::Error: std::error::Error + Send + Sync + 'static,
@@ -227,7 +234,7 @@ where
} => {
let kind = command_kind(&command);
info!(request_id = %request_id, command = %kind, "received command from control-plane");
match dispatch_command(command).await {
match dispatch_command(command, config).await {
Ok(result) => {
info!(request_id = %request_id, command = %kind, "command execution completed");
send_command_result(sink, request_id, kind, result).await?;