433 lines
15 KiB
Rust
433 lines
15 KiB
Rust
//! Refresh agents independently: one failed relay or snapshot does not cancel others.
|
|
|
|
use axum::{Json, extract::State, http::StatusCode, response::IntoResponse};
|
|
use futures_util::StreamExt;
|
|
use tracing::{info, warn};
|
|
use wakey_agent::protocol::{AgentCommand, InventoryRequest};
|
|
|
|
use super::types::{
|
|
RefreshFleetAgentResult, RefreshFleetDevicesRequest, RefreshFleetDevicesResponse,
|
|
};
|
|
use crate::api::{
|
|
ApiError,
|
|
commands::{RelayCommandResponse, relay_agent_command},
|
|
};
|
|
use crate::runtime::AppState;
|
|
use crate::state::Store;
|
|
|
|
const REFRESH_CONCURRENCY: usize = 8;
|
|
|
|
#[tracing::instrument(name = "fleet_refresh", skip_all, fields(agents = tracing::field::Empty, succeeded = tracing::field::Empty, failed = tracing::field::Empty, accepted = tracing::field::Empty))]
|
|
pub async fn refresh_fleet_devices(
|
|
State(state): State<AppState>,
|
|
Json(req): Json<RefreshFleetDevicesRequest>,
|
|
) -> Result<impl IntoResponse, ApiError> {
|
|
let agent_ids = select_agents(&state, req.agent_ids).await;
|
|
tracing::Span::current().record("agents", agent_ids.len());
|
|
let results = futures_util::stream::iter(
|
|
agent_ids
|
|
.into_iter()
|
|
.map(|agent_id| refresh_agent(&state, agent_id, req.timeout_ms)),
|
|
)
|
|
.buffer_unordered(REFRESH_CONCURRENCY)
|
|
.collect::<Vec<_>>()
|
|
.await;
|
|
let response = summarize(results);
|
|
let succeeded = response
|
|
.agents
|
|
.iter()
|
|
.filter(|result| result.status == "ok")
|
|
.count();
|
|
let span = tracing::Span::current();
|
|
span.record("succeeded", succeeded);
|
|
span.record("failed", response.agents.len() - succeeded);
|
|
span.record("accepted", response.total_accepted);
|
|
info!("fleet refresh completed");
|
|
Ok((StatusCode::OK, Json(response)))
|
|
}
|
|
|
|
/// Default to a point-in-time connected set. Explicit IDs remain eligible so
|
|
/// the relay can report disconnected agents as per-agent errors, not omit them.
|
|
async fn select_agents(state: &AppState, requested: Vec<String>) -> Vec<String> {
|
|
let ids = if requested.is_empty() {
|
|
state.sessions.read().await.keys().cloned().collect()
|
|
} else {
|
|
requested
|
|
};
|
|
normalize_refresh_agent_ids(ids)
|
|
}
|
|
|
|
pub(super) fn normalize_refresh_agent_ids(mut agent_ids: Vec<String>) -> Vec<String> {
|
|
agent_ids.sort();
|
|
agent_ids.dedup();
|
|
agent_ids
|
|
}
|
|
|
|
#[tracing::instrument(name = "fleet_refresh_agent", skip_all, fields(agent_id = %agent_id, outcome = tracing::field::Empty, accepted = tracing::field::Empty))]
|
|
async fn refresh_agent(
|
|
state: &AppState,
|
|
agent_id: String,
|
|
timeout_ms: Option<u64>,
|
|
) -> RefreshFleetAgentResult {
|
|
let response = relay_agent_command(
|
|
state,
|
|
&agent_id,
|
|
AgentCommand::Inventory(InventoryRequest {
|
|
query: None,
|
|
name: None,
|
|
ips: Vec::new(),
|
|
devs: Vec::new(),
|
|
nuds: Vec::new(),
|
|
macs: Vec::new(),
|
|
}),
|
|
timeout_ms,
|
|
)
|
|
.await;
|
|
let result = store_inventory(&state.store, &agent_id, response).await;
|
|
tracing::Span::current().record("outcome", result.status.as_str());
|
|
tracing::Span::current().record("accepted", result.accepted);
|
|
if result.status != "ok" {
|
|
// The response retains diagnostic detail; do not export payloads.
|
|
warn!("fleet agent refresh failed");
|
|
}
|
|
result
|
|
}
|
|
|
|
/// Validate the relay result before replacing anything. An empty *valid*
|
|
/// inventory is still a complete snapshot and intentionally clears old rows.
|
|
fn decode_inventory(
|
|
response: Result<RelayCommandResponse, ApiError>,
|
|
) -> Result<wakey_core::DeviceInventory, String> {
|
|
let response = response.map_err(|err| format!("{}: {}", err.status, err.message))?;
|
|
if response.status != "ok" {
|
|
return Err(response
|
|
.error
|
|
.map(|error| error.message)
|
|
.unwrap_or_else(|| "inventory command failed".into()));
|
|
}
|
|
let result = response
|
|
.result
|
|
.ok_or_else(|| "inventory command returned no result".to_string())?;
|
|
serde_json::from_value(result).map_err(|err| err.to_string())
|
|
}
|
|
|
|
async fn store_inventory(
|
|
store: &Store,
|
|
agent_id: &str,
|
|
response: Result<RelayCommandResponse, ApiError>,
|
|
) -> RefreshFleetAgentResult {
|
|
let accepted = async {
|
|
let inventory = decode_inventory(response)?;
|
|
store
|
|
.replace_agent_device_snapshot(agent_id, &inventory.devices)
|
|
.await
|
|
.map_err(|err| err.to_string())
|
|
}
|
|
.await;
|
|
match accepted {
|
|
Ok(accepted) => RefreshFleetAgentResult {
|
|
agent_id: agent_id.into(),
|
|
status: "ok".into(),
|
|
accepted,
|
|
error: None,
|
|
},
|
|
Err(error) => RefreshFleetAgentResult {
|
|
agent_id: agent_id.into(),
|
|
status: "error".into(),
|
|
accepted: 0,
|
|
error: Some(error),
|
|
},
|
|
}
|
|
}
|
|
|
|
/// Concurrent completion order is deliberately not part of the response contract.
|
|
fn summarize(mut agents: Vec<RefreshFleetAgentResult>) -> RefreshFleetDevicesResponse {
|
|
agents.sort_by(|a, b| a.agent_id.cmp(&b.agent_id));
|
|
RefreshFleetDevicesResponse {
|
|
total_accepted: agents.iter().map(|result| result.accepted).sum(),
|
|
agents,
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use serde_json::json;
|
|
|
|
fn response(result: Option<serde_json::Value>) -> RelayCommandResponse {
|
|
RelayCommandResponse {
|
|
request_id: "refresh-test".into(),
|
|
status: "ok".into(),
|
|
result,
|
|
error: None,
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn refresh_updates_cached_fleet_despite_a_disconnected_agent() {
|
|
use crate::runtime::{AgentReply, AgentSession, SessionEvent};
|
|
use std::{collections::HashMap, sync::Arc, time::Duration};
|
|
use tokio::sync::{Mutex, RwLock, mpsc};
|
|
use wakey_agent::protocol::{CommandResult, ServerMessage};
|
|
|
|
let dir =
|
|
std::env::temp_dir().join(format!("wakey-refresh-handler-{}", uuid::Uuid::new_v4()));
|
|
let state = AppState {
|
|
store: Arc::new(
|
|
Store::load_or_init(&dir.join("state.sqlite3"), vec![], Duration::from_secs(60))
|
|
.await
|
|
.unwrap(),
|
|
),
|
|
sessions: Arc::new(RwLock::new(HashMap::new())),
|
|
pending: Arc::new(Mutex::new(HashMap::new())),
|
|
public_url: "http://localhost".into(),
|
|
command_timeout: Duration::from_secs(30),
|
|
enroll_token_ttl: Duration::from_secs(60),
|
|
terminals: crate::runtime::terminals::TerminalRegistry::new(),
|
|
alerts: crate::api::AlertsHub::new(),
|
|
fleet: Arc::new(super::super::FleetCache::default()),
|
|
};
|
|
let token = state
|
|
.store
|
|
.issue_enroll_token(Duration::from_secs(60))
|
|
.await
|
|
.unwrap();
|
|
let agent_id = state
|
|
.store
|
|
.enroll(&token.enroll_token)
|
|
.await
|
|
.unwrap()
|
|
.agent_id;
|
|
let device = |ip: &str| {
|
|
wakey_core::Device::from_parts(
|
|
vec![wakey_core::NeighborEntry {
|
|
ip: ip.parse().unwrap(),
|
|
mac: Some("aa:bb:cc:dd:ee:01".parse().unwrap()),
|
|
dev: Some("br-lan".into()),
|
|
state: wakey_core::NeighborState::Reachable,
|
|
}],
|
|
vec![],
|
|
)
|
|
};
|
|
state
|
|
.store
|
|
.replace_agent_device_snapshot(&agent_id, &[device("192.168.1.10")])
|
|
.await
|
|
.unwrap();
|
|
let (tx, mut rx) = mpsc::unbounded_channel();
|
|
state.sessions.write().await.insert(
|
|
agent_id.clone(),
|
|
AgentSession {
|
|
connection_id: "test".into(),
|
|
tx,
|
|
capabilities: vec![],
|
|
capability_options: Default::default(),
|
|
},
|
|
);
|
|
assert_eq!(
|
|
select_agents(&state, vec![]).await.as_slice(),
|
|
std::slice::from_ref(&agent_id)
|
|
);
|
|
|
|
// Prime the real cache with old data before refresh touches the store.
|
|
let query = super::super::types::ListFleetDevicesQuery {
|
|
visibility: Some("all".into()),
|
|
..Default::default()
|
|
};
|
|
let before = super::super::load_fleet_devices(&state, &query)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(before.len(), 1);
|
|
assert_eq!(
|
|
before[0].ips,
|
|
["192.168.1.10".parse::<std::net::IpAddr>().unwrap()]
|
|
);
|
|
|
|
let refresh = refresh_fleet_devices(
|
|
State(state.clone()),
|
|
Json(RefreshFleetDevicesRequest {
|
|
agent_ids: vec!["disconnected".into(), agent_id.clone(), agent_id.clone()],
|
|
timeout_ms: None,
|
|
}),
|
|
);
|
|
// Answer through the production relay's pending-request channel. No
|
|
// sleep or polling: receiving the command proves its waiter exists.
|
|
let reply = async {
|
|
let SessionEvent::Message(ServerMessage::Command {
|
|
request_id,
|
|
command,
|
|
..
|
|
}) = rx.recv().await.unwrap()
|
|
else {
|
|
panic!("expected inventory command");
|
|
};
|
|
let AgentCommand::Inventory(request) = command else {
|
|
panic!("expected inventory command");
|
|
};
|
|
assert!(request.query.is_none() && request.name.is_none());
|
|
assert!(
|
|
request.ips.is_empty()
|
|
&& request.macs.is_empty()
|
|
&& request.devs.is_empty()
|
|
&& request.nuds.is_empty()
|
|
);
|
|
let result =
|
|
serde_json::to_value(CommandResult::Inventory(wakey_core::DeviceInventory {
|
|
devices: vec![device("192.168.1.20")],
|
|
}))
|
|
.unwrap();
|
|
state
|
|
.pending
|
|
.lock()
|
|
.await
|
|
.remove(request_id.as_str())
|
|
.unwrap()
|
|
.send(AgentReply::Result(result))
|
|
.unwrap_or_else(|_| panic!("relay waiter dropped"));
|
|
};
|
|
let (response, ()) = tokio::join!(refresh, reply);
|
|
let response = response.unwrap().into_response();
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
let body = axum::body::to_bytes(response.into_body(), 4096)
|
|
.await
|
|
.unwrap();
|
|
let result: RefreshFleetDevicesResponse = serde_json::from_slice(&body).unwrap();
|
|
assert_eq!(result.total_accepted, 1);
|
|
assert_eq!(result.agents.len(), 2);
|
|
assert!(
|
|
result
|
|
.agents
|
|
.windows(2)
|
|
.all(|pair| pair[0].agent_id < pair[1].agent_id)
|
|
);
|
|
let success = result
|
|
.agents
|
|
.iter()
|
|
.find(|row| row.agent_id == agent_id)
|
|
.unwrap();
|
|
assert_eq!(success.status, "ok");
|
|
assert_eq!(success.accepted, 1);
|
|
assert!(success.error.is_none());
|
|
let failed = result
|
|
.agents
|
|
.iter()
|
|
.find(|row| row.agent_id == "disconnected")
|
|
.unwrap();
|
|
assert_eq!(failed.status, "error");
|
|
assert_eq!(failed.accepted, 0);
|
|
assert_eq!(
|
|
failed.error.as_deref(),
|
|
Some("404 Not Found: agent is not connected")
|
|
);
|
|
assert!(state.pending.lock().await.is_empty());
|
|
assert!(matches!(
|
|
rx.try_recv(),
|
|
Err(mpsc::error::TryRecvError::Empty)
|
|
));
|
|
|
|
let after = super::super::load_fleet_devices(&state, &query)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(after.len(), 1);
|
|
assert_eq!(after[0].device_key, before[0].device_key);
|
|
assert_eq!(
|
|
after[0].ips,
|
|
["192.168.1.20".parse::<std::net::IpAddr>().unwrap()]
|
|
);
|
|
let stored = state
|
|
.store
|
|
.list_agent_device_rows_for_agent(&agent_id)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(stored.len(), 1);
|
|
assert_eq!(stored[0].endpoints.len(), 1);
|
|
assert_eq!(
|
|
stored[0].endpoints[0].key.ip,
|
|
Some("192.168.1.20".parse().unwrap())
|
|
);
|
|
drop(state);
|
|
std::fs::remove_dir_all(dir).unwrap();
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn refresh_stores_snapshots_but_invalid_results_leave_them_untouched() {
|
|
let dir = std::env::temp_dir().join(format!("wakey-refresh-{}", uuid::Uuid::new_v4()));
|
|
let store = Store::load_or_init(
|
|
&dir.join("state.sqlite3"),
|
|
vec![],
|
|
std::time::Duration::from_secs(60),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
let device = wakey_core::Device::from_parts(
|
|
vec![wakey_core::NeighborEntry {
|
|
ip: "192.168.1.10".parse().unwrap(),
|
|
mac: Some("aa:bb:cc:dd:ee:01".parse().unwrap()),
|
|
dev: Some("br-lan".into()),
|
|
state: wakey_core::NeighborState::Reachable,
|
|
}],
|
|
vec![],
|
|
);
|
|
let inventory = serde_json::to_value(wakey_core::DeviceInventory {
|
|
devices: vec![device],
|
|
})
|
|
.unwrap();
|
|
let stored = store_inventory(&store, "agent", Ok(response(Some(inventory)))).await;
|
|
assert_eq!(stored.status, "ok");
|
|
assert_eq!(stored.accepted, 1);
|
|
let revision = store.fleet_revision();
|
|
let mut command_failed = response(None);
|
|
command_failed.status = "error".into();
|
|
command_failed.error = Some(wakey_agent::protocol::ErrorPayload {
|
|
code: "inventory_failed".into(),
|
|
message: "source unavailable".into(),
|
|
retryable: None,
|
|
});
|
|
for failed in [
|
|
Ok(response(None)),
|
|
Ok(response(Some(json!({"devices": "invalid"})))),
|
|
Ok(command_failed),
|
|
Err(ApiError::new(
|
|
StatusCode::BAD_GATEWAY,
|
|
"agent_send_failed",
|
|
"send failed",
|
|
)),
|
|
] {
|
|
let invalid = store_inventory(&store, "agent", failed).await;
|
|
assert_eq!(invalid.status, "error");
|
|
assert_eq!(invalid.accepted, 0);
|
|
assert!(invalid.error.is_some());
|
|
assert_eq!(store.fleet_revision(), revision);
|
|
let rows = store
|
|
.list_agent_device_rows_for_agent("agent")
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(rows.len(), 1);
|
|
assert_eq!(rows[0].endpoints.len(), 1);
|
|
assert_eq!(
|
|
rows[0].endpoints[0].key.ip,
|
|
Some("192.168.1.10".parse().unwrap())
|
|
);
|
|
}
|
|
let rows = store
|
|
.list_agent_device_rows_for_agent("agent")
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(rows.len(), 1);
|
|
assert_eq!(rows[0].endpoints.len(), 1);
|
|
let empty =
|
|
store_inventory(&store, "agent", Ok(response(Some(json!({"devices": []}))))).await;
|
|
assert_eq!(empty.status, "ok");
|
|
assert_eq!(empty.accepted, 0);
|
|
assert!(
|
|
store
|
|
.list_agent_device_rows_for_agent("agent")
|
|
.await
|
|
.unwrap()
|
|
.is_empty()
|
|
);
|
|
drop(store);
|
|
std::fs::remove_dir_all(dir).unwrap();
|
|
}
|
|
}
|