diff options
| author | tv <tv@krebsco.de> | 2026-10-19 23:44:20 +0200 |
|---|---|---|
| committer | tv <tv@krebsco.de> | 2026-10-19 23:44:20 +0200 |
| commit | d3a9d5d3ebff98e342dff1df0ebad86f346719f5 (patch) | |
| tree | f24b52636c4d6aca66849d6deecc79e3d1f96eb5 | |
initial commit
| -rw-r--r-- | src/api/debug.rs | 41 | ||||
| -rw-r--r-- | src/api/mod.rs | 160 | ||||
| -rw-r--r-- | src/api/peers.rs | 161 | ||||
| -rw-r--r-- | src/api/store.rs | 359 | ||||
| -rw-r--r-- | src/auth.rs | 141 | ||||
| -rw-r--r-- | src/bootstrap.rs | 99 | ||||
| -rw-r--r-- | src/firewall.rs | 105 | ||||
| -rw-r--r-- | src/main.rs | 230 | ||||
| -rw-r--r-- | src/mycelium.rs | 882 | ||||
| -rw-r--r-- | src/peer.rs | 593 | ||||
| -rw-r--r-- | src/state.rs | 483 | ||||
| -rw-r--r-- | src/store.rs | 506 | ||||
| -rw-r--r-- | src/sync.rs | 447 | ||||
| -rw-r--r-- | src/transport.rs | 74 |
14 files changed, 4281 insertions, 0 deletions
diff --git a/src/api/debug.rs b/src/api/debug.rs new file mode 100644 index 0000000..3c24f62 --- /dev/null +++ b/src/api/debug.rs @@ -0,0 +1,41 @@ +//! /debug/* handlers (only routed with --debug-api). + +use std::collections::HashSet; + +use axum::{body::Bytes, extract::State, http::StatusCode, Json}; +use tracing::{info, warn}; + +use crate::firewall::FirewallRules; +use crate::state::AppState; + +pub(crate) async fn get_debug_firewall(State(state): State<AppState>) -> Json<FirewallRules> { + Json(state.firewall.rules()) +} + +pub(crate) async fn put_debug_firewall( + State(state): State<AppState>, + Json(rules): Json<FirewallRules>, +) -> StatusCode { + info!(inbound = ?rules.inbound, outbound = ?rules.outbound, "Firewall rules replaced"); + state.firewall.set_rules(rules); + StatusCode::NO_CONTENT +} + +/// Compatibility for the existing harness: block a set of peers in both directions. +/// Parses the body without requiring a Content-Type header, as before. +pub(crate) async fn put_debug_blacklist(State(state): State<AppState>, body: Bytes) -> StatusCode { + match serde_json::from_slice::<HashSet<String>>(&body) { + Ok(peers) => { + info!(peers = ?peers, "Blacklist replaced (both directions)"); + state.firewall.set_rules(FirewallRules { + inbound: peers.clone(), + outbound: peers, + }); + StatusCode::NO_CONTENT + } + Err(error) => { + warn!(error = %error, "Invalid blacklist body"); + StatusCode::BAD_REQUEST + } + } +} diff --git a/src/api/mod.rs b/src/api/mod.rs new file mode 100644 index 0000000..8038f16 --- /dev/null +++ b/src/api/mod.rs @@ -0,0 +1,160 @@ +//! HTTP API: router, OpenAPI document and shared response helpers. + +mod debug; +mod peers; +mod store; + +use axum::{ + http::{ + header::{CONTENT_TYPE, ETAG}, + HeaderValue, StatusCode, + }, + response::{IntoResponse, Response}, + routing::{delete, get, put}, + Router, +}; +use serde_json::Value as JsonValue; +use sha2::{Digest, Sha256}; +use tracing::{error, warn}; +use utoipa::OpenApi; +use utoipa_swagger_ui::SwaggerUi; + +use crate::auth::{basic_auth, AuthState}; +use crate::state::AppState; +use crate::store::StoreError; + +/// Builds the complete HTTP router. Everything except Swagger is authenticated. +pub(crate) fn router(state: AppState, auth_state: AuthState, debug_api: bool) -> Router { + let mut api: Router<AppState> = Router::new() + .route("/store", get(store::get_store_root)) + .route( + "/store/{*key}", + get(store::get_store_entry) + .put(store::put_store_entry) + .delete(store::delete_store_entry), + ) + .route("/stream/{*key}", get(store::stream_store_entry)) + .route("/peers", get(peers::list_peers)) + .route("/peers/{peer}", delete(peers::delete_peer)); + + if debug_api { + warn!("Debug API enabled"); + api = api + .route( + "/debug/firewall", + get(debug::get_debug_firewall).put(debug::put_debug_firewall), + ) + .route("/debug/blacklist", put(debug::put_debug_blacklist)); + } + + api.layer(axum::middleware::from_fn_with_state(auth_state, basic_auth)) + .merge(SwaggerUi::new("/swagger-ui").url("/openapi.json", ApiDoc::openapi())) + .with_state(state) +} + +// --- shared helpers ---------------------------------------------------------- + +pub(crate) fn store_status(error: StoreError, context: &str) -> StatusCode { + match error { + StoreError::Conflict => StatusCode::CONFLICT, + StoreError::NotFound => StatusCode::NOT_FOUND, + other => { + error!(error = %format!("{other:#}"), "{context}"); + StatusCode::INTERNAL_SERVER_ERROR + } + } +} + +fn etag_for(bytes: &[u8]) -> String { + let digest = Sha256::digest(bytes); + let hex = digest + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::<String>(); + + format!("\"{hex}\"") +} + +fn parse_key_path(key: &str) -> Result<Vec<String>, StatusCode> { + let segments: Vec<String> = key.split('/').map(str::to_owned).collect(); + + if segments.is_empty() || segments.iter().any(|segment| segment.is_empty()) { + return Err(StatusCode::BAD_REQUEST); + } + + Ok(segments) +} + +fn response_for_value( + value: &JsonValue, + accept: Option<&HeaderValue>, +) -> Result<Response, StatusCode> { + let wants_text = accept + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| { + value + .split(',') + .any(|item| item.trim().starts_with("text/plain")) + }); + + let (content_type, body) = if wants_text { + let body = match value { + JsonValue::String(value) => value.clone(), + other => serde_json::to_string(other).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?, + }; + + ("text/plain; charset=utf-8", body) + } else { + ( + "application/json", + serde_json::to_string(value).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?, + ) + }; + + let etag = etag_for(body.as_bytes()); + + let mut response = body.into_response(); + + response + .headers_mut() + .insert(CONTENT_TYPE, HeaderValue::from_static(content_type)); + + response.headers_mut().insert( + ETAG, + HeaderValue::from_str(&etag).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?, + ); + + Ok(response) +} + +// --- openapi spec ------------------------------------------------------------ +// TODO spec is mostly out of sync + +#[derive(OpenApi)] +#[openapi( + paths( + store::get_store_root, + store::get_store_entry, + store::put_store_entry, + store::delete_store_entry, + peers::list_peers, + peers::delete_peer, + ), + components( + schemas( + peers::PeerStatus, + peers::PeerListResponse, + peers::PeerLivenessStatus, + ) + ), + tags( + (name = "Store", description = "Key-value store operations"), + (name = "Peers", description = "Peer discovery and retirement") + ), + info( + title = "Automerge Replication API", + version = "1.0.0", + description = "HTTP API for a replicated Automerge key-value store" + ) +)] +struct ApiDoc; diff --git a/src/api/peers.rs b/src/api/peers.rs new file mode 100644 index 0000000..b2658a5 --- /dev/null +++ b/src/api/peers.rs @@ -0,0 +1,161 @@ +//! /peers handlers. + +use std::sync::Arc; + +use axum::{ + extract::{Path, Query, State}, + http::StatusCode, + Json, +}; +use serde::{Deserialize, Serialize}; +use tokio::sync::Mutex; +use tokio::time::Instant; +use tracing::warn; +use utoipa::ToSchema; + +use super::store_status; +use crate::peer::forgotten_path; +use crate::state::{unix_millis, AppState, HelloState, Liveness, PeerMembership, PeerState}; +use crate::store; + +#[derive(Debug, Serialize, ToSchema)] +#[serde(rename_all = "lowercase")] +pub(crate) enum PeerLivenessStatus { + Active, + Dormant, +} + +#[derive(Debug, Serialize, ToSchema)] +#[serde(rename_all = "camelCase")] +pub(crate) struct PeerStatus { + public_key: String, + validated_cluster: Option<String>, + status: PeerLivenessStatus, + handshake_complete: bool, + last_seen_secs_ago: Option<u64>, + dormant_for_secs: Option<u64>, + has_pending_delivery: bool, +} + +#[derive(Debug, Serialize, ToSchema)] +#[serde(rename_all = "camelCase")] +pub(crate) struct PeerListResponse { + peers: Vec<PeerStatus>, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct DeletePeerQuery { + #[serde(default)] + force: bool, +} + +#[utoipa::path( + get, + path = "/peers", + tag = "Peers", + responses( + ( + status = 200, + description = "Lists known peers and their local replication status", + body = PeerListResponse + ) + ) +)] +pub(crate) async fn list_peers( + State(state): State<AppState>, +) -> Result<Json<PeerListResponse>, StatusCode> { + let snapshot: Vec<(String, Arc<Mutex<PeerState>>)> = state + .peer_states + .read() + .await + .iter() + .map(|(pk, m)| (pk.clone(), Arc::clone(m))) + .collect(); + + let now = Instant::now(); + let mut peers = Vec::with_capacity(snapshot.len()); + + for (public_key, peer_state_mutex) in snapshot { + let ps = peer_state_mutex.lock().await; + + let (status, dormant_for_secs) = match ps.liveness { + Liveness::Active => (PeerLivenessStatus::Active, None), + Liveness::Dormant { since } => ( + PeerLivenessStatus::Dormant, + Some(now.saturating_duration_since(since).as_secs()), + ), + }; + + peers.push(PeerStatus { + public_key, + validated_cluster: match &ps.membership { + PeerMembership::Unknown => None, + PeerMembership::Member { cluster_id } => Some(cluster_id.clone()), + }, + status, + handshake_complete: ps.hello_state == HelloState::Established, + last_seen_secs_ago: ps + .last_seen + .map(|t| now.saturating_duration_since(t).as_secs()), + dormant_for_secs, + has_pending_delivery: ps.pending_outbound_sync.is_some(), + }); + } + + peers.sort_by(|a, b| a.public_key.cmp(&b.public_key)); + Ok(Json(PeerListResponse { peers })) +} + +#[utoipa::path( + delete, + path = "/peers/{peer}", + tag = "Peers", + params( + ("peer" = String, Path, description = "Peer public key"), + ("force" = Option<bool>, Query, description = "Also forget a peer that is currently alive; it returns on its next message") + ), + responses( + (status = 204, description = "Cluster-wide forget recorded"), + (status = 400, description = "Not a public key, or the local node's own key"), + (status = 409, description = "Peer is alive; pass force=true"), + (status = 503, description = "Node is not ready") + ) +)] +pub(crate) async fn delete_peer( + State(state): State<AppState>, + Path(peer): Path<String>, + Query(query): Query<DeletePeerQuery>, +) -> Result<StatusCode, StatusCode> { + if !state.is_ready().await { + return Err(StatusCode::SERVICE_UNAVAILABLE); + } + if peer.is_empty() || peer == state.local_id() { + return Err(StatusCode::BAD_REQUEST); + } + + // Unknown locally is fine: other nodes may still know the peer. + if let Some(peer_state_mutex) = state.get_peer_state(&peer).await { + if peer_state_mutex + .lock() + .await + .is_alive(&state.timing, Instant::now()) + && !query.force + { + return Err(StatusCode::CONFLICT); + } + } + + let record = serde_json::json!({ "at": unix_millis(), "by": state.local_id() }); + + let peer_clone = peer.clone(); + state + .mutate_doc(move |doc| store::write_path(doc, &forgotten_path(&peer_clone), &record)) + .await + .map_err(|e| store_status(e, "Failed to write forget record"))?; + + if state.peer_states.write().await.remove(&peer).is_some() { + warn!(peer = %peer, "Peer forgotten cluster-wide by operator"); + } + + Ok(StatusCode::NO_CONTENT) +} diff --git a/src/api/store.rs b/src/api/store.rs new file mode 100644 index 0000000..dd5be94 --- /dev/null +++ b/src/api/store.rs @@ -0,0 +1,359 @@ +//! /store and /stream handlers. + +use std::io::Error as IoError; + +use axum::{ + body::{Body, Bytes}, + extract::{Path, Query, State}, + http::{ + header::{ACCEPT, CONTENT_TYPE}, + HeaderMap, HeaderValue, StatusCode, + }, + response::{IntoResponse, Response}, +}; +use serde::Deserialize; +use serde_json::Value as JsonValue; +use tokio::sync::broadcast; +use tracing::{debug, error, instrument}; + +use super::{parse_key_path, response_for_value, store_status}; +use crate::state::AppState; +use crate::store::{self, StoreError}; +use crate::sync::paths_related; + +#[derive(Debug, Default, Deserialize)] +pub(crate) struct ReadQuery { + #[serde(default)] + conflicts: bool, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct DeleteQuery { + #[serde(default)] + prune: bool, +} + +fn conflicts_response( + value: JsonValue, + conflicts: Vec<store::ConflictReport>, +) -> Result<Response, StatusCode> { + let body = serde_json::json!({ "value": value, "conflicts": conflicts }); + response_for_value(&body, None) // always JSON +} + +#[utoipa::path( + get, + path = "/store", + tag = "Store", + params( + ("conflicts" = Option<bool>, Query, description = "Wrap the value and list every conflict resolved on the way"), + ), + responses( + (status = 200, description = "Returns the complete store as JSON or text, depending on Accept"), + (status = 503, description = "Node is not ready"), + (status = 500, description = "Failed to read or serialize the document") + ) +)] +pub(crate) async fn get_store_root( + State(state): State<AppState>, + headers: HeaderMap, + Query(query): Query<ReadQuery>, +) -> Result<Response, StatusCode> { + let accept = headers.get(ACCEPT).cloned(); + let want_conflicts = query.conflicts; + + state + .render( + move |doc| { + if want_conflicts { + store::read_path_with_conflicts(doc, &Vec::new()) + .map(|found| found.map(|(value, conflicts)| (value, Some(conflicts)))) + } else { + store::read_path(doc, &Vec::new()).map(|found| found.map(|value| (value, None))) + } + }, + move |result| { + let (value, conflicts) = result + .map_err(|e| store_status(e, "Failed to read store path"))? + .unwrap_or_else(|| (JsonValue::Object(Default::default()), Some(Vec::new()))); + match conflicts { + Some(conflicts) => conflicts_response(value, conflicts), + None => response_for_value(&value, accept.as_ref()), + } + }, + ) + .await +} + +#[utoipa::path( + get, + path = "/store/{key}", + tag = "Store", + params( + ("key" = String, Path, description = "Slash-separated store path, for example settings/theme"), + ("conflicts" = Option<bool>, Query, description = "Wrap the value and list every conflict resolved on the way"), + ), + responses( + (status = 200, description = "Returns the value as JSON or text, depending on Accept"), + (status = 400, description = "Invalid store path"), + (status = 404, description = "Value does not exist"), + (status = 500, description = "Failed to read or serialize the value"), + (status = 503, description = "Node is not ready") + ) +)] +pub(crate) async fn get_store_entry( + State(state): State<AppState>, + Path(key): Path<String>, + headers: HeaderMap, + Query(query): Query<ReadQuery>, +) -> Result<Response, StatusCode> { + if !state.is_ready().await { + return Err(StatusCode::SERVICE_UNAVAILABLE); + } + + let path = parse_key_path(&key)?; + + let accept = headers.get(ACCEPT).cloned(); + let want_conflicts = query.conflicts; + + state + .render( + move |doc| { + if want_conflicts { + store::read_path_with_conflicts(doc, &path) + .map(|found| found.map(|(value, conflicts)| (value, Some(conflicts)))) + } else { + store::read_path(doc, &path).map(|found| found.map(|value| (value, None))) + } + }, + move |result| { + let (value, conflicts) = result + .map_err(|e| store_status(e, "Failed to read store path"))? + .ok_or(StatusCode::NOT_FOUND)?; + match conflicts { + Some(conflicts) => conflicts_response(value, conflicts), + None => response_for_value(&value, accept.as_ref()), + } + }, + ) + .await +} + +#[utoipa::path( + put, + path = "/store/{key}", + tag = "Store", + params( + ("key" = String, Path, description = "Slash-separated store path") + ), + request_body( + content = String, + description = "Value to write. Use application/json for JSON values or text/plain for strings.", + content_type = "application/json" + ), + responses( + (status = 200, description = "Value written successfully"), + (status = 400, description = "Invalid path, malformed JSON, invalid UTF-8, or unsupported JSON value"), + (status = 409, description = "Path conflicts with an existing non-map value"), + (status = 415, description = "Unsupported Content-Type"), + (status = 500, description = "Internal error or synchronization notification failure"), + (status = 503, description = "Node is not ready") + ) +)] +pub(crate) async fn put_store_entry( + State(state): State<AppState>, + Path(key): Path<String>, + headers: HeaderMap, + body: Bytes, +) -> Result<Response, StatusCode> { + if !state.is_ready().await { + return Err(StatusCode::SERVICE_UNAVAILABLE); + } + + let path = parse_key_path(&key)?; + + let content_type = headers + .get(CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .unwrap_or("application/json") + .split(';') + .next() + .unwrap_or("") + .trim(); + + let value = match content_type { + "application/json" => { + serde_json::from_slice::<JsonValue>(&body).map_err(|_| StatusCode::BAD_REQUEST)? + } + "text/plain" => JsonValue::String( + String::from_utf8(body.to_vec()).map_err(|_| StatusCode::BAD_REQUEST)?, + ), + _ => return Err(StatusCode::UNSUPPORTED_MEDIA_TYPE), + }; + + let written = value.clone(); + state + .mutate_doc(move |doc| store::write_path(doc, &path, &written)) + .await + .map_err(|e| store_status(e, "Failed to write store value"))?; + + response_for_value(&value, headers.get(ACCEPT)) +} + +#[utoipa::path( + delete, + path = "/store/{key}", + tag = "Store", + params( + ("key" = String, Path, description = "Slash-separated store path"), + ("prune" = Option<bool>, Query, description = "Remove empty ancestor maps after deletion") + ), + responses( + (status = 204, description = "Value deleted successfully"), + (status = 400, description = "Invalid store path"), + (status = 404, description = "Value does not exist"), + (status = 500, description = "Deletion or synchronization notification failed"), + (status = 503, description = "Node is not ready") + ) +)] +#[instrument(skip(state), fields(path, prune = query.prune))] +pub(crate) async fn delete_store_entry( + State(state): State<AppState>, + Path(key): Path<String>, + Query(query): Query<DeleteQuery>, +) -> Result<StatusCode, StatusCode> { + let path = parse_key_path(&key)?; + + tracing::Span::current().record("path", path.join("/").as_str()); + + if !state.is_ready().await { + return Err(StatusCode::SERVICE_UNAVAILABLE); + } + + if path.is_empty() { + return Err(StatusCode::BAD_REQUEST); + } + + let prune = query.prune; + state + .mutate_doc(move |doc| store::delete_path(doc, &path, prune)) + .await + .map_err(|e| store_status(e, "Failed to delete store value"))?; + + Ok(StatusCode::NO_CONTENT) +} + +// --- streams ----------------------------------------------------------------- + +async fn render_traced( + state: &AppState, + path: Vec<String>, +) -> Result<(Option<JsonValue>, bool), StoreError> { + state + .render(move |doc| store::read_path_traced(doc, &path), |result| result) + .await +} + +pub(crate) async fn stream_store_entry( + State(state): State<AppState>, + Path(key): Path<String>, +) -> Result<Response, StatusCode> { + if !state.is_ready().await { + return Err(StatusCode::SERVICE_UNAVAILABLE); + } + + let path = parse_key_path(&key)?; + + // Subscribe first, so nothing between the initial read and the loop is missed. + let mut changes = state.change_events.subscribe(); + + let initial = render_traced(&state, path.clone()) + .await + .map_err(|e| store_status(e, "Failed to read initial stream value"))?; + + let stream_state = state.clone(); + let stream_key = key.clone(); + + let output = async_stream::stream! { + let (mut last_value, mut conflicted) = initial; + + match stream_event_line(&stream_key, last_value.as_ref()) { + Ok(line) => yield Ok::<Bytes, IoError>(line), + Err(error) => { + yield Err(error); + return; + } + } + + loop { + let relevant = match changes.recv().await { + // Conflicted paths can change without a patch on their path; re-read always. + Ok(set) => conflicted || set.paths.iter().any(|changed| paths_related(changed, &path)), + Err(broadcast::error::RecvError::Lagged(skipped)) => { + debug!(skipped, key = %stream_key, "Stream lagged; re-reading"); + true + } + Err(broadcast::error::RecvError::Closed) => break, + }; + + if !relevant { + continue; + } + + let (next_value, next_conflicted) = match render_traced(&stream_state, path.clone()).await { + Ok(result) => result, + Err(error) => { + error!(error = %format!("{error:#}"), key = %stream_key, "Failed to read updated stream value"); + yield Err(IoError::other("Failed to read updated store value")); + return; + } + }; + conflicted = next_conflicted; + + // A related change can still leave the resolved value unchanged. + if next_value == last_value { + continue; + } + + match stream_event_line(&stream_key, next_value.as_ref()) { + Ok(line) => yield Ok(line), + Err(error) => { + yield Err(error); + return; + } + } + last_value = next_value; + } + }; + + let mut response = Body::from_stream(output).into_response(); + response.headers_mut().insert( + CONTENT_TYPE, + HeaderValue::from_static("application/x-ndjson"), + ); + response.headers_mut().insert( + axum::http::header::CACHE_CONTROL, + HeaderValue::from_static("no-cache, no-transform"), + ); + Ok(response) +} + +fn stream_event_line(key: &str, value: Option<&JsonValue>) -> std::io::Result<Bytes> { + let mut event = serde_json::Map::new(); + + event.insert("key".to_owned(), JsonValue::String(key.to_owned())); + + match value { + Some(value) => { + event.insert("value".to_owned(), value.clone()); + } + None => { + event.insert("deleted".to_owned(), JsonValue::Bool(true)); + } + } + + let mut line = serde_json::to_vec(&JsonValue::Object(event)).map_err(IoError::other)?; + + line.push(b'\n'); + Ok(Bytes::from(line)) +} diff --git a/src/auth.rs b/src/auth.rs new file mode 100644 index 0000000..3926c70 --- /dev/null +++ b/src/auth.rs @@ -0,0 +1,141 @@ +use argon2::{password_hash::PasswordVerifier, Argon2}; +use axum::extract::State; +use axum::{ + body::Body, + http::{ + header::{AUTHORIZATION, WWW_AUTHENTICATE}, + HeaderMap, Request, StatusCode, + }, + response::{IntoResponse, Response}, +}; +use base64::{engine::general_purpose::STANDARD as BASE64, Engine}; +use dashmap::DashMap; +use password_hash::phc::PasswordHash; +use std::{ + collections::HashMap, + fs::File, + io::{BufRead, BufReader}, + sync::Arc, + time::{Duration, Instant}, +}; +use tokio::sync::Semaphore; +use tokio::time::sleep; + +fn load_users(path: &str) -> std::io::Result<HashMap<String, String>> { + let file = File::open(path)?; + let reader = BufReader::new(file); + + let mut users = HashMap::new(); + + for line in reader.lines() { + let line = line?; + + if let Some((user, hash)) = line.split_once(':') { + users.insert(user.to_owned(), hash.to_owned()); + } + } + + Ok(users) +} + +#[derive(Clone)] +pub struct AuthState { + idle_timeout: Duration, + + users: Arc<HashMap<String, String>>, // user -> PHC hash string + cache: Arc<DashMap<[u8; 32], Instant>>, + /// Bounds concurrent Argon2 verifications (CPU and ~19 MiB each). + verify_slots: Arc<Semaphore>, +} + +impl AuthState { + pub fn new(path: &str) -> Self { + let users = load_users(path).unwrap_or_default(); + + Self { + idle_timeout: Duration::from_secs(300), + cache: Arc::new(DashMap::new()), + users: Arc::new(users), + verify_slots: Arc::new(Semaphore::new(2)), + } + } +} + +fn parse_basic(headers: &HeaderMap) -> Option<(String, String)> { + let encoded = headers + .get(AUTHORIZATION)? + .to_str() + .ok()? + .strip_prefix("Basic ")?; + let decoded = String::from_utf8(BASE64.decode(encoded).ok()?).ok()?; + let (user, pass) = decoded.split_once(':')?; + Some((user.to_owned(), pass.to_owned())) +} + +async fn authenticate(state: &AuthState, headers: &HeaderMap) -> bool { + let Some((user, pass)) = parse_basic(headers) else { + return false; + }; + let Some(hash) = state.users.get(&user).cloned() else { + return false; + }; + + let key = *blake3::hash(format!("{user}\0{pass}\0{hash}").as_bytes()).as_bytes(); + + // Fast path. The shard guard lives only until the end of this `match`; + // never call into the map while holding it. + let fresh = match state.cache.get_mut(&key) { + Some(mut last_seen) if last_seen.elapsed() < state.idle_timeout => { + *last_seen = Instant::now(); // slide + true + } + _ => false, + }; + if fresh { + return true; + } + + // Slow path: Argon2 is deliberately expensive, so keep it off the runtime + // threads and bound how many run at once. + let Ok(_permit) = state.verify_slots.acquire().await else { + return false; + }; + + let verified = tokio::task::spawn_blocking(move || { + PasswordHash::new(&hash) + .map(|ph| { + Argon2::default() + .verify_password(pass.as_bytes(), &ph) + .is_ok() + }) + .unwrap_or(false) + }) + .await + .unwrap_or(false); + + if verified { + // Overwrites an expired entry; no separate remove needed. + state.cache.insert(key, Instant::now()); + } + + verified +} + +pub async fn basic_auth( + State(state): State<AuthState>, + headers: HeaderMap, + request: Request<Body>, + next: axum::middleware::Next, +) -> Response { + if !authenticate(&state, &headers).await { + sleep(Duration::from_millis(1500)).await; + return ( + StatusCode::UNAUTHORIZED, + [(WWW_AUTHENTICATE, r#"Basic realm="Restricted""#)], + "Unauthorized", + ) + .into_response(); + } + + next.run(request).await +} diff --git a/src/bootstrap.rs b/src/bootstrap.rs new file mode 100644 index 0000000..53f473f --- /dev/null +++ b/src/bootstrap.rs @@ -0,0 +1,99 @@ +//! Joining a cluster (or creating one as the genesis node). + +use std::sync::atomic::Ordering; +use std::time::Duration; + +use anyhow::{bail, Result}; +use tracing::{debug, error, info, instrument, warn}; +use uuid::Uuid; + +use crate::peer::{remember_peer, send_hello_once}; +use crate::state::{AppState, PeerMembership}; + +const BOOTSTRAP_TIMEOUT: Duration = Duration::from_secs(300); +const BOOTSTRAP_RETRY: Duration = Duration::from_millis(500); + +#[instrument(skip(state, seeds))] +pub(crate) async fn bootstrap(state: &AppState, seeds: &[String]) -> Result<()> { + if seeds.is_empty() { + info!("Initializing genesis node (no seed peers configured)"); + state.initial_sync_done.store(true, Ordering::Release); + state.mark_ready(Uuid::new_v4().to_string()).await?; + return Ok(()); + } + + info!( + timeout_seconds = BOOTSTRAP_TIMEOUT.as_secs(), + "Joining existing cluster" + ); + + let mut peers = Vec::new(); + + for seed in seeds { + match state.transport.resolve_peer(seed).await { + Ok(peer) => { + debug!(seed = %seed, peer = %peer, "Resolved seed to peer"); + peers.push(peer) + } + Err(e) => warn!(seed = %seed, error = %e, "Ignoring seed argument"), + } + } + + if peers.is_empty() { + bail!("None of the configured seeds could be resolved"); + } + + for peer in peers { + remember_peer(state, &peer).await; + if let Err(error) = send_hello_once(state, &peer).await { + warn!(peer = %peer, error = %format!("{error:#}"), "Initial HELLO failed; will retry"); + } + } + + let deadline = tokio::time::Instant::now() + BOOTSTRAP_TIMEOUT; + loop { + if let Some(cluster_id) = state.cluster_id().await { + if state.initial_sync_done.load(Ordering::Acquire) { + state.mark_ready(cluster_id.clone()).await?; + info!(cluster_id = %cluster_id, "Node is ready"); + return Ok(()); + } + } + + if tokio::time::Instant::now() >= deadline { + let mut peer_diagnostics: Vec<String> = Vec::new(); + + let peer_states = state.peer_states.read().await.clone(); + + for (peer, peer_state_mutex) in &peer_states { + let peer_state = peer_state_mutex.lock().await; + let validated_cluster = match &peer_state.membership { + PeerMembership::Unknown => None, + PeerMembership::Member { cluster_id } => Some(cluster_id.as_str()), + }; + peer_diagnostics.push(format!( + "{peer}: validated_cluster={:?}, pending_inbound_sync={}", + validated_cluster, + peer_state.pending_inbound_sync.len() + )) + } + let local_cluster = state.cluster_id().await; + let initial_sync_done = state.initial_sync_done.load(Ordering::Acquire); + error!( + local_cluster = ?local_cluster, + initial_sync_done, + peers = %peer_diagnostics.join("; "), + timeout_seconds = BOOTSTRAP_TIMEOUT.as_secs(), + "Bootstrap timed out" + ); + bail!( + "Bootstrap timed out after {:?}; local_cluster={:?}; initial_sync_done={}; peers=[{}]. Refusing to create a genesis document.", + BOOTSTRAP_TIMEOUT, + local_cluster, + initial_sync_done, + peer_diagnostics.join("; ") + ); + } + tokio::time::sleep(BOOTSTRAP_RETRY).await; + } +} diff --git a/src/firewall.rs b/src/firewall.rs new file mode 100644 index 0000000..917fe5f --- /dev/null +++ b/src/firewall.rs @@ -0,0 +1,105 @@ +use std::collections::HashSet; +use std::sync::{Arc, RwLock, RwLockReadGuard}; + +use anyhow::Result; +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; +use utoipa::ToSchema; + +use crate::transport::EncodedMessage; +use crate::transport::{MessageHandler, ParsedMessage, Transport}; + +#[derive(Debug, Clone, Default, Serialize, Deserialize, ToSchema)] +pub struct FirewallRules { + /// Messages from these peers are dropped before the application sees them. + #[serde(default)] + pub inbound: HashSet<String>, + /// Messages to these peers are silently dropped, like a lossy network. + #[serde(default)] + pub outbound: HashSet<String>, +} + +/// Uses a std lock: checks are tiny and the guard is never held across an await. +#[derive(Default)] +pub struct Firewall { + rules: RwLock<FirewallRules>, +} + +impl Firewall { + fn read(&self) -> RwLockReadGuard<'_, FirewallRules> { + self.rules + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + } + + pub fn rules(&self) -> FirewallRules { + self.read().clone() + } + + pub fn set_rules(&self, rules: FirewallRules) { + *self + .rules + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = rules; + } + + pub fn blocks_inbound(&self, peer: &str) -> bool { + self.read().inbound.contains(peer) + } + + pub fn blocks_outbound(&self, peer: &str) -> bool { + self.read().outbound.contains(peer) + } +} + +pub struct FirewallTransport { + inner: Arc<dyn Transport>, + firewall: Arc<Firewall>, +} + +impl FirewallTransport { + pub fn new(inner: Arc<dyn Transport>, firewall: Arc<Firewall>) -> Self { + Self { inner, firewall } + } +} + +#[async_trait] +impl Transport for FirewallTransport { + fn local_id(&self) -> &str { + self.inner.local_id() + } + + async fn resolve_peer(&self, hint: &str) -> Result<String> { + self.inner.resolve_peer(hint).await + } + + fn encode(&self, peer: &str, message: &ParsedMessage) -> Result<EncodedMessage> { + self.inner.encode(peer, message) + } + + async fn send_encoded(&self, message: &EncodedMessage) -> Result<()> { + if self.firewall.blocks_outbound(&message.peer) { + tracing::trace!(peer = %message.peer, "Firewall dropped outbound message"); + return Ok(()); + } + self.inner.send_encoded(message).await + } + + async fn start_consumer(&self, handler: MessageHandler) -> Result<()> { + let firewall = Arc::clone(&self.firewall); + + self.inner + .start_consumer(Box::new(move |sender, message| { + if firewall.blocks_inbound(&sender) { + tracing::trace!(peer = %sender, "Firewall dropped inbound message"); + return Box::pin(async { Ok(()) }); + } + handler(sender, message) + })) + .await + } + + async fn wait_for_exit(&self) -> std::io::Result<std::process::ExitStatus> { + self.inner.wait_for_exit().await + } +} diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..9774d78 --- /dev/null +++ b/src/main.rs @@ -0,0 +1,230 @@ +//! continuum: a replicated Automerge key-value store synchronized over Mycelium. +//! +//! main.rs only wires things together: arguments, state construction, background +//! tasks, bootstrap and the HTTP server. + +use std::collections::HashMap; +use std::future::IntoFuture; +use std::net::SocketAddr; +use std::path::PathBuf; +use std::sync::atomic::AtomicBool; +use std::sync::Arc; +use std::time::Duration; + +use anyhow::{Context, Result}; +use automerge::AutoCommit; +use clap::Parser; +use tokio::sync::{broadcast, mpsc, watch, RwLock}; +use tracing::{error, info}; +use url::Url; + +mod api; +mod auth; +mod bootstrap; +mod firewall; +mod mycelium; +mod peer; +mod state; +mod store; +mod sync; +mod transport; + +use auth::AuthState; +use firewall::{Firewall, FirewallTransport}; +use mycelium::{MyceliumSetup, MyceliumSpawnConfig}; +use state::{AppState, NodeLifecycle, Timing}; +use transport::Transport; + +#[derive(Parser, Debug)] +#[command(version, about = "continuum messaging system", long_about = None)] +struct Args { + #[arg(long)] + debug_api: bool, + + #[arg(long, value_name = "HOSTNAME:PORT", default_value = "127.0.0.1:3000")] + bind_addr: SocketAddr, + + #[arg( + long, + value_name = "SCHEME://HOSTNAME:PORT", + default_value = None, + conflicts_with_all = &["mycelium_key_file", "mycelium_network_key_file", "mycelium_network_name"], + )] + mycelium_rest_url: Option<Url>, + + #[arg(long)] + pub mycelium_key_file: Option<PathBuf>, + + #[arg(long)] + pub mycelium_network_key_file: Option<PathBuf>, + + #[arg(long)] + pub mycelium_network_name: Option<String>, + + #[arg(long)] + pub mycelium_tcp_listen_port: Option<u16>, + + #[arg( + long, + use_value_delimiter = true, + default_value = "tls://127.0.0.1:9651" + )] + pub mycelium_peers: Vec<String>, + + #[arg(long, value_name = "STRING,[STRING...]", value_delimiter = ',')] + seeds: Option<Vec<String>>, + + #[arg(long, value_name = "PATH")] + htpasswd: String, + + /// Keepalive interval; dormancy, resync and probe timings derive from it. + #[arg(long, value_name = "SECONDS", default_value_t = 60)] + keepalive_secs: u64, +} + +#[tokio::main] +async fn main() -> Result<()> { + console_subscriber::init(); + + let args = Args::parse(); + + let mycelium_setup = { + let mut spawn_config = MyceliumSpawnConfig { + rest_url_opt: args.mycelium_rest_url, + key_file_opt: args.mycelium_key_file, + network_key_file: args.mycelium_network_key_file, + peers: args.mycelium_peers, + ..Default::default() + }; + if let Some(net_name) = args.mycelium_network_name { + spawn_config.network_name = net_name; + } + MyceliumSetup::bootstrap(spawn_config).await? + }; + + let mycelium_rest_url = mycelium_setup.rest_url.clone(); + + info!( + bind_addr = args.bind_addr.to_string(), + mycelium_rest_url = mycelium_rest_url.to_string(), + "Starting Automerge replication service" + ); + + let seeds: Vec<String> = args.seeds.unwrap_or_default(); + info!(seed_count = seeds.len(), "Loaded bootstrap configuration"); + + let (sync_tx, sync_rx) = mpsc::unbounded_channel(); + let (change_tx, _) = watch::channel(0_u64); + let (change_events, _) = broadcast::channel(256); + + let firewall = Arc::new(Firewall::default()); + + let mycelium_transport = + Arc::new(mycelium::MyceliumTransport::connect(mycelium_setup).await?); + info!(local_id = mycelium_transport.local_id(), "Transport ready"); + + let state = AppState { + doc: Arc::new(RwLock::new(AutoCommit::new())), + lifecycle: Arc::new(RwLock::new(NodeLifecycle::Joining { + discovered_cluster_id: None, + })), + peer_states: Arc::new(RwLock::new(HashMap::new())), + transport: Arc::new(FirewallTransport::new( + mycelium_transport, + Arc::clone(&firewall), + )), + firewall, + sync_events: sync_tx, + change_tx, + initial_sync_done: Arc::new(AtomicBool::new(false)), + change_events, + timing: Timing::from_keepalive(Duration::from_secs(args.keepalive_secs)), + }; + + info!("Application state initialized"); + + let consumer_state = state.clone(); + tokio::spawn(async move { + let inner_state = consumer_state.clone(); + let _ = consumer_state + .transport + .start_consumer(Box::new(move |sender, msg| { + let state_clone = inner_state.clone(); + Box::pin(async move { peer::handle_message(&state_clone, sender, msg).await }) + })) + .await; + }); + info!("Mycelium consumer task spawned"); + + let sync_state = state.clone(); + tokio::spawn(async move { + sync::run_sync_worker(sync_state, sync_rx).await; + }); + info!("Sync worker task spawned"); + + let peer_worker_state = state.clone(); + tokio::spawn(async move { + peer::run_peer_worker(peer_worker_state).await; + }); + info!("Peer worker task spawned"); + + bootstrap::bootstrap(&state, &seeds) + .await + .context("Joining cluster")?; + + // Start the change publisher after bootstrapping: streams only exist from then on, + // and it avoids diffing the initial sync. + let publisher_state = state.clone(); + tokio::spawn(async move { sync::run_change_publisher(publisher_state).await }); + + info!( + cluster_id = ?state.cluster_id().await, + ready = state.is_ready().await, + "Bootstrap completed" + ); + + let auth_state = AuthState::new(&args.htpasswd); + let app = api::router(state.clone(), auth_state, args.debug_api); + + let listener = tokio::net::TcpListener::bind(args.bind_addr) + .await + .with_context(|| format!("Binding local API to {}", args.bind_addr))?; + + info!(address = %listener.local_addr()?, "HTTP API listening"); + + let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>(); + + let server_builder = axum::serve(listener, app).with_graceful_shutdown(async move { + let _ = shutdown_rx.await; + info!("Shutdown signal received. Draining API connections...") + }); + + let server_fut = server_builder.into_future(); + tokio::pin!(server_fut); + + tokio::select! { + exit_status = state.transport.wait_for_exit() => { + error!("Mycelium runtime process exited. Status: {:?}", exit_status); + let _ = shutdown_tx.send(()); + let _ = server_fut.as_mut().await; + } + + server_result = server_fut.as_mut() => { + if let Err(e) = server_result { + error!("Axum HTTP server encountered a runtime error: {:?}", e); + } else { + info!("Axum HTTP engine shutdown cleanly."); + } + std::mem::drop(state); + } + + _ = tokio::signal::ctrl_c() => { + info!("Received manual interrupt signal. Shutting down..."); + let _ = shutdown_tx.send(()); + let _ = server_fut.as_mut().await; + } + } + + info!("HTTP API server stopped"); + Ok(()) +} diff --git a/src/mycelium.rs b/src/mycelium.rs new file mode 100644 index 0000000..94c5a3d --- /dev/null +++ b/src/mycelium.rs @@ -0,0 +1,882 @@ +use anyhow::{anyhow, Context, Result}; +use async_trait::async_trait; +use base64::engine::{general_purpose::STANDARD as BASE64, Engine}; +use reqwest::StatusCode; +use serde::{Deserialize, Deserializer, Serialize, Serializer}; +use serde_bytes::{ByteBuf, Bytes}; +use std::borrow::Cow; +use std::boxed::Box; +use std::fs::File; +use std::io::Write; +use std::net::TcpListener; +use std::os::unix::io::FromRawFd; +use std::path::PathBuf; +use std::process::Stdio; +use std::time::Duration; +use thiserror::Error; +use tokio::process::{Child, Command}; +use url::Url; + +use crate::transport::EncodedMessage; +use crate::transport::Keepalive; +use crate::transport::{HelloReply, HelloRequest, SyncEnvelope}; +use crate::transport::{MessageHandler, ParsedMessage, Transport}; + +const MYCELIUM_POLL_TIMEOUT_SECS: u64 = 300; + +const TOPIC_PREFIX: &str = "continuum.v1."; +const TOPIC_HELLO: &str = "continuum.v1.hello"; +const TOPIC_HELLO_REPLY: &str = "continuum.v1.hello-reply"; +const TOPIC_SYNC: &str = "continuum.v1.sync"; +const TOPIC_PING: &str = "continuum.v1.ping"; +const TOPIC_PONG: &str = "continuum.v1.pong"; + +const EMPTY_QUEUE_DELAY: Duration = Duration::from_millis(25); + + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +struct InboundMessage { + id: String, + src_pk: String, + topic: String, + payload: String, +} + +#[derive(Debug, Serialize)] +struct OutboundMessage { + dst: Destination, + topic: String, + payload: String, +} + +#[derive(Debug, Serialize)] +struct Destination { + pk: String, +} + +/// Represents where a 32-byte cryptographic key originates. +#[derive(Debug, Clone)] +pub enum KeySource { + Raw([u8; 32]), + Path(PathBuf), +} + +/// Atomic structure ensuring private network credentials are always handled together. +#[derive(Debug, Clone)] +pub struct PrivateNetworkConfig { + pub name: String, + pub key: KeySource, +} + +/// Comprehensive parameters required to spawn or target an instance. +#[derive(Debug, Clone)] +pub struct MyceliumSpawnConfig { + pub rest_url_opt: Option<Url>, + pub key_file_opt: Option<PathBuf>, + pub network_key_file: Option<PathBuf>, + pub network_name: String, + pub tcp_port: Option<u16>, + pub api_port: Option<u16>, + pub rpc_port: Option<u16>, + pub discovery_port: Option<u16>, + pub peers: Vec<String>, +} + +impl Default for MyceliumSpawnConfig { + fn default() -> Self { + Self { + rest_url_opt: None, + key_file_opt: None, + network_key_file: None, + network_name: "something".to_string(), + tcp_port: None, + api_port: None, + rpc_port: None, + discovery_port: None, + peers: vec!["tls://127.0.0.1:9651".to_string()], + } + } +} + +pub enum MyceliumRuntime { + Supervised(Box<SupervisedMycelium>), + External, +} + +pub struct MyceliumSetup { + pub runtime: MyceliumRuntime, + pub rest_url: Url, +} + +impl MyceliumSetup { + /// Evaluates the configuration payload and initializes the chosen execution strategy. + pub async fn bootstrap(config: MyceliumSpawnConfig) -> Result<Self> { + // TODO warn for all ignored options if rest_url_opt is set + if let Some(url) = config.rest_url_opt { + tracing::info!( + "Configured to interface with an external standalone Mycelium instance at: {}", + url + ); + return Ok(Self { + runtime: MyceliumRuntime::External, + rest_url: url, + }); + } + + tracing::info!( + "Configured to initialize an embedded user-space Mycelium background worker" + ); + + let key = match config.key_file_opt { + Some(key_file) => KeySource::Path(key_file), + None => { + let mut bytes = [0u8; 32]; + getrandom::fill(&mut bytes)?; + KeySource::Raw(bytes) + } + }; + + let private_net_config = config + .network_key_file + .map(|net_key_path| PrivateNetworkConfig { + name: config.network_name.clone(), + key: KeySource::Path(net_key_path), + }); + + let supervised_node = SupervisedMycelium::spawn( + key, + private_net_config, + config.tcp_port, + config.api_port, + config.rpc_port, + config.discovery_port, + &config.peers, + ) + .await + .context("Failed to securely bootstrap embedded Mycelium background worker")?; + + let resolved_url = Url::parse(&format!( + "http://127.0.0.1:{}", + supervised_node.config.api_port + ))?; + + Ok(Self { + runtime: MyceliumRuntime::Supervised(Box::new(supervised_node)), + rest_url: resolved_url, + }) + } + + /// Blocks until the background daemon drops out. If running externally, hangs open indefinitely. + pub async fn wait_for_exit(&mut self) -> std::io::Result<std::process::ExitStatus> { + match &mut self.runtime { + MyceliumRuntime::Supervised(node) => node.wait_for_exit().await, + MyceliumRuntime::External => { + std::future::pending::<std::io::Result<std::process::ExitStatus>>().await + } + } + } +} + +#[derive(Deserialize, Debug, Clone)] +#[serde(rename_all = "camelCase")] +pub struct MyceliumAdminInfo { + pub node_ip: String, + pub node_pubkey: String, + pub node_start_time: u64, + pub node_subnet: String, +} + +pub async fn fetch_admin_info(rest_url: &Url) -> Result<MyceliumAdminInfo> { + let url = rest_url.join("/api/v1/admin")?; + let info = reqwest::Client::builder() + .timeout(Duration::from_secs(10)) + .build()? + .get(url) + .send() + .await? + .error_for_status()? + .json() + .await?; + Ok(info) +} + +#[derive(Debug, Clone)] +pub struct ClusterNodeConfig { + pub tcp_port: u16, + pub api_port: u16, + pub rpc_port: u16, + pub discovery_port: Option<u16>, + pub private_network: Option<String>, +} + +pub struct SupervisedMycelium { + child: Child, + pub config: ClusterNodeConfig, + _key_file_handle: Option<File>, + _net_key_file_handle: Option<File>, + key_file_path: PathBuf, + net_key_file_path: Option<PathBuf>, + is_ephemeral_key: bool, + is_ephemeral_net_key: bool, +} + +impl SupervisedMycelium { + async fn spawn( + node_key: KeySource, + private_network: Option<PrivateNetworkConfig>, + explicit_tcp: Option<u16>, + explicit_api: Option<u16>, + explicit_rpc: Option<u16>, + explicit_discovery_port: Option<u16>, + peers: &[String], + ) -> Result<Self> { + let tcp_port = match explicit_tcp { + Some(p) => p, + None => Self::find_available_port().context("Failed to allocate TCP wire port")?, + }; + let api_port = match explicit_api { + Some(p) => p, + None => Self::find_available_port().context("Failed to allocate API port")?, + }; + let rpc_port = match explicit_rpc { + Some(p) => p, + None => Self::find_available_port().context("Failed to allocate JSON-RPC port")?, + }; + let discovery_port = match explicit_discovery_port { + Some(0) => Some( + Self::find_available_port().context("Failed to allocate peer discovery port")?, + ), + p => p, + }; + + let config = ClusterNodeConfig { + tcp_port, + api_port, + rpc_port, + discovery_port, + private_network: private_network.as_ref().map(|pn| pn.name.clone()), + }; + + let parent_pid = std::process::id(); + + let (key_file_handle, key_file_path, node_fd, is_ephemeral_key) = + Self::resolve_key_strategy("mycelium_node_key", node_key, parent_pid)?; + + let (net_key_file_handle, net_key_file_path, net_fd, is_ephemeral_net_key) = + match private_network { + Some(net_cfg) => { + let (h, p, fd, eph) = + Self::resolve_key_strategy("mycelium_net_key", net_cfg.key, parent_pid)?; + (h, Some(p), fd, eph) + } + None => (None, None, None, false), + }; + + let api_addr = format!("127.0.0.1:{}", config.api_port); + let rpc_addr = format!("127.0.0.1:{}", config.rpc_port); + + let mut cmd = Command::new("mycelium"); + + let fds_to_inherit: Vec<i32> = [node_fd, net_fd].iter().filter_map(|&x| x).collect(); + if !fds_to_inherit.is_empty() { + unsafe { + cmd.pre_exec(move || { + for &fd in &fds_to_inherit { + libc::fcntl(fd, libc::F_SETFD, 0); + } + Ok(()) + }); + } + } + + cmd.arg("--key-file") + .arg(&key_file_path) + .arg("--tcp-listen-port") + .arg(config.tcp_port.to_string()) + .arg("--disable-quic") + .arg("--api-addr") + .arg(&api_addr) + .arg("--jsonrpc-addr") + .arg(&rpc_addr) + .arg("--no-tun"); + + if let Some(ref discovery_port) = config.discovery_port { + cmd.arg("--peer-discovery-port") + .arg(discovery_port.to_string()); + } else { + cmd.arg("--disable-peer-discovery"); + } + + if let Some(ref net_name) = config.private_network { + cmd.arg("--network-name").arg(net_name); + } + if let Some(ref net_key_path) = net_key_file_path { + cmd.arg("--network-key-file").arg(net_key_path); + } + + if !peers.is_empty() { + cmd.arg("--peers").arg(peers.join(" ")); + } + + cmd.env("RUST_LOG", "info"); + + cmd.stdout(Stdio::piped()); + cmd.stderr(Stdio::piped()); + + let mut child = cmd + .spawn() + .context("Failed to execute mycelium process daemon")?; + + // 1. FIXED: Pass-through streaming for child process STDOUT + if let Some(mut stdout) = child.stdout.take() { + tokio::spawn(async move { + let mut host_stdout = tokio::io::stdout(); + let _ = tokio::io::copy(&mut stdout, &mut host_stdout).await; + }); + } + + // 2. FIXED: Pass-through streaming for child process STDERR + if let Some(mut stderr) = child.stderr.take() { + tokio::spawn(async move { + let mut host_stderr = tokio::io::stderr(); + let _ = tokio::io::copy(&mut stderr, &mut host_stderr).await; + }); + } + + let node = Self { + child, + config, + _key_file_handle: key_file_handle, + _net_key_file_handle: net_key_file_handle, + key_file_path, + net_key_file_path, + is_ephemeral_key, + is_ephemeral_net_key, + }; + + //let admin_info = node.poll_admin_endpoint().await?; + //Ok((node, admin_info)) + Ok(node) + } + + fn resolve_key_strategy( + desc_tag: &str, + source: KeySource, + parent_pid: u32, + ) -> Result<(Option<File>, PathBuf, Option<i32>, bool)> { + match source { + KeySource::Raw(raw_bytes) => { + let (handle, fd) = Self::create_ephemeral_file(desc_tag, raw_bytes)?; + let pseudo_path = PathBuf::from(format!("/proc/{}/fd/{}", parent_pid, fd)); + Ok((Some(handle), pseudo_path, Some(fd), true)) + } + KeySource::Path(existing_path) => Ok((None, existing_path, None, false)), + } + } + + fn find_available_port() -> Result<u16> { + let listener = TcpListener::bind("127.0.0.1:0")?; + let addr = listener.local_addr()?; + Ok(addr.port()) + } + + fn create_ephemeral_file(desc_tag: &str, raw_key: [u8; 32]) -> Result<(File, i32)> { + #[cfg(target_os = "linux")] + { + let memfd_name = std::ffi::CString::new(desc_tag)?; + let fd = unsafe { libc::memfd_create(memfd_name.as_ptr(), 0) }; + if fd < 0 { + return Err(anyhow!( + "memfd_create failed natively inside Linux kernel for tag: {}", + desc_tag + )); + } + let mut file = unsafe { File::from_raw_fd(fd) }; + file.write_all(&raw_key)?; + file.flush()?; + Ok((file, fd)) + } + + #[cfg(not(target_os = "linux"))] + { + let mut path = std::env::tempdir(); + path.push(format!("{}_{}.key", desc_tag, uuid::Uuid::new_v4())); + let mut file = File::create(&path)?; + file.write_all(&raw_key)?; + file.flush()?; + Ok((file, -1)) + } + } + + //async fn poll_admin_endpoint(&mut self) -> Result<MyceliumAdminInfo> { + // let client = reqwest::Client::new(); + // let url = format!("http://127.0.0.1:{}/api/v1/admin", self.config.api_port); + // + // for _ in 0..30 { + // if let Some(status) = self.child.try_wait()? { + // return Err(anyhow!( + // "myceliumd process closed prematurely with status: {}", + // status + // )); + // } + // if let Ok(resp) = client.get(&url).send().await { + // if resp.status().is_success() { + // let info: MyceliumAdminInfo = resp.json().await?; + // return Ok(info); + // } + // } + // tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; + // } + // Err(anyhow!( + // "Timeout waiting for local mycelium administration socket to respond" + // )) + //} + + pub async fn wait_for_exit(&mut self) -> std::io::Result<std::process::ExitStatus> { + self.child.wait().await + } +} + +impl Drop for SupervisedMycelium { + fn drop(&mut self) { + let _ = self.child.start_kill(); + + #[cfg(not(target_os = "linux"))] + { + if self.is_ephemeral_key { + let _ = std::fs::remove_file(&self.key_file_path); + } + if self.is_ephemeral_net_key { + if let Some(ref path) = self.net_key_file_path { + let _ = std::fs::remove_file(path); + } + } + } + } +} + +fn is_pubkey(s: &str) -> bool { + s.len() == 64 && s.chars().all(|c| c.is_ascii_hexdigit()) +} + +pub struct MyceliumTransport { + // Wrap your zero-allocation lifecycle supervisor directly inside the driver plane + pub setup: tokio::sync::Mutex<MyceliumSetup>, + pub client: reqwest::Client, + pub poll_client: reqwest::Client, + pub messages_endpoint: Url, + pubkey_basepath: Url, + local_id: String, +} + +impl MyceliumTransport { + pub async fn connect(setup: MyceliumSetup) -> Result<Self> { + let rest_url = setup.rest_url.clone(); + + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(10)) + .build()?; + let poll_client = reqwest::Client::builder() + .timeout(Duration::from_secs(MYCELIUM_POLL_TIMEOUT_SECS + 30)) + .build()?; + + let admin: MyceliumAdminInfo = client + .get(rest_url.join("/api/v1/admin")?) + .send() + .await? + .error_for_status()? + .json() + .await + .context("Querying local Mycelium identity")?; + + Ok(Self { + setup: tokio::sync::Mutex::new(setup), + client, + poll_client, + messages_endpoint: rest_url.join("/api/v1/messages")?, + pubkey_basepath: rest_url.join("/api/v1/pubkey/")?, + local_id: admin.node_pubkey, + }) + } + + async fn lookup_pubkey(&self, ip: &std::net::Ipv6Addr) -> Result<String> { + #[derive(Deserialize)] + struct PublicKeyResponse { + #[serde(rename = "publicKey")] + public_key: String, + } + + let response: PublicKeyResponse = self + .client + .get(self.pubkey_basepath.join(&ip.to_string())?) + .send() + .await? + .error_for_status()? + .json() + .await?; + Ok(response.public_key) + } +} + +#[async_trait] +impl Transport for MyceliumTransport { + fn local_id(&self) -> &str { + &self.local_id + } + + async fn resolve_peer(&self, hint: &str) -> Result<String> { + if is_pubkey(hint) { + return Ok(hint.to_owned()); + } + + if let Ok(ip) = hint.parse::<std::net::Ipv6Addr>() { + return self.lookup_pubkey(&ip).await; + } + + let ip = tokio::time::timeout(Duration::from_secs(10), tokio::net::lookup_host((hint, 0))) + .await + .context("DNS lookup timed out")?? + .find_map(|addr| match addr.ip() { + std::net::IpAddr::V6(ip) => Some(ip), + _ => None, + }) + .ok_or_else(|| anyhow!("{hint}: no AAAA record found"))?; + + self.lookup_pubkey(&ip).await + } + + fn encode(&self, peer: &str, message: &ParsedMessage) -> Result<EncodedMessage> { + let outbound = format_message(peer, message)?; + Ok(EncodedMessage { + peer: peer.to_owned(), + body: serde_json::to_vec(&outbound)?.into(), + }) + } + + async fn send_encoded(&self, message: &EncodedMessage) -> Result<()> { + tracing::debug!( + payload_bytes = message.body.len(), + "Sending message" + ); + let response = self + .client + .post(self.messages_endpoint.clone()) + .header(reqwest::header::CONTENT_TYPE, "application/json") + .body(message.body.clone()) + .send() + .await?; + + let status = response.status(); + if !status.is_success() { + tracing::error!(status = %status, "Mycelium rejected outbound message"); + return Err(TransportError::MyceliumRejected { + peer: message.peer.clone(), + status, + } + .into()); + } + Ok(()) + } + + async fn start_consumer(&self, handler: MessageHandler) -> Result<()> { + tracing::info!( + endpoint = self.messages_endpoint.as_str(), + "Starting Mycelium message consumer" + ); + + loop { + tracing::trace!("Polling Mycelium inbox"); + + let response = match self + .poll_client + .get(self.messages_endpoint.clone()) + .query(&[("timeout", MYCELIUM_POLL_TIMEOUT_SECS)]) + .query(&[("topic", encode_topic(TOPIC_PREFIX))]) + .send() + .await + { + Ok(resp) => resp, + Err(error) => { + tracing::warn!(error = %error, "Mycelium inbox poll failed"); + tokio::time::sleep(Duration::from_millis(250)).await; + continue; + } + }; + + let status = response.status(); + + if status == StatusCode::NO_CONTENT { + tracing::trace!("Mycelium inbox is empty"); + tokio::time::sleep(EMPTY_QUEUE_DELAY).await; + continue; + } + + if !status.is_success() { + tracing::warn!(status = %status, "Mycelium inbox returned unsuccessful HTTP status"); + tokio::time::sleep(Duration::from_millis(500)).await; + continue; + } + + let raw_bytes = match response.bytes().await { + Ok(bytes) => bytes, + Err(error) => { + tracing::warn!(error = %error, "Failed to read raw Mycelium response stream"); + tokio::time::sleep(Duration::from_millis(100)).await; + continue; + } + }; + + let inbound_message = match serde_json::from_slice::<InboundMessage>(&raw_bytes) { + Ok(msg) => msg, + Err(error) => { + tracing::error!( + error = %error, + raw_payload = %String::from_utf8_lossy(&raw_bytes), + "Failed to decode inbound Mycelium message structure" + ); + tokio::time::sleep(Duration::from_millis(500)).await; + continue; + } + }; + + tracing::debug!( + id = %inbound_message.id, + payload_bytes = raw_bytes.len(), + "Received inbound Mycelium message" + ); + + match parse_message(&inbound_message) { + Ok((sender, parsed)) => { + if let Err(error) = handler(sender, parsed).await { + tracing::warn!(error = %format!("{error:#}"), "Failed to process inbound message"); + } + } + Err(transport_error) => { + tracing::error!( + error = %transport_error, + topic = %inbound_message.topic, + "Dropping poisoned packet: parsing failed gracefully" + ); + // No sleep required: we successfully popped a malformed packet and want to process the next one immediately + } + } + } + } + + async fn wait_for_exit(&self) -> std::io::Result<std::process::ExitStatus> { + let mut guard = self.setup.lock().await; + guard.wait_for_exit().await + } +} + +fn encode_topic(topic: &str) -> String { + BASE64.encode(topic.as_bytes()) +} + +fn decode_topic(topic: &str) -> Result<String, TransportError> { + let bytes = BASE64.decode(topic)?; + Ok(String::from_utf8(bytes)?) +} + +fn encode_payload(payload: &Vec<u8>) -> String { + BASE64.encode(payload) +} + +fn decode_payload(topic: &str, payload: &str) -> Result<ParsedMessage, TransportError> { + let bytes = BASE64.decode(payload)?; + match topic { + TOPIC_HELLO => { + let msg = bytes_to_hello_request(&bytes)?; + Ok(ParsedMessage::Hello(msg)) + } + TOPIC_HELLO_REPLY => { + let msg = bytes_to_hello_reply(&bytes)?; + Ok(ParsedMessage::HelloReply(msg)) + } + TOPIC_SYNC => { + let msg = bytes_to_sync_envelope(&bytes)?; + Ok(ParsedMessage::Sync(msg)) + } + TOPIC_PING => Ok(ParsedMessage::Ping(bytes_to_keepalive(&bytes)?)), + TOPIC_PONG => Ok(ParsedMessage::Pong(bytes_to_keepalive(&bytes)?)), + t => { + Err(TransportError::UnknownTopic(t.to_string())) + } + } +} + +fn parse_message(message: &InboundMessage) -> Result<(String, ParsedMessage), TransportError> { + let sender = message.src_pk.clone(); + + if sender.is_empty() { + return Err(TransportError::NoSender); + } + + let topic = decode_topic(&message.topic)?; + let parsed = decode_payload(&topic, &message.payload)?; + + Ok((sender, parsed)) +} + +fn format_message( + recipient: &str, + message: &ParsedMessage, +) -> Result<OutboundMessage, TransportError> { + let (topic, payload_bytes) = match message { + ParsedMessage::Hello(req) => (TOPIC_HELLO, hello_request_to_bytes(req)?), + ParsedMessage::HelloReply(res) => (TOPIC_HELLO_REPLY, hello_reply_to_bytes(res)?), + ParsedMessage::Sync(env) => (TOPIC_SYNC, sync_envelope_to_bytes(env)?), + ParsedMessage::Ping(k) => (TOPIC_PING, keepalive_to_bytes(k)?), + ParsedMessage::Pong(k) => (TOPIC_PONG, keepalive_to_bytes(k)?), + }; + Ok(OutboundMessage { + dst: Destination { + pk: recipient.to_owned(), + }, + topic: encode_topic(topic), + payload: encode_payload(&payload_bytes), + }) +} + +#[derive(Serialize, Deserialize)] +struct KeepaliveShadow { + heads_digest: ByteBuf, +} + +fn keepalive_to_bytes(msg: &Keepalive) -> Result<Vec<u8>, rmp_serde::encode::Error> { + rmp_serde::to_vec(&KeepaliveShadow { + heads_digest: ByteBuf::from(msg.heads_digest.to_vec()), + }) +} + +fn bytes_to_keepalive(bytes: &[u8]) -> Result<Keepalive, TransportError> { + let shadow: KeepaliveShadow = rmp_serde::from_slice(bytes)?; + let heads_digest = shadow + .heads_digest + .as_slice() + .try_into() + .map_err(|_| TransportError::InvalidPayload("heads digest must be 32 bytes".into()))?; + Ok(Keepalive { heads_digest }) +} + +#[derive(Serialize, Deserialize)] +pub struct HelloRequestShadow {} + +#[derive(Serialize, Deserialize)] +pub struct HelloReplyShadow<'a> { + // Cow allows zero-copy references when saving, but creates an owned String when loading + pub cluster_id: Cow<'a, str>, + #[serde(default)] + pub peers: Vec<String>, +} + +#[derive(Serialize, Deserialize)] +pub struct SyncEnvelopeShadow<'a> { + pub cluster_id: Cow<'a, str>, + #[serde(with = "automerge_msg_codec")] + pub message: automerge::sync::Message, +} + +// ===================================================================== +// 2. CUSTOM CODEC +// ===================================================================== +mod automerge_msg_codec { + use super::*; + use automerge::sync::Message; + + pub fn serialize<S>(msg: &Message, serializer: S) -> Result<S::Ok, S::Error> + where + S: Serializer, + { + let bytes = msg.clone().encode(); + Bytes::new(&bytes).serialize(serializer) + } + + pub fn deserialize<'de, D>(deserializer: D) -> Result<Message, D::Error> + where + D: Deserializer<'de>, + { + let byte_buf = ByteBuf::deserialize(deserializer)?; + Message::decode(&byte_buf) + .map_err(|e| serde::de::Error::custom(format!("Automerge decode failed: {:?}", e))) + } +} + +// ===================================================================== +// 3. CLEAN CONVERSION INTERFACE +// ===================================================================== + +// --- HelloRequest --- +pub fn hello_request_to_bytes(_msg: &HelloRequest) -> Result<Vec<u8>, rmp_serde::encode::Error> { + rmp_serde::to_vec(&HelloRequestShadow {}) +} +pub fn bytes_to_hello_request(bytes: &[u8]) -> Result<HelloRequest, rmp_serde::decode::Error> { + let _: HelloRequestShadow = rmp_serde::from_slice(bytes)?; + Ok(HelloRequest {}) +} + +// --- HelloReply --- +pub fn hello_reply_to_bytes(msg: &HelloReply) -> Result<Vec<u8>, rmp_serde::encode::Error> { + let shadow = HelloReplyShadow { + cluster_id: Cow::Borrowed(&msg.cluster_id), + peers: msg.peers.clone(), // Vec must be cloned for serialization wrapper matching + }; + rmp_serde::to_vec(&shadow) +} +pub fn bytes_to_hello_reply(bytes: &[u8]) -> Result<HelloReply, rmp_serde::decode::Error> { + let shadow: HelloReplyShadow = rmp_serde::from_slice(bytes)?; + Ok(HelloReply { + cluster_id: shadow.cluster_id.into_owned(), + peers: shadow.peers, + }) +} + +// --- SyncEnvelope --- +pub fn sync_envelope_to_bytes(msg: &SyncEnvelope) -> Result<Vec<u8>, rmp_serde::encode::Error> { + let shadow = SyncEnvelopeShadow { + cluster_id: Cow::Borrowed(&msg.cluster_id), + message: msg.message.clone(), // Clone required by automerge_msg_codec ownership rules + }; + rmp_serde::to_vec(&shadow) +} +pub fn bytes_to_sync_envelope(bytes: &[u8]) -> Result<SyncEnvelope, rmp_serde::decode::Error> { + let shadow: SyncEnvelopeShadow = rmp_serde::from_slice(bytes)?; + Ok(SyncEnvelope { + cluster_id: shadow.cluster_id.into_owned(), + message: shadow.message, + }) +} + +#[derive(Error, Debug)] +pub enum TransportError { + #[error("Mycelium delivered a message without a sender public key")] + NoSender, + + #[error("Base64 decoding failed: {0}")] + Base64Decode(#[from] base64::DecodeError), + + #[error("UTF-8 decoding failed: {0}")] + Utf8Decode(#[from] std::string::FromUtf8Error), + + #[error("MessagePack serialization failed: {0}")] + MessagePackEncode(#[from] rmp_serde::encode::Error), + + #[error("MessagePack deserialization failed: {0}")] + MessagePackDecode(#[from] rmp_serde::decode::Error), + + #[error("Unknown topic: {0}")] + UnknownTopic(String), + + #[error("Invalid payload: {0}")] + InvalidPayload(String), + + #[error("Network HTTP request failed: {0}")] + Http(#[from] reqwest::Error), + + #[error("Mycelium server rejected outbound message to {peer} (HTTP {status})")] + MyceliumRejected { + peer: String, + status: reqwest::StatusCode, + }, +} diff --git a/src/peer.rs b/src/peer.rs new file mode 100644 index 0000000..3dd344e --- /dev/null +++ b/src/peer.rs @@ -0,0 +1,593 @@ +//! Peers: handshake (HELLO / HELLO_REPLY), liveness and keepalives, self-healing, +//! cluster-wide forget records, and dispatch of inbound messages. + +use std::sync::Arc; +use std::time::Duration; + +use anyhow::{anyhow, bail, Context, Result}; +use automerge::AutoCommit; +use sha2::{Digest, Sha256}; +use tokio::sync::Mutex; +use tokio::time::{Instant, MissedTickBehavior}; +use tracing::{debug, info, instrument, trace, warn}; + +use crate::state::{AppState, HelloState, Liveness, PeerMembership, PeerState}; +use crate::store::{self, StoreError}; +use crate::sync::{handle_sync, SyncEvent}; +use crate::transport::{self, ParsedMessage}; + +const HELLO_RETRY_INITIAL: Duration = Duration::from_secs(1); +const HELLO_RETRY_MAX: Duration = Duration::from_secs(30); + +const PEER_TICK: Duration = Duration::from_millis(500); + +const META_ROOT: &str = "_continuum"; +const FORGET_CHECK_INTERVAL: Duration = Duration::from_secs(5); + +// --- forget records ---------------------------------------------------------- + +fn forgotten_dir() -> Vec<String> { + vec![META_ROOT.to_owned(), "forgotten".to_owned()] +} + +pub(crate) fn forgotten_path(pk: &str) -> Vec<String> { + let mut path = forgotten_dir(); + path.push(pk.to_owned()); + path +} + +fn forgotten_peers(doc: &AutoCommit) -> Vec<String> { + store::keys_at(doc, &forgotten_dir()).unwrap_or_default() +} + +async fn apply_forget_records(state: &AppState) { + // Before the initial sync we can't judge the records, and we have nothing to delete. + if !state.is_ready().await { + return; + } + + let forgotten = forgotten_peers(&*state.doc.read().await); + if forgotten.is_empty() { + return; + } + + let now = Instant::now(); + let mut cancel = Vec::new(); + + for pk in forgotten { + if pk == state.local_id() { + cancel.push(pk); // we are evidently alive + continue; + } + + let Some(peer_state_mutex) = state.get_peer_state(&pk).await else { + continue; // already forgotten here + }; + + let alive = peer_state_mutex.lock().await.is_alive(&state.timing, now); + + if alive { + cancel.push(pk); + } else if state.peer_states.write().await.remove(&pk).is_some() { + info!(peer = %pk, "Forgetting peer per cluster-wide record"); + } + } + + if cancel.is_empty() { + return; + } + + // Individual failures are logged inside; NotFound means another node already + // cancelled the record. + let _ = state + .mutate_doc(move |doc| { + for pk in &cancel { + match store::delete_path(&mut *doc, &forgotten_path(pk), false) { + Ok(()) => info!(peer = %pk, "Peer is alive; cancelling cluster-wide forget"), + Err(StoreError::NotFound) => {} + Err(error) => { + warn!(peer = %pk, error = %error, "Failed to cancel forget record") + } + } + } + Ok::<(), StoreError>(()) + }) + .await; +} + +// --- keepalive and self-healing ---------------------------------------------- + +fn heads_digest(doc: &mut AutoCommit) -> transport::HeadsDigest { + let mut heads = doc.get_heads(); + heads.sort_unstable(); + let mut hasher = Sha256::new(); + for head in &heads { + hasher.update(head.0); + } + hasher.finalize().into() +} + +async fn local_keepalive(state: &AppState) -> transport::Keepalive { + // AutoCommit::get_heads takes &mut self (it closes the pending transaction). + let heads_digest = heads_digest(&mut *state.doc.write().await); + transport::Keepalive { heads_digest } +} + +async fn heal_if_diverged( + state: &AppState, + peer: &str, + local: &transport::Keepalive, + remote: &transport::Keepalive, +) { + let Some(peer_state_mutex) = state.get_peer_state(peer).await else { + return; + }; + + let heal = peer_state_mutex.lock().await.needs_heal( + &local.heads_digest, + &remote.heads_digest, + &state.timing, + Instant::now(), + ); + + if heal { + warn!(peer = %peer, "Documents differ and sync is idle; re-handshaking to restart sync"); + ensure_handshake(state, peer).await; + } +} + +// --- peer discovery and HELLO ------------------------------------------------ + +#[instrument(skip(state, peer))] +pub(crate) async fn remember_peer(state: &AppState, peer: &str) { + if peer.is_empty() { + warn!("Ignoring empty peer public key"); + return; + } + + if peer == state.local_id() { + debug!("Ignoring own public key"); + return; + } + + let mut peer_states = state.peer_states.write().await; + + if !peer_states.contains_key(peer) { + peer_states.entry(peer.to_owned()).or_default(); + info!( + peer_count = peer_states.len(), + "Discovered and registered peer" + ); + } else { + trace!("Peer already registered"); + } +} + +#[instrument(skip(state), fields(peer = %peer))] +pub(crate) async fn send_hello_once(state: &AppState, peer: &str) -> Result<()> { + let Some(peer_state_mutex) = state.get_peer_state(peer).await else { + trace!("Peer was forgotten; not sending HELLO"); + return Ok(()); + }; + + let claimed = { + let mut peer_state = peer_state_mutex.lock().await; + if peer_state.hello_state == HelloState::NotSent { + peer_state.hello_state = HelloState::Sending { + next_delay: HELLO_RETRY_INITIAL, + }; + true + } else { + trace!(hello_state = ?peer_state.hello_state, "not sending"); + false + } + }; + + if !claimed { + return Ok(()); + } + + info!("Sending HELLO"); + transmit_hello(state, peer, HELLO_RETRY_INITIAL).await +} + +#[instrument(skip(state), fields(peer = %peer))] +async fn retry_hello_if_due(state: &AppState, peer: &str) -> Result<()> { + let Some(peer_state_mutex) = state.get_peer_state(peer).await else { + return Ok(()); + }; + + let delay = { + let mut peer_state = peer_state_mutex.lock().await; + match peer_state.hello_state { + HelloState::AwaitingReply { + next_attempt, + next_delay, + } if Instant::now() >= next_attempt => { + peer_state.hello_state = HelloState::Sending { next_delay }; + next_delay + } + _ => return Ok(()), + } + }; + + debug!( + next_retry_ms = delay.as_millis(), + "No HELLO_REPLY yet; re-sending HELLO" + ); + transmit_hello(state, peer, delay).await +} + +/// Sends a HELLO without holding any lock, then schedules the next retry. +/// Success and failure both end in AwaitingReply: Mycelium accepting the message +/// says nothing about whether the peer received it. +async fn transmit_hello(state: &AppState, peer: &str, delay: Duration) -> Result<()> { + let result = state + .transport + .send_message(peer, &ParsedMessage::Hello(transport::HelloRequest {})) + .await; + + if let Some(peer_state_mutex) = state.get_peer_state(peer).await { + let mut peer_state = peer_state_mutex.lock().await; + + if matches!(peer_state.hello_state, HelloState::Sending { .. }) { + let (wait, next_delay) = match peer_state.liveness { + Liveness::Dormant { .. } => { + (state.timing.dormant_probe, state.timing.dormant_probe) + } + Liveness::Active => (delay, delay.saturating_mul(2).min(HELLO_RETRY_MAX)), + }; + peer_state.hello_state = HelloState::AwaitingReply { + next_attempt: Instant::now() + wait, + next_delay, + }; + } + } + + result.context("Sending HELLO") +} + +async fn ensure_handshake(state: &AppState, peer: &str) { + if let Err(error) = send_hello_once(state, peer).await { + warn!(peer = %peer, error = %format!("{error:#}"), "HELLO failed; will retry"); + } +} + +/// Sends a HELLO_REPLY if one is owed and our cluster ID is known. +/// Otherwise it stays owed and the peer worker tries again. +async fn send_owed_hello_reply(state: &AppState, peer: &str) { + if state.cluster_id().await.is_none() { + return; + } + + let Some(peer_state_mutex) = state.get_peer_state(peer).await else { + return; + }; + + let owed = std::mem::take(&mut peer_state_mutex.lock().await.reply_owed); + if !owed { + return; + } + + if let Err(error) = send_hello_reply(state, peer).await { + warn!(peer = %peer, error = %format!("{error:#}"), "HELLO_REPLY failed; will retry"); + peer_state_mutex.lock().await.reply_owed = true; + } +} + +async fn flush_owed_hello_replies(state: &AppState) { + let peers: Vec<String> = state.peer_states.read().await.keys().cloned().collect(); + for peer in &peers { + send_owed_hello_reply(state, peer).await; + } +} + +#[instrument(skip(state), fields(peer = %peer))] +async fn send_hello_reply(state: &AppState, peer: &str) -> Result<()> { + let cluster_id = state + .cluster_id() + .await + .ok_or_else(|| anyhow!("Cannot send HELLO_REPLY without an established cluster ID"))?; + + let candidates: Vec<(String, Arc<Mutex<PeerState>>)> = state + .peer_states + .read() + .await + .iter() + .filter(|(known, _)| known.as_str() != peer) + .map(|(known, m)| (known.clone(), Arc::clone(m))) + .collect(); + + let mut peers = Vec::new(); + for (known, peer_state_mutex) in candidates { + let peer_state = peer_state_mutex.lock().await; + if peer_state.liveness == Liveness::Active + && peer_state.hello_state == HelloState::Established + { + peers.push(known); + } + } + + let reply = transport::HelloReply { cluster_id, peers }; + + let ready = state.is_ready().await; + + debug!(ready, "Prepared HELLO_REPLY"); + + state + .transport + .send_message(peer, &ParsedMessage::HelloReply(reply)) + .await + .context("Sending HELLO_REPLY")?; + + debug!("HELLO_REPLY sent"); + Ok(()) +} + +#[instrument(skip(state), fields(from_peer = %from_peer, advertised_peers = ?advertised_peers))] +async fn discover_peers(state: &AppState, from_peer: &str, advertised_peers: &[String]) { + for discovered_peer in advertised_peers { + if discovered_peer == from_peer { + continue; + } + + remember_peer(state, discovered_peer).await; + + if let Err(error) = send_hello_once(state, discovered_peer).await { + warn!( + peer = %discovered_peer, + error = %format!("{error:#}"), + "Failed to initiate HELLO with discovered peer" + ); + } + } +} + +// --- peer worker ------------------------------------------------------------- + +pub(crate) async fn run_peer_worker(state: AppState) { + let mut tick = tokio::time::interval(PEER_TICK); + tick.set_missed_tick_behavior(MissedTickBehavior::Delay); + + info!("Starting peer maintenance worker"); + + let mut last_forget_check: Option<Instant> = None; + + loop { + tick.tick().await; + + // While joining, keep handshakes fast; dormancy would slow probes to minutes. + let ready = state.is_ready().await; + + let peers: Vec<(String, Arc<Mutex<PeerState>>)> = state + .peer_states + .read() + .await + .iter() + .map(|(pk, m)| (pk.clone(), Arc::clone(m))) + .collect(); + + for (peer, peer_state_mutex) in &peers { + let now = Instant::now(); + + let (became_dormant, ping_due) = { + let mut peer_state = peer_state_mutex.lock().await; + let became_dormant = ready && peer_state.check_dormancy(&state.timing, now); + let ping_due = peer_state.take_keepalive_due(&state.timing, now); + (became_dormant, ping_due) + }; + + if became_dormant { + warn!( + peer = %peer, + silent_secs = state.timing.dormant_after.as_secs(), + probe_secs = state.timing.dormant_probe.as_secs(), + "Peer went dormant" + ); + } + + if ping_due { + let local = local_keepalive(&state).await; + if let Err(error) = state + .transport + .send_message(peer, &ParsedMessage::Ping(local)) + .await + { + debug!(peer = %peer, error = %format!("{error:#}"), "Keepalive PING failed"); + } + } + + if let Err(error) = retry_hello_if_due(&state, peer).await { + warn!(peer = %peer, error = %format!("{error:#}"), "HELLO retry failed"); + } + + send_owed_hello_reply(&state, peer).await; + } + + if last_forget_check.is_none_or(|t| t.elapsed() >= FORGET_CHECK_INTERVAL) { + apply_forget_records(&state).await; + last_forget_check = Some(Instant::now()); + } + } +} + +// --- peer membership --------------------------------------------------------- + +#[instrument(skip(state), fields(peer = %peer, received_cluster = %cluster_id))] +async fn establish_peer( + state: &AppState, + peer: &str, + cluster_id: String, +) -> Result<Vec<transport::SyncEnvelope>> { + let expected_cluster = state.cluster_id().await; + + if let Some(expected) = &expected_cluster { + if expected != &cluster_id { + bail!( + "Incompatible cluster at {peer}: expected {}, received {}", + expected, + cluster_id + ); + } + } + + let peer_state_mutex = state.get_or_create_peer_state(peer).await; + let mut peer_state = peer_state_mutex.lock().await; + + if let PeerMembership::Member { + cluster_id: existing, + } = &peer_state.membership + { + if existing != &cluster_id { + bail!( + "Peer {peer} changed cluster identity from {} to {}", + existing, + cluster_id + ); + } + } + + // Lock order: peer mutex -> lifecycle. + if expected_cluster.is_none() { + state.set_expected_cluster_id(cluster_id.clone()).await?; + } + + if matches!(peer_state.membership, PeerMembership::Unknown) { + peer_state.membership = PeerMembership::Member { + cluster_id: cluster_id.clone(), + }; + } + + // First reply completes the handshake: stop retries and sync from a clean slate. + // Duplicate replies (caused by our own retries) don't reset again. + if peer_state.hello_state != HelloState::Established { + peer_state.hello_state = HelloState::Established; + peer_state.reset_sync(); + info!("Handshake complete; sync state reset"); + } + + let pending = std::mem::take(&mut peer_state.pending_inbound_sync); + + info!(cluster_id = %cluster_id, pending_messages = pending.len(), "Peer membership established"); + Ok(pending) +} + +#[instrument(skip(state), fields(peer = %peer))] +pub(crate) async fn peer_cluster(state: &AppState, peer: &str) -> Option<String> { + let peer_state_mutex = state.get_peer_state(peer).await?; + let peer_state = peer_state_mutex.lock().await; + match &peer_state.membership { + PeerMembership::Unknown => None, + PeerMembership::Member { cluster_id } => Some(cluster_id.clone()), + } +} + +// --- inbound dispatch -------------------------------------------------------- + +#[instrument(skip(state, message))] +pub(crate) async fn handle_message( + state: &AppState, + peer: String, + message: transport::ParsedMessage, +) -> Result<()> { + if peer == state.local_id() { + debug!("Ignoring message from self"); + return Ok(()); + } + + remember_peer(state, &peer).await; + + // Any inbound message proves liveness. + let reactivated = state + .get_or_create_peer_state(&peer) + .await + .lock() + .await + .mark_seen(Instant::now()); + + if reactivated { + info!(peer = %peer, "Dormant peer is active again; re-handshaking"); + } + + debug!(message = ?message, "Received message"); + + match message { + transport::ParsedMessage::Hello(_) => { + // A HELLO means the peer is (re)starting its handshake with us, e.g. after a + // restart. Whatever our sync state believed about it is stale. + { + let peer_state_mutex = state.get_or_create_peer_state(&peer).await; + let mut peer_state = peer_state_mutex.lock().await; + peer_state.reset_sync(); + peer_state.reply_owed = true; + } + + // Sends now if we know our cluster ID; otherwise stays owed. + send_owed_hello_reply(state, &peer).await; + + // Our own HELLO goes out regardless of whether the reply could be sent. + if let Err(error) = send_hello_once(state, &peer).await { + warn!(peer = %peer, error = %format!("{error:#}"), "HELLO failed; will retry"); + } + + // Restart sync with the fresh state (no-op until the peer is a validated member). + let _ = state + .sync_events + .send(SyncEvent::PeerNeedsSync(peer.clone())); + } + transport::ParsedMessage::HelloReply(reply) => { + let pending = establish_peer(state, &peer, reply.cluster_id.clone()).await?; + + discover_peers(state, &peer, &reply.peers).await; + + // Our cluster ID may have just become known: pay any replies we owe. + flush_owed_hello_replies(state).await; + + if !pending.is_empty() { + debug!( + pending_messages = pending.len(), + "Replaying pending sync messages" + ); + for envelope in pending { + if let Err(error) = handle_sync(state, &peer, envelope).await { + warn!(peer = %peer, error = %format!("{error:#}"), "Discarding invalid pending sync envelope"); + } + } + } + + state + .sync_events + .send(SyncEvent::PeerValidated(peer.clone())) + .map_err(|error| anyhow!("Failed to notify sync worker: {error}"))?; + } + transport::ParsedMessage::Sync(envelope) => { + if reactivated { + ensure_handshake(state, &peer).await; + } + handle_sync(state, &peer, envelope).await?; + } + transport::ParsedMessage::Ping(remote) => { + if reactivated || peer_cluster(state, &peer).await.is_none() { + ensure_handshake(state, &peer).await; + } + let local = local_keepalive(state).await; + heal_if_diverged(state, &peer, &local, &remote).await; + if let Err(error) = state + .transport + .send_message(&peer, &ParsedMessage::Pong(local)) + .await + { + debug!(peer = %peer, error = %format!("{error:#}"), "PONG failed"); + } + } + transport::ParsedMessage::Pong(remote) => { + if reactivated || peer_cluster(state, &peer).await.is_none() { + ensure_handshake(state, &peer).await; + } + let local = local_keepalive(state).await; + heal_if_diverged(state, &peer, &local, &remote).await; + } + } + + Ok(()) +} diff --git a/src/state.rs b/src/state.rs new file mode 100644 index 0000000..3fe1790 --- /dev/null +++ b/src/state.rs @@ -0,0 +1,483 @@ +//! Shared node state: the document, lifecycle, per-peer state and timing. + +use std::collections::HashMap; +use std::sync::atomic::AtomicBool; +use std::sync::Arc; +use std::time::Duration; + +use anyhow::{bail, Result}; +use automerge::{sync, AutoCommit}; +use tokio::sync::{broadcast, mpsc, watch, Mutex, RwLock}; +use tokio::time::Instant; +use tracing::{debug, error, warn}; + +use crate::firewall::Firewall; +use crate::sync::{ChangeSet, SyncEvent}; +use crate::transport; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum NodeLifecycle { + // TODO maybe rename discovered_cluster_id to cluster_id if that doesn't impair readability. + Joining { + discovered_cluster_id: Option<String>, + }, + Ready { + cluster_id: String, + }, +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub(crate) enum PeerMembership { + #[default] + Unknown, + Member { + cluster_id: String, + }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum Liveness { + Active, + /// No inbound message for Timing::dormant_after. Kept indefinitely, probed slowly. + Dormant { since: Instant }, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub(crate) enum HelloState { + #[default] + NotSent, + /// A send is in flight. `next_delay` is how long to wait after it before retrying. + Sending { next_delay: Duration }, + /// HELLO was sent (or the send failed); no HELLO_REPLY yet. + AwaitingReply { + next_attempt: Instant, + next_delay: Duration, + }, + /// HELLO_REPLY received. No more retries. + Established, +} + +#[derive(Clone)] +pub(crate) struct PendingOutbound { + pub(crate) message: transport::EncodedMessage, + pub(crate) retry_delay: Duration, + pub(crate) next_attempt: Instant, +} + +#[derive(Clone)] +pub(crate) struct PeerState { + pub(crate) membership: PeerMembership, + pub(crate) sync_state: sync::State, + pub(crate) hello_state: HelloState, + + /// The peer sent HELLO; we still owe it a HELLO_REPLY (e.g. we had no cluster ID yet, + /// or sending failed). + pub(crate) reply_owed: bool, + + /// Received from this peer, waiting to be processed. + pub(crate) pending_inbound_sync: Vec<transport::SyncEnvelope>, + + /// Generated for this peer, waiting for delivery. + pub(crate) pending_outbound_sync: Option<PendingOutbound>, + + /// A sync pass was requested while an envelope was pending. Once that envelope + /// is delivered, the peer must be re-enqueued so the deferred changes go out. + pub(crate) needs_generate: bool, + + pub(crate) liveness: Liveness, + pub(crate) added_at: Instant, + pub(crate) last_seen: Option<Instant>, + pub(crate) last_ping_at: Option<Instant>, + + /// Last time a sync message from this peer was applied. + pub(crate) last_sync_rx: Option<Instant>, + + /// Last self-healing re-handshake (cooldown). + pub(crate) last_resync_at: Option<Instant>, +} + +impl Default for PeerState { + fn default() -> Self { + Self { + membership: PeerMembership::default(), + sync_state: sync::State::new(), + hello_state: HelloState::default(), + reply_owed: false, + pending_inbound_sync: Vec::new(), + pending_outbound_sync: None, + needs_generate: false, + liveness: Liveness::Active, + added_at: Instant::now(), + last_seen: None, + last_ping_at: None, + last_sync_rx: None, + last_resync_at: None, + } + } +} + +impl PeerState { + /// Start Automerge sync with this peer from scratch. Any envelope generated against + /// the old state is discarded; a stale RetryDelivery event is then a no-op. + pub(crate) fn reset_sync(&mut self) { + self.sync_state = sync::State::new(); + self.pending_outbound_sync = None; + self.needs_generate = false; + } + + fn quiet_for(&self, now: Instant) -> Duration { + now.saturating_duration_since(self.last_seen.unwrap_or(self.added_at)) + } + + /// Records inbound activity. Returns true if the peer was dormant. + pub(crate) fn mark_seen(&mut self, now: Instant) -> bool { + self.last_seen = Some(now); + + if matches!(self.liveness, Liveness::Dormant { .. }) { + self.liveness = Liveness::Active; + // Force a fresh handshake; its reply resets sync on both sides. + if !matches!(self.hello_state, HelloState::Sending { .. }) { + self.hello_state = HelloState::NotSent; + } + return true; + } + false + } + + /// Moves a silent peer to Dormant. Returns true on the transition. + pub(crate) fn check_dormancy(&mut self, timing: &Timing, now: Instant) -> bool { + if self.liveness != Liveness::Active || self.quiet_for(now) < timing.dormant_after { + return false; + } + + self.liveness = Liveness::Dormant { since: now }; + self.reset_sync(); + + if !matches!(self.hello_state, HelloState::Sending { .. }) { + self.hello_state = HelloState::AwaitingReply { + next_attempt: now + timing.dormant_probe, + next_delay: timing.dormant_probe, + }; + } + true + } + + /// True (and records the ping) if an established, quiet peer needs a keepalive. + pub(crate) fn take_keepalive_due(&mut self, timing: &Timing, now: Instant) -> bool { + let due = self.liveness == Liveness::Active + && self.hello_state == HelloState::Established + && self.quiet_for(now) >= timing.keepalive + && self + .last_ping_at + .is_none_or(|t| now.saturating_duration_since(t) >= timing.keepalive); + + if due { + self.last_ping_at = Some(now); + } + due + } + + /// Self-healing check on PING/PONG. Returns true if a re-handshake should start. + pub(crate) fn needs_heal( + &mut self, + local: &transport::HeadsDigest, + remote: &transport::HeadsDigest, + timing: &Timing, + now: Instant, + ) -> bool { + if local == remote + || self.liveness != Liveness::Active + || self.hello_state != HelloState::Established + { + return false; + } + + let sync_idle = now.saturating_duration_since(self.last_sync_rx.unwrap_or(self.added_at)) + >= timing.resync_after; + let cooled_down = self + .last_resync_at + .is_none_or(|t| now.saturating_duration_since(t) >= timing.resync_after); + + if !(sync_idle && cooled_down) { + return false; + } + + self.last_resync_at = Some(now); + self.hello_state = HelloState::NotSent; // establish_peer resets sync on the reply + true + } + + /// Heard from recently. A returning node counts as alive from its first message, + /// before the handshake finishes, so a cluster-wide forget can't make it flap. + pub(crate) fn is_alive(&self, timing: &Timing, now: Instant) -> bool { + self.liveness == Liveness::Active + && self + .last_seen + .is_some_and(|t| now.saturating_duration_since(t) < timing.dormant_after) + } +} + +#[derive(Clone, Copy, Debug)] +pub(crate) struct Timing { + pub(crate) keepalive: Duration, + pub(crate) dormant_after: Duration, + pub(crate) resync_after: Duration, + pub(crate) dormant_probe: Duration, +} + +impl Timing { + // Recommended values + // Environment | Keepalive | Notes + // --------------------------------------------+-----------+------------------------------------------------------------ + // Test harness | 1-2 s | Dormancy after about 4–8 s and healing within seconds keep tests short. + // LAN or small cluster (up to ~50 nodes) | 15-30 s | Under 3.5 messages/s per node at 50 nodes and 30 s. Current setting. + // WAN, mobile or metered links, or 100+ nodes | 60-120 s | At 100 nodes and 30 s it’s about 6–7 messages/s per node. Fine, but rarely worth it on battery or metered links. + // Several hundred nodes or more | - | The full mesh becomes the problem in itself, not the interval. Message counts grow O(N²) across the cluster, which calls for partial views (SWIM-style random probing and syncing with a subset of peers). + pub(crate) fn from_keepalive(keepalive: Duration) -> Self { + Self { + keepalive, + dormant_after: keepalive * 4, // dormant_after >= 4* keepalive to tolerate three lost PINGs in a row + resync_after: keepalive * 2, // resync_after should be >= 2* keepalive to guarantee at least one digest exchange while sync is idle + dormant_probe: keepalive * 10, // >= 10*: dormant peers cost almost nothing, but still rejoin within minutes + } + } +} + +#[derive(Clone)] +pub(crate) struct AppState { + pub(crate) doc: Arc<RwLock<AutoCommit>>, + + pub(crate) lifecycle: Arc<RwLock<NodeLifecycle>>, + + // Serializes cluster-ID establishment, peer membership, and readiness transitions. + //[concurrent]cluster_validation: Arc<Mutex<()>>, + pub(crate) peer_states: Arc<RwLock<HashMap<String, Arc<Mutex<PeerState>>>>>, + pub(crate) transport: Arc<dyn transport::Transport>, + pub(crate) firewall: Arc<Firewall>, + + pub(crate) sync_events: mpsc::UnboundedSender<SyncEvent>, + + // TODO see key_watchers.md ("Use a per-key notification registry instead of the global `watch<u64>`") + pub(crate) change_tx: watch::Sender<u64>, + + pub(crate) initial_sync_done: Arc<AtomicBool>, + + pub(crate) change_events: broadcast::Sender<Arc<ChangeSet>>, + + pub(crate) timing: Timing, +} + +impl AppState { + pub(crate) fn local_id(&self) -> &str { + self.transport.local_id() + } + + pub(crate) async fn is_ready(&self) -> bool { + matches!(*self.lifecycle.read().await, NodeLifecycle::Ready { .. }) + } + + pub(crate) async fn cluster_id(&self) -> Option<String> { + match &*self.lifecycle.read().await { + NodeLifecycle::Ready { cluster_id } => Some(cluster_id.clone()), + NodeLifecycle::Joining { + discovered_cluster_id, + } => discovered_cluster_id.clone(), + } + } + + // Record the expected cluster without making the node Ready. Bootstrap owns readiness. + pub(crate) async fn set_expected_cluster_id(&self, cluster_id: String) -> Result<()> { + //[concurrent]let _guard = self.cluster_validation.lock().await; + let mut lifecycle = self.lifecycle.write().await; + match &*lifecycle { + NodeLifecycle::Joining { + discovered_cluster_id: Some(existing), + } if existing != &cluster_id => { + bail!( + "Conflicting expected cluster IDs: expected {}, received {}", + existing, + cluster_id + ); + } + NodeLifecycle::Joining { + discovered_cluster_id: None, + } => { + let old_lifecycle = lifecycle.clone(); + *lifecycle = NodeLifecycle::Joining { + discovered_cluster_id: Some(cluster_id), + }; + debug!("Node transitioned from {old_lifecycle:?} {lifecycle:?}"); + Ok(()) + } + NodeLifecycle::Joining { + discovered_cluster_id: Some(_), + } => { + warn!("Node already Joining with the same identity"); + Ok(()) + } + NodeLifecycle::Ready { + cluster_id: existing, + } if existing == &cluster_id => { + warn!("Node already Ready with the same identity"); + Ok(()) + } + NodeLifecycle::Ready { + cluster_id: existing, + } => { + bail!( + "Cannot change established cluster ID: existing {}, received {}", + existing, + cluster_id + ); + } + } + } + + // Only bootstrap should call this, after validating document metadata. + pub(crate) async fn mark_ready(&self, cluster_id: String) -> Result<()> { + //[concurrent]let _guard = self.cluster_validation.lock().await; + let mut lifecycle = self.lifecycle.write().await; + match &*lifecycle { + NodeLifecycle::Joining { + discovered_cluster_id: Some(expected), + } if expected != &cluster_id => { + bail!( + "Cannot mark node ready: expected cluster {}, received {}", + expected, + cluster_id + ); + } + NodeLifecycle::Joining { .. } => { + let old_lifecycle = lifecycle.clone(); + *lifecycle = NodeLifecycle::Ready { cluster_id }; + debug!("Node transitioned from {old_lifecycle:?} {lifecycle:?}"); + } + NodeLifecycle::Ready { + cluster_id: existing, + } if existing == &cluster_id => { + debug!("Node already ready with the same identity"); + } + NodeLifecycle::Ready { + cluster_id: existing, + } => { + bail!( + "Cannot change established cluster ID: existing {}, received {}", + existing, + cluster_id + ); + } + } + Ok(()) + } + + /// Looks up a peer without creating it. Background work (sync, retries, + /// keepalives) must use this so a forgotten peer isn't resurrected. + pub(crate) async fn get_peer_state(&self, pk: &str) -> Option<Arc<Mutex<PeerState>>> { + self.peer_states.read().await.get(pk).cloned() + } + + pub(crate) async fn get_or_create_peer_state(&self, pk: &str) -> Arc<Mutex<PeerState>> { + let mut peers = self.peer_states.write().await; + + Arc::clone( + peers + .entry(pk.to_owned()) + .or_insert_with(|| Arc::new(Mutex::new(PeerState::default()))), + ) + } + + pub(crate) fn notify_store_changed(&self) { + self.change_tx.send_modify(|version| { + *version = version.wrapping_add(1); + }); + } + + /// Reads the document on the blocking pool; `finish` runs after the lock is released. + pub(crate) async fn render<T, R>( + &self, + read: impl FnOnce(&AutoCommit) -> T + Send + 'static, + finish: impl FnOnce(T) -> R + Send + 'static, + ) -> R + where + R: Send + 'static, + { + let doc = Arc::clone(&self.doc); + tokio::task::spawn_blocking(move || { + let data = { + let guard = doc.blocking_read(); + read(&guard) + }; + finish(data) + }) + .await + .expect("document read task panicked") + } + + /// Mutates the document on the blocking pool. + pub(crate) async fn write_doc<T>( + &self, + f: impl FnOnce(&mut AutoCommit) -> T + Send + 'static, + ) -> T + where + T: Send + 'static, + { + let doc = Arc::clone(&self.doc); + tokio::task::spawn_blocking(move || f(&mut doc.blocking_write())) + .await + .expect("document write task panicked") + } + + /// Mutates the document on the blocking pool as one Automerge change, then + /// publishes it. Write and notification run in the same blocking task, so a + /// cancelled request can't leave a committed write unannounced. + pub(crate) async fn mutate_doc<T, E>( + &self, + f: impl FnOnce(&mut AutoCommit) -> Result<T, E> + Send + 'static, + ) -> Result<T, E> + where + T: Send + 'static, + E: Send + 'static, + { + let state = self.clone(); + + tokio::task::spawn_blocking(move || { + let result = { + let mut doc = state.doc.blocking_write(); + let result = f(&mut doc); + match &result { + // One API write = one Automerge change. + Ok(_) => { + doc.commit(); + } + // Discard partial ops (e.g. old value deleted, new one not written). + Err(_) => { + doc.rollback(); + } + } + result + }; // lock released before notifying + + if result.is_ok() { + state.publish_change(); + } + result + }) + .await + .expect("document write task panicked") + } + + pub(crate) fn publish_change(&self) { + self.notify_store_changed(); + if let Err(error) = self.sync_events.send(SyncEvent::DocumentChanged) { + // The write is committed; failing the request would misreport it. + error!(error = %error, "Sync worker is gone; change will not be replicated"); + } + } +} + +pub(crate) fn unix_millis() -> u64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_millis() as u64) + .unwrap_or(0) +} diff --git a/src/store.rs b/src/store.rs new file mode 100644 index 0000000..08fe5c0 --- /dev/null +++ b/src/store.rs @@ -0,0 +1,506 @@ +//! JSON view over the Automerge document with conflict-aware merging. + +use anyhow::anyhow; +use automerge::{ + transaction::Transactable, ActorId, AutoCommit, ObjId, ObjType, ReadDoc, ScalarValue, Value, +}; +use base64::{engine::general_purpose::STANDARD as BASE64, Engine}; +use serde::Serialize; +use serde_json::Value as JsonValue; +use std::collections::BTreeSet; + +#[derive(Debug, thiserror::Error)] +pub enum StoreError { + #[error("path crosses a non-map value")] + Conflict, + #[error("value does not exist")] + NotFound, + #[error(transparent)] + Doc(#[from] automerge::AutomergeError), + #[error(transparent)] + Other(#[from] anyhow::Error), +} + +/// Automerge's total order on operations: Lamport counter, then actor. +fn op_order(id: &ObjId) -> (u64, Option<&ActorId>) { + match id { + ObjId::Root => (0, None), + ObjId::Id(counter, actor, _) => (*counter, Some(actor)), + } +} + +#[derive(Debug, Clone)] +enum Resolved { + /// Concurrent maps presented as one. Newest first; `[0]` is canonical for writes. + Maps(Vec<ObjId>), + /// Any other value: the newest concurrent candidate. + Other(Value<'static>, ObjId), +} + +/// All concurrent values of `key` across a merged set of parent maps, newest first. +fn candidates( + doc: &AutoCommit, + parents: &[ObjId], + key: &str, +) -> Result<Vec<(Value<'static>, ObjId)>, StoreError> { + let mut all = Vec::new(); + for parent in parents { + for (value, id) in doc.get_all(parent, key)? { + all.push((value.into_owned(), id)); + } + } + all.sort_by(|a, b| op_order(&b.1).cmp(&op_order(&a.1))); + Ok(all) +} + +fn resolve( + doc: &AutoCommit, + parents: &[ObjId], + key: &str, + log: Option<&mut ConflictLog>, +) -> Result<Option<Resolved>, StoreError> { + let mut all = candidates(doc, parents, key)?; + if all.is_empty() { + return Ok(None); + } + + let merged = matches!(all[0].0, Value::Object(ObjType::Map)); + + if all.len() > 1 { + if let Some(log) = log { + log.record(doc, key, &all, merged)?; + } + } + + if merged { + let maps = all + .into_iter() + .filter(|(value, _)| matches!(value, Value::Object(ObjType::Map))) + .map(|(_, id)| id) + .collect(); + Ok(Some(Resolved::Maps(maps))) + } else { + let (value, id) = all.swap_remove(0); + Ok(Some(Resolved::Other(value, id))) + } +} + +fn merged_keys(doc: &AutoCommit, maps: &[ObjId]) -> BTreeSet<String> { + maps.iter().flat_map(|map| doc.keys(map)).collect() +} + +fn resolved_to_json( + doc: &AutoCommit, + resolved: &Resolved, + mut log: Option<&mut ConflictLog>, +) -> Result<JsonValue, StoreError> { + let maps = match resolved { + Resolved::Maps(maps) => maps, + Resolved::Other(Value::Scalar(scalar), _) => { + return Ok(scalar_to_json(scalar.clone().into_owned())?) + } + Resolved::Other(Value::Object(_), id) => return object_to_json(doc, id, log), + }; + + let mut out = serde_json::Map::new(); + + if let [map] = maps.as_slice() { + // Fast path: one pass; get_all only for keys Automerge flags as conflicted. + for item in doc.map_range(map, ..) { + let key = item.key.to_string(); + + let child = if item.conflict { + match resolve(doc, maps, &key, log.as_deref_mut())? { + Some(child) => child, + None => continue, + } + } else { + let id = item.id(); + match Value::from(item.value) { + Value::Object(ObjType::Map) => Resolved::Maps(vec![id]), + value => Resolved::Other(value.into_owned(), id), + } + }; + + let json = nested(&mut log, &key, |log| resolved_to_json(doc, &child, log))?; + out.insert(key, json); + } + } else { + // Several concurrent maps: union of their keys. + for key in merged_keys(doc, maps) { + let Some(child) = resolve(doc, maps, &key, log.as_deref_mut())? else { + continue; + }; + let json = nested(&mut log, &key, |log| resolved_to_json(doc, &child, log))?; + out.insert(key, json); + } + } + + Ok(JsonValue::Object(out)) +} + +fn object_to_json( + doc: &AutoCommit, + id: &ObjId, + mut log: Option<&mut ConflictLog>, +) -> Result<JsonValue, StoreError> { + match doc.object_type(id)? { + ObjType::Map => resolved_to_json(doc, &Resolved::Maps(vec![id.clone()]), log), + ObjType::List => { + let mut out = Vec::new(); + for (index, item) in doc.list_range(id, ..).enumerate() { + let child_id = item.id(); + out.push(match Value::from(item.value) { + Value::Object(_) => nested(&mut log, &index.to_string(), |log| { + object_to_json(doc, &child_id, log) + })?, + Value::Scalar(scalar) => scalar_to_json(scalar.into_owned())?, + }); + } + Ok(JsonValue::Array(out)) + } + ObjType::Text => Ok(JsonValue::String(doc.text(id)?)), + ObjType::Table => Err(anyhow!("Automerge tables are not supported").into()), + } +} + +fn delete_in_all(doc: &mut AutoCommit, maps: &[ObjId], key: &str) -> Result<(), StoreError> { + for map in maps { + if doc.get(map, key)?.is_some() { + doc.delete(map, key)?; + } + } + Ok(()) +} + +fn read_path_inner( + doc: &AutoCommit, + path: &[String], + mut log: Option<&mut ConflictLog>, +) -> Result<Option<JsonValue>, StoreError> { + let mut current = Resolved::Maps(vec![automerge::ROOT]); + + for (index, segment) in path.iter().enumerate() { + let Resolved::Maps(maps) = ¤t else { + return Ok(None); + }; + if let Some(log) = log.as_deref_mut() { + // The last segment is the requested value itself, not an ancestor. + log.in_ancestors = index + 1 < path.len(); + } + match resolve(doc, maps, segment, log.as_deref_mut())? { + Some(next) => current = next, + None => return Ok(None), + } + if let Some(log) = log.as_deref_mut() { + log.path.push(segment.clone()); + } + } + + if let Some(log) = log.as_deref_mut() { + log.in_ancestors = false; + } + Ok(Some(resolved_to_json(doc, ¤t, log)?)) +} + +/// Reads a path; the empty path is the whole store. +pub fn read_path(doc: &AutoCommit, path: &[String]) -> Result<Option<JsonValue>, StoreError> { + read_path_inner(doc, path, None) +} + +pub fn read_path_with_conflicts( + doc: &AutoCommit, + path: &[String], +) -> Result<Option<(JsonValue, Vec<ConflictReport>)>, StoreError> { + let mut log = ConflictLog { want_values: true, ..Default::default() }; + let value = read_path_inner(doc, path, Some(&mut log))?; + Ok(value.map(|value| (value, log.reports))) +} + +/// For streams: the value, plus whether the path or subtree involves any conflict. +/// Reports are collected without values, so this costs little more than read_path. +pub fn read_path_traced( + doc: &AutoCommit, + path: &[String], +) -> Result<(Option<JsonValue>, bool), StoreError> { + let mut log = ConflictLog::default(); + let value = read_path_inner(doc, path, Some(&mut log))?; + Ok((value, !log.reports.is_empty())) +} + +/// Keys of the merged map at `path`; empty if it doesn't exist or isn't a map. +pub fn keys_at(doc: &AutoCommit, path: &[String]) -> Result<Vec<String>, StoreError> { + let mut maps = vec![automerge::ROOT]; + for segment in path { + maps = match resolve(doc, &maps, segment, None)? { + Some(Resolved::Maps(next)) => next, + _ => return Ok(Vec::new()), + }; + } + Ok(merged_keys(doc, &maps).into_iter().collect()) +} + +/// Writes (replaces) the value at `path`, creating intermediate maps. +pub fn write_path( + doc: &mut AutoCommit, + path: &[String], + value: &JsonValue, +) -> Result<(), StoreError> { + let (leaf, parents) = path.split_last().ok_or(StoreError::NotFound)?; + + let mut maps = vec![automerge::ROOT]; + for segment in parents { + maps = match resolve(doc, &maps, segment, None)? { + Some(Resolved::Maps(next)) => next, + Some(Resolved::Other(..)) => return Err(StoreError::Conflict), + // A concurrent creation elsewhere is merged on read. + None => vec![doc.put_object(&maps[0], segment.as_str(), ObjType::Map)?], + }; + } + + // Replace semantics: clear the key in every merged parent, or an old map + // would merge into the new value. + delete_in_all(doc, &maps, leaf)?; + json_to_automerge(doc, &maps[0], leaf, value)?; + Ok(()) +} + +/// Deletes the value at `path`, optionally pruning ancestors that become empty. +pub fn delete_path(doc: &mut AutoCommit, path: &[String], prune: bool) -> Result<(), StoreError> { + let (leaf, parents) = path.split_last().ok_or(StoreError::NotFound)?; + + // levels[d] = merged maps at path[..d] + let mut levels = vec![vec![automerge::ROOT]]; + for segment in parents { + let next = match resolve(doc, levels.last().expect("non-empty"), segment, None)? { + Some(Resolved::Maps(maps)) => maps, + _ => return Err(StoreError::NotFound), + }; + levels.push(next); + } + + let maps = levels.last().expect("non-empty"); + if resolve(doc, maps, leaf, None)?.is_none() { + return Err(StoreError::NotFound); + } + delete_in_all(doc, maps, leaf)?; + + if prune { + // Never prune the root (depth 0). + for depth in (1..levels.len()).rev() { + if !merged_keys(doc, &levels[depth]).is_empty() { + break; + } + delete_in_all(doc, &levels[depth - 1], &path[depth - 1])?; + } + } + + Ok(()) +} + +/// Where a value is placed inside its parent object. +#[derive(Clone, Copy)] +enum Slot<'a> { + /// A key in a map. + Key(&'a str), + /// An index in a list. The value is inserted here, not overwritten. + Index(usize), +} + +fn put_object( + doc: &mut AutoCommit, + parent: &ObjId, + slot: Slot, + obj_type: ObjType, +) -> Result<ObjId, StoreError> { + Ok(match slot { + Slot::Key(key) => doc.put_object(parent, key, obj_type)?, + Slot::Index(index) => doc.insert_object(parent, index, obj_type)?, + }) +} + +fn put_scalar( + doc: &mut AutoCommit, + parent: &ObjId, + slot: Slot, + value: impl Into<ScalarValue>, +) -> Result<(), StoreError> { + match slot { + Slot::Key(key) => doc.put(parent, key, value)?, + Slot::Index(index) => doc.insert(parent, index, value)?, + } + Ok(()) +} + +fn write_value( + doc: &mut AutoCommit, + parent: &ObjId, + slot: Slot, + value: &JsonValue, +) -> Result<(), StoreError> { + match value { + JsonValue::Object(object) => { + let id = put_object(doc, parent, slot, ObjType::Map)?; + for (child_key, child_value) in object { + write_value(doc, &id, Slot::Key(child_key), child_value)?; + } + } + JsonValue::Array(items) => { + let id = put_object(doc, parent, slot, ObjType::List)?; + for (index, item) in items.iter().enumerate() { + write_value(doc, &id, Slot::Index(index), item)?; + } + } + JsonValue::String(value) => put_scalar(doc, parent, slot, value.as_str())?, + JsonValue::Bool(value) => put_scalar(doc, parent, slot, *value)?, + JsonValue::Number(value) => { + if let Some(value) = value.as_i64() { + put_scalar(doc, parent, slot, value)?; + } else if let Some(value) = value.as_u64() { + put_scalar(doc, parent, slot, value)?; + } else if let Some(value) = value.as_f64() { + put_scalar(doc, parent, slot, value)?; + } else { + return Err(anyhow!("Unsupported JSON number").into()); + } + } + JsonValue::Null => put_scalar(doc, parent, slot, ScalarValue::Null)?, + } + + Ok(()) +} + +fn json_to_automerge( + doc: &mut AutoCommit, + parent: &ObjId, + key: &str, + value: &JsonValue, +) -> Result<(), StoreError> { + write_value(doc, parent, Slot::Key(key), value) +} + +fn scalar_to_json(scalar: ScalarValue) -> Result<JsonValue, StoreError> { + match scalar { + ScalarValue::Null => Ok(JsonValue::Null), + ScalarValue::Boolean(b) => Ok(JsonValue::Bool(b)), + ScalarValue::Int(i) => Ok(JsonValue::Number(i.into())), + ScalarValue::Uint(u) => Ok(JsonValue::Number(u.into())), + ScalarValue::F64(f) => serde_json::Number::from_f64(f) + .map(JsonValue::Number) + .ok_or(anyhow!("Non-finite float cannot be represented as JSON").into()), + ScalarValue::Str(s) => Ok(JsonValue::String(s.to_string())), + ScalarValue::Bytes(b) => Ok(JsonValue::String(BASE64.encode(b))), + ScalarValue::Counter(c) => Ok(JsonValue::Number(i64::from(c).into())), + ScalarValue::Timestamp(t) => Ok(JsonValue::Number(t.into())), + ScalarValue::Unknown { .. } => Err(anyhow!("Unsupported Automerge scalar type").into()), + } +} + +// --- conflict log ------------------------------------------------------------ + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub enum Resolution { + MergedMaps, + NewestWins, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct CandidateReport { + pub op_id: String, + pub selected: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub value: Option<JsonValue>, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct ConflictReport { + pub path: Vec<String>, + /// On the way to the requested path; reported without candidate values. + #[serde(skip_serializing_if = "std::ops::Not::not")] + pub ancestor: bool, + pub resolution: Resolution, + pub candidates: Vec<CandidateReport>, +} + +#[derive(Default)] +struct ConflictLog { + path: Vec<String>, + reports: Vec<ConflictReport>, + /// Render candidate values (only within the requested subtree). + want_values: bool, + /// Currently walking the ancestors of the requested path. + in_ancestors: bool, +} + +fn op_id_string(id: &ObjId) -> String { + match id { + ObjId::Root => "_root".to_owned(), + ObjId::Id(counter, actor, _) => format!("{counter}@{actor}"), + } +} + +fn value_to_json(doc: &AutoCommit, value: &Value<'_>, id: &ObjId) -> Result<JsonValue, StoreError> { + match value { + Value::Scalar(scalar) => Ok(scalar_to_json(scalar.clone().into_owned())?), + Value::Object(_) => object_to_json(doc, id, None), + } +} + +impl ConflictLog { + fn record( + &mut self, + doc: &AutoCommit, + key: &str, + all: &[(Value<'static>, ObjId)], + merged: bool, + ) -> Result<(), StoreError> { + let with_values = self.want_values && !self.in_ancestors; + + let candidates = all + .iter() + .enumerate() + .map(|(index, (value, id))| { + let selected = if merged { + matches!(value, Value::Object(ObjType::Map)) + } else { + index == 0 + }; + Ok(CandidateReport { + op_id: op_id_string(id), + selected, + value: if with_values { Some(value_to_json(doc, value, id)?) } else { None }, + }) + }) + .collect::<Result<Vec<_>, StoreError>>()?; + + let mut path = self.path.clone(); + path.push(key.to_owned()); + + self.reports.push(ConflictReport { + path, + ancestor: self.in_ancestors, + resolution: if merged { Resolution::MergedMaps } else { Resolution::NewestWins }, + candidates, + }); + Ok(()) + } +} + +/// Runs `f` with `segment` appended to the conflict path. +fn nested<T>( + log: &mut Option<&mut ConflictLog>, + segment: &str, + f: impl FnOnce(Option<&mut ConflictLog>) -> T, +) -> T { + if let Some(log) = log.as_deref_mut() { + log.path.push(segment.to_owned()); + } + let out = f(log.as_deref_mut()); + if let Some(log) = log.as_deref_mut() { + log.path.pop(); + } + out +} diff --git a/src/sync.rs b/src/sync.rs new file mode 100644 index 0000000..f1500cd --- /dev/null +++ b/src/sync.rs @@ -0,0 +1,447 @@ +//! Automerge synchronization: the sync worker, delivery with retries, inbound sync +//! messages, and the change publisher that feeds incremental streams. + +use std::collections::HashSet; +use std::sync::atomic::Ordering; +use std::sync::Arc; +use std::time::Duration; + +use anyhow::{anyhow, bail, Context, Result}; +use automerge::sync::SyncDoc; +use tokio::sync::mpsc; +use tokio::time::Instant; +use tracing::{debug, info, instrument, trace, warn}; + +use crate::peer::{peer_cluster, send_hello_once}; +use crate::state::{AppState, Liveness, PeerMembership, PendingOutbound}; +use crate::transport::{self, ParsedMessage}; + +const INITIAL_SYNC_RETRY: Duration = Duration::from_millis(500); +const MAX_SYNC_RETRY: Duration = Duration::from_secs(30); + +#[derive(Debug)] +pub(crate) enum SyncEvent { + PeerValidated(String), + DocumentChanged, + MessageReceived(String), + PeerNeedsSync(String), + RetryDelivery(String), +} + +/// Store paths touched by one batch of document changes. +#[derive(Debug)] +pub(crate) struct ChangeSet { + pub(crate) paths: Vec<Vec<String>>, +} + +#[instrument(skip(state), fields(peer = %peer, retry_delivery))] +async fn deliver_sync_to_peer(state: &AppState, peer: &str, retry_delivery: bool) -> Result<()> { + let Some(peer_state_mutex) = state.get_peer_state(peer).await else { + trace!("Peer was forgotten; skipping synchronization"); + return Ok(()); + }; + + let validated_cluster = { + let peer_state = peer_state_mutex.lock().await; + + if matches!(peer_state.liveness, Liveness::Dormant { .. }) { + // Reactivation re-handshakes, and PeerValidated then restarts sync. + trace!("Skipping synchronization with dormant peer"); + return Ok(()); + } + + match &peer_state.membership { + PeerMembership::Unknown => return Ok(()), + PeerMembership::Member { cluster_id } => cluster_id.clone(), + } + }; + + let cluster_id = state + .cluster_id() + .await + .ok_or_else(|| anyhow!("Cannot sync without an established cluster ID"))?; + + if validated_cluster != cluster_id { + bail!("Peer {peer} belongs to a different cluster"); + } + + let existing = { + let mut guard = peer_state_mutex.lock().await; + let peer_state = &mut *guard; + match &peer_state.pending_outbound_sync { + Some(entry) if retry_delivery => Some((entry.message.clone(), entry.next_attempt)), + Some(_) => { + peer_state.needs_generate = true; + return Ok(()); + } + None => None, + } + }; + + let encoded = match existing { + Some((message, next_attempt)) => { + if Instant::now() < next_attempt { + return Ok(()); + } + message + } + None if retry_delivery => return Ok(()), + None => { + let doc = Arc::clone(&state.doc); + let psm = Arc::clone(&peer_state_mutex); + let transport = Arc::clone(&state.transport); + let peer_owned = peer.to_owned(); + + let generated = tokio::task::spawn_blocking( + move || -> Result<Option<transport::EncodedMessage>> { + let message = { + let mut doc = doc.blocking_write(); + let mut peer_state = psm.blocking_lock(); + peer_state.needs_generate = false; + let sync_doc = doc.sync(); + sync_doc.generate_sync_message(&mut peer_state.sync_state) + }; + let Some(message) = message else { + return Ok(None); + }; + let envelope = transport::SyncEnvelope { + cluster_id, + message, + }; + transport + .encode(&peer_owned, &ParsedMessage::Sync(envelope)) + .map(Some) + }, + ) + .await + .context("Sync generation task panicked")??; + + let Some(encoded) = generated else { + return Ok(()); + }; + + peer_state_mutex.lock().await.pending_outbound_sync = Some(PendingOutbound { + message: encoded.clone(), + retry_delay: INITIAL_SYNC_RETRY, + next_attempt: Instant::now(), + }); + + encoded + } + }; + + // No state locks are held across the send. + match state.transport.send_encoded(&encoded).await { + Ok(()) => { + let regenerate = { + let mut peer_state = peer_state_mutex.lock().await; + peer_state.pending_outbound_sync = None; + std::mem::take(&mut peer_state.needs_generate) + }; + + debug!(regenerate, "Outbound sync message accepted by Mycelium"); + + if regenerate { + state + .sync_events + .send(SyncEvent::PeerNeedsSync(peer.to_owned())) + .map_err(|error| anyhow!("Failed to notify sync worker: {error}"))?; + } + + Ok(()) + } + Err(error) => { + let delay = { + let mut peer_state = peer_state_mutex.lock().await; + + let Some(ref mut entry) = peer_state.pending_outbound_sync else { + debug!("Pending envelope discarded by a sync reset during delivery"); + return Ok(()); + }; + + entry.retry_delay = entry.retry_delay.saturating_mul(2).min(MAX_SYNC_RETRY); + entry.next_attempt = Instant::now() + entry.retry_delay; + entry.retry_delay + }; + + warn!( + retry_ms = delay.as_millis(), + error = %format!("{error:#}"), + "Sync delivery failed; retaining the same outbound envelope" + ); + + let tx = state.sync_events.clone(); + let peer = peer.to_owned(); + tokio::spawn(async move { + tokio::time::sleep(delay).await; + let _ = tx.send(SyncEvent::RetryDelivery(peer)); + }); + + Ok(()) + } + } +} + +#[instrument(skip(state, envelope), fields(peer = %peer, cluster_id = %envelope.cluster_id))] +pub(crate) async fn handle_sync( + state: &AppState, + peer: &str, + envelope: transport::SyncEnvelope, +) -> Result<()> { + debug!("Processing sync envelope"); + + let peer_state_mutex = state.get_or_create_peer_state(peer).await; + + let needs_hello = { + let mut peer_state = peer_state_mutex.lock().await; + + match &peer_state.membership { + PeerMembership::Unknown => { + let duplicate = peer_state.pending_inbound_sync.iter().any(|pending| { + pending.message == envelope.message && pending.cluster_id == envelope.cluster_id + }); + if !duplicate { + peer_state.pending_inbound_sync.push(envelope.clone()); + info!( + pending_count = peer_state.pending_inbound_sync.len(), + "Queued sync envelope pending peer membership" + ); + } else { + trace!("Ignoring duplicate pending sync envelope"); + } + true + } + PeerMembership::Member { cluster_id } if cluster_id != &envelope.cluster_id => { + bail!( + "Sync cluster mismatch for peer {peer}: validated_cluster={cluster_id}, envelope_cluster={}", + envelope.cluster_id + ); + } + PeerMembership::Member { .. } => false, + } + }; + + if needs_hello { + // The peer treats us as a member but we don't know it, typically because we + // restarted. Our HELLO makes it reset its sync state and reply. + if let Err(error) = send_hello_once(state, peer).await { + warn!(peer = %peer, error = %format!("{error:#}"), "HELLO to unknown sync sender failed; will retry"); + } + return Ok(()); + } + + let expected_cluster = state + .cluster_id() + .await + .ok_or_else(|| anyhow!("Missing expected cluster ID after peer membership"))?; + if envelope.cluster_id != expected_cluster { + bail!( + "Incompatible cluster from {peer}: expected {}, received {}", + expected_cluster, + envelope.cluster_id + ); + } + + let caught_up = { + let doc = Arc::clone(&state.doc); + let peer_state_mutex = Arc::clone(&peer_state_mutex); + let check_caught_up = !state.initial_sync_done.load(Ordering::Acquire); + let message = envelope.message; + + tokio::task::spawn_blocking(move || -> Result<bool> { + let mut doc = doc.blocking_write(); + let mut peer_state = peer_state_mutex.blocking_lock(); + + doc.sync() + .receive_sync_message(&mut peer_state.sync_state, message) + .context("Automerge synchronization failed")?; + peer_state.last_sync_rx = Some(Instant::now()); + + Ok(check_caught_up + && peer_state + .sync_state + .their_heads + .as_ref() + .is_some_and(|heads| heads.iter().all(|h| doc.get_change_by_hash(h).is_some()))) + }) + .await + .context("Sync task panicked")?? + }; + + if caught_up && !state.initial_sync_done.swap(true, Ordering::AcqRel) { + info!(peer = %peer, "Initial sync complete"); + } + + state.notify_store_changed(); + + state + .sync_events + .send(SyncEvent::MessageReceived(peer.to_owned())) + .map_err(|error| anyhow!("Failed to notify sync worker: {error}"))?; + + Ok(()) +} + +pub(crate) async fn run_sync_worker(state: AppState, mut rx: mpsc::UnboundedReceiver<SyncEvent>) { + let mut dirty_peers = HashSet::<String>::new(); + let mut retry_peers = HashSet::<String>::new(); + let mut document_dirty = false; + + info!("Starting event-driven synchronization worker"); + + loop { + if dirty_peers.is_empty() && retry_peers.is_empty() && !document_dirty { + match rx.recv().await { + Some(event) => match event { + SyncEvent::PeerValidated(peer) + | SyncEvent::MessageReceived(peer) + | SyncEvent::PeerNeedsSync(peer) => { + dirty_peers.insert(peer); + } + SyncEvent::RetryDelivery(peer) => { + retry_peers.insert(peer); + } + SyncEvent::DocumentChanged => { + document_dirty = true; + } + }, + None => { + warn!("Synchronization event channel closed"); + return; + } + } + } + + // Coalesce event bursts. + while let Ok(event) = rx.try_recv() { + match event { + SyncEvent::PeerValidated(peer) + | SyncEvent::MessageReceived(peer) + | SyncEvent::PeerNeedsSync(peer) => { + dirty_peers.insert(peer); + } + SyncEvent::RetryDelivery(peer) => { + retry_peers.insert(peer); + } + SyncEvent::DocumentChanged => { + document_dirty = true; + } + } + } + + // Let pending HTTP task locks wake up and complete their lifecycle transitions gracefully. + tokio::task::yield_now().await; + + if document_dirty { + document_dirty = false; + + // Lock and extract only the cloned string keys into a small transient vector + let peer_keys: Vec<String> = { + let peers_guard = state.peer_states.read().await; + peers_guard.keys().cloned().collect() + }; + + for peer in peer_keys { + if peer_cluster(&state, &peer).await.is_some() { + dirty_peers.insert(peer); + } + } + } + + // Retry existing envelopes separately from generating new ones. + let retries: Vec<String> = retry_peers.drain().collect(); + + for peer in retries { + if let Err(error) = deliver_sync_to_peer(&state, &peer, true).await { + warn!( + peer = %peer, + error = %format!("{error:#}"), + "Sync delivery retry failed" + ); + } + } + + let peers: Vec<String> = dirty_peers.drain().collect(); + + for peer in peers { + if let Err(error) = deliver_sync_to_peer(&state, &peer, false).await { + warn!( + peer = %peer, + error = %format!("{error:#}"), + "Synchronization processing failed" + ); + } + } + } +} + +// --- change publisher -------------------------------------------------------- + +fn patch_store_path(patch: &automerge::Patch) -> Vec<String> { + let mut path: Vec<String> = patch + .path + .iter() + .map(|(_, prop)| match prop { + automerge::Prop::Map(key) => key.clone(), + automerge::Prop::Seq(index) => index.to_string(), + }) + .collect(); + + match &patch.action { + automerge::PatchAction::PutMap { key, .. } + | automerge::PatchAction::DeleteMap { key, .. } => path.push(key.clone()), + // List, text and counter changes: the containing object changed. + _ => {} + } + path +} + +pub(crate) fn paths_related(a: &[String], b: &[String]) -> bool { + a.starts_with(b) || b.starts_with(a) +} + +pub(crate) async fn run_change_publisher(state: AppState) { + let mut changed = state.change_tx.subscribe(); + let mut last_heads = state.write_doc(|doc| doc.get_heads()).await; + + info!("Starting change publisher"); + + while changed.changed().await.is_ok() { + let before = last_heads.clone(); + + let (heads, paths, conflicts) = state + .write_doc(move |doc| { + let heads = doc.get_heads(); + if heads == before { + return (heads, Vec::new(), Vec::new()); + } + + let mut paths = Vec::new(); + let mut conflicts = Vec::new(); + for patch in doc.diff(&before, &heads) { + let path = patch_store_path(&patch); + if matches!( + patch.action, + automerge::PatchAction::PutMap { conflict: true, .. } + ) { + conflicts.push(path.clone()); + } + paths.push(path); + } + (heads, paths, conflicts) + }) + .await; + + if heads == last_heads { + continue; + } + last_heads = heads; + + for path in &conflicts { + warn!(path = ?path, "Concurrent writes produced a conflict"); + } + + // Fails only when nobody is subscribed. + let _ = state.change_events.send(Arc::new(ChangeSet { paths })); + } +} diff --git a/src/transport.rs b/src/transport.rs new file mode 100644 index 0000000..95d3046 --- /dev/null +++ b/src/transport.rs @@ -0,0 +1,74 @@ +use anyhow::Result; +use async_trait::async_trait; +use automerge::sync::Message; +use std::boxed::Box; +use std::future::Future; +use std::pin::Pin; + +pub type HeadsDigest = [u8; 32]; + +#[derive(Debug, Clone)] +pub struct Keepalive { + pub heads_digest: HeadsDigest, +} + +#[derive(Debug)] +pub enum ParsedMessage { + Hello(HelloRequest), + HelloReply(HelloReply), + Sync(SyncEnvelope), + Ping(Keepalive), + Pong(Keepalive), +} + +#[derive(Debug)] +pub struct HelloRequest {} + +#[derive(Debug, Clone)] +pub struct HelloReply { + pub cluster_id: String, + pub peers: Vec<String>, +} + +#[derive(Debug, Clone)] +pub struct SyncEnvelope { + pub cluster_id: String, + pub message: Message, +} + +/// A message serialized for the wire. Cloning is cheap (shared buffer). +#[derive(Clone, Debug)] +pub struct EncodedMessage { + pub peer: String, + pub body: bytes::Bytes, +} + +pub type HandlerFn = dyn Fn(String, ParsedMessage) -> Pin<Box<dyn Future<Output = Result<()>> + Send>> + + Send + + Sync + + 'static; +pub type MessageHandler = Box<HandlerFn>; + +#[async_trait] +pub trait Transport: Send + Sync { + /// This node's identity on the transport (for Mycelium: its public key). + fn local_id(&self) -> &str; + + /// Resolves a user-supplied peer hint (e.g. a seed) to a peer identity. + async fn resolve_peer(&self, hint: &str) -> Result<String>; + + fn encode(&self, peer: &str, message: &ParsedMessage) -> Result<EncodedMessage>; + async fn send_encoded(&self, message: &EncodedMessage) -> Result<()>; + + /// Sends an outbound message payload to a targeted peer on a specific topic. + async fn send_message(&self, peer: &str, message: &ParsedMessage) -> Result<()> { + let encoded = self.encode(peer, message)?; + self.send_encoded(&encoded).await + } + + /// Starts the long-polling inbound event loops, driving messages to the provided handler closure. + async fn start_consumer(&self, handler: MessageHandler) -> Result<()>; + + /// Asynchronously monitors for fatal network client or sub-process failures. + async fn wait_for_exit(&self) -> std::io::Result<std::process::ExitStatus>; +} |
