The start of greatness is hella rocky

wakey-agent/Cargo.toml

@@ -0,0 +1,22 @@
+[package]
+name = "wakey-agent"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+anyhow = "1"
+clap = { version = "4", features = ["derive"] }
+futures-util = "0.3"
+http = "1"
+macaddr = { version = "1", features = ["serde", "serde_std"] }
+reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] }
+serde = { version = "1", features = ["derive"] }
+serde_json = "1"
+tokio = { version = "1", features = ["fs", "macros", "rt-multi-thread", "time", "signal", "sync"] }
+tokio-tungstenite = { version = "0.24", features = ["rustls-tls-webpki-roots"] }
+toml = "0.8"
+tracing = "0.1"
+tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt"] }
+url = "2"
+wakey = { path = ".." }
+wakey-core = { path = "../wakey-core" }

wakey-agent/src/config.rs

@@ -0,0 +1,75 @@
+use anyhow::{Context, Result};
+use serde::{Deserialize, Serialize};
+use std::path::Path;
+
+pub const DEFAULT_CONFIG_PATH: &str = "/etc/wakey-agent/config.toml";
+
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
+pub struct AgentConfig {
+    pub server_url: String,
+    pub agent_id: String,
+    pub agent_token: String,
+    #[serde(default = "default_reconnect_base_ms")]
+    pub reconnect_base_ms: u64,
+    #[serde(default = "default_reconnect_max_ms")]
+    pub reconnect_max_ms: u64,
+}
+
+const fn default_reconnect_base_ms() -> u64 {
+    1_000
+}
+
+const fn default_reconnect_max_ms() -> u64 {
+    30_000
+}
+
+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()))?;
+    toml::from_str(&content)
+        .with_context(|| format!("failed to parse agent config {}", path.display()))
+}
+
+pub fn save_config(path: &Path, config: &AgentConfig) -> Result<()> {
+    if let Some(parent) = path.parent() {
+        std::fs::create_dir_all(parent)
+            .with_context(|| format!("failed to create config dir {}", parent.display()))?;
+    }
+
+    let content = toml::to_string_pretty(config).context("failed to serialize agent config")?;
+    let tmp = path.with_extension("toml.tmp");
+    std::fs::write(&tmp, content)
+        .with_context(|| format!("failed to write temp config {}", tmp.display()))?;
+    std::fs::rename(&tmp, path).with_context(|| {
+        format!(
+            "failed to move temp config {} into {}",
+            tmp.display(),
+            path.display()
+        )
+    })?;
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn config_roundtrip() {
+        let dir = std::env::temp_dir().join(format!("wakey-agent-config-{}", std::process::id()));
+        let path = dir.join("config.toml");
+        let config = AgentConfig {
+            server_url: "https://example.com".into(),
+            agent_id: "agent-1".into(),
+            agent_token: "secret".into(),
+            reconnect_base_ms: 123,
+            reconnect_max_ms: 456,
+        };
+
+        save_config(&path, &config).expect("save");
+        let loaded = load_config(&path).expect("load");
+        assert_eq!(loaded, config);
+        let _ = std::fs::remove_file(&path);
+        let _ = std::fs::remove_dir_all(&dir);
+    }
+}

wakey-agent/src/dispatch.rs

@@ -0,0 +1,114 @@
+use anyhow::Result;
+use tracing::{debug, instrument};
+
+use crate::protocol::{
+    AgentCommand, CommandResult, DevsRequest, InventoryRequest, LeasesRequest, StatusRequest,
+    WakeRequest,
+};
+
+#[instrument(skip_all)]
+pub async fn dispatch_command(command: AgentCommand) -> Result<CommandResult> {
+    match command {
+        AgentCommand::Status(req) => dispatch_status(req).await,
+        AgentCommand::Leases(req) => dispatch_leases(req).await,
+        AgentCommand::Devs(req) => dispatch_devs(req).await,
+        AgentCommand::Inventory(req) => dispatch_inventory(req).await,
+        AgentCommand::Wake(req) => dispatch_wake(req).await,
+    }
+}
+
+async fn dispatch_status(req: StatusRequest) -> Result<CommandResult> {
+    let query = req.into_device_query();
+    let status = if query.name.is_some()
+        && query.filter.ips.is_empty()
+        && query.filter.devs.is_empty()
+        && query.filter.nuds.is_empty()
+        && query.filter.macs.is_empty()
+    {
+        wakey::get_status_for_input(query.name.clone().unwrap_or_default()).await?
+    } else {
+        wakey::get_status(query).await?
+    };
+    debug!(rows = status.table.len(), "dispatched status command");
+    Ok(CommandResult::Status(status))
+}
+
+async fn dispatch_leases(req: LeasesRequest) -> Result<CommandResult> {
+    let leases = wakey::get_leases(wakey_core::LeaseQuery {
+        include_state: req.include_state,
+    })
+    .await?;
+    Ok(CommandResult::Leases(leases))
+}
+
+async fn dispatch_devs(req: DevsRequest) -> Result<CommandResult> {
+    let mut devs = if let Some(name) = &req.dev {
+        wakey::get_interface_summary(name).await?.into_iter().collect()
+    } else {
+        wakey::get_interface_summaries().await?
+    };
+    if req.up_only {
+        devs.retain(|dev| dev.operstate == "up");
+    }
+    Ok(CommandResult::Devs(devs))
+}
+
+async fn dispatch_inventory(req: InventoryRequest) -> Result<CommandResult> {
+    let inventory = wakey::inventory(req.into_device_query()).await?;
+    Ok(CommandResult::Inventory(inventory))
+}
+
+async fn dispatch_wake(req: WakeRequest) -> Result<CommandResult> {
+    validate_wake_request(&req)?;
+    let result = match (req.query, req.mac, req.ip) {
+        (Some(query), None, None) => wakey::wake_from_query(query).await?,
+        (None, Some(mac), ip) => wakey::wake_explicit(mac, ip).await?,
+        _ => unreachable!("wake request validated before dispatch"),
+    };
+    Ok(CommandResult::Wake(result))
+}
+
+pub fn validate_wake_request(req: &WakeRequest) -> Result<()> {
+    let has_query = req.query.is_some();
+    let has_mac = req.mac.is_some();
+    let has_ip = req.ip.is_some();
+
+    if has_ip && !has_mac {
+        anyhow::bail!("wake request `ip` requires `mac`");
+    }
+    if has_query && (has_mac || has_ip) {
+        anyhow::bail!("wake query mode and explicit mode are mutually exclusive");
+    }
+    if !has_query && !has_mac {
+        anyhow::bail!("wake request requires either `query` or `mac`");
+    }
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::protocol::WakeRequest;
+
+    #[test]
+    fn wake_rejects_ip_without_mac() {
+        let err = validate_wake_request(&WakeRequest {
+            query: None,
+            mac: None,
+            ip: Some("192.168.1.1".parse().expect("ip")),
+        })
+        .expect_err("should reject");
+        assert!(err.to_string().contains("requires `mac`"));
+    }
+
+    #[test]
+    fn wake_rejects_mixed_mode() {
+        let err = validate_wake_request(&WakeRequest {
+            query: Some("pc".into()),
+            mac: Some("aa:bb:cc:dd:ee:ff".parse().expect("mac")),
+            ip: None,
+        })
+        .expect_err("should reject");
+        assert!(err.to_string().contains("mutually exclusive"));
+    }
+}

wakey-agent/src/enroll.rs

@@ -0,0 +1,70 @@
+use anyhow::{Context, Result};
+use reqwest::StatusCode;
+use serde::{Deserialize, Serialize};
+use std::path::Path;
+use tracing::info;
+
+use crate::config::{AgentConfig, save_config};
+
+#[derive(Debug, Serialize)]
+struct EnrollRequest<'a> {
+    enroll_token: &'a str,
+}
+
+#[derive(Debug, Deserialize)]
+struct EnrollResponse {
+    agent_id: String,
+    agent_token: String,
+    server_url: Option<String>,
+}
+
+pub async fn enroll(server_url: &str, enroll_token: &str, config_path: &Path) -> Result<AgentConfig> {
+    let server_url = normalize_server_url(server_url);
+    let endpoint = format!("{server_url}/api/v1/agents/enroll");
+    let client = reqwest::Client::new();
+    let response = client
+        .post(endpoint)
+        .json(&EnrollRequest { enroll_token })
+        .send()
+        .await
+        .context("failed to call enrollment endpoint")?;
+
+    if response.status() != StatusCode::OK {
+        let status = response.status();
+        let body = response
+            .text()
+            .await
+            .unwrap_or_else(|_| "<unable to read enrollment error body>".into());
+        anyhow::bail!("enrollment failed with {status}: {body}");
+    }
+
+    let payload: EnrollResponse = response
+        .json()
+        .await
+        .context("failed to decode enrollment response")?;
+
+    let config = AgentConfig {
+        server_url: payload.server_url.unwrap_or(server_url),
+        agent_id: payload.agent_id,
+        agent_token: payload.agent_token,
+        reconnect_base_ms: 1_000,
+        reconnect_max_ms: 30_000,
+    };
+    save_config(config_path, &config)?;
+    info!(config_path = %config_path.display(), "wrote agent config");
+    Ok(config)
+}
+
+pub fn normalize_server_url(server_url: &str) -> String {
+    server_url.trim_end_matches('/').to_string()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::normalize_server_url;
+
+    #[test]
+    fn normalize_server_url_trims_slash() {
+        assert_eq!(normalize_server_url("https://example.com/"), "https://example.com");
+    }
+}

wakey-agent/src/main.rs

@@ -0,0 +1,72 @@
+mod config;
+mod dispatch;
+mod enroll;
+mod protocol;
+mod session;
+mod tracing;
+
+use std::path::PathBuf;
+
+use clap::{ArgAction, Args, Parser, Subcommand};
+use ::tracing::info;
+
+#[derive(Parser)]
+#[command(name = "wakey-agent")]
+#[command(version, about = "Outbound control-plane agent for Wakey")]
+struct Cli {
+    /// Increase log verbosity. Use `-v` for debug and `-vv` for trace.
+    #[arg(short = 'v', long = "verbose", action = ArgAction::Count, global = true)]
+    verbose: u8,
+
+    #[command(subcommand)]
+    command: Command,
+}
+
+#[derive(Subcommand)]
+enum Command {
+    /// Start the agent daemon and maintain the outbound control-plane session.
+    Serve(ServeArgs),
+    /// Enroll this router with a control plane and write agent config.
+    Enroll(EnrollArgs),
+}
+
+#[derive(Args)]
+struct ServeArgs {
+    /// Path to the agent config file.
+    #[arg(long, default_value = config::DEFAULT_CONFIG_PATH)]
+    config: PathBuf,
+}
+
+#[derive(Args)]
+struct EnrollArgs {
+    /// Base HTTPS URL of the control plane.
+    #[arg(long)]
+    server_url: String,
+    /// One-time or short-lived enroll token provided by the control plane.
+    #[arg(long)]
+    enroll_token: String,
+    /// Path to the agent config file to write.
+    #[arg(long, default_value = config::DEFAULT_CONFIG_PATH)]
+    config: PathBuf,
+}
+
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+    let cli = Cli::parse();
+    tracing::init(cli.verbose);
+
+    match cli.command {
+        Command::Serve(args) => {
+            let config = config::load_config(&args.config)?;
+            info!(config_path = %args.config.display(), agent_id = %config.agent_id, "starting wakey-agent");
+            session::run(config).await?;
+        }
+        Command::Enroll(args) => {
+            let config = enroll::enroll(&args.server_url, &args.enroll_token, &args.config).await?;
+            println!("agent_id={}", config.agent_id);
+            println!("config={}", args.config.display());
+        }
+    }
+
+    Ok(())
+}

wakey-agent/src/protocol.rs

@@ -0,0 +1,173 @@
+use macaddr::MacAddr;
+use serde::{Deserialize, Serialize};
+use std::net::IpAddr;
+use wakey_core::{
+    DeviceFilters, DeviceInventory, DeviceQuery, DhcpLeaseWithState, InterfaceSummary, NeighborEntry,
+    Status, WakeResult,
+};
+use wakey_core::parse::mac;
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct StatusRequest {
+    pub query: Option<String>,
+    pub name: Option<String>,
+    #[serde(default)]
+    pub ips: Vec<IpAddr>,
+    #[serde(default)]
+    pub devs: Vec<String>,
+    #[serde(default)]
+    pub nuds: Vec<wakey_core::NeighborState>,
+    #[serde(default)]
+    #[serde(with = "mac::vec_mac")]
+    pub macs: Vec<MacAddr>,
+}
+
+impl StatusRequest {
+    pub fn into_device_query(self) -> DeviceQuery {
+        if let Some(query) = self.query.as_ref()
+            && self.name.is_none()
+            && self.ips.is_empty()
+            && self.devs.is_empty()
+            && self.nuds.is_empty()
+            && self.macs.is_empty()
+        {
+            return DeviceQuery {
+                name: Some(query.clone()),
+                ..Default::default()
+            };
+        }
+
+        DeviceQuery {
+            name: self.name.or(self.query),
+            filter: DeviceFilters {
+                ips: self.ips,
+                devs: self.devs,
+                nuds: self.nuds,
+                macs: self.macs,
+            },
+        }
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct LeasesRequest {
+    #[serde(default)]
+    pub include_state: bool,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct DevsRequest {
+    pub dev: Option<String>,
+    #[serde(default)]
+    pub up_only: bool,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct InventoryRequest {
+    pub query: Option<String>,
+    pub name: Option<String>,
+    #[serde(default)]
+    pub ips: Vec<IpAddr>,
+    #[serde(default)]
+    pub devs: Vec<String>,
+    #[serde(default)]
+    pub nuds: Vec<wakey_core::NeighborState>,
+    #[serde(default)]
+    #[serde(with = "mac::vec_mac")]
+    pub macs: Vec<MacAddr>,
+}
+
+impl InventoryRequest {
+    pub fn into_device_query(self) -> DeviceQuery {
+        StatusRequest {
+            query: self.query,
+            name: self.name,
+            ips: self.ips,
+            devs: self.devs,
+            nuds: self.nuds,
+            macs: self.macs,
+        }
+        .into_device_query()
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct WakeRequest {
+    pub query: Option<String>,
+    #[serde(default)]
+    #[serde(with = "mac::option_mac")]
+    pub mac: Option<MacAddr>,
+    #[serde(default)]
+    pub ip: Option<IpAddr>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(tag = "kind", rename_all = "snake_case")]
+pub enum AgentCommand {
+    Status(StatusRequest),
+    Leases(LeasesRequest),
+    Devs(DevsRequest),
+    Inventory(InventoryRequest),
+    Wake(WakeRequest),
+}
+
+#[derive(Debug, Serialize)]
+#[serde(tag = "kind", rename_all = "snake_case")]
+pub enum CommandResult {
+    Status(Status<NeighborEntry>),
+    Leases(Vec<DhcpLeaseWithState>),
+    Devs(Vec<InterfaceSummary>),
+    Inventory(DeviceInventory),
+    Wake(WakeResult),
+}
+
+#[derive(Debug, Serialize)]
+#[serde(tag = "type", rename_all = "snake_case")]
+pub enum ClientMessage {
+    Hello {
+        agent_id: String,
+    },
+    Auth {
+        agent_id: String,
+        agent_token: String,
+    },
+    Heartbeat {
+        agent_id: String,
+    },
+    Result {
+        request_id: String,
+        result: CommandResult,
+    },
+    Error {
+        request_id: String,
+        error: String,
+    },
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+#[serde(tag = "type", rename_all = "snake_case")]
+pub enum ServerMessage {
+    Command {
+        request_id: String,
+        command: AgentCommand,
+    },
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn command_serialization_is_stable() {
+        let msg = ServerMessage::Command {
+            request_id: "req-1".into(),
+            command: AgentCommand::Leases(LeasesRequest {
+                include_state: true,
+            }),
+        };
+
+        let json = serde_json::to_string(&msg).expect("serialize");
+        assert!(json.contains("\"request_id\":\"req-1\""));
+        assert!(json.contains("\"kind\":\"leases\""));
+    }
+}

wakey-agent/src/session.rs

@@ -0,0 +1,156 @@
+use anyhow::{Context, Result};
+use futures_util::{SinkExt, StreamExt};
+use tokio::time::{Duration, MissedTickBehavior, interval, sleep};
+use tokio_tungstenite::{connect_async, tungstenite::Message};
+use tracing::{debug, error, info, warn};
+
+use crate::config::AgentConfig;
+use crate::dispatch::dispatch_command;
+use crate::protocol::{ClientMessage, ServerMessage};
+
+pub async fn run(config: AgentConfig) -> Result<()> {
+    let mut backoff = config.reconnect_base_ms.max(100);
+    loop {
+        match run_once(&config).await {
+            Ok(()) => {
+                backoff = config.reconnect_base_ms.max(100);
+            }
+            Err(err) => {
+                warn!(error = %err, backoff_ms = backoff, "agent session ended; reconnecting");
+                sleep(Duration::from_millis(backoff)).await;
+                backoff = (backoff.saturating_mul(2)).min(config.reconnect_max_ms.max(backoff));
+            }
+        }
+    }
+}
+
+async fn run_once(config: &AgentConfig) -> Result<()> {
+    let ws_url = websocket_url(&config.server_url)?;
+    info!(%ws_url, agent_id = %config.agent_id, "connecting agent websocket");
+    let (stream, _) = connect_async(ws_url.as_str())
+        .await
+        .context("failed to connect websocket")?;
+    let (mut sink, mut source) = stream.split();
+
+    send_json(&mut sink, &ClientMessage::Hello {
+        agent_id: config.agent_id.clone(),
+    })
+    .await?;
+    send_json(
+        &mut sink,
+        &ClientMessage::Auth {
+            agent_id: config.agent_id.clone(),
+            agent_token: config.agent_token.clone(),
+        },
+    )
+    .await?;
+
+    let mut heartbeat = interval(Duration::from_secs(30));
+    heartbeat.set_missed_tick_behavior(MissedTickBehavior::Skip);
+
+    loop {
+        tokio::select! {
+            _ = heartbeat.tick() => {
+                send_json(&mut sink, &ClientMessage::Heartbeat {
+                    agent_id: config.agent_id.clone(),
+                }).await?;
+            }
+            maybe_msg = source.next() => {
+                let msg = match maybe_msg {
+                    Some(msg) => msg.context("websocket frame failed")?,
+                    None => anyhow::bail!("websocket closed by server"),
+                };
+
+                match msg {
+                    Message::Text(text) => {
+                        let message: ServerMessage = serde_json::from_str(&text)
+                            .context("failed to decode server message")?;
+                        handle_server_message(&mut sink, message).await?;
+                    }
+                    Message::Ping(payload) => {
+                        sink.send(Message::Pong(payload)).await.context("failed to send pong")?;
+                    }
+                    Message::Pong(_) => {}
+                    Message::Close(frame) => {
+                        anyhow::bail!("websocket closed: {:?}", frame);
+                    }
+                    Message::Binary(_) => {
+                        debug!("ignoring unexpected binary websocket frame");
+                    }
+                    Message::Frame(_) => {}
+                }
+            }
+        }
+    }
+}
+
+async fn handle_server_message<S>(sink: &mut S, message: ServerMessage) -> Result<()>
+where
+    S: SinkExt<Message> + Unpin,
+    <S as futures_util::Sink<Message>>::Error: std::error::Error + Send + Sync + 'static,
+{
+    match message {
+        ServerMessage::Command {
+            request_id,
+            command,
+        } => match dispatch_command(command).await {
+            Ok(result) => {
+                send_json(sink, &ClientMessage::Result { request_id, result }).await?;
+            }
+            Err(err) => {
+                error!(request_id, error = %err, "command dispatch failed");
+                send_json(
+                    sink,
+                    &ClientMessage::Error {
+                        request_id,
+                        error: err.to_string(),
+                    },
+                )
+                .await?;
+            }
+        },
+    }
+    Ok(())
+}
+
+async fn send_json<S>(sink: &mut S, message: &ClientMessage) -> Result<()>
+where
+    S: SinkExt<Message> + Unpin,
+    <S as futures_util::Sink<Message>>::Error: std::error::Error + Send + Sync + 'static,
+{
+    let payload = serde_json::to_string(message).context("failed to serialize websocket message")?;
+    sink.send(Message::Text(payload))
+        .await
+        .context("failed to send websocket message")?;
+    Ok(())
+}
+
+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() {
+        "http" => "ws",
+        "https" => "wss",
+        "ws" => "ws",
+        "wss" => "wss",
+        other => anyhow::bail!("unsupported server_url scheme `{other}`"),
+    };
+
+    let mut url = base;
+    url.set_scheme(scheme)
+        .map_err(|_| anyhow::anyhow!("failed to convert server_url scheme"))?;
+    url.set_path("/api/v1/agent/ws");
+    url.set_query(None);
+    url.set_fragment(None);
+    Ok(url)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::websocket_url;
+
+    #[test]
+    fn websocket_url_uses_expected_path() {
+        let url = websocket_url("https://example.com/control").expect("url");
+        assert_eq!(url.as_str(), "wss://example.com/api/v1/agent/ws");
+    }
+}

wakey-agent/src/tracing.rs

@@ -0,0 +1,20 @@
+use tracing_subscriber::{EnvFilter, fmt, layer::SubscriberExt, util::SubscriberInitExt};
+
+pub fn init(verbose: u8) {
+    let filter = EnvFilter::try_from_default_env()
+        .or_else(|_| EnvFilter::try_new(default_filter(verbose)))
+        .expect("static tracing filter should parse");
+
+    tracing_subscriber::registry()
+        .with(filter)
+        .with(fmt::layer())
+        .init();
+}
+
+fn default_filter(verbose: u8) -> &'static str {
+    match verbose {
+        0 => "wakey_agent=info,wakey=info",
+        1 => "wakey_agent=debug,wakey=debug",
+        _ => "wakey_agent=trace,wakey=debug",
+    }
+}