use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::{mpsc, Notify};
use tokio::task::JoinHandle;
use tracing::warn;
use super::lockfile::{LockKind, Lockfile};
use super::{map, reader};
use crate::http::RestClient;
use crate::protocol::Clip;
use crate::store::{queries, Store};
use crate::ws::{self, DecryptFailReason, WsConfig, WsEvent, WsStatus};
pub type OnNewClipCallback = Arc<dyn Fn(&Clip) + Send + Sync>;
pub type OnConnectedCallback = Arc<dyn Fn() + Send + Sync>;
pub struct Writer {
handle: Option<JoinHandle<()>>,
stop: Arc<Notify>,
_lock: Lockfile, on_connected: Option<OnConnectedCallback>,
}
impl Writer {
pub async fn start(
store: Arc<Store>,
client: Arc<RestClient>,
ws_cfg: WsConfig,
lock_path: PathBuf,
kind: LockKind,
on_new_clip: Option<OnNewClipCallback>,
on_connected: Option<OnConnectedCallback>,
) -> std::io::Result<Option<Self>> {
let lock = match Lockfile::try_acquire(&lock_path, kind)? {
Some(l) => l,
None => return Ok(None),
};
let stop = Arc::new(Notify::new());
let stop_clone = stop.clone();
let store_clone = store.clone();
let client_clone = client.clone();
let enc_key = ws_cfg.encryption_key;
let on_connected_for_task = on_connected.clone();
let handle = tokio::spawn(async move {
let _ = reader::backfill_once(
&store_clone,
&client_clone,
reader::BackfillBudget::default(),
enc_key.as_ref(),
)
.await;
let (tx, mut rx) = mpsc::channel::<WsEvent>(64);
let _ws_handle = tokio::spawn(ws::run(ws_cfg, tx));
loop {
tokio::select! {
_ = stop_clone.notified() => return,
maybe = rx.recv() => {
match maybe {
None => {
return;
}
Some(event) => {
handle_ws_event(
&store_clone,
&client_clone,
event,
enc_key,
on_new_clip.as_ref(),
on_connected_for_task.as_ref(),
)
.await;
}
}
}
}
}
});
Ok(Some(Self {
handle: Some(handle),
stop,
_lock: lock,
on_connected,
}))
}
pub fn with_on_connected(mut self, cb: OnConnectedCallback) -> Self {
self.on_connected = Some(cb);
self
}
pub async fn shutdown(mut self) {
self.stop.notify_waiters();
if let Some(h) = self.handle.take() {
let _ = h.await;
}
}
}
pub(crate) fn dispatch_on_connected(cb: &OnConnectedCallback) {
let cb = cb.clone();
tokio::spawn(async move { cb() });
}
async fn handle_ws_event(
store: &Store,
client: &RestClient,
event: WsEvent,
enc_key: Option<[u8; 32]>,
on_new_clip: Option<&OnNewClipCallback>,
on_connected: Option<&OnConnectedCallback>,
) {
match event {
WsEvent::NewClip { clip, plaintext: _ } => {
let _ = enc_key; let clip = *clip; if clip.encrypted {
warn!(
clip_id = %clip.clip_id,
"writer: received NewClip with encrypted=true — skipping"
);
return;
}
match map::clip_wire_to_stored(&clip) {
Ok(Some(stored)) => {
let inserted = match queries::insert_clip(store, &stored) {
Ok(()) => true,
Err(e) => {
warn!(clip_id = %stored.id, error = %e, "writer: insert_clip failed");
false
}
};
if inserted {
if let Err(e) = queries::set_watermark(store, &stored.id) {
warn!(clip_id = %stored.id, error = %e, "writer: set_watermark failed");
}
if let Some(cb) = on_new_clip {
cb(&clip);
}
}
}
Ok(None) => {
}
Err(e) => {
warn!(clip_id = %clip.clip_id, error = %e, "writer: map_wire_to_stored failed");
}
}
}
WsEvent::ClipDecryptFailed { clip_id, reason } => {
let reason_str = match reason {
DecryptFailReason::MissingKey => "no encryption key available".into(),
DecryptFailReason::TagFailed(e) => format!("key mismatch: {e}"),
};
warn!(clip_id = %clip_id, reason = %reason_str, "writer: skipping encrypted clip — decrypt failed");
}
WsEvent::ClipDeleted { clip_id: _ } => {
}
WsEvent::Revoked { reason } => {
warn!(
reason = ?reason,
"writer: device revoked by relay — writer will stop receiving events"
);
}
WsEvent::TokenRotated {
token: _,
device_id: _,
} => {
}
WsEvent::KeyExchangeRequested { device_id } => {
log::info!(
"writer: received KeyExchangeRequested device_id={:?}",
device_id
);
let Some(did) = device_id else {
log::info!("writer: key_exchange_requested missing device_id — skipping");
return;
};
let Some(key) = enc_key else {
log::info!(
"writer: cannot bear key — no encryption key on this device (target_device_id={})",
did
);
return;
};
log::info!("writer: bearing key for target_device_id={}", did);
if let Err(e) = crate::key_exchange::handle_event(client, &did, &key).await {
log::warn!(
"writer: key_exchange responder failed: target_device_id={}, error={}",
did,
e
);
} else {
log::info!(
"writer: posted encrypted key bundle to peer target_device_id={}",
did
);
}
}
WsEvent::Status(WsStatus::Connected) => {
tracing::debug!("writer: WS connected");
if let Some(cb) = on_connected {
dispatch_on_connected(cb);
}
}
WsEvent::Status(WsStatus::Disconnected) => {
tracing::debug!("writer: WS disconnected — ws::run will reconnect");
}
WsEvent::Status(WsStatus::Connecting) => {
tracing::debug!("writer: WS connecting");
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::path::Path;
use std::sync::atomic::{AtomicUsize, Ordering};
fn make_test_store() -> Arc<crate::store::Store> {
Arc::new(crate::store::Store::open(Path::new(":memory:")).expect("in-memory store"))
}
fn make_test_client() -> Arc<crate::http::RestClient> {
Arc::new(crate::http::RestClient::for_test_offline())
}
#[tokio::test]
async fn writer_invokes_on_connected_on_status_connected() {
let counter = Arc::new(AtomicUsize::new(0));
let counter_c = counter.clone();
let cb: OnConnectedCallback = Arc::new(move || {
counter_c.fetch_add(1, Ordering::SeqCst);
});
let store = make_test_store();
let client = make_test_client();
handle_ws_event(
&store,
&client,
WsEvent::Status(WsStatus::Connected),
None,
None,
Some(&cb),
)
.await;
handle_ws_event(
&store,
&client,
WsEvent::Status(WsStatus::Disconnected),
None,
None,
Some(&cb),
)
.await;
handle_ws_event(
&store,
&client,
WsEvent::Status(WsStatus::Connected),
None,
None,
Some(&cb),
)
.await;
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
assert_eq!(
counter.load(Ordering::SeqCst),
2,
"callback should fire exactly once per Connected event"
);
}
#[tokio::test]
async fn writer_skips_callback_when_none_set() {
let store = make_test_store();
let client = make_test_client();
handle_ws_event(
&store,
&client,
WsEvent::Status(WsStatus::Connected),
None,
None,
None,
)
.await;
}
}