diff --git a/wakey-control-plane/src/state/store/agent_devices.rs b/wakey-control-plane/src/state/store/agent_devices.rs index 39e6a73..3cf4672 100644 --- a/wakey-control-plane/src/state/store/agent_devices.rs +++ b/wakey-control-plane/src/state/store/agent_devices.rs @@ -5,6 +5,7 @@ use super::Store; use super::helpers::core::*; use crate::state::types::*; +mod read; mod snapshot; impl Store { @@ -95,242 +96,23 @@ impl Store { Ok(incoming_keys.len()) } - /// List all agent device rows with their child MAC/IP/hostname/fact rows. + /// List all agent devices and their child rows from one consistent read view. #[tracing::instrument(name = "db_list_fleet_devices", skip_all)] pub async fn list_agent_device_rows(&self) -> Result> { - let devices = sqlx::query_as!( - AgentDeviceRow, - r#"SELECT agent_id, device_key, presence, display_name, - first_seen_unix, last_seen_unix - FROM agent_devices - ORDER BY agent_id, device_key"#, - ) - .fetch_all(&self.pool) - .await - .context("failed listing agent devices")?; - - let macs = sqlx::query_as!( - AgentDeviceMacRow, - r#"SELECT agent_id as "agent_id!", device_key as "device_key!", mac as "mac!" - FROM agent_device_macs"#, - ) - .fetch_all(&self.pool) - .await - .context("failed listing agent device macs")?; - - let ips = sqlx::query_as!( - AgentDeviceIpRow, - r#"SELECT agent_id as "agent_id!", device_key as "device_key!", ip as "ip!" - FROM agent_device_ips"#, - ) - .fetch_all(&self.pool) - .await - .context("failed listing agent device ips")?; - - let hostnames = sqlx::query_as!( - AgentDeviceHostnameRow, - r#"SELECT agent_id as "agent_id!", device_key as "device_key!", hostname as "hostname!" - FROM agent_device_hostnames"#, - ) - .fetch_all(&self.pool) - .await - .context("failed listing agent device hostnames")?; - - let endpoints = sqlx::query_as::<_, AgentDeviceEndpointRow>( - r#"SELECT agent_id, device_key, endpoint_key, source, mac, ip, - hostname, interface, presence, first_seen_unix, last_seen_unix - FROM agent_device_endpoints"#, - ) - .fetch_all(&self.pool) - .await - .context("failed listing agent device endpoints")?; - - let facts = sqlx::query_as!( - AgentDeviceFactRow, - r#"SELECT agent_id as "agent_id!", device_key as "device_key!", fact_json as "fact_json!" - FROM agent_device_facts"#, - ) - .fetch_all(&self.pool) - .await - .context("failed listing agent device facts")?; - - Ok(assemble_device_rows( - devices, macs, ips, hostnames, endpoints, facts, - )) + read::load(&self.pool, None).await } - /// List agent device rows for a single agent. + /// List one agent's devices using the same assembly path as fleet reads. #[allow(dead_code)] #[tracing::instrument(name = "db_list_agent_devices", skip_all)] pub async fn list_agent_device_rows_for_agent( &self, agent_id: &str, ) -> Result> { - let devices = sqlx::query_as!( - AgentDeviceRow, - r#"SELECT agent_id, device_key, presence, display_name, - first_seen_unix, last_seen_unix - FROM agent_devices - WHERE agent_id = ?1 - ORDER BY device_key"#, - agent_id, - ) - .fetch_all(&self.pool) - .await - .context("failed listing agent devices for agent")?; - - let macs = sqlx::query_as!( - AgentDeviceMacRow, - r#"SELECT agent_id as "agent_id!", device_key as "device_key!", mac as "mac!" - FROM agent_device_macs - WHERE agent_id = ?1"#, - agent_id, - ) - .fetch_all(&self.pool) - .await - .context("failed listing agent device macs for agent")?; - - let ips = sqlx::query_as!( - AgentDeviceIpRow, - r#"SELECT agent_id as "agent_id!", device_key as "device_key!", ip as "ip!" - FROM agent_device_ips - WHERE agent_id = ?1"#, - agent_id, - ) - .fetch_all(&self.pool) - .await - .context("failed listing agent device ips for agent")?; - - let hostnames = sqlx::query_as!( - AgentDeviceHostnameRow, - r#"SELECT agent_id as "agent_id!", device_key as "device_key!", hostname as "hostname!" - FROM agent_device_hostnames - WHERE agent_id = ?1"#, - agent_id, - ) - .fetch_all(&self.pool) - .await - .context("failed listing agent device hostnames for agent")?; - - let endpoints = sqlx::query_as::<_, AgentDeviceEndpointRow>( - r#"SELECT agent_id, device_key, endpoint_key, source, mac, ip, - hostname, interface, presence, first_seen_unix, last_seen_unix - FROM agent_device_endpoints - WHERE agent_id = ?1"#, - ) - .bind(agent_id) - .fetch_all(&self.pool) - .await - .context("failed listing agent device endpoints for agent")?; - - let facts = sqlx::query_as!( - AgentDeviceFactRow, - r#"SELECT agent_id as "agent_id!", device_key as "device_key!", fact_json as "fact_json!" - FROM agent_device_facts - WHERE agent_id = ?1"#, - agent_id, - ) - .fetch_all(&self.pool) - .await - .context("failed listing agent device facts for agent")?; - - Ok(assemble_device_rows( - devices, macs, ips, hostnames, endpoints, facts, - )) + read::load(&self.pool, Some(agent_id)).await } } -fn assemble_device_rows( - devices: Vec, - macs: Vec, - ips: Vec, - hostnames: Vec, - endpoints: Vec, - facts: Vec, -) -> Vec { - use std::collections::BTreeMap; - - let mut mac_map: BTreeMap<(&str, &str), Vec<&AgentDeviceMacRow>> = BTreeMap::new(); - for row in &macs { - mac_map - .entry((row.agent_id.as_str(), row.device_key.as_str())) - .or_default() - .push(row); - } - let mut ip_map: BTreeMap<(&str, &str), Vec<&AgentDeviceIpRow>> = BTreeMap::new(); - for row in &ips { - ip_map - .entry((row.agent_id.as_str(), row.device_key.as_str())) - .or_default() - .push(row); - } - let mut hostname_map: BTreeMap<(&str, &str), Vec> = BTreeMap::new(); - for row in &hostnames { - hostname_map - .entry((row.agent_id.as_str(), row.device_key.as_str())) - .or_default() - .push(row.hostname.clone()); - } - let mut endpoint_map: BTreeMap<(&str, &str), Vec> = BTreeMap::new(); - for row in &endpoints { - if let Some(endpoint) = row.to_endpoint() { - endpoint_map - .entry((row.agent_id.as_str(), row.device_key.as_str())) - .or_default() - .push(endpoint); - } - } - let mut fact_map: BTreeMap<(&str, &str), Vec> = BTreeMap::new(); - for row in &facts { - fact_map - .entry((row.agent_id.as_str(), row.device_key.as_str())) - .or_default() - .push(row.fact_json.clone()); - } - - devices - .into_iter() - .map(|device| { - let key = (device.agent_id.as_str(), device.device_key.as_str()); - let macs: Vec = mac_map - .get(&key) - .into_iter() - .flatten() - .filter_map(|row| match macaddr::MacAddr::try_from(*row) { - Ok(mac) => Some(mac), - Err(e) => { - ::tracing::warn!(error = %e, agent_id = %device.agent_id, device_key = %device.device_key, raw_mac = %row.mac, "failed to parse mac from agent_device_macs row"); - None - } - }) - .collect(); - let ips: Vec = ip_map - .get(&key) - .into_iter() - .flatten() - .filter_map(|row| match std::net::IpAddr::try_from(*row) { - Ok(ip) => Some(ip), - Err(e) => { - ::tracing::warn!(error = %e, agent_id = %device.agent_id, device_key = %device.device_key, raw_ip = %row.ip, "failed to parse ip from agent_device_ips row"); - None - } - }) - .collect(); - let hostnames = hostname_map.get(&key).cloned().unwrap_or_default(); - let endpoints = endpoint_map.get(&key).cloned().unwrap_or_default(); - let facts = fact_map.get(&key).cloned().unwrap_or_default(); - AgentDeviceWithChildren { - macs, - ips, - hostnames, - endpoints, - facts, - device, - } - }) - .collect() -} - pub fn device_key_from_id(device_id: &DeviceId) -> String { match device_id { DeviceId::Mac(mac) => format!("mac:{}", mac.to_string().to_ascii_lowercase()), diff --git a/wakey-control-plane/src/state/store/agent_devices/read.rs b/wakey-control-plane/src/state/store/agent_devices/read.rs new file mode 100644 index 0000000..eddd135 --- /dev/null +++ b/wakey-control-plane/src/state/store/agent_devices/read.rs @@ -0,0 +1,270 @@ +//! Fleet and single-agent reads share one sequence and a single SQLite read view. + +use crate::state::types::*; +use anyhow::{Context, Result}; +use sqlx::{SqliteConnection, SqlitePool}; + +mod assembly; +use assembly::DeviceRows; + +/// A read transaction prevents child rows from belonging to a newer snapshot +/// than their parents. Keep scoped SQL explicit so agent reads use indexed filters. +pub(super) async fn load( + pool: &SqlitePool, + agent_id: Option<&str>, +) -> Result> { + let mut tx = pool + .begin() + .await + .context("failed starting device read transaction")?; + let rows = DeviceRows { + devices: load_devices(&mut tx, agent_id).await?, + macs: load_macs(&mut tx, agent_id).await?, + ips: load_ips(&mut tx, agent_id).await?, + hostnames: load_hostnames(&mut tx, agent_id).await?, + endpoints: load_endpoints(&mut tx, agent_id).await?, + facts: load_facts(&mut tx, agent_id).await?, + }; + tx.commit() + .await + .context("failed completing device read transaction")?; + Ok(rows.into_devices()) +} + +async fn load_devices( + connection: &mut SqliteConnection, + agent_id: Option<&str>, +) -> Result> { + match agent_id { + Some(agent_id) => sqlx::query_as!( + AgentDeviceRow, + r#"SELECT agent_id, device_key, presence, display_name, + first_seen_unix, last_seen_unix + FROM agent_devices + WHERE agent_id = ?1 + ORDER BY device_key"#, + agent_id, + ) + .fetch_all(&mut *connection) + .await + .context("failed listing agent devices for agent"), + None => sqlx::query_as!( + AgentDeviceRow, + r#"SELECT agent_id, device_key, presence, display_name, + first_seen_unix, last_seen_unix + FROM agent_devices + ORDER BY agent_id, device_key"#, + ) + .fetch_all(&mut *connection) + .await + .context("failed listing agent devices"), + } +} + +async fn load_macs( + connection: &mut SqliteConnection, + agent_id: Option<&str>, +) -> Result> { + match agent_id { + Some(agent_id) => sqlx::query_as!( + AgentDeviceMacRow, + r#"SELECT agent_id as "agent_id!", device_key as "device_key!", mac as "mac!" + FROM agent_device_macs + WHERE agent_id = ?1"#, + agent_id, + ) + .fetch_all(&mut *connection) + .await + .context("failed listing agent device macs for agent"), + None => sqlx::query_as!( + AgentDeviceMacRow, + r#"SELECT agent_id as "agent_id!", device_key as "device_key!", mac as "mac!" + FROM agent_device_macs"#, + ) + .fetch_all(&mut *connection) + .await + .context("failed listing agent device macs"), + } +} + +async fn load_ips( + connection: &mut SqliteConnection, + agent_id: Option<&str>, +) -> Result> { + match agent_id { + Some(agent_id) => sqlx::query_as!( + AgentDeviceIpRow, + r#"SELECT agent_id as "agent_id!", device_key as "device_key!", ip as "ip!" + FROM agent_device_ips + WHERE agent_id = ?1"#, + agent_id, + ) + .fetch_all(&mut *connection) + .await + .context("failed listing agent device ips for agent"), + None => sqlx::query_as!( + AgentDeviceIpRow, + r#"SELECT agent_id as "agent_id!", device_key as "device_key!", ip as "ip!" + FROM agent_device_ips"#, + ) + .fetch_all(&mut *connection) + .await + .context("failed listing agent device ips"), + } +} + +async fn load_hostnames( + connection: &mut SqliteConnection, + agent_id: Option<&str>, +) -> Result> { + match agent_id { + Some(agent_id) => sqlx::query_as!( + AgentDeviceHostnameRow, + r#"SELECT agent_id as "agent_id!", device_key as "device_key!", hostname as "hostname!" + FROM agent_device_hostnames + WHERE agent_id = ?1"#, + agent_id, + ) + .fetch_all(&mut *connection) + .await + .context("failed listing agent device hostnames for agent"), + None => sqlx::query_as!( + AgentDeviceHostnameRow, + r#"SELECT agent_id as "agent_id!", device_key as "device_key!", hostname as "hostname!" + FROM agent_device_hostnames"#, + ) + .fetch_all(&mut *connection) + .await + .context("failed listing agent device hostnames"), + } +} + +async fn load_endpoints( + connection: &mut SqliteConnection, + agent_id: Option<&str>, +) -> Result> { + match agent_id { + Some(agent_id) => sqlx::query_as::<_, AgentDeviceEndpointRow>( + r#"SELECT agent_id, device_key, endpoint_key, source, mac, ip, + hostname, interface, presence, first_seen_unix, last_seen_unix + FROM agent_device_endpoints + WHERE agent_id = ?1"#, + ) + .bind(agent_id) + .fetch_all(&mut *connection) + .await + .context("failed listing agent device endpoints for agent"), + None => sqlx::query_as::<_, AgentDeviceEndpointRow>( + r#"SELECT agent_id, device_key, endpoint_key, source, mac, ip, + hostname, interface, presence, first_seen_unix, last_seen_unix + FROM agent_device_endpoints"#, + ) + .fetch_all(&mut *connection) + .await + .context("failed listing agent device endpoints"), + } +} + +async fn load_facts( + connection: &mut SqliteConnection, + agent_id: Option<&str>, +) -> Result> { + match agent_id { + Some(agent_id) => { +sqlx::query_as!( + AgentDeviceFactRow, + r#"SELECT agent_id as "agent_id!", device_key as "device_key!", fact_json as "fact_json!" + FROM agent_device_facts + WHERE agent_id = ?1"#, + agent_id, + ) + .fetch_all(&mut *connection) + .await + .context("failed listing agent device facts for agent") + } + None => { +sqlx::query_as!( + AgentDeviceFactRow, + r#"SELECT agent_id as "agent_id!", device_key as "device_key!", fact_json as "fact_json!" + FROM agent_device_facts"#, + ) + .fetch_all(&mut *connection) + .await + .context("failed listing agent device facts") + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::state::store::helpers::test_helpers::TestStore; + use wakey_core::{Device, NeighborEntry, NeighborState}; + + fn device(ip: &str) -> Device { + Device::from_parts( + vec![NeighborEntry { + ip: ip.parse().unwrap(), + mac: Some("aa:bb:cc:dd:ee:01".parse().unwrap()), + dev: Some("br-lan".into()), + state: NeighborState::Reachable, + }], + vec![], + ) + } + + #[tokio::test] + async fn scoped_reads_match_fleet_rows_without_leaking_other_agents() { + let ts = TestStore::new().await; + let store = ts.store(); + for (agent, ip) in [("a", "192.168.1.10"), ("b", "192.168.1.20")] { + store + .replace_agent_device_snapshot(agent, &[device(ip)]) + .await + .unwrap(); + } + let fleet = store.list_agent_device_rows().await.unwrap(); + assert_eq!(fleet.len(), 2); + for agent in ["a", "b", "missing"] { + let scoped = store.list_agent_device_rows_for_agent(agent).await.unwrap(); + let expected = fleet + .iter() + .filter(|row| row.device.agent_id == agent) + .cloned() + .collect::>(); + // Compare every parent/child field, including timestamps and endpoint metadata. + assert_eq!(format!("{scoped:?}"), format!("{expected:?}")); + } + } + + #[tokio::test] + async fn child_reads_share_the_parent_snapshot_during_a_concurrent_write() { + let ts = TestStore::new().await; + let store = ts.store(); + store + .replace_agent_device_snapshot("a", &[device("192.168.1.10")]) + .await + .unwrap(); + let mut tx = store.pool.begin().await.unwrap(); + let parents = load_devices(&mut tx, Some("a")).await.unwrap(); + assert_eq!(parents.len(), 1); + + // In WAL mode the writer can commit while the reader still sees the + // original snapshot. No sleeps or timing assumptions are needed. + store + .replace_agent_device_snapshot("a", &[device("192.168.1.20")]) + .await + .unwrap(); + let ips = load_ips(&mut tx, Some("a")).await.unwrap(); + let endpoints = load_endpoints(&mut tx, Some("a")).await.unwrap(); + assert_eq!(ips[0].ip, "192.168.1.10"); + assert_eq!(endpoints[0].ip.as_deref(), Some("192.168.1.10")); + tx.commit().await.unwrap(); + + let current = store.list_agent_device_rows_for_agent("a").await.unwrap(); + assert_eq!( + current[0].ips, + ["192.168.1.20".parse::().unwrap()] + ); + } +} diff --git a/wakey-control-plane/src/state/store/agent_devices/read/assembly.rs b/wakey-control-plane/src/state/store/agent_devices/read/assembly.rs new file mode 100644 index 0000000..def6a65 --- /dev/null +++ b/wakey-control-plane/src/state/store/agent_devices/read/assembly.rs @@ -0,0 +1,198 @@ +//! Attach decoded children to their parent using the full agent/device identity. + +use std::collections::BTreeMap; + +use crate::state::types::*; + +#[derive(Default)] +pub(super) struct DeviceRows { + pub(super) devices: Vec, + pub(super) macs: Vec, + pub(super) ips: Vec, + pub(super) hostnames: Vec, + pub(super) endpoints: Vec, + pub(super) facts: Vec, +} + +type DeviceIndex = BTreeMap<(String, String), AgentDeviceWithChildren>; + +impl DeviceRows { + pub(super) fn into_devices(self) -> Vec { + let mut index = self + .devices + .into_iter() + .map(|device| { + let key = (device.agent_id.clone(), device.device_key.clone()); + ( + key, + AgentDeviceWithChildren { + device, + macs: Vec::new(), + ips: Vec::new(), + hostnames: Vec::new(), + endpoints: Vec::new(), + facts: Vec::new(), + }, + ) + }) + .collect::(); + add_macs(&mut index, self.macs); + add_ips(&mut index, self.ips); + add_hostnames(&mut index, self.hostnames); + add_endpoints(&mut index, self.endpoints); + add_facts(&mut index, self.facts); + // Matches the SQL order: agent_id/device_key for fleet reads, and + // device_key within one agent. Child row order is left unchanged. + index.into_values().collect() + } +} + +fn add_macs(index: &mut DeviceIndex, rows: Vec) { + for row in rows { + let mac = match macaddr::MacAddr::try_from(&row) { + Ok(mac) => mac, + Err(error) => { + tracing::warn!(%error, agent_id = %row.agent_id, device_key = %row.device_key, + raw_mac = %row.mac, "failed to parse mac from agent_device_macs row"); + continue; + } + }; + if let Some(device) = index.get_mut(&(row.agent_id, row.device_key)) { + device.macs.push(mac); + } + } +} + +fn add_ips(index: &mut DeviceIndex, rows: Vec) { + for row in rows { + let ip = match std::net::IpAddr::try_from(&row) { + Ok(ip) => ip, + Err(error) => { + tracing::warn!(%error, agent_id = %row.agent_id, device_key = %row.device_key, + raw_ip = %row.ip, "failed to parse ip from agent_device_ips row"); + continue; + } + }; + if let Some(device) = index.get_mut(&(row.agent_id, row.device_key)) { + device.ips.push(ip); + } + } +} + +fn add_hostnames(index: &mut DeviceIndex, rows: Vec) { + for row in rows { + if let Some(device) = index.get_mut(&(row.agent_id, row.device_key)) { + device.hostnames.push(row.hostname); + } + } +} + +fn add_endpoints(index: &mut DeviceIndex, rows: Vec) { + for row in rows { + if let Some(endpoint) = row.to_endpoint() + && let Some(device) = index.get_mut(&(row.agent_id, row.device_key)) + { + device.endpoints.push(endpoint); + } + } +} + +fn add_facts(index: &mut DeviceIndex, rows: Vec) { + for row in rows { + if let Some(device) = index.get_mut(&(row.agent_id, row.device_key)) { + device.facts.push(row.fact_json); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn parent(agent: &str, key: &str) -> AgentDeviceRow { + AgentDeviceRow { + agent_id: agent.into(), + device_key: key.into(), + presence: "unknown".into(), + display_name: None, + first_seen_unix: 1, + last_seen_unix: 2, + } + } + + #[test] + fn assembly_keeps_agent_boundaries_order_and_empty_parents() { + let rows = DeviceRows { + devices: vec![ + parent("b", "shared"), + parent("a", "shared"), + parent("a", "empty"), + ], + hostnames: vec![ + AgentDeviceHostnameRow { + agent_id: "b".into(), + device_key: "shared".into(), + hostname: "second".into(), + }, + AgentDeviceHostnameRow { + agent_id: "a".into(), + device_key: "shared".into(), + hostname: "first".into(), + }, + AgentDeviceHostnameRow { + agent_id: "missing".into(), + device_key: "shared".into(), + hostname: "orphan".into(), + }, + ], + facts: vec![AgentDeviceFactRow { + agent_id: "a".into(), + device_key: "shared".into(), + fact_json: "{\"action\":\"remove\"}".into(), + }], + ..Default::default() + } + .into_devices(); + assert_eq!(rows.len(), 3); + assert_eq!(rows[0].device.device_key, "empty"); + assert!(rows[0].hostnames.is_empty()); + assert_eq!(rows[1].hostnames, ["first"]); + assert_eq!(rows[1].facts, ["{\"action\":\"remove\"}"]); + assert_eq!(rows[2].hostnames, ["second"]); + assert!(rows[2].facts.is_empty()); + } + + #[test] + fn malformed_children_do_not_hide_valid_children_or_the_parent() { + let rows = DeviceRows { + devices: vec![parent("a", "pc")], + macs: ["not-a-mac", "aa:bb:cc:dd:ee:01"] + .into_iter() + .map(|mac| AgentDeviceMacRow { + agent_id: "a".into(), + device_key: "pc".into(), + mac: mac.into(), + }) + .collect(), + ips: ["not-an-ip", "192.168.1.10"] + .into_iter() + .map(|ip| AgentDeviceIpRow { + agent_id: "a".into(), + device_key: "pc".into(), + ip: ip.into(), + }) + .collect(), + ..Default::default() + } + .into_devices(); + assert_eq!(rows.len(), 1); + assert_eq!( + rows[0].macs, + ["aa:bb:cc:dd:ee:01".parse::().unwrap()] + ); + assert_eq!( + rows[0].ips, + ["192.168.1.10".parse::().unwrap()] + ); + } +}