request sync + over WS + merge devices + fix deployment scripts

This commit is contained in:
lda
2026-04-28 23:19:05 +07:00 Verified
parent 132e15d134
commit 91e49a5806
17 changed files with 420 additions and 55 deletions
@@ -0,0 +1,12 @@
{
"db_name": "SQLite",
"query": "UPDATE device_identifiers SET device_id = ?1 WHERE device_id = ?2",
"describe": {
"columns": [],
"parameters": {
"Right": 2
},
"nullable": []
},
"hash": "7247cc488921de2313a602697590502771a231bde7d03cd9ccac55a70971da45"
}
View File
View File
+5
View File
@@ -104,8 +104,10 @@ main() {
# Ensure execute bits on staged files we know should be executable # Ensure execute bits on staged files we know should be executable
for f in \ for f in \
"$STAGING/etc/init.d/"* \ "$STAGING/etc/init.d/"* \
"$STAGING/etc/hotplug.d/"*/* \
"$STAGING/etc/ldlda_help/"*.sh \ "$STAGING/etc/ldlda_help/"*.sh \
"$STAGING/root/.bin/wakey" \ "$STAGING/root/.bin/wakey" \
"$STAGING/root/.bin/wakey-agent" \
"$STAGING/root/.bin/kill_wakey.sh" \ "$STAGING/root/.bin/kill_wakey.sh" \
"$STAGING/root/.bin/remote_deploy_wakey.sh"; do "$STAGING/root/.bin/remote_deploy_wakey.sh"; do
[ -e "$f" ] && chmod +x "$f" 2>/dev/null || true [ -e "$f" ] && chmod +x "$f" 2>/dev/null || true
@@ -114,11 +116,14 @@ main() {
# Normalize line endings for shell scripts (avoid CRLF issues on OpenWrt) # Normalize line endings for shell scripts (avoid CRLF issues on OpenWrt)
for f in \ for f in \
"$STAGING/etc/init.d/"* \ "$STAGING/etc/init.d/"* \
"$STAGING/etc/hotplug.d/"*/* \
"$STAGING/etc/ldlda_help/"*.sh \ "$STAGING/etc/ldlda_help/"*.sh \
"$STAGING/root/.bin/"*.sh; do "$STAGING/root/.bin/"*.sh; do
[ -f "$f" ] && sed -i 's/\r$//' "$f" 2>/dev/null || true [ -f "$f" ] && sed -i 's/\r$//' "$f" 2>/dev/null || true
done done
chown -R root:root "$STAGING" 2>/dev/null || true
# Copy staged tree into / # Copy staged tree into /
tar -C "$STAGING" -cf - . | tar -C / -xpf - || fail "install copy failed" tar -C "$STAGING" -cf - . | tar -C / -xpf - || fail "install copy failed"
+3
View File
@@ -82,6 +82,9 @@ if (Test-Path $hotplugSrc) {
$content = Get-Content -Raw -LiteralPath $_.FullName $content = Get-Content -Raw -LiteralPath $_.FullName
$content = $content -replace "`r`n", "`n" $content = $content -replace "`r`n", "`n"
Set-Content -NoNewline -LiteralPath $dest -Value $content -Encoding UTF8 Set-Content -NoNewline -LiteralPath $dest -Value $content -Encoding UTF8
if (Get-Command chmod -ErrorAction SilentlyContinue) {
chmod +x $dest
}
} }
} }
} }
+44 -16
View File
@@ -64,6 +64,40 @@ sudo_cmd() {
fail 'root privileges are required for install/restart steps; set SUDO or run as root' fail 'root privileges are required for install/restart steps; set SUDO or run as root'
} }
priv() {
if [ "$(id -u)" -ne 0 ]; then
$SUDO_BIN "$@"
else
"$@"
fi
}
restore_label() {
path=$1
if command -v restorecon >/dev/null 2>&1; then
priv restorecon "$path" || true
fi
}
install_executable_force() {
src=$1
dst=$2
dir=$(dirname "$dst")
tmp="$dst.new.$$"
old="$dst.old.$$"
priv mkdir -p "$dir"
priv rm -f "$tmp"
priv cp -f "$src" "$tmp"
priv chmod 0755 "$tmp"
if [ -e "$dst" ]; then
priv mv -f "$dst" "$old" || true
fi
priv mv -f "$tmp" "$dst"
priv rm -f "$old"
restore_label "$dst"
}
main() { main() {
REPO_URL="${WAKEY_CC_REPO_URL:-https://git.ldlda.com/lda/wakey.git}" REPO_URL="${WAKEY_CC_REPO_URL:-https://git.ldlda.com/lda/wakey.git}"
REF="${WAKEY_CC_REF:-main}" REF="${WAKEY_CC_REF:-main}"
@@ -124,26 +158,20 @@ main() {
[ -f "$STAGING/bin/wakey-control-plane" ] || fail 'bundle missing bin/wakey-control-plane' [ -f "$STAGING/bin/wakey-control-plane" ] || fail 'bundle missing bin/wakey-control-plane'
[ -f "$STAGING/ui/dist/index.html" ] || fail 'bundle missing ui/dist/index.html' [ -f "$STAGING/ui/dist/index.html" ] || fail 'bundle missing ui/dist/index.html'
if [ "$(id -u)" -ne 0 ] && [ -n "$SUDO_BIN" ]; then priv mkdir -p "$ROOT"
$SUDO_BIN mkdir -p "$ROOT" install_executable_force \
$SUDO_BIN cp -a "$STAGING/." "$ROOT/" "$STAGING/bin/wakey-control-plane" \
else "$ROOT/bin/wakey-control-plane"
mkdir -p "$ROOT" priv cp -a "$STAGING/ui" "$ROOT/"
cp -a "$STAGING/." "$ROOT/" priv cp -a "$STAGING/scripts" "$ROOT/"
fi priv cp -a "$STAGING/deploy" "$ROOT/"
if [ -z "${WAKEY_CC_NO_RESTART:-}" ] && command -v systemctl >/dev/null 2>&1; then if [ -z "${WAKEY_CC_NO_RESTART:-}" ] && command -v systemctl >/dev/null 2>&1; then
if systemctl list-unit-files "$SERVICE" >/dev/null 2>&1; then if systemctl list-unit-files "$SERVICE" >/dev/null 2>&1; then
log "restarting $SERVICE" log "restarting $SERVICE"
if [ "$(id -u)" -ne 0 ] && [ -n "$SUDO_BIN" ]; then priv systemctl daemon-reload
$SUDO_BIN systemctl daemon-reload priv systemctl restart "$SERVICE"
$SUDO_BIN systemctl restart "$SERVICE" priv systemctl --no-pager --full status "$SERVICE" | sed -n '1,16p'
$SUDO_BIN systemctl --no-pager --full status "$SERVICE" | sed -n '1,16p'
else
systemctl daemon-reload
systemctl restart "$SERVICE"
systemctl --no-pager --full status "$SERVICE" | sed -n '1,16p'
fi
else else
log "service $SERVICE not installed; skipped restart" log "service $SERVICE not installed; skipped restart"
fi fi
+36
View File
@@ -130,6 +130,17 @@ pub struct WakeRequest {
pub ip: Option<IpAddr>, pub ip: Option<IpAddr>,
} }
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentObservation {
pub kind: String,
pub action: String,
pub mac: Option<String>,
pub ip: Option<IpAddr>,
pub hostname: Option<String>,
pub first_seen_unix: u64,
pub last_seen_unix: u64,
}
#[derive(Debug, Clone, Serialize, Deserialize)] #[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")] #[serde(tag = "kind", rename_all = "snake_case")]
pub enum AgentCommand { pub enum AgentCommand {
@@ -161,6 +172,10 @@ pub enum ClientMessage {
Heartbeat { Heartbeat {
agent_id: String, agent_id: String,
}, },
Observations {
agent_id: String,
observations: Vec<AgentObservation>,
},
Result { Result {
request_id: RequestId, request_id: RequestId,
result: CommandResult, result: CommandResult,
@@ -178,6 +193,7 @@ pub enum ServerMessage {
request_id: RequestId, request_id: RequestId,
command: AgentCommand, command: AgentCommand,
}, },
SyncObservations,
} }
#[cfg(test)] #[cfg(test)]
@@ -215,4 +231,24 @@ mod tests {
assert!(json.contains("\"type\":\"result\"")); assert!(json.contains("\"type\":\"result\""));
assert!(json.contains("\"kind\":\"devs\"")); assert!(json.contains("\"kind\":\"devs\""));
} }
#[test]
fn observations_message_serializes() {
let msg = ClientMessage::Observations {
agent_id: "agent-a".into(),
observations: vec![AgentObservation {
kind: "dhcp".into(),
action: "update".into(),
mac: Some("aa:bb:cc:dd:ee:ff".into()),
ip: Some("192.168.1.10".parse().expect("ip")),
hostname: Some("lda".into()),
first_seen_unix: 10,
last_seen_unix: 20,
}],
};
let json = serde_json::to_string(&msg).expect("serialize");
assert!(json.contains("\"type\":\"observations\""));
assert!(json.contains("\"kind\":\"dhcp\""));
}
} }
+55 -32
View File
@@ -9,7 +9,7 @@ use tracing::{debug, error, info, info_span, warn};
use crate::config::AgentConfig; use crate::config::AgentConfig;
use crate::dispatch::dispatch_command; use crate::dispatch::dispatch_command;
use crate::protocol::{AgentCommand, ClientMessage, ErrorPayload, ServerMessage}; use crate::protocol::{AgentCommand, AgentObservation, ClientMessage, ErrorPayload, ServerMessage};
pub async fn run(config: AgentConfig) -> Result<()> { pub async fn run(config: AgentConfig) -> Result<()> {
let mut backoff = config.reconnect_base_ms.max(100); let mut backoff = config.reconnect_base_ms.max(100);
@@ -85,13 +85,13 @@ async fn run_once(config: &AgentConfig) -> Result<()> {
.await?; .await?;
info!(agent_id = %config.agent_id, "agent websocket session authenticated"); info!(agent_id = %config.agent_id, "agent websocket session authenticated");
let http_client = reqwest::Client::new();
let mut heartbeat = interval(Duration::from_secs(30)); let mut heartbeat = interval(Duration::from_secs(30));
heartbeat.set_missed_tick_behavior(MissedTickBehavior::Skip); heartbeat.set_missed_tick_behavior(MissedTickBehavior::Skip);
let mut observation_sync = interval(Duration::from_secs( let mut observation_sync = interval(Duration::from_secs(
config.observation_sync_interval_seconds.max(1), config.observation_sync_interval_seconds.max(1),
)); ));
observation_sync.set_missed_tick_behavior(MissedTickBehavior::Skip); observation_sync.set_missed_tick_behavior(MissedTickBehavior::Skip);
observation_sync.reset();
loop { loop {
tokio::select! { tokio::select! {
@@ -102,7 +102,7 @@ async fn run_once(config: &AgentConfig) -> Result<()> {
debug!(agent_id = %config.agent_id, "heartbeat sent"); debug!(agent_id = %config.agent_id, "heartbeat sent");
} }
_ = observation_sync.tick() => { _ = observation_sync.tick() => {
if let Err(err) = send_agent_observations(&http_client, config).await { if let Err(err) = send_agent_observations_ws(&mut sink, config).await {
warn!(agent_id = %config.agent_id, error = %err, "failed to sync local observations"); warn!(agent_id = %config.agent_id, error = %err, "failed to sync local observations");
} }
} }
@@ -119,7 +119,7 @@ async fn run_once(config: &AgentConfig) -> Result<()> {
Message::Text(text) => { Message::Text(text) => {
match serde_json::from_str::<ServerMessage>(&text) { match serde_json::from_str::<ServerMessage>(&text) {
Ok(message) => { Ok(message) => {
handle_server_message(config, &mut sink, message).await?; handle_server_message(config, &mut sink, &mut observation_sync, message).await?;
} }
Err(err) => { Err(err) => {
// Allow the server to introduce extra frame types without // Allow the server to introduce extra frame types without
@@ -150,18 +150,7 @@ async fn run_once(config: &AgentConfig) -> Result<()> {
struct UploadAgentObservationsRequest { struct UploadAgentObservationsRequest {
agent_id: String, agent_id: String,
agent_token: String, agent_token: String,
observations: Vec<AgentObservationRequest>, observations: Vec<AgentObservation>,
}
#[derive(Debug, Serialize)]
struct AgentObservationRequest {
kind: String,
action: String,
mac: Option<String>,
ip: Option<IpAddr>,
hostname: Option<String>,
first_seen_unix: u64,
last_seen_unix: u64,
} }
pub async fn sync_observations_once(config: &AgentConfig) -> Result<usize> { pub async fn sync_observations_once(config: &AgentConfig) -> Result<usize> {
@@ -170,10 +159,7 @@ pub async fn sync_observations_once(config: &AgentConfig) -> Result<usize> {
} }
async fn send_agent_observations(client: &reqwest::Client, config: &AgentConfig) -> Result<usize> { async fn send_agent_observations(client: &reqwest::Client, config: &AgentConfig) -> Result<usize> {
let observations = let observations = load_agent_observations(config).await?;
wakey::wakey_linux::dhcp::list_local_observations_from_path(&config.observation_store_path)
.await
.context("failed to read local observations")?;
if observations.is_empty() { if observations.is_empty() {
return Ok(0); return Ok(0);
} }
@@ -182,18 +168,7 @@ async fn send_agent_observations(client: &reqwest::Client, config: &AgentConfig)
let payload = UploadAgentObservationsRequest { let payload = UploadAgentObservationsRequest {
agent_id: config.agent_id.clone(), agent_id: config.agent_id.clone(),
agent_token: config.agent_token.clone(), agent_token: config.agent_token.clone(),
observations: observations observations,
.into_iter()
.map(|observation| AgentObservationRequest {
kind: observation.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,
})
.collect(),
}; };
let response = client let response = client
@@ -218,6 +193,47 @@ async fn send_agent_observations(client: &reqwest::Client, config: &AgentConfig)
Ok(payload.observations.len()) Ok(payload.observations.len())
} }
async fn send_agent_observations_ws<S>(sink: &mut S, config: &AgentConfig) -> Result<usize>
where
S: SinkExt<Message> + Unpin,
<S as futures_util::Sink<Message>>::Error: std::error::Error + Send + Sync + 'static,
{
let observations = load_agent_observations(config).await?;
if observations.is_empty() {
return Ok(0);
}
let count = observations.len();
send_json(
sink,
&ClientMessage::Observations {
agent_id: config.agent_id.clone(),
observations,
},
)
.await?;
debug!(agent_id = %config.agent_id, observations = count, "sent observations over websocket");
Ok(count)
}
async fn load_agent_observations(config: &AgentConfig) -> Result<Vec<AgentObservation>> {
let observations =
wakey::wakey_linux::dhcp::list_local_observations_from_path(&config.observation_store_path)
.await
.context("failed to read local observations")?;
Ok(observations
.into_iter()
.map(|observation| AgentObservation {
kind: observation.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,
})
.collect())
}
pub fn next_backoff_ms(current_ms: u64, max_ms: u64) -> u64 { pub fn next_backoff_ms(current_ms: u64, max_ms: u64) -> u64 {
let cap = max_ms.max(current_ms); let cap = max_ms.max(current_ms);
current_ms.saturating_mul(2).min(cap) current_ms.saturating_mul(2).min(cap)
@@ -226,6 +242,7 @@ pub fn next_backoff_ms(current_ms: u64, max_ms: u64) -> u64 {
async fn handle_server_message<S>( async fn handle_server_message<S>(
config: &AgentConfig, config: &AgentConfig,
sink: &mut S, sink: &mut S,
observation_sync: &mut tokio::time::Interval,
message: ServerMessage, message: ServerMessage,
) -> Result<()> ) -> Result<()>
where where
@@ -261,6 +278,11 @@ where
} }
} }
} }
ServerMessage::SyncObservations => {
info!("received observation sync request from control-plane");
send_agent_observations_ws(sink, config).await?;
observation_sync.reset();
}
} }
Ok(()) Ok(())
} }
@@ -336,6 +358,7 @@ fn client_message_kind(message: &ClientMessage) -> &'static str {
ClientMessage::Hello { .. } => "hello", ClientMessage::Hello { .. } => "hello",
ClientMessage::Auth { .. } => "auth", ClientMessage::Auth { .. } => "auth",
ClientMessage::Heartbeat { .. } => "heartbeat", ClientMessage::Heartbeat { .. } => "heartbeat",
ClientMessage::Observations { .. } => "observations",
ClientMessage::Result { .. } => "result", ClientMessage::Result { .. } => "result",
ClientMessage::Error { .. } => "error", ClientMessage::Error { .. } => "error",
} }
@@ -0,0 +1,12 @@
{
"db_name": "SQLite",
"query": "UPDATE device_identifiers SET device_id = ?1 WHERE device_id = ?2",
"describe": {
"columns": [],
"parameters": {
"Right": 2
},
"nullable": []
},
"hash": "7247cc488921de2313a602697590502771a231bde7d03cd9ccac55a70971da45"
}
+3 -2
View File
@@ -5,7 +5,7 @@ mod stats;
pub use devices::{ pub use devices::{
attach_device_identifier, attach_observation_identifier, create_known_device, attach_device_identifier, attach_observation_identifier, create_known_device,
forget_known_device, list_known_devices, forget_known_device, list_known_devices, merge_known_device,
}; };
pub use enroll::{ pub use enroll::{
EnrollTokenStatus, IssueEnrollTokenResponse, RevokeAgentResponse, RevokeEnrollTokenResponse, EnrollTokenStatus, IssueEnrollTokenResponse, RevokeAgentResponse, RevokeEnrollTokenResponse,
@@ -13,6 +13,7 @@ pub use enroll::{
set_agent_nickname, set_agent_nickname,
}; };
pub use observations::{ pub use observations::{
list_agent_observation_history, list_agent_observations, upload_agent_observations, list_agent_observation_history, list_agent_observations, request_agent_observation_sync,
upload_agent_observations,
}; };
pub use stats::{StateStatsResponse, state_stats}; pub use stats::{StateStatsResponse, state_stats};
@@ -56,6 +56,11 @@ pub struct ForgetKnownDeviceResponse {
pub forgotten: bool, pub forgotten: bool,
} }
#[derive(Debug, Deserialize)]
pub struct MergeKnownDeviceRequest {
pub source_device_id: String,
}
pub async fn create_known_device( pub async fn create_known_device(
State(state): State<AppState>, State(state): State<AppState>,
Json(req): Json<CreateKnownDeviceRequest>, Json(req): Json<CreateKnownDeviceRequest>,
@@ -193,6 +198,33 @@ pub async fn attach_observation_identifier(
} }
} }
pub async fn merge_known_device(
State(state): State<AppState>,
AxumPath(device_id): AxumPath<String>,
Json(req): Json<MergeKnownDeviceRequest>,
) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
match state
.store
.merge_known_devices(&device_id, &req.source_device_id)
.await
{
Ok(Some(device)) => Ok((StatusCode::OK, Json(known_device_response(device)))),
Ok(None) => Err(json_error(
StatusCode::NOT_FOUND,
"known_device_not_found",
"target or source known device not found",
)),
Err(err) => {
warn!(error = %err, "failed to merge known devices");
Err(json_error(
StatusCode::BAD_REQUEST,
"merge_known_device_failed",
&err.to_string(),
))
}
}
}
fn known_device_response(device: crate::state::KnownDevice) -> KnownDeviceResponse { fn known_device_response(device: crate::state::KnownDevice) -> KnownDeviceResponse {
KnownDeviceResponse { KnownDeviceResponse {
device_id: device.device_id, device_id: device.device_id,
@@ -1,12 +1,13 @@
use axum::Json; use axum::Json;
use axum::extract::{Query, State}; use axum::extract::{Path as AxumPath, Query, State};
use axum::http::StatusCode; use axum::http::StatusCode;
use axum::response::IntoResponse; use axum::response::IntoResponse;
use serde::Deserialize; use serde::Deserialize;
use tracing::warn; use tracing::warn;
use wakey_agent::protocol::ServerMessage;
use crate::api::json_error; use crate::api::json_error;
use crate::runtime::AppState; use crate::runtime::{AppState, SessionEvent};
use crate::state::{ use crate::state::{
AgentDeviceObservationEvent, AgentDeviceObservationInput, AgentDeviceObservationView, AgentDeviceObservationEvent, AgentDeviceObservationInput, AgentDeviceObservationView,
}; };
@@ -35,6 +36,12 @@ pub struct UploadAgentObservationsResponse {
pub accepted: usize, pub accepted: usize,
} }
#[derive(Debug, serde::Serialize, serde::Deserialize)]
pub struct RequestAgentObservationSyncResponse {
pub agent_id: String,
pub requested: bool,
}
#[derive(Debug, Deserialize)] #[derive(Debug, Deserialize)]
pub struct ListObservationsQuery { pub struct ListObservationsQuery {
pub agent_id: Option<String>, pub agent_id: Option<String>,
@@ -130,6 +137,43 @@ pub async fn list_agent_observations(
} }
} }
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( pub async fn list_agent_observation_history(
State(state): State<AppState>, State(state): State<AppState>,
Query(query): Query<ListObservationHistoryQuery>, Query(query): Query<ListObservationHistoryQuery>,
+2 -2
View File
@@ -14,8 +14,8 @@ pub use control::{
StateStatsResponse, attach_device_identifier, attach_observation_identifier, StateStatsResponse, attach_device_identifier, attach_observation_identifier,
create_known_device, enroll, forget_known_device, healthz, issue_enroll_token, create_known_device, enroll, forget_known_device, healthz, issue_enroll_token,
list_agent_observation_history, list_agent_observations, list_enroll_tokens, list_agent_observation_history, list_agent_observations, list_enroll_tokens,
list_known_devices, revoke_agent, revoke_enroll_token, set_agent_nickname, state_stats, list_known_devices, merge_known_device, request_agent_observation_sync, revoke_agent,
upload_agent_observations, revoke_enroll_token, set_agent_nickname, state_stats, upload_agent_observations,
}; };
pub fn json_error( pub fn json_error(
+8
View File
@@ -101,6 +101,10 @@ fn control_api_routes() -> Router<AppState> {
"/api/v1/control/observations/history", "/api/v1/control/observations/history",
get(api::list_agent_observation_history), get(api::list_agent_observation_history),
) )
.route(
"/api/v1/control/agents/{agent_id}/observations/sync",
post(api::request_agent_observation_sync),
)
.route( .route(
"/api/v1/control/devices", "/api/v1/control/devices",
get(api::list_known_devices).post(api::create_known_device), get(api::list_known_devices).post(api::create_known_device),
@@ -109,6 +113,10 @@ fn control_api_routes() -> Router<AppState> {
"/api/v1/control/devices/{device_id}", "/api/v1/control/devices/{device_id}",
axum::routing::delete(api::forget_known_device), axum::routing::delete(api::forget_known_device),
) )
.route(
"/api/v1/control/devices/{device_id}/merge",
post(api::merge_known_device),
)
.route( .route(
"/api/v1/control/devices/{device_id}/identifiers", "/api/v1/control/devices/{device_id}/identifiers",
post(api::attach_device_identifier), post(api::attach_device_identifier),
+53
View File
@@ -337,6 +337,59 @@ mod tests {
cleanup_dir(&dir); cleanup_dir(&dir);
} }
#[tokio::test]
async fn merge_known_devices_moves_identifiers_and_deletes_source() {
let (store, dir) = make_store().await;
let target = 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:01".into(),
}],
})
.await
.expect("target should create");
let source = store
.create_known_device(KnownDeviceInput {
display_name: "lda duplicate".into(),
pinned: false,
notes: None,
identifiers: vec![DeviceIdentifierInput {
kind: "mac".into(),
value: "aa:bb:cc:dd:ee:02".into(),
}],
})
.await
.expect("source should create");
let merged = store
.merge_known_devices(&target.device_id, &source.device_id)
.await
.expect("merge should succeed")
.expect("target should remain");
assert_eq!(merged.device_id, target.device_id);
assert_eq!(merged.identifiers.len(), 2);
assert!(
merged
.identifiers
.iter()
.any(|identifier| identifier.value == "aa:bb:cc:dd:ee:02")
);
assert!(
store
.get_known_device(&source.device_id)
.await
.expect("source lookup should work")
.is_none()
);
cleanup_dir(&dir);
}
#[tokio::test] #[tokio::test]
async fn agent_observations_upsert_current_state_and_events() { async fn agent_observations_upsert_current_state_and_events() {
let (store, dir) = make_store().await; let (store, dir) = make_store().await;
@@ -73,6 +73,75 @@ impl Store {
Ok(result.rows_affected() > 0) Ok(result.rows_affected() > 0)
} }
pub async fn merge_known_devices(
&self,
target_device_id: &str,
source_device_id: &str,
) -> Result<Option<KnownDevice>> {
if target_device_id == source_device_id {
return self.get_known_device(target_device_id).await;
}
let now = now_unix();
let now_i64 = i64::try_from(now).context("known device timestamp overflow")?;
let mut tx = self
.pool
.begin()
.await
.context("failed starting known device merge transaction")?;
let target_exists = sqlx::query_scalar!(
r#"SELECT COUNT(*) as "count!: i64" FROM known_devices WHERE device_id = ?1"#,
target_device_id
)
.fetch_one(&mut *tx)
.await
.context("failed checking target known device existence")?;
if target_exists == 0 {
return Ok(None);
}
let source_exists = sqlx::query_scalar!(
r#"SELECT COUNT(*) as "count!: i64" FROM known_devices WHERE device_id = ?1"#,
source_device_id
)
.fetch_one(&mut *tx)
.await
.context("failed checking source known device existence")?;
if source_exists == 0 {
return Ok(None);
}
sqlx::query!(
"UPDATE device_identifiers SET device_id = ?1 WHERE device_id = ?2",
target_device_id,
source_device_id
)
.execute(&mut *tx)
.await
.context("failed moving source identifiers to target device")?;
sqlx::query!(
"DELETE FROM known_devices WHERE device_id = ?1",
source_device_id
)
.execute(&mut *tx)
.await
.context("failed deleting merged source known device")?;
sqlx::query!(
"UPDATE known_devices SET updated_at_unix = ?1 WHERE device_id = ?2",
now_i64,
target_device_id
)
.execute(&mut *tx)
.await
.context("failed updating merged target known device timestamp")?;
tx.commit()
.await
.context("failed committing known device merge transaction")?;
self.get_known_device(target_device_id).await
}
pub async fn attach_device_identifier( pub async fn attach_device_identifier(
&self, &self,
device_id: &str, device_id: &str,
+40 -1
View File
@@ -8,7 +8,7 @@ use std::time::Instant;
use tokio::sync::mpsc; use tokio::sync::mpsc;
use tracing::{debug, info, info_span, warn}; use tracing::{debug, info, info_span, warn};
use uuid::Uuid; use uuid::Uuid;
use wakey_agent::protocol::{ErrorPayload, RequestId}; use wakey_agent::protocol::{AgentObservation, ErrorPayload, RequestId, ServerMessage};
use crate::runtime::{AgentReply, AgentSession, AppState, SessionEvent}; use crate::runtime::{AgentReply, AgentSession, AppState, SessionEvent};
use crate::state::AuditEventInput; use crate::state::AuditEventInput;
@@ -26,6 +26,10 @@ enum IncomingClientMessage {
Heartbeat { Heartbeat {
agent_id: String, agent_id: String,
}, },
Observations {
agent_id: String,
observations: Vec<AgentObservation>,
},
Result { Result {
request_id: RequestId, request_id: RequestId,
result: serde_json::Value, result: serde_json::Value,
@@ -238,6 +242,7 @@ async fn process_agent_text(
{ {
warn!(error = %err, "failed to append audit event for auth success"); warn!(error = %err, "failed to append audit event for auth success");
} }
let _ = tx.send(SessionEvent::Message(ServerMessage::SyncObservations));
} }
IncomingClientMessage::Heartbeat { agent_id } => { IncomingClientMessage::Heartbeat { agent_id } => {
if authed_agent_id.as_deref() != Some(agent_id.as_str()) { if authed_agent_id.as_deref() != Some(agent_id.as_str()) {
@@ -246,6 +251,40 @@ async fn process_agent_text(
ensure_current_session(state, &agent_id, connection_id).await?; ensure_current_session(state, &agent_id, connection_id).await?;
debug!(agent_id = %agent_id, "heartbeat received"); debug!(agent_id = %agent_id, "heartbeat received");
} }
IncomingClientMessage::Observations {
agent_id,
observations,
} => {
if authed_agent_id.as_deref() != Some(agent_id.as_str()) {
anyhow::bail!("observations for unauthenticated or mismatched agent");
}
ensure_current_session(state, &agent_id, connection_id).await?;
let inputs = observations
.into_iter()
.map(|observation| crate::state::AgentDeviceObservationInput {
kind: observation.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,
})
.collect();
match state
.store
.upsert_agent_observations(&agent_id, inputs)
.await
{
Ok(accepted) => {
debug!(agent_id = %agent_id, accepted, "agent websocket observations accepted");
}
Err(err) => {
warn!(agent_id = %agent_id, error = %err, "failed to store websocket observations");
anyhow::bail!("failed to store observations: {err}");
}
}
}
IncomingClientMessage::Result { request_id, result } => { IncomingClientMessage::Result { request_id, result } => {
let agent_id = authed_agent_id let agent_id = authed_agent_id
.as_deref() .as_deref()