lda/wakey
1
0
Fork 0
Code Issues 5 Pull Requests Actions Packages Projects Releases 18 Wiki Activity

commit full of PRs

Browse Source 
PRs full of commit

basically i used the ENTIRE mimo grant to convert a system to another. hope this one is good. may fail.
This commit is contained in:
lda
2026-05-03 00:48:11 +07:00 Verified
parent b5a28c801c
commit b850c218ec
112 changed files with 2407 additions and 4237 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wakey-control-plane/src/api/control.rs
+2 -7
View File
@@ -1,12 +1,11 @@
mod devices;
mod enroll;
mod fleet;
mod observations;
mod stats;
pub use devices::{
attach_device_identifier, attach_observation_identifier, create_known_device,
detach_device_identifier, forget_known_device, list_known_devices, merge_known_device,
attach_device_identifier, create_known_device, detach_device_identifier, forget_known_device,
list_known_devices, merge_known_device,
};
pub use enroll::{
EnrollTokenStatus, IssueEnrollTokenResponse, RevokeAgentResponse, RevokeEnrollTokenResponse,
@@ -14,8 +13,4 @@ pub use enroll::{
set_agent_nickname,
};
pub use fleet::{list_fleet_devices, refresh_fleet_devices, wake_fleet_device};
pub use observations::{
list_agent_observation_history, list_agent_observations, request_agent_observation_sync,
upload_agent_observations,
};
pub use stats::{StateStatsResponse, state_stats};
wakey-control-plane/src/api/control/devices.rs
-32
View File
@@ -25,11 +25,6 @@ pub struct DeviceIdentifierRequest {
pub value: String,
}
#[derive(Debug, Deserialize)]
pub struct AttachObservationIdentifierRequest {
pub observation_key: String,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct KnownDeviceResponse {
pub device_id: String,
@@ -171,33 +166,6 @@ pub async fn attach_device_identifier(
}
}
pub async fn attach_observation_identifier(
State(state): State<AppState>,
AxumPath(device_id): AxumPath<String>,
Json(req): Json<AttachObservationIdentifierRequest>,
) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
match state
.store
.attach_observation_identifier(&device_id, &req.observation_key)
.await
{
Ok(Some(device)) => Ok((StatusCode::OK, Json(known_device_response(device)))),
Ok(None) => Err(json_error(
StatusCode::NOT_FOUND,
"known_device_not_found",
"known device not found",
)),
Err(err) => {
warn!(error = %err, "failed to attach observation identifier");
Err(json_error(
StatusCode::BAD_REQUEST,
"attach_observation_identifier_failed",
&err.to_string(),
))
}
}
}
pub async fn detach_device_identifier(
State(state): State<AppState>,
AxumPath((device_id, identifier_key)): AxumPath<(String, String)>,
wakey-control-plane/src/api/control/fleet/build.rs
+131 -131
View File
@@ -1,6 +1,10 @@
use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::net::IpAddr;
use crate::state::{AgentDeviceObservation, DeviceIdentifier, KnownDevice, KnownDeviceSummary};
use macaddr::MacAddr;
use wakey_core::Presence;
use crate::state::{AgentDeviceWithChildren, DeviceIdentifier, KnownDevice, KnownDeviceSummary};
use super::types::{FleetDevice, FleetDeviceAgent, FleetWakeRoute, ListFleetDevicesQuery};
@@ -22,20 +26,20 @@ struct FleetAccumulator {
display_name: Option<String>,
known_device: Option<KnownDeviceSummary>,
pinned: bool,
ips: BTreeSet<String>,
macs: BTreeSet<String>,
ips: BTreeSet<IpAddr>,
macs: BTreeSet<MacAddr>,
hostnames: BTreeSet<String>,
sources: BTreeSet<String>,
agents: BTreeMap<String, FleetDeviceAgent>,
first_seen_unix: Option<u64>,
last_seen_unix: Option<u64>,
presence_rank: u8,
presence: Presence,
routes: BTreeMap<String, FleetWakeRoute>,
}
pub(crate) fn build_fleet_devices(
known_devices: Vec<KnownDevice>,
observations: Vec<AgentDeviceObservation>,
agent_devices: Vec<AgentDeviceWithChildren>,
context: &FleetBuildContext,
) -> Vec<FleetDevice> {
let mut by_key = BTreeMap::<String, FleetAccumulator>::new();
@@ -49,23 +53,22 @@ pub(crate) fn build_fleet_devices(
display_name: Some(device.display_name.clone()),
known_device: Some(known_device_summary(&device)),
pinned: device.pinned,
presence_rank: 1,
..Default::default()
});
for identifier in device.identifiers {
add_identifier_to_entry(entry, &identifier);
for identifier in &device.identifiers {
add_identifier_to_entry(entry, identifier);
}
}
for observation in observations {
let key = observation_group_key(&observation, context);
for agent_device in agent_devices {
let key = device_group_key(&agent_device, context);
let entry = by_key
.entry(key.clone())
.or_insert_with(|| FleetAccumulator {
device_key: key,
..Default::default()
});
add_observation_to_entry(entry, observation, context);
add_agent_device_to_entry(entry, agent_device, context);
}
let mut devices = by_key
@@ -76,7 +79,7 @@ pub(crate) fn build_fleet_devices(
b.pinned
.cmp(&a.pinned)
.then_with(|| b.known_device.is_some().cmp(&a.known_device.is_some()))
.then_with(|| presence_rank(&b.presence).cmp(&presence_rank(&a.presence)))
.then_with(|| b.presence.cmp(&a.presence))
.then_with(|| b.last_seen_unix.cmp(&a.last_seen_unix))
.then_with(|| a.display_name.cmp(&b.display_name))
});
@@ -98,7 +101,7 @@ pub(crate) fn filter_fleet_devices(devices: &mut Vec<FleetDevice>, query: &ListF
}
if let Some(presence) = presence.as_deref()
&& presence != "all"
&& device.presence != presence
&& device.presence.as_str() != presence
{
return false;
}
@@ -118,18 +121,18 @@ pub(crate) fn filter_fleet_devices(devices: &mut Vec<FleetDevice>, query: &ListF
return false;
}
if let Some(search) = search.as_deref() {
let mut haystack = vec![
device.device_key.as_str(),
device.display_name.as_str(),
device.presence.as_str(),
let mut haystack: Vec<String> = vec![
device.device_key.clone(),
device.display_name.clone(),
device.presence.as_str().to_string(),
];
haystack.extend(device.ips.iter().map(String::as_str));
haystack.extend(device.macs.iter().map(String::as_str));
haystack.extend(device.hostnames.iter().map(String::as_str));
haystack.extend(device.sources.iter().map(String::as_str));
haystack.extend(device.agents.iter().map(|agent| agent.agent_id.as_str()));
haystack.extend(device.ips.iter().map(|ip| ip.to_string()));
haystack.extend(device.macs.iter().map(|mac| mac.to_string()));
haystack.extend(device.hostnames.clone());
haystack.extend(device.sources.clone());
haystack.extend(device.agents.iter().map(|a| a.agent_id.clone()));
if !haystack
.into_iter()
.iter()
.any(|value| value.to_ascii_lowercase().contains(search))
{
return false;
@@ -145,70 +148,73 @@ fn fleet_device_is_operator_noise(device: &FleetDevice) -> bool {
&& device.hostnames.is_empty()
&& device.recommended_route.is_none()
&& device.ips.is_empty()
&& device.presence == "offline"
&& device.presence == Presence::Offline
}
fn observation_group_key(
observation: &AgentDeviceObservation,
context: &FleetBuildContext,
) -> String {
if let Some(summary) = observation_known_device(observation, context) {
fn device_group_key(agent_device: &AgentDeviceWithChildren, context: &FleetBuildContext) -> String {
if let Some(summary) = device_known_device(agent_device, context) {
return format!("known:{}", summary.device_id);
}
if let Some(mac) = observation.mac.as_deref() {
if let Some(mac) = agent_device.macs.first() {
return format!("mac:{mac}");
}
if let Some(ip) = observation.ip.as_deref() {
if let Some(ip) = agent_device.ips.first() {
return format!("ip:{ip}");
}
observation.observation_key.clone()
agent_device.device.device_key.clone()
}
fn add_observation_to_entry(
fn add_agent_device_to_entry(
entry: &mut FleetAccumulator,
observation: AgentDeviceObservation,
agent_device: AgentDeviceWithChildren,
context: &FleetBuildContext,
) {
let observation_offline = observation_is_offline(&observation);
if let Some(summary) = observation_known_device(&observation, context)
let device_offline = agent_device.device.presence() == Presence::Offline;
if let Some(summary) = device_known_device(&agent_device, context)
&& entry.known_device.is_none()
{
entry.display_name = Some(summary.display_name.clone());
if let Some(ref name) = agent_device.device.display_name {
entry.display_name = Some(name.clone());
}
entry.pinned = summary.pinned;
entry.known_device = Some(summary);
}
if let Some(mac) = observation.mac.as_deref() {
entry.macs.insert(mac.to_string());
for mac in &agent_device.macs {
entry.macs.insert(*mac);
}
if !observation_offline && let Some(ip) = observation.ip.as_deref() {
entry.ips.insert(ip.to_string());
}
if let Some(hostname) = observation.hostname.as_deref() {
if entry.display_name.is_none() {
entry.display_name = Some(hostname.to_string());
if !device_offline {
for ip in &agent_device.ips {
entry.ips.insert(*ip);
}
entry.hostnames.insert(hostname.to_string());
}
entry.sources.insert(observation.kind.clone());
for hostname in &agent_device.hostnames {
if entry.display_name.is_none() {
entry.display_name = Some(hostname.clone());
}
entry.hostnames.insert(hostname.clone());
}
entry.sources.insert("device".to_string());
let first_seen = agent_device.device.first_seen();
let last_seen = agent_device.device.last_seen();
entry.first_seen_unix = Some(
entry
.first_seen_unix
.map(|current| current.min(observation.first_seen_unix))
.unwrap_or(observation.first_seen_unix),
.map(|current| current.min(first_seen))
.unwrap_or(first_seen),
);
entry.last_seen_unix = Some(
entry
.last_seen_unix
.map(|current| current.max(observation.last_seen_unix))
.unwrap_or(observation.last_seen_unix),
.map(|current| current.max(last_seen))
.unwrap_or(last_seen),
);
entry.presence_rank = entry
.presence_rank
.max(observation_presence_rank(&observation));
let device_presence = agent_device.device.presence();
entry.presence = std::cmp::max(entry.presence, device_presence);
let agent_id = agent_device.device.agent_id.clone();
let status = context
.agent_status
.get(&observation.agent_id)
.get(&agent_id)
.cloned()
.unwrap_or(AgentRuntimeStatus {
nickname: None,
@@ -216,71 +222,95 @@ fn add_observation_to_entry(
});
entry
.agents
.entry(observation.agent_id.clone())
.entry(agent_id.clone())
.and_modify(|agent| {
agent.last_seen_unix = agent.last_seen_unix.max(observation.last_seen_unix);
agent.last_seen_unix = agent.last_seen_unix.max(last_seen);
agent.connected = status.connected;
agent.nickname = status.nickname.clone();
})
.or_insert(FleetDeviceAgent {
agent_id: observation.agent_id.clone(),
agent_id: agent_id.clone(),
nickname: status.nickname.clone(),
connected: status.connected,
last_seen_unix: observation.last_seen_unix,
last_seen_unix: last_seen,
});
let route_id = route_id(
&observation.agent_id,
observation.mac.as_deref(),
observation.ip.as_deref(),
&observation.kind,
);
let wakeable = status.connected && observation.mac.is_some() && !observation_offline;
entry.routes.insert(
route_id.clone(),
FleetWakeRoute {
route_id: route_id.clone(),
agent_id: observation.agent_id,
nickname: status.nickname,
connected: status.connected,
mac: observation.mac,
ip: observation.ip,
hostname: observation.hostname,
source: observation.kind,
last_seen_unix: observation.last_seen_unix,
wakeable,
},
);
if let Some(route) = entry.routes.get_mut(&route_id) {
route.wakeable = route.connected && route.mac.is_some() && !observation_offline;
for mac in &agent_device.macs {
let ip_for_mac = agent_device.ips.first().copied();
let hostname_for_mac = agent_device.hostnames.first().cloned();
let rid = route_id(&agent_id, Some(mac), ip_for_mac.as_ref(), "device");
let wakeable = status.connected && !device_offline;
entry.routes.insert(
rid.clone(),
FleetWakeRoute {
route_id: rid,
agent_id: agent_id.clone(),
nickname: status.nickname.clone(),
connected: status.connected,
mac: Some(*mac),
ip: ip_for_mac,
hostname: hostname_for_mac,
source: "device".to_string(),
last_seen_unix: last_seen,
wakeable,
},
);
}
if agent_device.macs.is_empty()
&& let Some(ip) = agent_device.ips.first()
{
let hostname = agent_device.hostnames.first().cloned();
let rid = route_id(&agent_id, None, Some(ip), "device");
entry.routes.insert(
rid.clone(),
FleetWakeRoute {
route_id: rid,
agent_id: agent_id.clone(),
nickname: status.nickname.clone(),
connected: status.connected,
mac: None,
ip: Some(*ip),
hostname,
source: "device".to_string(),
last_seen_unix: last_seen,
wakeable: false,
},
);
}
}
fn add_identifier_to_entry(entry: &mut FleetAccumulator, identifier: &DeviceIdentifier) {
match identifier.kind.as_str() {
"mac" => {
entry.macs.insert(identifier.value.clone());
if let Ok(mac) = identifier.value.parse::<MacAddr>() {
entry.macs.insert(mac);
}
}
"ip" => {
entry.ips.insert(identifier.value.clone());
if let Ok(ip) = identifier.value.parse::<IpAddr>() {
entry.ips.insert(ip);
}
}
_ => {}
}
}
fn observation_known_device(
observation: &AgentDeviceObservation,
fn device_known_device(
agent_device: &AgentDeviceWithChildren,
context: &FleetBuildContext,
) -> Option<KnownDeviceSummary> {
observation
.mac
.as_deref()
.and_then(|mac| context.identifier_map.get(&format!("mac:{mac}")).cloned())
agent_device
.macs
.first()
.map(|mac| format!("mac:{}", mac.to_string().to_ascii_lowercase()))
.and_then(|key| context.identifier_map.get(&key).cloned())
.or_else(|| {
observation
.ip
.as_deref()
.and_then(|ip| context.identifier_map.get(&format!("ip:{ip}")).cloned())
agent_device
.ips
.first()
.map(|ip| format!("ip:{ip}"))
.and_then(|key| context.identifier_map.get(&key).cloned())
})
}
@@ -309,8 +339,8 @@ impl FleetAccumulator {
let display_name = self
.display_name
.or_else(|| self.hostnames.iter().next().cloned())
.or_else(|| self.macs.iter().next().cloned())
.or_else(|| self.ips.iter().next().cloned())
.or_else(|| self.macs.iter().next().map(|mac| mac.to_string()))
.or_else(|| self.ips.iter().next().map(|ip| ip.to_string()))
.unwrap_or_else(|| "(unknown device)".to_string());
FleetDevice {
@@ -325,50 +355,20 @@ impl FleetAccumulator {
sources: self.sources.into_iter().collect(),
first_seen_unix: self.first_seen_unix,
last_seen_unix: self.last_seen_unix,
presence: rank_presence(self.presence_rank).to_string(),
presence: self.presence,
route_candidates,
recommended_route,
}
}
}
fn observation_presence_rank(observation: &AgentDeviceObservation) -> u8 {
match observation.last_action.as_str() {
"remove" => 0,
"add" | "old" | "update" => 2,
_ => 1,
}
}
fn observation_is_offline(observation: &AgentDeviceObservation) -> bool {
observation.last_action == "remove"
}
fn rank_presence(rank: u8) -> &'static str {
match rank {
3 => "online",
2 => "likely_online",
0 => "offline",
_ => "unknown",
}
}
fn presence_rank(presence: &str) -> u8 {
match presence {
"online" => 3,
"likely_online" => 2,
"offline" => 0,
_ => 1,
}
}
fn route_id(agent_id: &str, mac: Option<&str>, ip: Option<&str>, source: &str) -> String {
fn route_id(agent_id: &str, mac: Option<&MacAddr>, ip: Option<&IpAddr>, source: &str) -> String {
format!(
"{}|{}|{}|{}",
agent_id,
source,
mac.unwrap_or(""),
ip.unwrap_or("")
mac.map(|m| m.to_string()).unwrap_or_default(),
ip.map(|i| i.to_string()).unwrap_or_default()
)
}
wakey-control-plane/src/api/control/fleet/inventory.rs
-208
View File
@@ -1,208 +0,0 @@
use std::collections::BTreeSet;
use std::net::IpAddr;
use serde::Deserialize;
use crate::state::AgentDeviceObservationInput;
#[derive(Debug, Deserialize)]
struct InventoryEnvelope {
kind: String,
devices: Vec<InventoryDevice>,
}
#[derive(Debug, Deserialize)]
struct InventoryDevice {
#[serde(default)]
names: Vec<String>,
#[serde(default)]
ips: Vec<IpAddr>,
#[serde(default)]
macs: Vec<String>,
#[serde(default)]
neighbors: Vec<InventoryNeighbor>,
#[serde(default)]
leases: Vec<InventoryLease>,
#[serde(default)]
observations: Vec<InventoryObservationFact>,
#[serde(default)]
presence: String,
}
#[derive(Debug, Deserialize)]
struct InventoryNeighbor {
ip: IpAddr,
#[serde(default)]
mac: Option<String>,
#[serde(default)]
state: Option<String>,
}
#[derive(Debug, Deserialize)]
struct InventoryLease {
ip: IpAddr,
mac: String,
#[serde(default)]
name: Option<String>,
}
#[derive(Debug, Deserialize)]
struct InventoryObservationFact {
kind: String,
action: String,
#[serde(default)]
mac: Option<String>,
#[serde(default)]
ip: Option<IpAddr>,
#[serde(default)]
hostname: Option<String>,
}
pub(crate) fn inventory_result_to_observations(
result: serde_json::Value,
) -> anyhow::Result<Vec<AgentDeviceObservationInput>> {
let envelope: InventoryEnvelope = serde_json::from_value(result)?;
if envelope.kind != "inventory" {
anyhow::bail!("expected inventory result, got {}", envelope.kind);
}
let now = now_unix();
let mut out = Vec::new();
for device in envelope.devices {
let hostname = device
.names
.iter()
.find(|name| !name.trim().is_empty())
.cloned();
let mut wrote_source_observation = false;
for neighbor in &device.neighbors {
if neighbor.mac.is_none() && device.macs.is_empty() {
out.push(inventory_observation(
inventory_action_for_neighbor(neighbor),
None,
Some(neighbor.ip.to_string()),
hostname.clone(),
now,
));
wrote_source_observation = true;
continue;
}
let macs = neighbor
.mac
.iter()
.chain(device.macs.iter())
.collect::<BTreeSet<_>>();
for mac in macs {
out.push(inventory_observation(
inventory_action_for_neighbor(neighbor),
Some(mac.clone()),
Some(neighbor.ip.to_string()),
hostname.clone(),
now,
));
wrote_source_observation = true;
}
}
for observation in &device.observations {
if observation.mac.is_none() && observation.ip.is_none() {
continue;
}
out.push(inventory_observation(
inventory_action_for_observation(observation),
observation.mac.clone(),
observation.ip.map(|ip| ip.to_string()),
observation.hostname.clone().or_else(|| hostname.clone()),
now,
));
wrote_source_observation = true;
}
for lease in &device.leases {
out.push(inventory_observation(
"update",
Some(lease.mac.clone()),
Some(lease.ip.to_string()),
lease.name.clone().or_else(|| hostname.clone()),
now,
));
wrote_source_observation = true;
}
if wrote_source_observation {
continue;
}
let action = if device.presence == "offline" {
"remove"
} else {
"update"
};
if !device.macs.is_empty() {
for mac in device.macs {
out.push(inventory_observation(
action,
Some(mac),
device.ips.first().map(ToString::to_string),
hostname.clone(),
now,
));
}
} else {
for ip in &device.ips {
out.push(inventory_observation(
action,
None,
Some(ip.to_string()),
hostname.clone(),
now,
));
}
}
}
Ok(out)
}
fn inventory_observation(
action: &str,
mac: Option<String>,
ip: Option<String>,
hostname: Option<String>,
now: u64,
) -> AgentDeviceObservationInput {
AgentDeviceObservationInput {
kind: "inventory".into(),
action: action.into(),
mac,
ip,
hostname,
first_seen_unix: now,
last_seen_unix: now,
}
}
fn inventory_action_for_neighbor(neighbor: &InventoryNeighbor) -> &'static str {
if neighbor
.state
.as_deref()
.is_some_and(|state| state.eq_ignore_ascii_case("FAILED"))
{
"remove"
} else {
"update"
}
}
fn inventory_action_for_observation(observation: &InventoryObservationFact) -> &str {
match observation.action.as_str() {
"add" | "old" | "update" => "update",
"remove" | "del" => "remove",
_ if observation.kind == "neigh" => "update",
_ => "update",
}
}
fn now_unix() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|duration| duration.as_secs())
.unwrap_or_default()
}
wakey-control-plane/src/api/control/fleet/mod.rs
+31 -49
View File
@@ -12,7 +12,6 @@ use crate::api::json_error;
use crate::runtime::AppState;
mod build;
mod inventory;
mod types;
#[cfg(test)]
@@ -22,7 +21,6 @@ use build::{
AgentRuntimeStatus, FleetBuildContext, build_fleet_devices, filter_fleet_devices,
known_device_summary,
};
use inventory::inventory_result_to_observations;
use types::{
FleetDevice, ListFleetDevicesQuery, RefreshFleetAgentResult, RefreshFleetDevicesRequest,
RefreshFleetDevicesResponse, WakeFleetDeviceRequest, WakeFleetDeviceResponse,
@@ -87,28 +85,35 @@ pub async fn refresh_fleet_devices(
});
continue;
};
match inventory_result_to_observations(result) {
Ok(observations) => match state
.store
.upsert_agent_observations_snapshot(&agent_id, "inventory", observations)
.await
{
Ok(accepted) => {
total_accepted = total_accepted.saturating_add(accepted);
results.push(RefreshFleetAgentResult {
match serde_json::from_value::<Vec<wakey_core::Device>>(
result
.get("devices")
.cloned()
.unwrap_or(serde_json::Value::Array(vec![])),
) {
Ok(devices) => {
match state
.store
.replace_agent_device_snapshot(&agent_id, &devices)
.await
{
Ok(accepted) => {
total_accepted = total_accepted.saturating_add(accepted);
results.push(RefreshFleetAgentResult {
agent_id,
status: "ok".into(),
accepted,
error: None,
});
}
Err(err) => results.push(RefreshFleetAgentResult {
agent_id,
status: "ok".into(),
accepted,
error: None,
});
status: "error".into(),
accepted: 0,
error: Some(err.to_string()),
}),
}
Err(err) => results.push(RefreshFleetAgentResult {
agent_id,
status: "error".into(),
accepted: 0,
error: Some(err.to_string()),
}),
},
}
Err(err) => results.push(RefreshFleetAgentResult {
agent_id,
status: "error".into(),
@@ -191,7 +196,7 @@ pub async fn wake_fleet_device(
)
})?;
let Some(mac) = route.mac.as_deref() else {
let Some(mac) = route.mac else {
return Err(json_error(
StatusCode::BAD_REQUEST,
"wake_route_unavailable",
@@ -213,33 +218,13 @@ pub async fn wake_fleet_device(
));
}
let mac = mac.parse().map_err(|err| {
json_error(
StatusCode::BAD_REQUEST,
"invalid_wake_route",
&format!("invalid route MAC: {err}"),
)
})?;
let ip = route
.ip
.as_deref()
.map(str::parse)
.transpose()
.map_err(|err| {
json_error(
StatusCode::BAD_REQUEST,
"invalid_wake_route",
&format!("invalid route IP: {err}"),
)
})?;
let command = relay_agent_command(
&state,
&route.agent_id,
AgentCommand::Wake(WakeRequest {
query: None,
mac: Some(mac),
ip,
ip: route.ip,
}),
req.timeout_ms,
)
@@ -256,10 +241,7 @@ async fn load_fleet_devices(
query: &ListFleetDevicesQuery,
) -> anyhow::Result<Vec<FleetDevice>> {
let known_devices = state.store.list_known_devices().await?;
let observations = state
.store
.list_agent_observations(None, query.limit.unwrap_or(1000).max(1))
.await?;
let agent_devices = state.store.list_agent_device_rows().await?;
let connected = {
let sessions = state.sessions.read().await;
sessions.keys().cloned().collect::<BTreeSet<_>>()
@@ -291,7 +273,7 @@ async fn load_fleet_devices(
agent_status,
identifier_map,
};
let mut devices = build_fleet_devices(known_devices, observations, &context);
let mut devices = build_fleet_devices(known_devices, agent_devices, &context);
filter_fleet_devices(&mut devices, query);
let limit = query.limit.unwrap_or(500).clamp(1, 1000);
devices.truncate(limit);
wakey-control-plane/src/api/control/fleet/tests.rs
+164 -168
View File
@@ -1,12 +1,15 @@
use std::collections::HashMap;
use std::net::IpAddr;
use macaddr::MacAddr;
use wakey_core::Presence;
use super::build::{
AgentRuntimeStatus, FleetBuildContext, build_fleet_devices, filter_fleet_devices,
known_device_summary,
};
use super::inventory::inventory_result_to_observations;
use super::types::ListFleetDevicesQuery;
use crate::state::{AgentDeviceObservation, DeviceIdentifier, KnownDevice};
use crate::state::{AgentDeviceRow, AgentDeviceWithChildren, DeviceIdentifier, KnownDevice};
fn context(connected: &[&str]) -> FleetBuildContext {
FleetBuildContext {
@@ -26,27 +29,83 @@ fn context(connected: &[&str]) -> FleetBuildContext {
}
}
fn observation(
fn agent_device(
agent_id: &str,
device_key: &str,
mac: Option<&str>,
ip: Option<&str>,
last_seen_unix: u64,
) -> AgentDeviceObservation {
AgentDeviceObservation {
observation_key: format!(
"agent:{agent_id}:dhcp:{}",
mac.map(|mac| format!("mac:{mac}"))
.or_else(|| ip.map(|ip| format!("ip:{ip}")))
.unwrap_or_default()
),
agent_id: agent_id.into(),
kind: "dhcp".into(),
mac: mac.map(str::to_string),
ip: ip.map(str::to_string),
hostname: Some("lda".into()),
first_seen_unix: 1,
last_seen_unix,
last_action: "update".into(),
last_seen_unix: i64,
) -> AgentDeviceWithChildren {
AgentDeviceWithChildren {
device: AgentDeviceRow {
agent_id: agent_id.into(),
device_key: device_key.into(),
presence: "likely_online".into(),
display_name: Some("lda".into()),
first_seen_unix: 1,
last_seen_unix,
},
macs: mac
.map(|m| m.parse::<MacAddr>().unwrap())
.into_iter()
.collect(),
ips: ip
.map(|i| i.parse::<IpAddr>().unwrap())
.into_iter()
.collect(),
hostnames: vec!["lda".to_string()],
facts: vec![],
}
}
fn offline_agent_device(
agent_id: &str,
device_key: &str,
mac: Option<&str>,
ip: Option<&str>,
last_seen_unix: i64,
) -> AgentDeviceWithChildren {
AgentDeviceWithChildren {
device: AgentDeviceRow {
agent_id: agent_id.into(),
device_key: device_key.into(),
presence: "offline".into(),
display_name: Some("lda".into()),
first_seen_unix: 1,
last_seen_unix,
},
macs: mac
.map(|m| m.parse::<MacAddr>().unwrap())
.into_iter()
.collect(),
ips: ip
.map(|i| i.parse::<IpAddr>().unwrap())
.into_iter()
.collect(),
hostnames: vec!["lda".to_string()],
facts: vec![],
}
}
fn offline_ip_only_unknown(
agent_id: &str,
device_key: &str,
ip: &str,
last_seen_unix: i64,
) -> AgentDeviceWithChildren {
AgentDeviceWithChildren {
device: AgentDeviceRow {
agent_id: agent_id.into(),
device_key: device_key.into(),
presence: "offline".into(),
display_name: None,
first_seen_unix: 1,
last_seen_unix,
},
macs: vec![],
ips: vec![ip.parse().unwrap()],
hostnames: vec![],
facts: vec![],
}
}
@@ -55,14 +114,16 @@ fn fleet_grouping_combines_same_mac_across_agents() {
let devices = build_fleet_devices(
Vec::new(),
vec![
observation(
agent_device(
"agent-a",
"mac:aa:bb:cc:dd:ee:ff",
Some("aa:bb:cc:dd:ee:ff"),
Some("192.168.1.2"),
10,
),
observation(
agent_device(
"agent-b",
"mac:aa:bb:cc:dd:ee:ff",
Some("aa:bb:cc:dd:ee:ff"),
Some("192.168.2.2"),
20,
@@ -71,8 +132,9 @@ fn fleet_grouping_combines_same_mac_across_agents() {
&context(&["agent-a", "agent-b"]),
);
let expected_mac: MacAddr = "aa:bb:cc:dd:ee:ff".parse().unwrap();
assert_eq!(devices.len(), 1);
assert_eq!(devices[0].macs, vec!["aa:bb:cc:dd:ee:ff"]);
assert_eq!(devices[0].macs, vec![expected_mac]);
assert_eq!(devices[0].agents.len(), 2);
assert_eq!(
devices[0]
@@ -84,7 +146,7 @@ fn fleet_grouping_combines_same_mac_across_agents() {
}
#[test]
fn known_device_with_two_macs_absorbs_both_observation_groups() {
fn known_device_with_two_macs_absorbs_both_device_groups() {
let known = KnownDevice {
device_id: "dev-1".into(),
display_name: "lda".into(),
@@ -120,8 +182,20 @@ fn known_device_with_two_macs_absorbs_both_observation_groups() {
let devices = build_fleet_devices(
vec![known],
vec![
observation("agent-a", Some("aa:bb:cc:dd:ee:01"), None, 10),
observation("agent-a", Some("aa:bb:cc:dd:ee:02"), None, 20),
agent_device(
"agent-a",
"mac:aa:bb:cc:dd:ee:01",
Some("aa:bb:cc:dd:ee:01"),
None,
10,
),
agent_device(
"agent-a",
"mac:aa:bb:cc:dd:ee:02",
Some("aa:bb:cc:dd:ee:02"),
None,
20,
),
],
&ctx,
);
@@ -135,44 +209,54 @@ fn known_device_with_two_macs_absorbs_both_observation_groups() {
fn ip_only_unknown_is_visible_but_not_wakeable() {
let devices = build_fleet_devices(
Vec::new(),
vec![observation("agent-a", None, Some("192.168.1.2"), 10)],
vec![agent_device(
"agent-a",
"ip:192.168.1.2",
None,
Some("192.168.1.2"),
10,
)],
&context(&["agent-a"]),
);
let expected_ip: IpAddr = "192.168.1.2".parse().unwrap();
assert_eq!(devices.len(), 1);
assert_eq!(devices[0].ips, vec![expected_ip]);
assert!(devices[0].recommended_route.is_none());
}
#[test]
fn offline_device_has_offline_presence() {
let devices = build_fleet_devices(
Vec::new(),
vec![offline_agent_device(
"agent-a",
"mac:aa:bb:cc:dd:ee:ff",
Some("aa:bb:cc:dd:ee:ff"),
Some("192.168.1.2"),
20,
)],
&context(&["agent-a"]),
);
assert_eq!(devices.len(), 1);
assert_eq!(devices[0].ips, vec!["192.168.1.2"]);
assert!(devices[0].recommended_route.is_none());
assert!(!devices[0].route_candidates[0].wakeable);
}
#[test]
fn offline_observation_does_not_advertise_current_ip_or_wake_route() {
let mut offline = observation(
"agent-a",
Some("aa:bb:cc:dd:ee:ff"),
Some("192.168.1.2"),
20,
);
offline.kind = "neigh".into();
offline.last_action = "remove".into();
let devices = build_fleet_devices(Vec::new(), vec![offline], &context(&["agent-a"]));
assert_eq!(devices.len(), 1);
assert!(devices[0].ips.is_empty());
assert_eq!(devices[0].presence, "offline");
assert_eq!(devices[0].presence, Presence::Offline);
assert!(devices[0].recommended_route.is_none());
assert!(!devices[0].route_candidates[0].wakeable);
}
#[test]
fn unknown_ip_only_remove_observation_is_hidden() {
let mut offline = observation("agent-a", None, Some("192.168.1.2"), 20);
offline.kind = "neigh".into();
offline.hostname = None;
offline.last_action = "remove".into();
let mut devices = build_fleet_devices(Vec::new(), vec![offline], &context(&["agent-a"]));
fn unknown_ip_only_offline_device_is_hidden_by_default() {
let mut devices = build_fleet_devices(
Vec::new(),
vec![offline_ip_only_unknown(
"agent-a",
"ip:192.168.1.2",
"192.168.1.2",
20,
)],
&context(&["agent-a"]),
);
assert_eq!(devices.len(), 1);
assert_eq!(devices[0].display_name, "(unknown device)");
@@ -181,13 +265,17 @@ fn unknown_ip_only_remove_observation_is_hidden() {
}
#[test]
fn visibility_all_keeps_unknown_ip_only_remove_observation() {
let mut offline = observation("agent-a", None, Some("192.168.1.2"), 20);
offline.kind = "neigh".into();
offline.hostname = None;
offline.last_action = "remove".into();
let mut devices = build_fleet_devices(Vec::new(), vec![offline], &context(&["agent-a"]));
fn visibility_all_keeps_unknown_ip_only_offline_device() {
let mut devices = build_fleet_devices(
Vec::new(),
vec![offline_ip_only_unknown(
"agent-a",
"ip:192.168.1.2",
"192.168.1.2",
20,
)],
&context(&["agent-a"]),
);
filter_fleet_devices(
&mut devices,
@@ -201,7 +289,7 @@ fn visibility_all_keeps_unknown_ip_only_remove_observation() {
}
#[test]
fn known_ip_only_remove_observation_is_kept() {
fn known_ip_only_offline_device_is_kept() {
let known = KnownDevice {
device_id: "dev-1".into(),
display_name: "lda".into(),
@@ -220,114 +308,22 @@ fn known_ip_only_remove_observation_is_kept() {
let mut ctx = context(&["agent-a"]);
ctx.identifier_map
.insert("ip:192.168.1.2".into(), known_device_summary(&known));
let mut offline = observation("agent-a", None, Some("192.168.1.2"), 20);
offline.kind = "neigh".into();
offline.hostname = None;
offline.last_action = "remove".into();
let devices = build_fleet_devices(vec![known], vec![offline], &ctx);
let devices = build_fleet_devices(
vec![known],
vec![offline_agent_device(
"agent-a",
"ip:192.168.1.2",
None,
Some("192.168.1.2"),
20,
)],
&ctx,
);
let expected_ip: IpAddr = "192.168.1.2".parse().unwrap();
assert_eq!(devices.len(), 1);
assert_eq!(devices[0].device_key, "known:dev-1");
assert_eq!(devices[0].display_name, "lda");
assert!(devices[0].ips.contains(&"192.168.1.2".to_string()));
}
#[test]
fn inventory_result_maps_to_stored_observations() {
let observations = inventory_result_to_observations(serde_json::json!({
"kind": "inventory",
"devices": [{
"names": ["lda"],
"ips": ["192.168.1.2"],
"macs": ["aa:bb:cc:dd:ee:ff"],
"presence": "likely_online"
}]
}))
.expect("inventory should map");
assert_eq!(observations.len(), 1);
assert_eq!(observations[0].kind, "inventory");
assert_eq!(observations[0].action, "update");
assert_eq!(observations[0].mac.as_deref(), Some("aa:bb:cc:dd:ee:ff"));
}
#[test]
fn inventory_result_preserves_neighbor_failed_ip_as_remove() {
let observations = inventory_result_to_observations(serde_json::json!({
"kind": "inventory",
"devices": [{
"names": ["lda"],
"ips": ["192.168.1.2", "192.168.1.3"],
"macs": ["aa:bb:cc:dd:ee:ff"],
"neighbors": [
{
"ip": "192.168.1.2",
"mac": "aa:bb:cc:dd:ee:ff",
"state": "FAILED"
},
{
"ip": "192.168.1.3",
"mac": "aa:bb:cc:dd:ee:ff",
"state": "REACHABLE"
}
],
"presence": "online"
}]
}))
.expect("inventory should map");
assert_eq!(observations.len(), 2);
let removed = observations
.iter()
.find(|observation| observation.ip.as_deref() == Some("192.168.1.2"))
.expect("failed neighbor observation should exist");
assert_eq!(removed.action, "remove");
assert_eq!(removed.mac.as_deref(), Some("aa:bb:cc:dd:ee:ff"));
let current = observations
.iter()
.find(|observation| observation.ip.as_deref() == Some("192.168.1.3"))
.expect("reachable neighbor observation should exist");
assert_eq!(current.action, "update");
assert_eq!(current.mac.as_deref(), Some("aa:bb:cc:dd:ee:ff"));
}
#[test]
fn inventory_result_preserves_hook_observations_and_leases() {
let observations = inventory_result_to_observations(serde_json::json!({
"kind": "inventory",
"devices": [{
"names": ["lda"],
"ips": ["192.168.1.2", "192.168.1.3"],
"macs": ["aa:bb:cc:dd:ee:ff"],
"observations": [{
"kind": "neigh",
"action": "remove",
"mac": "aa:bb:cc:dd:ee:ff",
"ip": "192.168.1.2"
}],
"leases": [{
"expires_epoch": 1893456000_u64,
"ip": "192.168.1.3",
"mac": "aa:bb:cc:dd:ee:ff",
"name": "lda"
}],
"presence": "likely_online"
}]
}))
.expect("inventory should map");
assert_eq!(observations.len(), 2);
assert!(observations.iter().any(|observation| {
observation.action == "remove"
&& observation.mac.as_deref() == Some("aa:bb:cc:dd:ee:ff")
&& observation.ip.as_deref() == Some("192.168.1.2")
}));
assert!(observations.iter().any(|observation| {
observation.action == "update"
&& observation.mac.as_deref() == Some("aa:bb:cc:dd:ee:ff")
&& observation.ip.as_deref() == Some("192.168.1.3")
&& observation.hostname.as_deref() == Some("lda")
}));
assert!(devices[0].ips.contains(&expected_ip));
}
wakey-control-plane/src/api/control/fleet/types.rs
+12 -5
View File
@@ -1,7 +1,12 @@
use macaddr::MacAddr;
use serde::{Deserialize, Serialize};
use std::net::IpAddr;
use wakey_core::Presence;
use crate::api::commands::RelayCommandResponse;
use crate::state::KnownDeviceSummary;
use wakey_core::parse::mac;
#[derive(Debug, Default, Deserialize)]
pub struct ListFleetDevicesQuery {
@@ -53,14 +58,15 @@ pub struct FleetDevice {
pub display_name: String,
pub known_device: Option<KnownDeviceSummary>,
pub pinned: bool,
pub ips: Vec<String>,
pub macs: Vec<String>,
pub ips: Vec<IpAddr>,
#[serde(with = "mac::vec_mac")]
pub macs: Vec<MacAddr>,
pub hostnames: Vec<String>,
pub agents: Vec<FleetDeviceAgent>,
pub sources: Vec<String>,
pub first_seen_unix: Option<u64>,
pub last_seen_unix: Option<u64>,
pub presence: String,
pub presence: Presence,
pub route_candidates: Vec<FleetWakeRoute>,
pub recommended_route: Option<FleetWakeRoute>,
}
@@ -79,8 +85,9 @@ pub struct FleetWakeRoute {
pub agent_id: String,
pub nickname: Option<String>,
pub connected: bool,
pub mac: Option<String>,
pub ip: Option<String>,
#[serde(with = "mac::option_mac")]
pub mac: Option<MacAddr>,
pub ip: Option<IpAddr>,
pub hostname: Option<String>,
pub source: String,
pub last_seen_unix: u64,
wakey-control-plane/src/api/control/observations.rs
-243
View File
@@ -1,243 +0,0 @@
use std::collections::BTreeMap;
use axum::Json;
use axum::extract::{Path as AxumPath, Query, State};
use axum::http::StatusCode;
use axum::response::IntoResponse;
use serde::Deserialize;
use tracing::warn;
use wakey_agent::protocol::ServerMessage;
use crate::api::json_error;
use crate::runtime::{AppState, SessionEvent};
use crate::state::{
AgentDeviceObservationEvent, AgentDeviceObservationInput, AgentDeviceObservationView,
};
#[derive(Debug, Deserialize)]
pub struct UploadAgentObservationsRequest {
pub agent_id: String,
pub agent_token: String,
#[serde(default)]
pub observations: Vec<AgentObservationRequest>,
}
#[derive(Debug, Deserialize)]
pub struct AgentObservationRequest {
pub kind: String,
pub action: String,
pub mac: Option<String>,
pub ip: Option<String>,
pub hostname: Option<String>,
pub first_seen_unix: u64,
pub last_seen_unix: u64,
}
#[derive(Debug, serde::Serialize, serde::Deserialize)]
pub struct UploadAgentObservationsResponse {
pub accepted: usize,
}
#[derive(Debug, serde::Serialize, serde::Deserialize)]
pub struct RequestAgentObservationSyncResponse {
pub agent_id: String,
pub requested: bool,
}
#[derive(Debug, Deserialize)]
pub struct ListObservationsQuery {
pub agent_id: Option<String>,
pub limit: Option<usize>,
}
#[derive(Debug, Deserialize)]
pub struct ListObservationHistoryQuery {
pub agent_id: Option<String>,
pub kind: Option<String>,
pub mac: Option<String>,
pub ip: Option<String>,
pub observation_key: Option<String>,
pub limit: Option<usize>,
}
pub async fn upload_agent_observations(
State(state): State<AppState>,
Json(req): Json<UploadAgentObservationsRequest>,
) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
if !state
.store
.verify_agent_token(&req.agent_id, &req.agent_token)
.await
{
return Err(json_error(
StatusCode::UNAUTHORIZED,
"agent_auth_rejected",
"agent credentials rejected",
));
}
let mut by_kind: BTreeMap<String, Vec<AgentDeviceObservationInput>> = BTreeMap::new();
for observation in req.observations {
let kind = observation.kind.trim().to_ascii_lowercase();
let entry = by_kind.entry(kind.clone()).or_default();
entry.push(AgentDeviceObservationInput {
kind,
action: observation.action,
mac: observation.mac,
ip: observation.ip,
hostname: observation.hostname,
first_seen_unix: observation.first_seen_unix,
last_seen_unix: observation.last_seen_unix,
});
}
let mut accepted = 0usize;
for (kind, observations) in by_kind {
match state
.store
.upsert_agent_observations_snapshot(&req.agent_id, &kind, observations)
.await
{
Ok(written) => {
accepted = accepted.saturating_add(written);
}
Err(err) => {
warn!(error = %err, agent_id = %req.agent_id, kind = %kind, "failed to upload agent observations");
return Err(json_error(
StatusCode::BAD_REQUEST,
"upload_observations_failed",
&err.to_string(),
));
}
}
}
Ok((
StatusCode::OK,
Json(UploadAgentObservationsResponse { accepted }),
))
}
pub async fn list_agent_observations(
State(state): State<AppState>,
Query(query): Query<ListObservationsQuery>,
) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
match state
.store
.list_agent_observation_views(query.agent_id.as_deref(), query.limit.unwrap_or(500))
.await
{
Ok(observations) => Ok((
StatusCode::OK,
Json(
observations
.into_iter()
.map(agent_observation_response)
.collect::<Vec<_>>(),
),
)),
Err(err) => {
warn!(error = %err, "failed to list agent observations");
Err(json_error(
StatusCode::INTERNAL_SERVER_ERROR,
"list_observations_failed",
&err.to_string(),
))
}
}
}
pub async fn request_agent_observation_sync(
State(state): State<AppState>,
AxumPath(agent_id): AxumPath<String>,
) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
let tx = {
let sessions = state.sessions.read().await;
sessions.get(&agent_id).map(|session| session.tx.clone())
};
let Some(tx) = tx else {
return Ok((
StatusCode::NOT_FOUND,
Json(RequestAgentObservationSyncResponse {
agent_id,
requested: false,
}),
));
};
match tx.send(SessionEvent::Message(ServerMessage::SyncObservations)) {
Ok(()) => Ok((
StatusCode::OK,
Json(RequestAgentObservationSyncResponse {
agent_id,
requested: true,
}),
)),
Err(err) => {
warn!(error = %err, "failed to request agent observation sync");
Err(json_error(
StatusCode::BAD_GATEWAY,
"agent_observation_sync_request_failed",
&err.to_string(),
))
}
}
}
pub async fn list_agent_observation_history(
State(state): State<AppState>,
Query(query): Query<ListObservationHistoryQuery>,
) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
let kind = normalized_query_value(query.kind).map(|value| value.to_ascii_lowercase());
let mac = normalized_query_value(query.mac).map(|value| value.to_ascii_lowercase());
let ip = normalized_query_value(query.ip);
let observation_key = normalized_query_value(query.observation_key);
match state
.store
.list_agent_observation_events(
query.agent_id.as_deref(),
kind.as_deref(),
mac.as_deref(),
ip.as_deref(),
observation_key.as_deref(),
query.limit.unwrap_or(500),
)
.await
{
Ok(events) => Ok((
StatusCode::OK,
Json(
events
.into_iter()
.map(agent_observation_event_response)
.collect::<Vec<_>>(),
),
)),
Err(err) => {
warn!(error = %err, "failed to list agent observation history");
Err(json_error(
StatusCode::INTERNAL_SERVER_ERROR,
"list_observation_history_failed",
&err.to_string(),
))
}
}
}
fn agent_observation_response(
observation: AgentDeviceObservationView,
) -> AgentDeviceObservationView {
observation
}
fn agent_observation_event_response(
event: AgentDeviceObservationEvent,
) -> AgentDeviceObservationEvent {
event
}
fn normalized_query_value(value: Option<String>) -> Option<String> {
value
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
}
wakey-control-plane/src/api/mod.rs
+4 -6
View File
@@ -11,12 +11,10 @@ pub use audit::list_audit_events;
pub use commands::{list_agents, run_command};
pub use control::{
EnrollTokenStatus, IssueEnrollTokenResponse, RevokeAgentResponse, RevokeEnrollTokenResponse,
StateStatsResponse, attach_device_identifier, attach_observation_identifier,
create_known_device, detach_device_identifier, enroll, forget_known_device, healthz,
issue_enroll_token, list_agent_observation_history, list_agent_observations,
list_enroll_tokens, list_fleet_devices, list_known_devices, merge_known_device,
refresh_fleet_devices, request_agent_observation_sync, revoke_agent, revoke_enroll_token,
set_agent_nickname, state_stats, upload_agent_observations, wake_fleet_device,
StateStatsResponse, attach_device_identifier, create_known_device, detach_device_identifier,
enroll, forget_known_device, healthz, issue_enroll_token, list_enroll_tokens,
list_fleet_devices, list_known_devices, merge_known_device, refresh_fleet_devices,
revoke_agent, revoke_enroll_token, set_agent_nickname, state_stats, wake_fleet_device,
};
pub fn json_error(
wakey-control-plane/src/config/types.rs
+1
View File
@@ -12,6 +12,7 @@ pub struct DaemonConfig {
pub state_file: PathBuf,
pub command_timeout: Duration,
pub enroll_token_ttl: Duration,
#[allow(dead_code)]
pub observation_retention: Duration,
pub pid_file: PathBuf,
pub ui_dist_dir: PathBuf,
wakey-control-plane/src/runtime/mod.rs
-32
View File
@@ -71,10 +71,6 @@ fn public_api_routes(ui_dist_dir: std::path::PathBuf) -> Router<AppState> {
)
.route("/healthz", get(api::healthz))
.route("/api/v1/agents/enroll", post(api::enroll))
.route(
"/api/v1/agents/observations",
post(api::upload_agent_observations),
)
.route("/api/v1/agent/ws", get(ws::agent_ws))
}
@@ -102,18 +98,6 @@ fn control_api_routes() -> Router<AppState> {
post(api::refresh_fleet_devices),
)
.route("/api/v1/control/fleet/wake", post(api::wake_fleet_device))
.route(
"/api/v1/control/observations",
get(api::list_agent_observations),
)
.route(
"/api/v1/control/observations/history",
get(api::list_agent_observation_history),
)
.route(
"/api/v1/control/agents/{agent_id}/observations/sync",
post(api::request_agent_observation_sync),
)
.route(
"/api/v1/control/devices",
get(api::list_known_devices).post(api::create_known_device),
@@ -134,10 +118,6 @@ fn control_api_routes() -> Router<AppState> {
"/api/v1/control/devices/{device_id}/identifiers/{identifier_key}",
axum::routing::delete(api::detach_device_identifier),
)
.route(
"/api/v1/control/devices/{device_id}/identifiers/from-observation",
post(api::attach_observation_identifier),
)
.route("/api/v1/control/audit/events", get(api::list_audit_events))
.route("/api/v1/control/alerts", get(api::active_alerts))
.route("/api/v1/control/alerts/history", get(api::alert_history))
@@ -226,18 +206,6 @@ pub async fn serve(daemon: config::DaemonConfig) -> Result<()> {
}
Err(err) => warn!(error = %err, "periodic gc failed"),
}
match app_state
.store
.gc_stale_observations(daemon.observation_retention)
.await
{
Ok(removed) => {
if removed > 0 {
info!(removed, "periodic gc removed stale observations");
}
}
Err(err) => warn!(error = %err, "periodic observation gc failed"),
}
}
join = &mut server => {
let _ = remove_pid_file(&daemon.pid_file);
wakey-control-plane/src/state/mod.rs
+3 -2
View File
@@ -2,8 +2,9 @@ mod store;
mod types;
pub use store::Store;
#[cfg(test)]
pub use types::AgentDeviceRow;
pub use types::{
AgentDeviceObservation, AgentDeviceObservationEvent, AgentDeviceObservationInput,
AgentDeviceObservationView, AlertState, AuditEvent, AuditEventFilter, AuditEventInput,
AgentDeviceWithChildren, AlertState, AuditEvent, AuditEventFilter, AuditEventInput,
DeviceIdentifier, DeviceIdentifierInput, KnownDevice, KnownDeviceInput, KnownDeviceSummary,
};
wakey-control-plane/src/state/store.rs
+97 -437
View File
@@ -8,10 +8,9 @@ use tracing::{info, warn};
use uuid::Uuid;
use crate::state::types::{
AgentDeviceObservation, AgentDeviceObservationEvent, AgentDeviceObservationInput,
AgentDeviceObservationView, AlertState, AlertTransition, AuditEvent, AuditEventFilter,
AuditEventInput, DeviceIdentifier, DeviceIdentifierInput, EnrollTokenInfo, IssuedAgent,
IssuedEnrollToken, KnownDevice, KnownDeviceInput, KnownDeviceSummary, StateStats,
AlertState, AlertTransition, AuditEvent, AuditEventFilter, AuditEventInput, DeviceIdentifier,
DeviceIdentifierInput, EnrollTokenInfo, IssuedAgent, IssuedEnrollToken, KnownDevice,
KnownDeviceInput, StateStats,
};
pub struct Store {
@@ -21,8 +20,9 @@ pub struct Store {
const SCHEMA_VERSION_KEY: &str = "schema_version";
const SEEDED_ENROLL_TOKEN_PREFIX: &str = "seeded_enroll_token:";
const SCHEMA_VERSION: u32 = 1;
const SCHEMA_VERSION: u32 = 2;
pub(crate) mod agent_devices;
mod alerts;
mod audit;
mod db;
@@ -30,7 +30,6 @@ mod devices;
mod enrollment;
mod helpers;
mod import_sled;
mod observations;
use helpers::alerts_audit::*;
use helpers::core::*;
@@ -45,20 +44,7 @@ mod tests {
use crate::state::{DeviceIdentifierInput, KnownDeviceInput};
use super::Store;
async fn make_store() -> (Store, std::path::PathBuf) {
let dir =
std::env::temp_dir().join(format!("wakey-cp-store-test-{}", uuid::Uuid::new_v4()));
let db_path = dir.join("state.sqlite3");
let store = Store::load_or_init(&db_path, Vec::new(), Duration::from_secs(60))
.await
.expect("store should initialize");
(store, dir)
}
fn cleanup_dir(path: &std::path::Path) {
let _ = fs::remove_dir_all(path);
}
use super::helpers::test_helpers::TestStore;
async fn insert_token(store: &Store, token: &str, expires_at_unix: u64) {
sqlx::query(
@@ -81,15 +67,16 @@ mod tests {
Err(err) => err,
};
assert!(err.to_string().contains("legacy sled store"));
cleanup_dir(&dir);
let _ = fs::remove_dir_all(&dir);
}
#[tokio::test]
async fn gc_removes_expired_tokens() {
let (store, dir) = make_store().await;
insert_token(&store, "enr-expired-gc-test", 1).await;
let ts = TestStore::new().await;
insert_token(ts.store(), "enr-expired-gc-test", 1).await;
let removed = store
let removed = ts
.store()
.gc_expired_enroll_tokens()
.await
.expect("gc should succeed");
@@ -98,19 +85,19 @@ mod tests {
let exists =
sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM enroll_tokens WHERE token = ?1")
.bind("enr-expired-gc-test")
.fetch_one(&store.pool)
.fetch_one(&ts.store().pool)
.await
.expect("read should succeed");
assert_eq!(exists, 0);
cleanup_dir(&dir);
}
#[tokio::test]
async fn enroll_rejects_expired_token() {
let (store, dir) = make_store().await;
insert_token(&store, "enr-expired-enroll-test", 1).await;
let ts = TestStore::new().await;
insert_token(ts.store(), "enr-expired-enroll-test", 1).await;
let err = store
let err = ts
.store()
.enroll("enr-expired-enroll-test")
.await
.expect_err("expired token should be rejected");
@@ -119,119 +106,122 @@ mod tests {
let exists =
sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM enroll_tokens WHERE token = ?1")
.bind("enr-expired-enroll-test")
.fetch_one(&store.pool)
.fetch_one(&ts.store().pool)
.await
.expect("read should succeed");
assert_eq!(exists, 0);
cleanup_dir(&dir);
}
#[tokio::test]
async fn stats_counts_agents_and_expired_tokens() {
let (store, dir) = make_store().await;
insert_token(&store, "enr-valid-test", i64::MAX as u64).await;
let ts = TestStore::new().await;
insert_token(ts.store(), "enr-valid-test", i64::MAX as u64).await;
let _issued = store
let _issued = ts
.store()
.issue_enroll_token(Duration::from_secs(60))
.await
.expect("issue should succeed");
insert_token(&store, "enr-expired-stats-test", 1).await;
insert_token(ts.store(), "enr-expired-stats-test", 1).await;
let issued_agent = store
let issued_agent = ts
.store()
.enroll("enr-valid-test")
.await
.expect("enroll should succeed for valid token");
assert!(!issued_agent.agent_id.is_empty());
let stats = store.stats().await.expect("stats should succeed");
let stats = ts.store().stats().await.expect("stats should succeed");
assert_eq!(stats.agent_count, 1);
assert_eq!(stats.enroll_token_count, 2);
assert_eq!(stats.expired_enroll_token_count, 1);
cleanup_dir(&dir);
}
#[tokio::test]
async fn revoke_agent_removes_credentials() {
let (store, dir) = make_store().await;
let ts = TestStore::new().await;
insert_token(&store, "enr-revoke-agent-test", i64::MAX as u64).await;
insert_token(ts.store(), "enr-revoke-agent-test", i64::MAX as u64).await;
let issued = store
let issued = ts
.store()
.enroll("enr-revoke-agent-test")
.await
.expect("enroll should succeed");
assert!(
store
ts.store()
.verify_agent_token(&issued.agent_id, &issued.agent_token)
.await
);
let removed = store
let removed = ts
.store()
.revoke_agent(&issued.agent_id)
.await
.expect("revoke should succeed");
assert!(removed);
assert!(
!store
!ts.store()
.verify_agent_token(&issued.agent_id, &issued.agent_token)
.await
);
let removed_again = store
let removed_again = ts
.store()
.revoke_agent(&issued.agent_id)
.await
.expect("second revoke should succeed");
assert!(!removed_again);
cleanup_dir(&dir);
}
#[tokio::test]
async fn nickname_set_and_clear_roundtrip() {
let (store, dir) = make_store().await;
let ts = TestStore::new().await;
insert_token(&store, "enr-nickname-test", i64::MAX as u64).await;
insert_token(ts.store(), "enr-nickname-test", i64::MAX as u64).await;
let issued = store
let issued = ts
.store()
.enroll("enr-nickname-test")
.await
.expect("enroll should succeed");
let updated = store
let updated = ts
.store()
.set_agent_nickname(&issued.agent_id, Some("kitchen-router"))
.await
.expect("nickname set should succeed");
assert!(updated);
let listed = store.list_agents_with_nicknames().await;
let listed = ts.store().list_agents_with_nicknames().await;
assert!(listed.iter().any(|(id, name)| {
id == &issued.agent_id && name.as_deref() == Some("kitchen-router")
}));
let cleared = store
let cleared = ts
.store()
.set_agent_nickname(&issued.agent_id, None)
.await
.expect("nickname clear should succeed");
assert!(cleared);
let listed = store.list_agents_with_nicknames().await;
let listed = ts.store().list_agents_with_nicknames().await;
assert!(
listed
.iter()
.any(|(id, name)| id == &issued.agent_id && name.is_none())
);
cleanup_dir(&dir);
}
#[tokio::test]
async fn known_device_can_hold_multiple_manual_mac_identifiers() {
let (store, dir) = make_store().await;
let ts = TestStore::new().await;
let created = store
let created = ts
.store()
.create_known_device(KnownDeviceInput {
display_name: "lda".into(),
pinned: true,
@@ -249,7 +239,8 @@ mod tests {
assert_eq!(created.identifiers.len(), 1);
assert_eq!(created.identifiers[0].value, "aa:bb:cc:dd:ee:01");
let updated = store
let updated = ts
.store()
.attach_device_identifier(
&created.device_id,
DeviceIdentifierInput {
@@ -269,7 +260,8 @@ mod tests {
.any(|identifier| identifier.value == "aa:bb:cc:dd:ee:02")
);
let matched = store
let matched = ts
.store()
.lookup_known_device_by_identifier(DeviceIdentifierInput {
kind: "mac".into(),
value: "aa:bb:cc:dd:ee:02".into(),
@@ -278,14 +270,13 @@ mod tests {
.expect("lookup should succeed")
.expect("identifier should match");
assert_eq!(matched.device_id, created.device_id);
cleanup_dir(&dir);
}
#[tokio::test]
async fn known_device_identifier_is_unique_across_devices() {
let (store, dir) = make_store().await;
let first = store
let ts = TestStore::new().await;
let first = ts
.store()
.create_known_device(KnownDeviceInput {
display_name: "lda".into(),
pinned: true,
@@ -297,7 +288,8 @@ mod tests {
})
.await
.expect("first device should create");
let second = store
let second = ts
.store()
.create_known_device(KnownDeviceInput {
display_name: "other".into(),
pinned: false,
@@ -307,7 +299,8 @@ mod tests {
.await
.expect("second device should create");
let err = store
let err = ts
.store()
.attach_device_identifier(
&second.device_id,
DeviceIdentifierInput {
@@ -322,7 +315,11 @@ mod tests {
.contains("failed attaching device identifier")
);
let listed = store.list_known_devices().await.expect("list should work");
let listed = ts
.store()
.list_known_devices()
.await
.expect("list should work");
assert_eq!(listed.len(), 2);
assert!(
listed
@@ -333,14 +330,13 @@ mod tests {
.len()
== 1
);
cleanup_dir(&dir);
}
#[tokio::test]
async fn device_identifier_can_be_detached_manually() {
let (store, dir) = make_store().await;
let created = store
let ts = TestStore::new().await;
let created = ts
.store()
.create_known_device(KnownDeviceInput {
display_name: "lda".into(),
pinned: true,
@@ -359,7 +355,8 @@ mod tests {
.await
.expect("known device should create");
let updated = store
let updated = ts
.store()
.detach_device_identifier(&created.device_id, "ip:192.168.1.2")
.await
.expect("identifier detach should succeed")
@@ -371,7 +368,8 @@ mod tests {
"mac:aa:bb:cc:dd:ee:ff"
);
let unmatched = store
let unmatched = ts
.store()
.lookup_known_device_by_identifier(DeviceIdentifierInput {
kind: "ip".into(),
value: "192.168.1.2".into(),
@@ -379,14 +377,13 @@ mod tests {
.await
.expect("lookup should succeed");
assert!(unmatched.is_none());
cleanup_dir(&dir);
}
#[tokio::test]
async fn merge_known_devices_moves_identifiers_and_deletes_source() {
let (store, dir) = make_store().await;
let target = store
let ts = TestStore::new().await;
let target = ts
.store()
.create_known_device(KnownDeviceInput {
display_name: "lda".into(),
pinned: true,
@@ -398,7 +395,8 @@ mod tests {
})
.await
.expect("target should create");
let source = store
let source = ts
.store()
.create_known_device(KnownDeviceInput {
display_name: "lda duplicate".into(),
pinned: false,
@@ -411,7 +409,8 @@ mod tests {
.await
.expect("source should create");
let merged = store
let merged = ts
.store()
.merge_known_devices(&target.device_id, &source.device_id)
.await
.expect("merge should succeed")
@@ -426,362 +425,19 @@ mod tests {
.any(|identifier| identifier.value == "aa:bb:cc:dd:ee:02")
);
assert!(
store
ts.store()
.get_known_device(&source.device_id)
.await
.expect("source lookup should work")
.is_none()
);
cleanup_dir(&dir);
}
#[tokio::test]
async fn agent_observations_upsert_current_state_and_events() {
let (store, dir) = make_store().await;
let accepted = store
.upsert_agent_observations(
"agent-a",
vec![crate::state::AgentDeviceObservationInput {
kind: "dhcp".into(),
action: "update".into(),
mac: Some("AA:BB:CC:DD:EE:FF".into()),
ip: Some("192.168.1.10".into()),
hostname: Some("lda".into()),
first_seen_unix: 10,
last_seen_unix: 20,
}],
)
.await
.expect("observation upsert should succeed");
assert_eq!(accepted, 1);
let rows = store
.list_agent_observations(Some("agent-a"), 10)
.await
.expect("observations should list");
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].mac.as_deref(), Some("aa:bb:cc:dd:ee:ff"));
assert_eq!(rows[0].hostname.as_deref(), Some("lda"));
let event_count =
sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM agent_device_observation_events")
.fetch_one(&store.pool)
.await
.expect("event count should read");
assert_eq!(event_count, 1);
let accepted = store
.upsert_agent_observations(
"agent-a",
vec![crate::state::AgentDeviceObservationInput {
kind: "dhcp".into(),
action: "update".into(),
mac: Some("AA:BB:CC:DD:EE:FF".into()),
ip: Some("192.168.1.10".into()),
hostname: Some("lda".into()),
first_seen_unix: 10,
last_seen_unix: 20,
}],
)
.await
.expect("duplicate observation upsert should succeed");
assert_eq!(accepted, 1);
let event_count =
sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM agent_device_observation_events")
.fetch_one(&store.pool)
.await
.expect("event count should read");
assert_eq!(event_count, 1);
store
.upsert_agent_observations(
"agent-a",
vec![crate::state::AgentDeviceObservationInput {
kind: "dhcp".into(),
action: "remove".into(),
mac: Some("AA:BB:CC:DD:EE:FF".into()),
ip: Some("192.168.1.10".into()),
hostname: Some("lda".into()),
first_seen_unix: 10,
last_seen_unix: 30,
}],
)
.await
.expect("changed observation upsert should succeed");
let events = store
.list_agent_observation_events(
Some("agent-a"),
None,
Some("aa:bb:cc:dd:ee:ff"),
None,
None,
10,
)
.await
.expect("observation events should list");
assert_eq!(events.len(), 2);
assert_eq!(events[0].action, "remove");
assert_eq!(
events[0].observation_key,
"agent:agent-a:dhcp:mac:aa:bb:cc:dd:ee:ff"
);
cleanup_dir(&dir);
}
#[tokio::test]
async fn observation_snapshot_prunes_missing_keys() {
let (store, dir) = make_store().await;
store
.upsert_agent_observations(
"agent-a",
vec![
crate::state::AgentDeviceObservationInput {
kind: "dhcp".into(),
action: "update".into(),
mac: Some("AA:BB:CC:DD:EE:01".into()),
ip: Some("192.168.1.10".into()),
hostname: Some("first".into()),
first_seen_unix: 10,
last_seen_unix: 20,
},
crate::state::AgentDeviceObservationInput {
kind: "dhcp".into(),
action: "update".into(),
mac: Some("AA:BB:CC:DD:EE:02".into()),
ip: Some("192.168.1.11".into()),
hostname: Some("second".into()),
first_seen_unix: 10,
last_seen_unix: 20,
},
],
)
.await
.expect("initial observations should upsert");
store
.upsert_agent_observations_snapshot(
"agent-a",
"dhcp",
vec![crate::state::AgentDeviceObservationInput {
kind: "dhcp".into(),
action: "update".into(),
mac: Some("AA:BB:CC:DD:EE:01".into()),
ip: Some("192.168.1.10".into()),
hostname: Some("first".into()),
first_seen_unix: 10,
last_seen_unix: 30,
}],
)
.await
.expect("snapshot upsert should succeed");
let rows = store
.list_agent_observations(Some("agent-a"), 10)
.await
.expect("observations should list");
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].mac.as_deref(), Some("aa:bb:cc:dd:ee:01"));
cleanup_dir(&dir);
}
#[tokio::test]
async fn observation_gc_removes_stale_rows() {
let (store, dir) = make_store().await;
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("time should be monotonic")
.as_secs();
let old = now.saturating_sub(10);
store
.upsert_agent_observations(
"agent-a",
vec![
crate::state::AgentDeviceObservationInput {
kind: "dhcp".into(),
action: "update".into(),
mac: Some("AA:BB:CC:DD:EE:10".into()),
ip: Some("192.168.1.20".into()),
hostname: Some("old".into()),
first_seen_unix: old,
last_seen_unix: old,
},
crate::state::AgentDeviceObservationInput {
kind: "dhcp".into(),
action: "update".into(),
mac: Some("AA:BB:CC:DD:EE:11".into()),
ip: Some("192.168.1.21".into()),
hostname: Some("fresh".into()),
first_seen_unix: now,
last_seen_unix: now,
},
],
)
.await
.expect("observations should upsert");
let removed = store
.gc_stale_observations(Duration::from_secs(5))
.await
.expect("gc should succeed");
assert!(removed >= 1);
let rows = store
.list_agent_observations(Some("agent-a"), 10)
.await
.expect("observations should list");
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].hostname.as_deref(), Some("fresh"));
cleanup_dir(&dir);
}
#[tokio::test]
async fn agent_observation_views_include_matching_known_device() {
let (store, dir) = make_store().await;
let device = store
.create_known_device(KnownDeviceInput {
display_name: "lda".into(),
pinned: true,
notes: None,
identifiers: vec![DeviceIdentifierInput {
kind: "mac".into(),
value: "aa:bb:cc:dd:ee:ff".into(),
}],
})
.await
.expect("known device should create");
store
.upsert_agent_observations(
"agent-a",
vec![
crate::state::AgentDeviceObservationInput {
kind: "dhcp".into(),
action: "update".into(),
mac: Some("AA:BB:CC:DD:EE:FF".into()),
ip: Some("192.168.1.10".into()),
hostname: Some("lda".into()),
first_seen_unix: 10,
last_seen_unix: 20,
},
crate::state::AgentDeviceObservationInput {
kind: "dhcp".into(),
action: "update".into(),
mac: Some("00:11:22:33:44:55".into()),
ip: Some("192.168.1.11".into()),
hostname: Some("guest".into()),
first_seen_unix: 11,
last_seen_unix: 21,
},
],
)
.await
.expect("observation upsert should succeed");
let rows = store
.list_agent_observation_views(Some("agent-a"), 10)
.await
.expect("observation views should list");
assert_eq!(rows.len(), 2);
let known = rows
.iter()
.find(|row| row.mac.as_deref() == Some("aa:bb:cc:dd:ee:ff"))
.expect("known observation should be present");
let known_device = known
.known_device
.as_ref()
.expect("known observation should join device");
assert_eq!(known_device.device_id, device.device_id);
assert_eq!(known_device.display_name, "lda");
assert!(known_device.pinned);
let unknown = rows
.iter()
.find(|row| row.mac.as_deref() == Some("00:11:22:33:44:55"))
.expect("unknown observation should be present");
assert!(unknown.known_device.is_none());
cleanup_dir(&dir);
}
#[tokio::test]
async fn observation_identifier_can_be_attached_to_known_device() {
let (store, dir) = make_store().await;
let device = store
.create_known_device(KnownDeviceInput {
display_name: "lda".into(),
pinned: true,
notes: None,
identifiers: Vec::new(),
})
.await
.expect("known device should create");
store
.upsert_agent_observations(
"agent-a",
vec![crate::state::AgentDeviceObservationInput {
kind: "dhcp".into(),
action: "update".into(),
mac: Some("AA:BB:CC:DD:EE:FF".into()),
ip: Some("192.168.1.10".into()),
hostname: Some("lda".into()),
first_seen_unix: 10,
last_seen_unix: 20,
}],
)
.await
.expect("observation upsert should succeed");
let observation = store
.list_agent_observations(Some("agent-a"), 10)
.await
.expect("observations should list")
.pop()
.expect("observation should exist");
let updated = store
.attach_observation_identifier(&device.device_id, &observation.observation_key)
.await
.expect("observation identifier should attach")
.expect("device should exist");
assert_eq!(updated.identifiers.len(), 1);
assert_eq!(updated.identifiers[0].kind, "mac");
assert_eq!(updated.identifiers[0].value, "aa:bb:cc:dd:ee:ff");
let views = store
.list_agent_observation_views(Some("agent-a"), 10)
.await
.expect("observation views should list");
assert_eq!(
views[0]
.known_device
.as_ref()
.map(|device| device.device_id.as_str()),
Some(device.device_id.as_str())
);
cleanup_dir(&dir);
}
#[tokio::test]
async fn audit_events_append_and_filter() {
let (store, dir) = make_store().await;
let ts = TestStore::new().await;
store
ts.store()
.append_audit_event(crate::state::AuditEventInput {
actor_type: "admin_api".into(),
actor_id: None,
@@ -796,7 +452,7 @@ mod tests {
.await
.expect("append first event should succeed");
store
ts.store()
.append_audit_event(crate::state::AuditEventInput {
actor_type: "agent".into(),
actor_id: Some("agent-2".into()),
@@ -811,7 +467,8 @@ mod tests {
.await
.expect("append second event should succeed");
let all = store
let all = ts
.store()
.list_audit_events(crate::state::AuditEventFilter {
limit: 10,
..Default::default()
@@ -820,7 +477,8 @@ mod tests {
.expect("list all should succeed");
assert_eq!(all.len(), 2);
let filtered = store
let filtered = ts
.store()
.list_audit_events(crate::state::AuditEventFilter {
agent_id: Some("agent-1".into()),
event_type: Some("command_result".into()),
@@ -833,7 +491,8 @@ mod tests {
assert_eq!(filtered.len(), 1);
assert_eq!(filtered[0].request_id.as_deref(), Some("req-1"));
let rejected = store
let rejected = ts
.store()
.list_audit_events(crate::state::AuditEventFilter {
outcome: Some("rejected".into()),
limit: 10,
@@ -843,12 +502,11 @@ mod tests {
.expect("rejected list should succeed");
assert_eq!(rejected.len(), 1);
assert_eq!(rejected[0].latency_ms, None);
cleanup_dir(&dir);
}
#[tokio::test]
async fn alert_transitions_track_open_and_resolve() {
let (store, dir) = make_store().await;
let ts = TestStore::new().await;
let alert = crate::state::AlertState {
alert_id: "agent_offline:agent-a".into(),
kind: "agent_offline".into(),
@@ -862,26 +520,28 @@ mod tests {
metadata: serde_json::json!({}),
};
let opened = store
let opened = ts
.store()
.sync_alert_transitions(std::slice::from_ref(&alert))
.await
.expect("open transition should succeed");
assert_eq!(opened.len(), 1);
assert_eq!(opened[0].to_status, "active");
let resolved = store
let resolved = ts
.store()
.sync_alert_transitions(&[])
.await
.expect("resolve transition should succeed");
assert_eq!(resolved.len(), 1);
assert_eq!(resolved[0].to_status, "resolved");
let history = store
let history = ts
.store()
.list_alert_transitions(None, 10)
.await
.expect("history should load");
assert!(history.len() >= 2);
cleanup_dir(&dir);
}
#[tokio::test]
@@ -923,7 +583,7 @@ mod tests {
.contains("invalid or already-used enroll token")
);
cleanup_dir(&dir);
let _ = fs::remove_dir_all(&dir);
}
#[tokio::test]
@@ -984,6 +644,6 @@ mod tests {
.any(|token| token.enroll_token == "enr-import-test")
);
cleanup_dir(&dir);
let _ = fs::remove_dir_all(&dir);
}
}
wakey-control-plane/src/state/store/agent_devices.rs
+608
View File
@@ -0,0 +1,608 @@
use anyhow::{Context, Result};
use wakey_core::{Device, DeviceId};
use super::Store;
use super::helpers::core::*;
use crate::state::types::*;
impl Store {
/// Replace the complete device snapshot for an agent.
///
/// This is the only write path for agent device state. Every sync source
/// (WebSocket snapshot, fleet refresh, HTTP upload) must go through here.
pub async fn replace_agent_device_snapshot(
&self,
agent_id: &str,
devices: &[Device],
) -> Result<usize> {
let mut tx = self
.pool
.begin()
.await
.context("failed starting device snapshot transaction")?;
let mut incoming_keys = std::collections::HashSet::with_capacity(devices.len());
let snapshot_time = now_unix();
let snapshot_time_i64 = i64::try_from(snapshot_time).context("snapshot time overflow")?;
let existing_keys: Vec<String> =
sqlx::query_scalar("SELECT device_key FROM agent_devices WHERE agent_id = ?1")
.bind(agent_id)
.fetch_all(&mut *tx)
.await
.context("failed fetching existing keys")?;
for device in devices {
let Some(device_id) = &device.id else {
continue;
};
let device_key = device_key_from_id(device_id);
incoming_keys.insert(device_key.clone());
let presence = presence_to_str(device.presence);
let display_name: Option<String> = None;
sqlx::query(
"INSERT INTO agent_devices (agent_id, device_key, presence, display_name, first_seen_unix, last_seen_unix)
VALUES (?1, ?2, ?3, ?4, ?5, ?6)
ON CONFLICT (agent_id, device_key) DO UPDATE SET
presence = excluded.presence,
last_seen_unix = excluded.last_seen_unix",
)
.bind(agent_id)
.bind(&device_key)
.bind(presence)
.bind(&display_name)
.bind(snapshot_time_i64)
.bind(snapshot_time_i64)
.execute(&mut *tx)
.await
.context("failed upserting agent device")?;
// Replace child rows.
sqlx::query("DELETE FROM agent_device_macs WHERE agent_id = ?1 AND device_key = ?2")
.bind(agent_id)
.bind(&device_key)
.execute(&mut *tx)
.await
.context("failed deleting device macs")?;
for mac in &device.macs {
let mac_str = mac.to_string().to_ascii_lowercase();
sqlx::query(
"INSERT INTO agent_device_macs (agent_id, device_key, mac) VALUES (?1, ?2, ?3)",
)
.bind(agent_id)
.bind(&device_key)
.bind(&mac_str)
.execute(&mut *tx)
.await
.context("failed inserting device mac")?;
}
sqlx::query("DELETE FROM agent_device_ips WHERE agent_id = ?1 AND device_key = ?2")
.bind(agent_id)
.bind(&device_key)
.execute(&mut *tx)
.await
.context("failed deleting device ips")?;
for ip in &device.ips {
let ip_str = ip.to_string();
sqlx::query(
"INSERT INTO agent_device_ips (agent_id, device_key, ip) VALUES (?1, ?2, ?3)",
)
.bind(agent_id)
.bind(&device_key)
.bind(&ip_str)
.execute(&mut *tx)
.await
.context("failed inserting device ip")?;
}
sqlx::query(
"DELETE FROM agent_device_hostnames WHERE agent_id = ?1 AND device_key = ?2",
)
.bind(agent_id)
.bind(&device_key)
.execute(&mut *tx)
.await
.context("failed deleting device hostnames")?;
for hostname in &device.names {
sqlx::query(
"INSERT INTO agent_device_hostnames (agent_id, device_key, hostname) VALUES (?1, ?2, ?3)",
)
.bind(agent_id)
.bind(&device_key)
.bind(hostname)
.execute(&mut *tx)
.await
.context("failed inserting device hostname")?;
}
sqlx::query("DELETE FROM agent_device_facts WHERE agent_id = ?1 AND device_key = ?2")
.bind(agent_id)
.bind(&device_key)
.execute(&mut *tx)
.await
.context("failed deleting device facts")?;
for observation in &device.observations {
let fact_json =
serde_json::to_string(observation).context("failed serializing fact")?;
sqlx::query(
"INSERT INTO agent_device_facts (agent_id, device_key, fact_json) VALUES (?1, ?2, ?3)",
)
.bind(agent_id)
.bind(&device_key)
.bind(&fact_json)
.execute(&mut *tx)
.await
.context("failed inserting device fact")?;
}
}
for old_key in existing_keys {
if !incoming_keys.contains(&old_key) {
sqlx::query("DELETE FROM agent_devices WHERE agent_id = ?1 AND device_key = ?2")
.bind(agent_id)
.bind(&old_key)
.execute(&mut *tx)
.await
.context("failed pruning old agent device")?;
}
}
tx.commit()
.await
.context("failed committing device snapshot")?;
Ok(incoming_keys.len())
}
/// List all agent device rows with their child MAC/IP/hostname/fact rows.
pub async fn list_agent_device_rows(&self) -> Result<Vec<AgentDeviceWithChildren>> {
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 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, facts))
}
/// List agent device rows for a single agent.
#[allow(dead_code)]
pub async fn list_agent_device_rows_for_agent(
&self,
agent_id: &str,
) -> Result<Vec<AgentDeviceWithChildren>> {
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 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, facts))
}
}
fn assemble_device_rows(
devices: Vec<AgentDeviceRow>,
macs: Vec<AgentDeviceMacRow>,
ips: Vec<AgentDeviceIpRow>,
hostnames: Vec<AgentDeviceHostnameRow>,
facts: Vec<AgentDeviceFactRow>,
) -> Vec<AgentDeviceWithChildren> {
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<String>> = 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 fact_map: BTreeMap<(&str, &str), Vec<String>> = 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<macaddr::MacAddr> = mac_map
.get(&key)
.into_iter()
.flatten()
.filter_map(|row| macaddr::MacAddr::try_from(*row).ok())
.collect();
let ips: Vec<std::net::IpAddr> = ip_map
.get(&key)
.into_iter()
.flatten()
.filter_map(|row| std::net::IpAddr::try_from(*row).ok())
.collect();
let hostnames = hostname_map.get(&key).cloned().unwrap_or_default();
let facts = fact_map.get(&key).cloned().unwrap_or_default();
AgentDeviceWithChildren {
macs,
ips,
hostnames,
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()),
DeviceId::Ip(ip) => format!("ip:{ip}"),
}
}
fn presence_to_str(presence: wakey_core::Presence) -> &'static str {
match presence {
wakey_core::Presence::Online => "online",
wakey_core::Presence::LikelyOnline => "likely_online",
wakey_core::Presence::Unknown => "unknown",
wakey_core::Presence::Offline => "offline",
}
}
#[cfg(test)]
mod tests {
use super::super::helpers::test_helpers::TestStore;
use std::time::Duration;
use wakey_core::{
Device, DeviceId, DeviceObservationFact, NeighborEntry, NeighborState, Presence,
};
fn sample_device(mac: &str, ip: &str, name: &str) -> Device {
Device {
id: Some(DeviceId::Mac(mac.parse().expect("mac"))),
names: vec![name.to_string()],
ips: vec![ip.parse().expect("ip")],
macs: vec![mac.parse().expect("mac")],
interfaces: vec!["br-lan".to_string()],
neighbors: vec![NeighborEntry {
ip: ip.parse().expect("ip"),
dev: Some("br-lan".to_string()),
mac: Some(mac.parse().expect("mac")),
state: NeighborState::Reachable,
}],
leases: vec![],
observations: vec![],
presence: Presence::Online,
}
}
#[tokio::test]
async fn snapshot_inserts_devices_and_children() {
let ts = TestStore::new().await;
let devices = vec![sample_device(
"aa:bb:cc:dd:ee:01",
"192.168.1.10",
"first-pc",
)];
let count = ts
.store()
.replace_agent_device_snapshot("agent-a", &devices)
.await
.expect("snapshot should succeed");
assert_eq!(count, 1);
let rows = ts
.store()
.list_agent_device_rows_for_agent("agent-a")
.await
.expect("list should succeed");
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].device.device_key, "mac:aa:bb:cc:dd:ee:01");
assert_eq!(rows[0].device.presence, "online");
assert_eq!(
rows[0].macs,
vec!["aa:bb:cc:dd:ee:01".parse::<macaddr::MacAddr>().unwrap()]
);
assert_eq!(
rows[0].ips,
vec!["192.168.1.10".parse::<std::net::IpAddr>().unwrap()]
);
assert_eq!(rows[0].hostnames, vec!["first-pc"]);
}
#[tokio::test]
async fn second_snapshot_prunes_missing_devices() {
let ts = TestStore::new().await;
let first = vec![
sample_device("aa:bb:cc:dd:ee:01", "192.168.1.10", "first"),
sample_device("aa:bb:cc:dd:ee:02", "192.168.1.11", "second"),
];
ts.store()
.replace_agent_device_snapshot("agent-a", &first)
.await
.expect("first snapshot should succeed");
let second = vec![sample_device("aa:bb:cc:dd:ee:01", "192.168.1.10", "first")];
ts.store()
.replace_agent_device_snapshot("agent-a", &second)
.await
.expect("second snapshot should succeed");
let rows = ts
.store()
.list_agent_device_rows_for_agent("agent-a")
.await
.expect("list should succeed");
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].device.device_key, "mac:aa:bb:cc:dd:ee:01");
}
#[tokio::test]
async fn empty_snapshot_clears_devices() {
let ts = TestStore::new().await;
let devices = vec![sample_device("aa:bb:cc:dd:ee:01", "192.168.1.10", "first")];
ts.store()
.replace_agent_device_snapshot("agent-a", &devices)
.await
.expect("snapshot should succeed");
ts.store()
.replace_agent_device_snapshot("agent-a", &[])
.await
.expect("empty snapshot should succeed");
let rows = ts
.store()
.list_agent_device_rows_for_agent("agent-a")
.await
.expect("list should succeed");
assert_eq!(rows.len(), 0);
}
#[tokio::test]
async fn first_seen_survives_snapshot_update() {
let ts = TestStore::new().await;
let devices = vec![sample_device("aa:bb:cc:dd:ee:01", "192.168.1.10", "first")];
ts.store()
.replace_agent_device_snapshot("agent-a", &devices)
.await
.expect("first snapshot should succeed");
let first_seen = ts
.store()
.list_agent_device_rows_for_agent("agent-a")
.await
.expect("list should succeed")[0]
.device
.first_seen_unix;
tokio::time::sleep(Duration::from_secs(1)).await;
ts.store()
.replace_agent_device_snapshot("agent-a", &devices)
.await
.expect("second snapshot should succeed");
let rows = ts
.store()
.list_agent_device_rows_for_agent("agent-a")
.await
.expect("list should succeed");
assert_eq!(rows[0].device.first_seen_unix, first_seen);
}
#[tokio::test]
async fn snapshot_does_not_prune_other_agents() {
let ts = TestStore::new().await;
let devices_a = vec![sample_device("aa:bb:cc:dd:ee:01", "192.168.1.10", "first")];
let devices_b = vec![sample_device("aa:bb:cc:dd:ee:02", "192.168.1.11", "second")];
ts.store()
.replace_agent_device_snapshot("agent-a", &devices_a)
.await
.expect("snapshot a should succeed");
ts.store()
.replace_agent_device_snapshot("agent-b", &devices_b)
.await
.expect("snapshot b should succeed");
ts.store()
.replace_agent_device_snapshot("agent-a", &[])
.await
.expect("clear a should succeed");
let rows_b = ts
.store()
.list_agent_device_rows_for_agent("agent-b")
.await
.expect("list b should succeed");
assert_eq!(rows_b.len(), 1);
assert_eq!(rows_b[0].device.device_key, "mac:aa:bb:cc:dd:ee:02");
}
#[tokio::test]
async fn device_without_id_is_skipped() {
let ts = TestStore::new().await;
let devices = vec![Device {
id: None,
names: vec!["no-id".to_string()],
ips: vec![],
macs: vec![],
interfaces: vec![],
neighbors: vec![],
leases: vec![],
observations: vec![],
presence: Presence::Unknown,
}];
let count = ts
.store()
.replace_agent_device_snapshot("agent-a", &devices)
.await
.expect("snapshot should succeed");
assert_eq!(count, 0);
let rows = ts
.store()
.list_agent_device_rows_for_agent("agent-a")
.await
.expect("list should succeed");
assert_eq!(rows.len(), 0);
}
#[tokio::test]
async fn facts_are_stored_as_json() {
let ts = TestStore::new().await;
let devices = vec![Device {
id: Some(DeviceId::Mac("aa:bb:cc:dd:ee:01".parse().expect("mac"))),
names: vec!["pc".to_string()],
ips: vec!["192.168.1.10".parse().expect("ip")],
macs: vec!["aa:bb:cc:dd:ee:01".parse().expect("mac")],
interfaces: vec![],
neighbors: vec![],
leases: vec![],
observations: vec![DeviceObservationFact {
kind: "dhcp".to_string(),
action: "update".to_string(),
mac: Some("aa:bb:cc:dd:ee:01".parse().expect("mac")),
ip: Some("192.168.1.10".parse().expect("ip")),
hostname: Some("pc".to_string()),
first_seen_unix: Some(10),
last_seen_unix: Some(20),
}],
presence: Presence::LikelyOnline,
}];
ts.store()
.replace_agent_device_snapshot("agent-a", &devices)
.await
.expect("snapshot should succeed");
let rows = ts
.store()
.list_agent_device_rows_for_agent("agent-a")
.await
.expect("list should succeed");
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].facts.len(), 1);
assert!(rows[0].facts[0].contains("dhcp"));
}
}
wakey-control-plane/src/state/store/devices.rs
-15
View File
@@ -181,21 +181,6 @@ impl Store {
self.get_known_device(device_id).await
}
pub async fn attach_observation_identifier(
&self,
device_id: &str,
observation_key: &str,
) -> Result<Option<KnownDevice>> {
let observation = get_observation_identifier_row(&self.pool, observation_key).await?;
let Some(observation) = observation else {
anyhow::bail!("observation not found");
};
let input = observation_identifier_to_input(observation)?;
self.attach_device_identifier(device_id, input).await
}
pub async fn detach_device_identifier(
&self,
device_id: &str,
wakey-control-plane/src/state/store/helpers.rs
+2
View File
@@ -4,3 +4,5 @@ pub(super) mod alerts_audit;
pub(super) mod core;
pub(super) mod legacy;
pub(super) mod rows;
#[cfg(test)]
pub(crate) mod test_helpers;
wakey-control-plane/src/state/store/helpers/alerts_audit.rs
+83 -82
View File
@@ -1,159 +1,162 @@
use super::*;
pub(in crate::state::store) async fn insert_audit_event(
pool: &SqlitePool,
key: &str,
event: &AuditEvent,
) -> Result<()> {
pub async fn insert_audit_event(pool: &SqlitePool, key: &str, event: &AuditEvent) -> Result<()> {
let metadata_json =
serde_json::to_string(&event.metadata).context("failed to encode audit metadata")?;
sqlx::query(
let ts_unix = i64::try_from(event.ts_unix).context("audit timestamp overflow")?;
let latency_ms = event
.latency_ms
.map(i64::try_from)
.transpose()
.context("audit latency overflow")?;
sqlx::query!(
"INSERT OR REPLACE INTO audit_events
(event_key, event_id, ts_unix, actor_type, actor_id, agent_id, request_id,
event_type, outcome, latency_ms, message, metadata_json)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12)",
key,
event.event_id,
ts_unix,
event.actor_type,
event.actor_id,
event.agent_id,
event.request_id,
event.event_type,
event.outcome,
latency_ms,
event.message,
metadata_json
)
.bind(key)
.bind(&event.event_id)
.bind(i64::try_from(event.ts_unix).context("audit timestamp overflow")?)
.bind(&event.actor_type)
.bind(&event.actor_id)
.bind(&event.agent_id)
.bind(&event.request_id)
.bind(&event.event_type)
.bind(&event.outcome)
.bind(
event
.latency_ms
.map(i64::try_from)
.transpose()
.context("audit latency overflow")?,
)
.bind(&event.message)
.bind(metadata_json)
.execute(pool)
.await
.context("failed persisting audit event")?;
Ok(())
}
pub(in crate::state::store) async fn insert_active_alert(
pub async fn insert_active_alert(
tx: &mut Transaction<'_, Sqlite>,
alert: &AlertState,
) -> Result<()> {
let metadata_json =
serde_json::to_string(&alert.metadata).context("failed to encode active alert metadata")?;
sqlx::query(
let alert_value = i64::try_from(alert.value).context("active alert value overflow")?;
let alert_threshold =
i64::try_from(alert.threshold).context("active alert threshold overflow")?;
let last_seen_unix =
i64::try_from(alert.last_seen_unix).context("active alert timestamp overflow")?;
sqlx::query!(
"INSERT INTO active_alerts
(alert_id, kind, severity, status, agent_id, message, value, threshold,
last_seen_unix, metadata_json)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
alert.alert_id,
alert.kind,
alert.severity,
alert.status,
alert.agent_id,
alert.message,
alert_value,
alert_threshold,
last_seen_unix,
metadata_json
)
.bind(&alert.alert_id)
.bind(&alert.kind)
.bind(&alert.severity)
.bind(&alert.status)
.bind(&alert.agent_id)
.bind(&alert.message)
.bind(i64::try_from(alert.value).context("active alert value overflow")?)
.bind(i64::try_from(alert.threshold).context("active alert threshold overflow")?)
.bind(i64::try_from(alert.last_seen_unix).context("active alert timestamp overflow")?)
.bind(metadata_json)
.execute(&mut **tx)
.await
.context("failed writing active alert snapshot")?;
Ok(())
}
pub(in crate::state::store) async fn insert_active_alert_pool(
pool: &SqlitePool,
alert: &AlertState,
) -> Result<()> {
pub async fn insert_active_alert_pool(pool: &SqlitePool, alert: &AlertState) -> Result<()> {
let metadata_json =
serde_json::to_string(&alert.metadata).context("failed to encode active alert metadata")?;
sqlx::query(
let alert_value = i64::try_from(alert.value).context("active alert value overflow")?;
let alert_threshold =
i64::try_from(alert.threshold).context("active alert threshold overflow")?;
let last_seen_unix =
i64::try_from(alert.last_seen_unix).context("active alert timestamp overflow")?;
sqlx::query!(
"INSERT OR REPLACE INTO active_alerts
(alert_id, kind, severity, status, agent_id, message, value, threshold,
last_seen_unix, metadata_json)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
alert.alert_id,
alert.kind,
alert.severity,
alert.status,
alert.agent_id,
alert.message,
alert_value,
alert_threshold,
last_seen_unix,
metadata_json
)
.bind(&alert.alert_id)
.bind(&alert.kind)
.bind(&alert.severity)
.bind(&alert.status)
.bind(&alert.agent_id)
.bind(&alert.message)
.bind(i64::try_from(alert.value).context("active alert value overflow")?)
.bind(i64::try_from(alert.threshold).context("active alert threshold overflow")?)
.bind(i64::try_from(alert.last_seen_unix).context("active alert timestamp overflow")?)
.bind(metadata_json)
.execute(pool)
.await
.context("failed writing active alert snapshot")?;
Ok(())
}
pub(in crate::state::store) async fn insert_alert_transition(
pub async fn insert_alert_transition(
tx: &mut Transaction<'_, Sqlite>,
key: &str,
transition: &AlertTransition,
) -> Result<()> {
let metadata_json = serde_json::to_string(&transition.metadata)
.context("failed to encode alert transition metadata")?;
sqlx::query(
let ts_unix = i64::try_from(transition.ts_unix).context("alert timestamp overflow")?;
sqlx::query!(
"INSERT INTO alert_transitions
(transition_key, transition_id, ts_unix, alert_id, kind, agent_id,
from_status, to_status, message, metadata_json)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
key,
transition.transition_id,
ts_unix,
transition.alert_id,
transition.kind,
transition.agent_id,
transition.from_status,
transition.to_status,
transition.message,
metadata_json
)
.bind(key)
.bind(&transition.transition_id)
.bind(i64::try_from(transition.ts_unix).context("alert timestamp overflow")?)
.bind(&transition.alert_id)
.bind(&transition.kind)
.bind(&transition.agent_id)
.bind(&transition.from_status)
.bind(&transition.to_status)
.bind(&transition.message)
.bind(metadata_json)
.execute(&mut **tx)
.await
.context("failed persisting alert transition")?;
Ok(())
}
pub(in crate::state::store) async fn insert_alert_transition_pool(
pub async fn insert_alert_transition_pool(
pool: &SqlitePool,
key: &str,
transition: &AlertTransition,
) -> Result<()> {
let metadata_json = serde_json::to_string(&transition.metadata)
.context("failed to encode alert transition metadata")?;
sqlx::query(
let ts_unix = i64::try_from(transition.ts_unix).context("alert timestamp overflow")?;
sqlx::query!(
"INSERT OR REPLACE INTO alert_transitions
(transition_key, transition_id, ts_unix, alert_id, kind, agent_id,
from_status, to_status, message, metadata_json)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
key,
transition.transition_id,
ts_unix,
transition.alert_id,
transition.kind,
transition.agent_id,
transition.from_status,
transition.to_status,
transition.message,
metadata_json
)
.bind(key)
.bind(&transition.transition_id)
.bind(i64::try_from(transition.ts_unix).context("alert timestamp overflow")?)
.bind(&transition.alert_id)
.bind(&transition.kind)
.bind(&transition.agent_id)
.bind(&transition.from_status)
.bind(&transition.to_status)
.bind(&transition.message)
.bind(metadata_json)
.execute(pool)
.await
.context("failed persisting alert transition")?;
Ok(())
}
pub(in crate::state::store) fn audit_event_from_row(
row: sqlx::sqlite::SqliteRow,
) -> Result<AuditEvent> {
pub fn audit_event_from_row(row: sqlx::sqlite::SqliteRow) -> Result<AuditEvent> {
let ts_unix: i64 = row.try_get("ts_unix")?;
let latency_ms: Option<i64> = row.try_get("latency_ms")?;
let metadata_json: String = row.try_get("metadata_json")?;
@@ -175,7 +178,7 @@ pub(in crate::state::store) fn audit_event_from_row(
})
}
pub(in crate::state::store) fn alert_state_from_row(row: AlertStateRow) -> Result<AlertState> {
pub fn alert_state_from_row(row: AlertStateRow) -> Result<AlertState> {
Ok(AlertState {
alert_id: row.alert_id,
kind: row.kind,
@@ -193,9 +196,7 @@ pub(in crate::state::store) fn alert_state_from_row(row: AlertStateRow) -> Resul
})
}
pub(in crate::state::store) fn alert_transition_from_row(
row: AlertTransitionRow,
) -> Result<AlertTransition> {
pub fn alert_transition_from_row(row: AlertTransitionRow) -> Result<AlertTransition> {
Ok(AlertTransition {
transition_id: row.transition_id,
ts_unix: u64::try_from(row.ts_unix).context("negative alert timestamp in state db")?,
wakey-control-plane/src/state/store/helpers/core.rs
+19 -204
View File
@@ -1,6 +1,6 @@
use super::*;
pub(in crate::state::store) async fn open_sqlite_pool(path: &Path) -> Result<SqlitePool> {
pub async fn open_sqlite_pool(path: &Path) -> Result<SqlitePool> {
let options = SqliteConnectOptions::new()
.filename(path)
.create_if_missing(true)
@@ -14,7 +14,7 @@ pub(in crate::state::store) async fn open_sqlite_pool(path: &Path) -> Result<Sql
.with_context(|| format!("failed to open SQLite state db {}", path.display()))
}
pub(in crate::state::store) async fn sql_count(pool: &SqlitePool, table: &str) -> Result<i64> {
pub async fn sql_count(pool: &SqlitePool, table: &str) -> Result<i64> {
let sql = format!("SELECT COUNT(*) FROM {table}");
sqlx::query_scalar::<_, i64>(&sql)
.fetch_one(pool)
@@ -22,14 +22,14 @@ pub(in crate::state::store) async fn sql_count(pool: &SqlitePool, table: &str) -
.with_context(|| format!("failed counting {table}"))
}
pub(in crate::state::store) fn now_unix() -> u64 {
pub fn now_unix() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0)
}
pub(in crate::state::store) fn decode_expiry(raw: &[u8]) -> Result<u64> {
pub fn decode_expiry(raw: &[u8]) -> Result<u64> {
if raw.len() != 8 {
anyhow::bail!("invalid token expiry length {}", raw.len());
}
@@ -38,7 +38,7 @@ pub(in crate::state::store) fn decode_expiry(raw: &[u8]) -> Result<u64> {
Ok(u64::from_le_bytes(arr))
}
pub(in crate::state::store) fn decode_schema(raw: &[u8]) -> Result<u32> {
pub fn decode_schema(raw: &[u8]) -> Result<u32> {
if raw.len() != 4 {
anyhow::bail!("invalid schema version length {}", raw.len());
}
@@ -47,11 +47,11 @@ pub(in crate::state::store) fn decode_schema(raw: &[u8]) -> Result<u32> {
Ok(u32::from_le_bytes(arr))
}
pub(in crate::state::store) fn seeded_enroll_token_key(token: &str) -> String {
pub fn seeded_enroll_token_key(token: &str) -> String {
format!("{SEEDED_ENROLL_TOKEN_PREFIX}{token}")
}
pub(in crate::state::store) fn normalize_required_text(value: &str, field: &str) -> Result<String> {
pub fn normalize_required_text(value: &str, field: &str) -> Result<String> {
let normalized = value.trim();
if normalized.is_empty() {
anyhow::bail!("{field} must not be empty");
@@ -59,14 +59,14 @@ pub(in crate::state::store) fn normalize_required_text(value: &str, field: &str)
Ok(normalized.to_string())
}
pub(in crate::state::store) fn normalize_optional_text(value: Option<&str>) -> Option<String> {
pub fn normalize_optional_text(value: Option<&str>) -> Option<String> {
value
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
pub(in crate::state::store) fn normalize_device_identifier(
pub fn normalize_device_identifier(
input: DeviceIdentifierInput,
) -> Result<NormalizedDeviceIdentifier> {
let kind = normalize_required_text(&input.kind, "identifier kind")?.to_ascii_lowercase();
@@ -80,15 +80,11 @@ pub(in crate::state::store) fn normalize_device_identifier(
}
#[cfg_attr(not(test), allow(dead_code))]
pub(in crate::state::store) fn normalized_identifier_key_owned(
identifier: NormalizedDeviceIdentifier,
) -> String {
pub fn normalized_identifier_key_owned(identifier: NormalizedDeviceIdentifier) -> String {
identifier.identifier_key
}
pub(in crate::state::store) async fn list_enroll_token_rows(
pool: &SqlitePool,
) -> Result<Vec<EnrollTokenRow>> {
pub async fn list_enroll_token_rows(pool: &SqlitePool) -> Result<Vec<EnrollTokenRow>> {
sqlx::query_as!(
EnrollTokenRow,
r#"SELECT token as "token!", expires_at_unix FROM enroll_tokens ORDER BY expires_at_unix, token"#,
@@ -98,10 +94,7 @@ pub(in crate::state::store) async fn list_enroll_token_rows(
.context("failed listing enroll tokens")
}
pub(in crate::state::store) fn enroll_token_info_from_row(
row: EnrollTokenRow,
now: u64,
) -> Result<EnrollTokenInfo> {
pub fn enroll_token_info_from_row(row: EnrollTokenRow, now: u64) -> Result<EnrollTokenInfo> {
let expires_at_unix =
u64::try_from(row.expires_at_unix).context("negative token expiry in state db")?;
Ok(EnrollTokenInfo {
@@ -111,9 +104,7 @@ pub(in crate::state::store) fn enroll_token_info_from_row(
})
}
pub(in crate::state::store) async fn list_known_device_rows(
pool: &SqlitePool,
) -> Result<Vec<KnownDeviceRow>> {
pub async fn list_known_device_rows(pool: &SqlitePool) -> Result<Vec<KnownDeviceRow>> {
sqlx::query_as!(
KnownDeviceRow,
r#"SELECT device_id as "device_id!", display_name as "display_name!",
@@ -126,7 +117,7 @@ pub(in crate::state::store) async fn list_known_device_rows(
.context("failed listing known devices")
}
pub(in crate::state::store) async fn get_known_device_row(
pub async fn get_known_device_row(
pool: &SqlitePool,
device_id: &str,
) -> Result<Option<KnownDeviceRow>> {
@@ -143,11 +134,11 @@ pub(in crate::state::store) async fn get_known_device_row(
.context("failed reading known device")
}
pub(in crate::state::store) fn known_device_row_device_id(row: &KnownDeviceRow) -> &str {
pub fn known_device_row_device_id(row: &KnownDeviceRow) -> &str {
&row.device_id
}
pub(in crate::state::store) fn known_device_from_row_and_identifiers(
pub fn known_device_from_row_and_identifiers(
row: KnownDeviceRow,
identifiers: Vec<DeviceIdentifier>,
) -> Result<KnownDevice> {
@@ -164,7 +155,7 @@ pub(in crate::state::store) fn known_device_from_row_and_identifiers(
})
}
pub(in crate::state::store) async fn list_device_identifier_rows(
pub async fn list_device_identifier_rows(
pool: &SqlitePool,
device_id: &str,
) -> Result<Vec<DeviceIdentifierRow>> {
@@ -182,109 +173,7 @@ pub(in crate::state::store) async fn list_device_identifier_rows(
.context("failed listing device identifiers")
}
pub(in crate::state::store) async fn get_observation_identifier_row(
pool: &SqlitePool,
observation_key: &str,
) -> Result<Option<ObservationIdentifierRow>> {
sqlx::query_as!(
ObservationIdentifierRow,
r#"SELECT mac, ip
FROM agent_device_observations
WHERE observation_key = ?1"#,
observation_key
)
.fetch_optional(pool)
.await
.context("failed reading observation identifier")
}
pub(in crate::state::store) async fn get_observation_current_row(
tx: &mut Transaction<'_, Sqlite>,
observation_key: &str,
) -> Result<Option<ObservationCurrentRow>> {
sqlx::query_as!(
ObservationCurrentRow,
r#"SELECT mac, ip, hostname, first_seen_unix, last_seen_unix,
last_action as "last_action!"
FROM agent_device_observations
WHERE observation_key = ?1"#,
observation_key
)
.fetch_optional(&mut **tx)
.await
.context("failed reading current observation")
}
pub(in crate::state::store) fn observation_current_changed(
current: Option<&ObservationCurrentRow>,
next: &AgentDeviceObservation,
first_seen_unix: i64,
last_seen_unix: i64,
) -> bool {
let Some(current) = current else {
return true;
};
current.mac != next.mac
|| current.ip != next.ip
|| current.hostname != next.hostname
|| current.first_seen_unix != first_seen_unix
|| current.last_seen_unix != last_seen_unix
|| current.last_action != next.last_action
}
pub(in crate::state::store) fn observation_identifier_to_input(
observation: ObservationIdentifierRow,
) -> Result<DeviceIdentifierInput> {
observation
.mac
.map(|value| DeviceIdentifierInput {
kind: "mac".into(),
value,
})
.or_else(|| {
observation.ip.map(|value| DeviceIdentifierInput {
kind: "ip".into(),
value,
})
})
.ok_or_else(|| anyhow::anyhow!("observation has no attachable mac or ip"))
}
pub(in crate::state::store) fn normalize_agent_observation(
agent_id: &str,
input: AgentDeviceObservationInput,
) -> Result<AgentDeviceObservation> {
let kind = normalize_required_text(&input.kind, "observation kind")?.to_ascii_lowercase();
let action = normalize_required_text(&input.action, "observation action")?.to_ascii_lowercase();
let mac = normalize_optional_text(input.mac.as_deref()).map(|value| value.to_ascii_lowercase());
let ip = normalize_optional_text(input.ip.as_deref());
let hostname = normalize_optional_text(input.hostname.as_deref());
let identifier = observation_identifier(&kind, mac.as_deref(), ip.as_deref())
.ok_or_else(|| anyhow::anyhow!("observation requires mac or ip"))?;
let observation_key = format!("agent:{agent_id}:{kind}:{identifier}");
Ok(AgentDeviceObservation {
observation_key,
agent_id: agent_id.to_string(),
kind,
mac,
ip,
hostname,
first_seen_unix: input.first_seen_unix,
last_seen_unix: input.last_seen_unix,
last_action: action,
})
}
fn observation_identifier(kind: &str, mac: Option<&str>, ip: Option<&str>) -> Option<String> {
match (kind, mac, ip) {
("neigh" | "inventory", Some(mac), Some(ip)) => Some(format!("mac:{mac}:ip:{ip}")),
(_, Some(mac), _) => Some(format!("mac:{mac}")),
(_, None, Some(ip)) => Some(format!("ip:{ip}")),
(_, None, None) => None,
}
}
pub(in crate::state::store) async fn insert_device_identifier_tx(
pub async fn insert_device_identifier_tx(
tx: &mut Transaction<'_, Sqlite>,
device_id: &str,
identifier: &NormalizedDeviceIdentifier,
@@ -313,9 +202,7 @@ pub(in crate::state::store) async fn insert_device_identifier_tx(
Ok(())
}
pub(in crate::state::store) fn device_identifier_from_row(
row: DeviceIdentifierRow,
) -> Result<DeviceIdentifier> {
pub fn device_identifier_from_row(row: DeviceIdentifierRow) -> Result<DeviceIdentifier> {
Ok(DeviceIdentifier {
identifier_key: row.identifier_key,
device_id: row.device_id,
@@ -325,75 +212,3 @@ pub(in crate::state::store) fn device_identifier_from_row(
.context("negative device identifier timestamp in state db")?,
})
}
#[cfg_attr(not(test), allow(dead_code))]
pub(in crate::state::store) fn agent_observation_from_row(
row: AgentObservationRow,
) -> Result<AgentDeviceObservation> {
Ok(AgentDeviceObservation {
observation_key: row.observation_key,
agent_id: row.agent_id,
kind: row.kind,
mac: row.mac,
ip: row.ip,
hostname: row.hostname,
first_seen_unix: u64::try_from(row.first_seen_unix)
.context("negative observation first_seen timestamp in state db")?,
last_seen_unix: u64::try_from(row.last_seen_unix)
.context("negative observation last_seen timestamp in state db")?,
last_action: row.last_action,
})
}
pub(in crate::state::store) fn agent_observation_view_from_row(
row: AgentObservationViewRow,
) -> Result<AgentDeviceObservationView> {
let known_device = match (row.device_id, row.display_name, row.pinned) {
(Some(device_id), Some(display_name), Some(pinned)) => Some(KnownDeviceSummary {
device_id,
display_name,
pinned: pinned != 0,
}),
_ => None,
};
Ok(AgentDeviceObservationView {
observation_key: row.observation_key,
agent_id: row.agent_id,
kind: row.kind,
mac: row.mac,
ip: row.ip,
hostname: row.hostname,
first_seen_unix: u64::try_from(row.first_seen_unix)
.context("negative observation first_seen timestamp in state db")?,
last_seen_unix: u64::try_from(row.last_seen_unix)
.context("negative observation last_seen timestamp in state db")?,
last_action: row.last_action,
known_device,
})
}
pub(in crate::state::store) fn agent_observation_event_from_row(
row: AgentObservationEventRow,
) -> Result<AgentDeviceObservationEvent> {
let known_device = match (row.device_id, row.display_name, row.pinned) {
(Some(device_id), Some(display_name), Some(pinned)) => Some(KnownDeviceSummary {
device_id,
display_name,
pinned: pinned != 0,
}),
_ => None,
};
Ok(AgentDeviceObservationEvent {
event_id: row.event_id,
observation_key: row.observation_key,
agent_id: row.agent_id,
kind: row.kind,
action: row.action,
mac: row.mac,
ip: row.ip,
hostname: row.hostname,
ts_unix: u64::try_from(row.ts_unix)
.context("negative observation event timestamp in state db")?,
known_device,
})
}
wakey-control-plane/src/state/store/helpers/legacy.rs
+28 -41
View File
@@ -1,6 +1,6 @@
use super::*;
pub(in crate::state::store) async fn import_tree_raw(
pub async fn import_tree_raw(
pool: &SqlitePool,
legacy: &sled::Db,
sled_tree: &str,
@@ -28,10 +28,7 @@ pub(in crate::state::store) async fn import_tree_raw(
Ok(())
}
pub(in crate::state::store) async fn import_enroll_tokens(
pool: &SqlitePool,
legacy: &sled::Db,
) -> Result<()> {
pub async fn import_enroll_tokens(pool: &SqlitePool, legacy: &sled::Db) -> Result<()> {
let tree = legacy
.open_tree("enroll_tokens")
.context("failed to open legacy enroll_tokens tree")?;
@@ -39,12 +36,13 @@ pub(in crate::state::store) async fn import_enroll_tokens(
let (token, expiry) = item.context("failed reading legacy enroll token")?;
let token =
String::from_utf8(token.to_vec()).context("legacy enroll token key is not utf-8")?;
let expires_at_unix = decode_expiry(expiry.as_ref())?;
sqlx::query(
let expires_at_unix = i64::try_from(decode_expiry(expiry.as_ref())?)
.context("legacy token expiry overflow")?;
sqlx::query!(
"INSERT OR REPLACE INTO enroll_tokens (token, expires_at_unix) VALUES (?1, ?2)",
token,
expires_at_unix
)
.bind(token)
.bind(i64::try_from(expires_at_unix).context("legacy token expiry overflow")?)
.execute(pool)
.await
.context("failed importing legacy enroll token")?;
@@ -52,10 +50,7 @@ pub(in crate::state::store) async fn import_enroll_tokens(
Ok(())
}
pub(in crate::state::store) async fn import_agents(
pool: &SqlitePool,
legacy: &sled::Db,
) -> Result<()> {
pub async fn import_agents(pool: &SqlitePool, legacy: &sled::Db) -> Result<()> {
let tree = legacy
.open_tree("agents")
.context("failed to open legacy agents tree")?;
@@ -65,20 +60,19 @@ pub(in crate::state::store) async fn import_agents(
String::from_utf8(agent_id.to_vec()).context("legacy agent id is not utf-8")?;
let agent_token =
String::from_utf8(agent_token.to_vec()).context("legacy agent token is not utf-8")?;
sqlx::query("INSERT OR REPLACE INTO agents (agent_id, agent_token) VALUES (?1, ?2)")
.bind(agent_id)
.bind(agent_token)
.execute(pool)
.await
.context("failed importing legacy agent")?;
sqlx::query!(
"INSERT OR REPLACE INTO agents (agent_id, agent_token) VALUES (?1, ?2)",
agent_id,
agent_token
)
.execute(pool)
.await
.context("failed importing legacy agent")?;
}
Ok(())
}
pub(in crate::state::store) async fn import_agent_meta(
pool: &SqlitePool,
legacy: &sled::Db,
) -> Result<()> {
pub async fn import_agent_meta(pool: &SqlitePool, legacy: &sled::Db) -> Result<()> {
let tree = legacy
.open_tree("agent_meta")
.context("failed to open legacy agent_meta tree")?;
@@ -88,20 +82,19 @@ pub(in crate::state::store) async fn import_agent_meta(
String::from_utf8(agent_id.to_vec()).context("legacy agent id is not utf-8")?;
let nickname =
String::from_utf8(nickname.to_vec()).context("legacy nickname is not utf-8")?;
sqlx::query("INSERT OR REPLACE INTO agent_meta (agent_id, nickname) VALUES (?1, ?2)")
.bind(agent_id)
.bind(nickname)
.execute(pool)
.await
.context("failed importing legacy agent metadata")?;
sqlx::query!(
"INSERT OR REPLACE INTO agent_meta (agent_id, nickname) VALUES (?1, ?2)",
agent_id,
nickname
)
.execute(pool)
.await
.context("failed importing legacy agent metadata")?;
}
Ok(())
}
pub(in crate::state::store) async fn import_audit_events(
pool: &SqlitePool,
legacy: &sled::Db,
) -> Result<()> {
pub async fn import_audit_events(pool: &SqlitePool, legacy: &sled::Db) -> Result<()> {
let tree = legacy
.open_tree("audit_events")
.context("failed to open legacy audit_events tree")?;
@@ -117,10 +110,7 @@ pub(in crate::state::store) async fn import_audit_events(
Ok(())
}
pub(in crate::state::store) async fn import_active_alerts(
pool: &SqlitePool,
legacy: &sled::Db,
) -> Result<()> {
pub async fn import_active_alerts(pool: &SqlitePool, legacy: &sled::Db) -> Result<()> {
let tree = legacy
.open_tree("active_alerts")
.context("failed to open legacy active_alerts tree")?;
@@ -135,10 +125,7 @@ pub(in crate::state::store) async fn import_active_alerts(
Ok(())
}
pub(in crate::state::store) async fn import_alert_transitions(
pool: &SqlitePool,
legacy: &sled::Db,
) -> Result<()> {
pub async fn import_alert_transitions(pool: &SqlitePool, legacy: &sled::Db) -> Result<()> {
let tree = legacy
.open_tree("alert_transitions")
.context("failed to open legacy alert_transitions tree")?;
wakey-control-plane/src/state/store/helpers/rows.rs
+41 -97
View File
@@ -1,108 +1,52 @@
pub(in crate::state::store) struct NormalizedDeviceIdentifier {
pub(in crate::state::store) identifier_key: String,
pub(in crate::state::store) kind: String,
pub(in crate::state::store) value: String,
pub struct NormalizedDeviceIdentifier {
pub identifier_key: String,
pub kind: String,
pub value: String,
}
pub(in crate::state::store) struct EnrollTokenRow {
pub(in crate::state::store) token: String,
pub(in crate::state::store) expires_at_unix: i64,
pub struct EnrollTokenRow {
pub token: String,
pub expires_at_unix: i64,
}
pub(in crate::state::store) struct KnownDeviceRow {
pub(in crate::state::store) device_id: String,
pub(in crate::state::store) display_name: String,
pub(in crate::state::store) pinned: i64,
pub(in crate::state::store) created_at_unix: i64,
pub(in crate::state::store) updated_at_unix: i64,
pub(in crate::state::store) notes: Option<String>,
pub struct KnownDeviceRow {
pub device_id: String,
pub display_name: String,
pub pinned: i64,
pub created_at_unix: i64,
pub updated_at_unix: i64,
pub notes: Option<String>,
}
pub(in crate::state::store) struct DeviceIdentifierRow {
pub(in crate::state::store) identifier_key: String,
pub(in crate::state::store) device_id: String,
pub(in crate::state::store) kind: String,
pub(in crate::state::store) value: String,
pub(in crate::state::store) created_at_unix: i64,
pub struct DeviceIdentifierRow {
pub identifier_key: String,
pub device_id: String,
pub kind: String,
pub value: String,
pub created_at_unix: i64,
}
pub(in crate::state::store) struct AgentObservationRow {
pub(in crate::state::store) observation_key: String,
pub(in crate::state::store) agent_id: String,
pub(in crate::state::store) kind: String,
pub(in crate::state::store) mac: Option<String>,
pub(in crate::state::store) ip: Option<String>,
pub(in crate::state::store) hostname: Option<String>,
pub(in crate::state::store) first_seen_unix: i64,
pub(in crate::state::store) last_seen_unix: i64,
pub(in crate::state::store) last_action: String,
pub struct AlertStateRow {
pub alert_id: String,
pub kind: String,
pub severity: String,
pub status: String,
pub agent_id: Option<String>,
pub message: String,
pub value: i64,
pub threshold: i64,
pub last_seen_unix: i64,
pub metadata_json: String,
}
pub(in crate::state::store) struct AgentObservationViewRow {
pub(in crate::state::store) observation_key: String,
pub(in crate::state::store) agent_id: String,
pub(in crate::state::store) kind: String,
pub(in crate::state::store) mac: Option<String>,
pub(in crate::state::store) ip: Option<String>,
pub(in crate::state::store) hostname: Option<String>,
pub(in crate::state::store) first_seen_unix: i64,
pub(in crate::state::store) last_seen_unix: i64,
pub(in crate::state::store) last_action: String,
pub(in crate::state::store) device_id: Option<String>,
pub(in crate::state::store) display_name: Option<String>,
pub(in crate::state::store) pinned: Option<i64>,
}
pub(in crate::state::store) struct AgentObservationEventRow {
pub(in crate::state::store) event_id: String,
pub(in crate::state::store) observation_key: String,
pub(in crate::state::store) agent_id: String,
pub(in crate::state::store) kind: String,
pub(in crate::state::store) action: String,
pub(in crate::state::store) mac: Option<String>,
pub(in crate::state::store) ip: Option<String>,
pub(in crate::state::store) hostname: Option<String>,
pub(in crate::state::store) ts_unix: i64,
pub(in crate::state::store) device_id: Option<String>,
pub(in crate::state::store) display_name: Option<String>,
pub(in crate::state::store) pinned: Option<i64>,
}
pub(in crate::state::store) struct ObservationIdentifierRow {
pub(in crate::state::store) mac: Option<String>,
pub(in crate::state::store) ip: Option<String>,
}
pub(in crate::state::store) struct ObservationCurrentRow {
pub(in crate::state::store) mac: Option<String>,
pub(in crate::state::store) ip: Option<String>,
pub(in crate::state::store) hostname: Option<String>,
pub(in crate::state::store) first_seen_unix: i64,
pub(in crate::state::store) last_seen_unix: i64,
pub(in crate::state::store) last_action: String,
}
pub(in crate::state::store) struct AlertStateRow {
pub(in crate::state::store) alert_id: String,
pub(in crate::state::store) kind: String,
pub(in crate::state::store) severity: String,
pub(in crate::state::store) status: String,
pub(in crate::state::store) agent_id: Option<String>,
pub(in crate::state::store) message: String,
pub(in crate::state::store) value: i64,
pub(in crate::state::store) threshold: i64,
pub(in crate::state::store) last_seen_unix: i64,
pub(in crate::state::store) metadata_json: String,
}
pub(in crate::state::store) struct AlertTransitionRow {
pub(in crate::state::store) transition_id: String,
pub(in crate::state::store) ts_unix: i64,
pub(in crate::state::store) alert_id: String,
pub(in crate::state::store) kind: String,
pub(in crate::state::store) agent_id: Option<String>,
pub(in crate::state::store) from_status: Option<String>,
pub(in crate::state::store) to_status: String,
pub(in crate::state::store) message: String,
pub(in crate::state::store) metadata_json: String,
pub struct AlertTransitionRow {
pub transition_id: String,
pub ts_unix: i64,
pub alert_id: String,
pub kind: String,
pub agent_id: Option<String>,
pub from_status: Option<String>,
pub to_status: String,
pub message: String,
pub metadata_json: String,
}
wakey-control-plane/src/state/store/helpers/test_helpers.rs
+36
View File
@@ -0,0 +1,36 @@
use std::fs;
use std::time::Duration;
use super::super::Store;
#[cfg(test)]
pub struct TestStore {
store: Option<Store>,
pub dir: std::path::PathBuf,
}
#[cfg(test)]
impl TestStore {
pub fn store(&self) -> &Store {
self.store.as_ref().expect("store already taken")
}
pub async fn new() -> Self {
let dir = std::env::temp_dir().join(format!("wakey-cp-test-{}", uuid::Uuid::new_v4()));
let db_path = dir.join("state.sqlite3");
let store = Store::load_or_init(&db_path, Vec::new(), Duration::from_secs(60))
.await
.expect("store should initialize");
TestStore {
store: Some(store),
dir,
}
}
}
impl Drop for TestStore {
fn drop(&mut self) {
self.store.take();
let _ = fs::remove_dir_all(&self.dir);
}
}
wakey-control-plane/src/state/store/observations.rs
-399
View File
@@ -1,399 +0,0 @@
use std::{collections::BTreeSet, ops::DerefMut};
use super::*;
async fn upsert_observation_tx(
tx: &mut Transaction<'_, Sqlite>,
observation: &AgentDeviceObservation,
) -> Result<()> {
let first_seen_unix = i64::try_from(observation.first_seen_unix)
.context("observation first_seen overflow")?;
let last_seen_unix = i64::try_from(observation.last_seen_unix)
.context("observation last_seen overflow")?;
let current = get_observation_current_row(tx, &observation.observation_key)
.await
.context("failed checking existing observation")?;
let append_event = observation_current_changed(
current.as_ref(),
observation,
first_seen_unix,
last_seen_unix,
);
sqlx::query!(
"INSERT INTO agent_device_observations
(observation_key, agent_id, kind, mac, ip, hostname,
first_seen_unix, last_seen_unix, last_action)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)
ON CONFLICT(observation_key) DO UPDATE SET
mac = excluded.mac,
ip = excluded.ip,
hostname = excluded.hostname,
first_seen_unix = MIN(agent_device_observations.first_seen_unix, excluded.first_seen_unix),
last_seen_unix = MAX(agent_device_observations.last_seen_unix, excluded.last_seen_unix),
last_action = excluded.last_action",
observation.observation_key,
observation.agent_id,
observation.kind,
observation.mac,
observation.ip,
observation.hostname,
first_seen_unix,
last_seen_unix,
observation.last_action
)
.execute(tx.deref_mut())
.await
.context("failed upserting agent device observation")?;
if append_event {
let event_id = format!("ode-{}", Uuid::new_v4());
sqlx::query!(
"INSERT INTO agent_device_observation_events
(event_id, agent_id, kind, action, mac, ip, hostname, ts_unix)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
event_id,
observation.agent_id,
observation.kind,
observation.last_action,
observation.mac,
observation.ip,
observation.hostname,
last_seen_unix
)
.execute(tx.deref_mut())
.await
.context("failed appending agent device observation event")?;
}
Ok(())
}
async fn store_observation_snapshot(
tx: &mut Transaction<'_, Sqlite>,
agent_id: &str,
kind: &str,
snapshot_unix: u64,
keys: &BTreeSet<String>,
) -> Result<()> {
let snapshot_unix = i64::try_from(snapshot_unix).context("snapshot timestamp overflow")?;
sqlx::query!(
"INSERT INTO agent_observation_snapshots
(agent_id, kind, last_dump_unix)
VALUES (?1, ?2, ?3)
ON CONFLICT(agent_id, kind) DO UPDATE SET
last_dump_unix = excluded.last_dump_unix",
agent_id,
kind,
snapshot_unix
)
.execute(tx.deref_mut())
.await
.context("failed upserting observation snapshot")?;
sqlx::query!(
"DELETE FROM agent_observation_snapshot_keys WHERE agent_id = ?1 AND kind = ?2",
agent_id,
kind
)
.execute(tx.deref_mut())
.await
.context("failed clearing observation snapshot keys")?;
for key in keys {
sqlx::query!(
"INSERT INTO agent_observation_snapshot_keys (agent_id, kind, observation_key)
VALUES (?1, ?2, ?3)",
agent_id,
kind,
key
)
.execute(tx.deref_mut())
.await
.context("failed inserting observation snapshot key")?;
}
sqlx::query!(
"DELETE FROM agent_device_observations
WHERE agent_id = ?1
AND kind = ?2
AND observation_key NOT IN (
SELECT observation_key
FROM agent_observation_snapshot_keys
WHERE agent_id = ?1 AND kind = ?2
)",
agent_id,
kind
)
.execute(tx.deref_mut())
.await
.context("failed removing stale observation keys")?;
Ok(())
}
impl Store {
pub async fn upsert_agent_observations(
&self,
agent_id: &str,
observations: Vec<AgentDeviceObservationInput>,
) -> Result<usize> {
let mut tx = self
.pool
.begin()
.await
.context("failed starting observation transaction")?;
let mut written = 0usize;
for observation in observations {
let observation = normalize_agent_observation(agent_id, observation)?;
upsert_observation_tx(&mut tx, &observation).await?;
written = written.saturating_add(1);
}
tx.commit()
.await
.context("failed committing observation transaction")?;
Ok(written)
}
pub async fn upsert_agent_observations_snapshot(
&self,
agent_id: &str,
kind: &str,
observations: Vec<AgentDeviceObservationInput>,
) -> Result<usize> {
let mut tx = self
.pool
.begin()
.await
.context("failed starting observation snapshot transaction")?;
let mut written = 0usize;
let mut snapshot_keys = BTreeSet::new();
let kind = normalize_required_text(kind, "observation kind")?.to_ascii_lowercase();
for observation in observations {
let observation = normalize_agent_observation(agent_id, observation)?;
if observation.kind != kind {
anyhow::bail!(
"observation kind mismatch: expected {kind} got {}",
observation.kind
);
}
snapshot_keys.insert(observation.observation_key.clone());
upsert_observation_tx(&mut tx, &observation).await?;
written = written.saturating_add(1);
}
let snapshot_unix = now_unix();
store_observation_snapshot(&mut tx, agent_id, &kind, snapshot_unix, &snapshot_keys)
.await?;
tx.commit()
.await
.context("failed committing observation snapshot transaction")?;
Ok(written)
}
pub async fn gc_stale_observations(&self, retention: Duration) -> Result<u64> {
if retention.as_secs() == 0 {
return Ok(0);
}
let cutoff = now_unix().saturating_sub(retention.as_secs());
let cutoff = i64::try_from(cutoff).context("observation retention overflow")?;
let removed_observations = sqlx::query!(
"DELETE FROM agent_device_observations WHERE last_seen_unix < ?1",
cutoff
)
.execute(&self.pool)
.await
.context("failed removing stale observations")?
.rows_affected();
let removed_events = sqlx::query!(
"DELETE FROM agent_device_observation_events WHERE ts_unix < ?1",
cutoff
)
.execute(&self.pool)
.await
.context("failed removing stale observation events")?
.rows_affected();
Ok(removed_observations.saturating_add(removed_events))
}
pub async fn list_agent_observations(
&self,
agent_id: Option<&str>,
limit: usize,
) -> Result<Vec<AgentDeviceObservation>> {
let limit = limit.clamp(1, 1000);
let limit = i64::try_from(limit).context("observation limit overflow")?;
let rows = if let Some(agent_id) = agent_id {
sqlx::query_as!(
AgentObservationRow,
r#"SELECT observation_key as "observation_key!", agent_id as "agent_id!",
kind as "kind!", mac, ip, hostname,
first_seen_unix, last_seen_unix, last_action as "last_action!"
FROM agent_device_observations
WHERE agent_id = ?1
ORDER BY last_seen_unix DESC
LIMIT ?2"#,
agent_id,
limit
)
.fetch_all(&self.pool)
.await
} else {
sqlx::query_as!(
AgentObservationRow,
r#"SELECT observation_key as "observation_key!", agent_id as "agent_id!",
kind as "kind!", mac, ip, hostname,
first_seen_unix, last_seen_unix, last_action as "last_action!"
FROM agent_device_observations
ORDER BY last_seen_unix DESC
LIMIT ?1"#,
limit
)
.fetch_all(&self.pool)
.await
}
.context("failed listing agent observations")?;
rows.into_iter().map(agent_observation_from_row).collect()
}
pub async fn list_agent_observation_views(
&self,
agent_id: Option<&str>,
limit: usize,
) -> Result<Vec<AgentDeviceObservationView>> {
let limit = limit.clamp(1, 1000);
let limit = i64::try_from(limit).context("observation limit overflow")?;
let rows = if let Some(agent_id) = agent_id {
sqlx::query_as!(
AgentObservationViewRow,
r#"SELECT observations.observation_key as "observation_key!",
observations.agent_id as "agent_id!",
observations.kind as "kind!",
observations.mac,
observations.ip,
observations.hostname,
observations.first_seen_unix,
observations.last_seen_unix,
observations.last_action as "last_action!",
known_devices.device_id,
known_devices.display_name,
known_devices.pinned
FROM agent_device_observations observations
LEFT JOIN device_identifiers identifiers
ON identifiers.identifier_key =
CASE
WHEN observations.mac IS NOT NULL THEN 'mac:' || observations.mac
WHEN observations.ip IS NOT NULL THEN 'ip:' || observations.ip
END
LEFT JOIN known_devices ON known_devices.device_id = identifiers.device_id
WHERE observations.agent_id = ?1
ORDER BY observations.last_seen_unix DESC
LIMIT ?2"#,
agent_id,
limit
)
.fetch_all(&self.pool)
.await
} else {
sqlx::query_as!(
AgentObservationViewRow,
r#"SELECT observations.observation_key as "observation_key!",
observations.agent_id as "agent_id!",
observations.kind as "kind!",
observations.mac,
observations.ip,
observations.hostname,
observations.first_seen_unix,
observations.last_seen_unix,
observations.last_action as "last_action!",
known_devices.device_id,
known_devices.display_name,
known_devices.pinned
FROM agent_device_observations observations
LEFT JOIN device_identifiers identifiers
ON identifiers.identifier_key =
CASE
WHEN observations.mac IS NOT NULL THEN 'mac:' || observations.mac
WHEN observations.ip IS NOT NULL THEN 'ip:' || observations.ip
END
LEFT JOIN known_devices ON known_devices.device_id = identifiers.device_id
ORDER BY observations.last_seen_unix DESC
LIMIT ?1"#,
limit
)
.fetch_all(&self.pool)
.await
}
.context("failed listing agent observation views")?;
rows.into_iter()
.map(agent_observation_view_from_row)
.collect()
}
pub async fn list_agent_observation_events(
&self,
agent_id: Option<&str>,
kind: Option<&str>,
mac: Option<&str>,
ip: Option<&str>,
observation_key: Option<&str>,
limit: usize,
) -> Result<Vec<AgentDeviceObservationEvent>> {
let limit = limit.clamp(1, 1000);
let limit = i64::try_from(limit).context("observation event limit overflow")?;
let rows = sqlx::query_as!(
AgentObservationEventRow,
r#"SELECT events.event_id as "event_id!",
('agent:' || events.agent_id || ':' || events.kind || ':' ||
CASE
WHEN events.kind IN ('neigh', 'inventory') AND events.mac IS NOT NULL AND events.ip IS NOT NULL
THEN 'mac:' || events.mac || ':ip:' || events.ip
WHEN events.mac IS NOT NULL THEN 'mac:' || events.mac
WHEN events.ip IS NOT NULL THEN 'ip:' || events.ip
ELSE ''
END) as "observation_key!",
events.agent_id as "agent_id!",
events.kind as "kind!",
events.action as "action!",
events.mac,
events.ip,
events.hostname,
events.ts_unix,
known_devices.device_id,
known_devices.display_name,
known_devices.pinned
FROM agent_device_observation_events events
LEFT JOIN device_identifiers identifiers
ON identifiers.identifier_key =
CASE
WHEN events.mac IS NOT NULL THEN 'mac:' || events.mac
WHEN events.ip IS NOT NULL THEN 'ip:' || events.ip
END
LEFT JOIN known_devices ON known_devices.device_id = identifiers.device_id
WHERE (?1 IS NULL OR events.agent_id = ?1)
AND (?2 IS NULL OR events.kind = ?2)
AND (?3 IS NULL OR events.mac = ?3)
AND (?4 IS NULL OR events.ip = ?4)
AND (?5 IS NULL OR ('agent:' || events.agent_id || ':' || events.kind || ':' ||
CASE
WHEN events.kind IN ('neigh', 'inventory') AND events.mac IS NOT NULL AND events.ip IS NOT NULL
THEN 'mac:' || events.mac || ':ip:' || events.ip
WHEN events.mac IS NOT NULL THEN 'mac:' || events.mac
WHEN events.ip IS NOT NULL THEN 'ip:' || events.ip
ELSE ''
END) = ?5)
ORDER BY events.ts_unix DESC
LIMIT ?6"#,
agent_id,
kind,
mac,
ip,
observation_key,
limit
)
.fetch_all(&self.pool)
.await
.context("failed listing agent observation events")?;
rows.into_iter()
.map(agent_observation_event_from_row)
.collect()
}
}
wakey-control-plane/src/state/types.rs
+77 -49
View File
@@ -1,6 +1,8 @@
use std::path::PathBuf;
use macaddr::MacAddr;
use serde::{Deserialize, Serialize};
use wakey_core::Presence;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IssuedAgent {
@@ -64,47 +66,6 @@ pub struct DeviceIdentifierInput {
pub value: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentDeviceObservation {
pub observation_key: String,
pub agent_id: String,
pub kind: String,
pub mac: Option<String>,
pub ip: Option<String>,
pub hostname: Option<String>,
pub first_seen_unix: u64,
pub last_seen_unix: u64,
pub last_action: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentDeviceObservationView {
pub observation_key: String,
pub agent_id: String,
pub kind: String,
pub mac: Option<String>,
pub ip: Option<String>,
pub hostname: Option<String>,
pub first_seen_unix: u64,
pub last_seen_unix: u64,
pub last_action: String,
pub known_device: Option<KnownDeviceSummary>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentDeviceObservationEvent {
pub event_id: String,
pub observation_key: String,
pub agent_id: String,
pub kind: String,
pub action: String,
pub mac: Option<String>,
pub ip: Option<String>,
pub hostname: Option<String>,
pub ts_unix: u64,
pub known_device: Option<KnownDeviceSummary>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KnownDeviceSummary {
pub device_id: String,
@@ -113,14 +74,81 @@ pub struct KnownDeviceSummary {
}
#[derive(Debug, Clone)]
pub struct AgentDeviceObservationInput {
pub kind: String,
pub action: String,
pub mac: Option<String>,
pub ip: Option<String>,
pub hostname: Option<String>,
pub first_seen_unix: u64,
pub last_seen_unix: u64,
pub struct AgentDeviceRow {
pub agent_id: String,
pub device_key: String,
pub presence: String,
pub display_name: Option<String>,
pub first_seen_unix: i64,
pub last_seen_unix: i64,
}
impl AgentDeviceRow {
pub fn presence(&self) -> Presence {
Presence::from(self.presence.as_str())
}
pub fn first_seen(&self) -> u64 {
self.first_seen_unix.max(0) as u64
}
pub fn last_seen(&self) -> u64 {
self.last_seen_unix.max(0) as u64
}
}
#[derive(Debug, Clone)]
pub struct AgentDeviceMacRow {
pub agent_id: String,
pub device_key: String,
pub mac: String,
}
impl TryFrom<&AgentDeviceMacRow> for MacAddr {
type Error = macaddr::ParseError;
fn try_from(row: &AgentDeviceMacRow) -> Result<Self, Self::Error> {
row.mac.parse()
}
}
#[derive(Debug, Clone)]
pub struct AgentDeviceIpRow {
pub agent_id: String,
pub device_key: String,
pub ip: String,
}
impl TryFrom<&AgentDeviceIpRow> for std::net::IpAddr {
type Error = std::net::AddrParseError;
fn try_from(row: &AgentDeviceIpRow) -> Result<Self, Self::Error> {
row.ip.parse()
}
}
#[derive(Debug, Clone)]
pub struct AgentDeviceHostnameRow {
pub agent_id: String,
pub device_key: String,
pub hostname: String,
}
#[derive(Debug, Clone)]
pub struct AgentDeviceFactRow {
pub agent_id: String,
pub device_key: String,
pub fact_json: String,
}
#[derive(Debug, Clone)]
pub struct AgentDeviceWithChildren {
pub device: AgentDeviceRow,
pub macs: Vec<macaddr::MacAddr>,
pub ips: Vec<std::net::IpAddr>,
pub hostnames: Vec<String>,
#[allow(dead_code)]
pub facts: Vec<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
wakey-control-plane/src/ws.rs
+18 -42
View File
@@ -4,12 +4,12 @@ use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use axum::response::IntoResponse;
use futures_util::{SinkExt, StreamExt};
use serde::Deserialize;
use std::collections::BTreeMap;
use std::time::Instant;
use tokio::sync::mpsc;
use tracing::{debug, info, info_span, warn};
use uuid::Uuid;
use wakey_agent::protocol::{AgentObservation, ErrorPayload, RequestId, ServerMessage};
use wakey_agent::protocol::{ErrorPayload, RequestId, ServerMessage};
use wakey_core::Device;
use crate::runtime::{AgentReply, AgentSession, AppState, SessionEvent};
use crate::state::AuditEventInput;
@@ -27,9 +27,9 @@ enum IncomingClientMessage {
Heartbeat {
agent_id: String,
},
Observations {
DeviceSnapshot {
agent_id: String,
observations: Vec<AgentObservation>,
devices: Vec<Device>,
},
Result {
request_id: RequestId,
@@ -243,7 +243,7 @@ async fn process_agent_text(
{
warn!(error = %err, "failed to append audit event for auth success");
}
let _ = tx.send(SessionEvent::Message(ServerMessage::SyncObservations));
let _ = tx.send(SessionEvent::Message(ServerMessage::SyncDeviceSnapshot));
}
IncomingClientMessage::Heartbeat { agent_id } => {
if authed_agent_id.as_deref() != Some(agent_id.as_str()) {
@@ -252,48 +252,24 @@ async fn process_agent_text(
ensure_current_session(state, &agent_id, connection_id).await?;
debug!(agent_id = %agent_id, "heartbeat received");
}
IncomingClientMessage::Observations {
agent_id,
observations,
} => {
IncomingClientMessage::DeviceSnapshot { agent_id, devices } => {
if authed_agent_id.as_deref() != Some(agent_id.as_str()) {
anyhow::bail!("observations for unauthenticated or mismatched agent");
anyhow::bail!("device_snapshot for unauthenticated or mismatched agent");
}
ensure_current_session(state, &agent_id, connection_id).await?;
let mut by_kind: BTreeMap<String, Vec<crate::state::AgentDeviceObservationInput>> =
BTreeMap::new();
for observation in observations {
let kind = observation.kind.trim().to_ascii_lowercase();
by_kind
.entry(kind.clone())
.or_default()
.push(crate::state::AgentDeviceObservationInput {
kind,
action: observation.action,
mac: observation.mac,
ip: observation.ip.map(|ip| ip.to_string()),
hostname: observation.hostname,
first_seen_unix: observation.first_seen_unix,
last_seen_unix: observation.last_seen_unix,
});
}
let mut accepted = 0usize;
for (kind, inputs) in by_kind {
match state
.store
.upsert_agent_observations_snapshot(&agent_id, &kind, inputs)
.await
{
Ok(written) => {
accepted = accepted.saturating_add(written);
}
Err(err) => {
warn!(agent_id = %agent_id, error = %err, kind = %kind, "failed to store websocket observations");
anyhow::bail!("failed to store observations: {err}");
}
match state
.store
.replace_agent_device_snapshot(&agent_id, &devices)
.await
{
Ok(accepted) => {
debug!(agent_id = %agent_id, accepted, "agent websocket device snapshot accepted");
}
Err(err) => {
warn!(agent_id = %agent_id, error = %err, "failed to store websocket device snapshot");
anyhow::bail!("failed to store device snapshot: {err}");
}
}
debug!(agent_id = %agent_id, accepted, "agent websocket observations accepted");
}
IncomingClientMessage::Result { request_id, result } => {
let agent_id = authed_agent_id