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

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

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
wakey-control-plane/.sqlx/query-7247cc488921de2313a602697590502771a231bde7d03cd9ccac55a70971da45.json
Generated
+12
View File
@@ -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"
}
wakey-control-plane/src/api/control.rs
+3 -2
View File
@@ -5,7 +5,7 @@ mod stats;
pub use devices::{
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::{
EnrollTokenStatus, IssueEnrollTokenResponse, RevokeAgentResponse, RevokeEnrollTokenResponse,
@@ -13,6 +13,7 @@ pub use enroll::{
set_agent_nickname,
};
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};
wakey-control-plane/src/api/control/devices.rs
+32
View File
@@ -56,6 +56,11 @@ pub struct ForgetKnownDeviceResponse {
pub forgotten: bool,
}
#[derive(Debug, Deserialize)]
pub struct MergeKnownDeviceRequest {
pub source_device_id: String,
}
pub async fn create_known_device(
State(state): State<AppState>,
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 {
KnownDeviceResponse {
device_id: device.device_id,
wakey-control-plane/src/api/control/observations.rs
+46 -2
View File
@@ -1,12 +1,13 @@
use axum::Json;
use axum::extract::{Query, State};
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;
use crate::runtime::{AppState, SessionEvent};
use crate::state::{
AgentDeviceObservationEvent, AgentDeviceObservationInput, AgentDeviceObservationView,
};
@@ -35,6 +36,12 @@ 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>,
@@ -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(
State(state): State<AppState>,
Query(query): Query<ListObservationHistoryQuery>,
wakey-control-plane/src/api/mod.rs
+2 -2
View File
@@ -14,8 +14,8 @@ pub use control::{
StateStatsResponse, attach_device_identifier, attach_observation_identifier,
create_known_device, enroll, forget_known_device, healthz, issue_enroll_token,
list_agent_observation_history, list_agent_observations, list_enroll_tokens,
list_known_devices, revoke_agent, revoke_enroll_token, set_agent_nickname, state_stats,
upload_agent_observations,
list_known_devices, merge_known_device, request_agent_observation_sync, revoke_agent,
revoke_enroll_token, set_agent_nickname, state_stats, upload_agent_observations,
};
pub fn json_error(
wakey-control-plane/src/runtime/mod.rs
+8
View File
@@ -101,6 +101,10 @@ fn control_api_routes() -> Router<AppState> {
"/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),
@@ -109,6 +113,10 @@ fn control_api_routes() -> Router<AppState> {
"/api/v1/control/devices/{device_id}",
axum::routing::delete(api::forget_known_device),
)
.route(
"/api/v1/control/devices/{device_id}/merge",
post(api::merge_known_device),
)
.route(
"/api/v1/control/devices/{device_id}/identifiers",
post(api::attach_device_identifier),
wakey-control-plane/src/state/store.rs
+53
View File
@@ -337,6 +337,59 @@ mod tests {
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]
async fn agent_observations_upsert_current_state_and_events() {
let (store, dir) = make_store().await;
wakey-control-plane/src/state/store/devices.rs
+69
View File
@@ -73,6 +73,75 @@ impl Store {
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(
&self,
device_id: &str,
wakey-control-plane/src/ws.rs
+40 -1
View File
@@ -8,7 +8,7 @@ use std::time::Instant;
use tokio::sync::mpsc;
use tracing::{debug, info, info_span, warn};
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::state::AuditEventInput;
@@ -26,6 +26,10 @@ enum IncomingClientMessage {
Heartbeat {
agent_id: String,
},
Observations {
agent_id: String,
observations: Vec<AgentObservation>,
},
Result {
request_id: RequestId,
result: serde_json::Value,
@@ -238,6 +242,7 @@ async fn process_agent_text(
{
warn!(error = %err, "failed to append audit event for auth success");
}
let _ = tx.send(SessionEvent::Message(ServerMessage::SyncObservations));
}
IncomingClientMessage::Heartbeat { agent_id } => {
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?;
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 } => {
let agent_id = authed_agent_id
.as_deref()