what a plan! what an impl! this is ahh

2026-04-11 15:25:18 +07:00
parent a36e9d058d
commit e2feaeab8c
43 changed files with 1090 additions and 1629 deletions
@@ -1,10 +1,11 @@
 use anyhow::{Context, Result};
 use serde::{Deserialize, Serialize};
+use std::fmt;
 use std::path::Path;

 pub const DEFAULT_CONFIG_PATH: &str = "/etc/wakey-agent/config.toml";

-#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
+#[derive(Clone, Serialize, Deserialize, PartialEq, Eq)]
 pub struct AgentConfig {
    pub server_url: String,
    pub agent_id: String,
@@ -15,6 +16,18 @@ pub struct AgentConfig {
    pub reconnect_max_ms: u64,
 }

+impl fmt::Debug for AgentConfig {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("AgentConfig")
+            .field("server_url", &self.server_url)
+            .field("agent_id", &self.agent_id)
+            .field("agent_token", &"<redacted>")
+            .field("reconnect_base_ms", &self.reconnect_base_ms)
+            .field("reconnect_max_ms", &self.reconnect_max_ms)
+            .finish()
+    }
+}
+
 const fn default_reconnect_base_ms() -> u64 {
    1_000
 }
@@ -69,6 +82,7 @@ mod tests {
        save_config(&path, &config).expect("save");
        let loaded = load_config(&path).expect("load");
        assert_eq!(loaded, config);
+        assert!(format!("{:?}", loaded).contains("<redacted>"));
        let _ = std::fs::remove_file(&path);
        let _ = std::fs::remove_dir_all(&dir);
    }
@@ -61,10 +61,77 @@ pub fn normalize_server_url(server_url: &str) -> String {

 #[cfg(test)]
 mod tests {
-    use super::normalize_server_url;
+    use super::*;
+    use std::io::{Read, Write};
+    use std::net::{SocketAddr, TcpListener};
+    use std::thread;
+
+    fn spawn_enroll_server(response_body: &'static str, status: &'static str) -> (String, thread::JoinHandle<()>) {
+        let listener = TcpListener::bind("127.0.0.1:0").expect("bind test listener");
+        let addr: SocketAddr = listener.local_addr().expect("local addr");
+        let handle = thread::spawn(move || {
+            let (mut stream, _) = listener.accept().expect("accept connection");
+
+            // Read until end-of-headers; body content is irrelevant for this test.
+            let mut buf = [0u8; 4096];
+            let mut req = Vec::new();
+            loop {
+                let n = stream.read(&mut buf).expect("read request");
+                if n == 0 {
+                    break;
+                }
+                req.extend_from_slice(&buf[..n]);
+                if req.windows(4).any(|w| w == b"\r\n\r\n") {
+                    break;
+                }
+            }
+
+            let response = format!(
+                "HTTP/1.1 {status}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
+                response_body.len(),
+                response_body
+            );
+            stream
+                .write_all(response.as_bytes())
+                .expect("write response");
+        });
+
+        (format!("http://{}", addr), handle)
+    }

    #[test]
    fn normalize_server_url_trims_slash() {
        assert_eq!(normalize_server_url("https://example.com/"), "https://example.com");
    }
+
+    #[tokio::test]
+    async fn enroll_response_persists_config() {
+        let response = r#"{"agent_id":"agent-123","agent_token":"token-xyz","server_url":"https://control.example.com"}"#;
+        let (server_url, handle) = spawn_enroll_server(response, "200 OK");
+
+        let dir = std::env::temp_dir().join(format!(
+            "wakey-agent-enroll-{}-{}",
+            std::process::id(),
+            std::time::SystemTime::now()
+                .duration_since(std::time::UNIX_EPOCH)
+                .expect("time")
+                .as_nanos()
+        ));
+        let path = dir.join("config.toml");
+
+        let config = enroll(&server_url, "enroll-abc", &path)
+            .await
+            .expect("enroll should succeed");
+
+        assert_eq!(config.agent_id, "agent-123");
+        assert_eq!(config.agent_token, "token-xyz");
+        assert_eq!(config.server_url, "https://control.example.com");
+
+        let persisted = crate::config::load_config(&path).expect("load persisted config");
+        assert_eq!(persisted, config);
+
+        handle.join().expect("server thread joined");
+        let _ = std::fs::remove_file(&path);
+        let _ = std::fs::remove_dir_all(&dir);
+    }
 }
@@ -0,0 +1 @@
+pub mod protocol;
@@ -1,5 +1,6 @@
 use macaddr::MacAddr;
 use serde::{Deserialize, Serialize};
+use std::fmt;
 use std::net::IpAddr;
 use wakey_core::{
    DeviceFilters, DeviceInventory, DeviceQuery, DhcpLeaseWithState, InterfaceSummary, NeighborEntry,
@@ -7,6 +8,65 @@ use wakey_core::{
 };
 use wakey_core::parse::mac;

+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct RequestId(String);
+
+impl RequestId {
+    pub fn as_str(&self) -> &str {
+        &self.0
+    }
+}
+
+impl TryFrom<String> for RequestId {
+    type Error = String;
+
+    fn try_from(value: String) -> Result<Self, Self::Error> {
+        if value.trim().is_empty() {
+            return Err("request_id must not be empty".into());
+        }
+        Ok(Self(value))
+    }
+}
+
+impl From<RequestId> for String {
+    fn from(value: RequestId) -> Self {
+        value.0
+    }
+}
+
+impl fmt::Display for RequestId {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
+impl Serialize for RequestId {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        serializer.serialize_str(&self.0)
+    }
+}
+
+impl<'de> Deserialize<'de> for RequestId {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        let raw = String::deserialize(deserializer)?;
+        RequestId::try_from(raw).map_err(serde::de::Error::custom)
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ErrorPayload {
+    pub code: String,
+    pub message: String,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub retryable: Option<bool>,
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct StatusRequest {
    pub query: Option<String>,
@@ -135,12 +195,12 @@ pub enum ClientMessage {
        agent_id: String,
    },
    Result {
-        request_id: String,
+        request_id: RequestId,
        result: CommandResult,
    },
    Error {
-        request_id: String,
-        error: String,
+        request_id: RequestId,
+        error: ErrorPayload,
    },
 }

@@ -148,7 +208,7 @@ pub enum ClientMessage {
 #[serde(tag = "type", rename_all = "snake_case")]
 pub enum ServerMessage {
    Command {
-        request_id: String,
+        request_id: RequestId,
        command: AgentCommand,
    },
 }
@@ -160,7 +220,7 @@ mod tests {
    #[test]
    fn command_serialization_is_stable() {
        let msg = ServerMessage::Command {
-            request_id: "req-1".into(),
+            request_id: RequestId::try_from("req-1".to_string()).expect("request id"),
            command: AgentCommand::Leases(LeasesRequest {
                include_state: true,
            }),
@@ -170,4 +230,10 @@ mod tests {
        assert!(json.contains("\"request_id\":\"req-1\""));
        assert!(json.contains("\"kind\":\"leases\""));
    }
+
+    #[test]
+    fn request_id_rejects_empty() {
+        let err = RequestId::try_from("   ".to_string()).expect_err("must fail");
+        assert!(err.contains("must not be empty"));
+    }
 }
@@ -6,7 +6,7 @@ use tracing::{debug, error, info, warn};

 use crate::config::AgentConfig;
 use crate::dispatch::dispatch_command;
-use crate::protocol::{ClientMessage, ServerMessage};
+use crate::protocol::{ClientMessage, ErrorPayload, ServerMessage};

 pub async fn run(config: AgentConfig) -> Result<()> {
    let mut backoff = config.reconnect_base_ms.max(100);
@@ -18,7 +18,7 @@ pub async fn run(config: AgentConfig) -> Result<()> {
            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));
+                backoff = next_backoff_ms(backoff, config.reconnect_max_ms);
            }
        }
    }
@@ -63,9 +63,16 @@ async fn run_once(config: &AgentConfig) -> Result<()> {

                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?;
+                        match serde_json::from_str::<ServerMessage>(&text) {
+                            Ok(message) => {
+                                handle_server_message(&mut sink, message).await?;
+                            }
+                            Err(err) => {
+                                // Allow the server to introduce extra frame types without
+                                // forcing reconnects for older agents.
+                                warn!(error = %err, payload = %text, "ignoring unknown server message");
+                            }
+                        }
                    }
                    Message::Ping(payload) => {
                        sink.send(Message::Pong(payload)).await.context("failed to send pong")?;
@@ -84,6 +91,11 @@ async fn run_once(config: &AgentConfig) -> Result<()> {
    }
 }

+pub fn next_backoff_ms(current_ms: u64, max_ms: u64) -> u64 {
+    let cap = max_ms.max(current_ms);
+    current_ms.saturating_mul(2).min(cap)
+}
+
 async fn handle_server_message<S>(sink: &mut S, message: ServerMessage) -> Result<()>
 where
    S: SinkExt<Message> + Unpin,
@@ -98,12 +110,16 @@ where
                send_json(sink, &ClientMessage::Result { request_id, result }).await?;
            }
            Err(err) => {
-                error!(request_id, error = %err, "command dispatch failed");
+                error!(request_id = %request_id, error = %err, "command dispatch failed");
                send_json(
                    sink,
                    &ClientMessage::Error {
                        request_id,
-                        error: err.to_string(),
+                        error: ErrorPayload {
+                            code: "command_dispatch_failed".into(),
+                            message: err.to_string(),
+                            retryable: None,
+                        },
                    },
                )
                .await?;
@@ -146,11 +162,23 @@ pub fn websocket_url(server_url: &str) -> Result<url::Url> {

 #[cfg(test)]
 mod tests {
-    use super::websocket_url;
+    use super::{next_backoff_ms, 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");
    }
+
+    #[test]
+    fn backoff_doubles_until_cap() {
+        assert_eq!(next_backoff_ms(1_000, 30_000), 2_000);
+        assert_eq!(next_backoff_ms(16_000, 30_000), 30_000);
+        assert_eq!(next_backoff_ms(30_000, 30_000), 30_000);
+    }
+
+    #[test]
+    fn backoff_never_shrinks_when_max_is_lower() {
+        assert_eq!(next_backoff_ms(8_000, 1_000), 8_000);
+    }
 }