diff --git a/.gitignore b/.gitignore index 85358f1..321fe5d 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,4 @@ src/version.rs db_v2.sqlite3 test.* .idea +.codeartsdoer/ diff --git a/Cargo.lock b/Cargo.lock index 73fdf29..a5a10b8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -294,6 +294,7 @@ checksum = "acee9fd5073ab6b045a275b3e709c163dd36c90685219cb21804a147b58dba43" dependencies = [ "async-trait", "axum-core", + "base64 0.13.1", "bitflags 1.3.2", "bytes", "futures-util", @@ -310,8 +311,10 @@ dependencies = [ "serde", "serde_json", "serde_urlencoded", + "sha-1", "sync_wrapper 0.1.2", "tokio", + "tokio-tungstenite 0.17.2", "tower 0.4.13", "tower-http 0.3.5", "tower-layer", @@ -1580,6 +1583,7 @@ dependencies = [ "serde_json", "sodiumoxide", "sqlx", + "tokio-stream", "tokio-tungstenite 0.17.2", "tower-http 0.3.5", "tungstenite 0.17.3", diff --git a/Cargo.toml b/Cargo.toml index 2821178..7d8eefa 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -29,7 +29,7 @@ machine-uid = "0.2" mac_address = "1.1.5" whoami = "1.2" base64 = "0.13" -axum = { version = "0.5", features = ["headers"] } +axum = { version = "0.5", features = ["headers", "ws"] } sqlx = { version = "0.6", features = [ "runtime-tokio-rustls", "sqlite", "macros", "chrono", "json" ] } deadpool = "0.8" async-trait = "0.1" @@ -40,6 +40,7 @@ chrono = "0.4" jsonwebtoken = "8" headers = "0.3" once_cell = "1.8" +tokio-stream = "0.1" sodiumoxide = "0.2" tokio-tungstenite = "0.17" tungstenite = "0.17" diff --git a/docs/management.md b/docs/management.md index 00453b3..71c98fa 100644 --- a/docs/management.md +++ b/docs/management.md @@ -1,14 +1,18 @@ # Console management integration -The opt-in `/v1` HTTP API is independent of RustDesk's client protocol ports. Set `RD_MANAGEMENT_BIND` to a private socket address, `RD_MANAGEMENT_TOKEN` to an independent random Bearer credential of at least 32 characters, and `RD_MANAGEMENT_DIR` to a persistent directory shared with the node agent. Do not publish this listener to the client-facing network. +The opt-in `/v1` HTTP API is independent of RustDesk's client protocol ports. Set `RD_MANAGEMENT_BIND` to a private socket address, `RD_MANAGEMENT_TOKEN` to an independent random Bearer credential of at least 32 characters, and `RD_MANAGEMENT_DIR` to a persistent directory for management state. Do not publish this listener to the client-facing network. -hbbs exposes `GET /v1/status`, `GET /v1/config`, `PUT /v1/config`, `GET /v1/logs`, `GET /v1/peers` and `GET/PUT /v1/bans`. hbbr additionally exposes `GET /v1/sessions` and `DELETE /v1/sessions/:uuid`; it has no peers endpoint. All endpoints require authorization. API version is 1. Session closure is asynchronous (`202`, `state: closing`). IP bans check both relay endpoints. Reported target IDs are not authenticated source identities. +hbbs exposes `GET /v1/status`, `GET /v1/config`, `PUT /v1/config`, `GET /v1/logs`, `GET /v1/peers`, `GET/PUT /v1/bans`, `POST /v1/restart`, `GET /v1/events` (SSE), `GET /v1/relays`, `GET /v1/relays/:id/sessions`, `DELETE /v1/relays/:id/sessions/:uuid`, `GET/PUT /v1/relays/:id/config`, `GET /v1/relays/:id/logs`, `POST /v1/relays/:id/restart` and `GET /v1/control` (WebSocket upgrade). hbbr additionally exposes `GET /v1/sessions`, `DELETE /v1/sessions/:uuid` and `POST /v1/restart`; it has no peers endpoint. All endpoints require authorization. API version is 1. Session closure is asynchronous (`202`, `state: closing`). IP bans check both relay endpoints. Reported target IDs are not authenticated source identities. + +Relay aggregation: when `RD_HBBS_API_URL` is set on hbbr, it opens an outbound WebSocket to the hbbs management API (`/v1/control`), registers with its relay identity and reports heartbeats, session lifecycle events and session snapshots. hbbs fans the active ban policy out to every registered relay (on update and on registration), aggregates relay sessions behind `GET /v1/relays/:id/sessions`, forwards disconnect requests through `DELETE /v1/relays/:id/sessions/:uuid`, proxies relay-scoped management requests (`GET/PUT /v1/relays/:id/config`, `GET /v1/relays/:id/logs`, `POST /v1/relays/:id/restart`) over the control channel using request ids and HTTP-equivalent responses (a relay that stays silent fails the proxy request after 10 seconds), and folds each relay address reported at registration into the live relay server list (static `relay-servers` plus registered relays; dynamic `relay-servers` updates replace the static baseline and keep the registered addresses). hbbr identity comes from `RD_RELAY_ID` (or `HOSTNAME`); `RD_RELAY_ADDRESS` optionally advertises the address clients use to reach the relay and is folded into the live relay server list only when set — the relay id (often a container hostname) is never broadcast as an address. A re-registering relay replaces its previous connection, and registrations silent for 45 seconds are dropped. `GET /v1/events` streams aggregated events (`peer_online`, `peer_offline`, `session_opened`, `session_closed`, `relay_registered`, `relay_disconnected`, `bans_replaced`) as server-sent events with a 15 second keep-alive. Files are service-specific: `hbbs-config.json`, `hbbr-config.json`, `hbbs-bans.json`, and `hbbr-bans.json`. Configurations use `{ "values": { "port": "21116" } }`, and policies use `{ "device_ids": [], "ips": [] }`. Saved bans are validated and restored before protocol startup even when `RD_MANAGEMENT_BIND` is empty. Disabling the HTTP listener does not clear admission policy; explicitly replace the bans with empty lists before disabling management if clearing is intended. Ban replacements serialize persistence outside the protocol admission lock and activate only after the write succeeds. Config updates follow the same contract: the merged desired configuration is persisted atomically first, and running state changes only after the write succeeds. `PUT /v1/config` accepts a partial `{"values": {...}}` merge against the current desired configuration. Schema entries marked `"dynamic": true` take effect immediately without a restart: hbbr applies `downgrade-threshold`, `downgrade-start-check`, `limit-speed`, `total-bandwidth` and `single-bandwidth` live, and hbbs applies `always-use-relay` and `relay-servers` live. The response reports `applied` entries and `pending_restart` entries (with `current` and `pending` values); pending entries take effect on the next restart, when the saved file is loaded before the protocol starts. `GET /v1/config` returns the live effective values, the schema and the current `pending_restart` diff. Secret entries are stored in the clear in the shared directory but are always reported as `__REDACTED__`; submitting the placeholder keeps the stored value unchanged. A missing dynamic handler falls the entry back to restart-required handling. -`GET /v1/logs?tail=N` returns up to 2000 recent in-memory log lines (default 200) captured from the standard log format; it requires no container log access. +`GET /v1/logs?tail=N` returns up to 2000 recent in-memory log lines (default 200) captured from the standard log format; it requires no container log access. `GET /v1/relays/:id/logs` proxies the same read for a registered relay. + +`POST /v1/restart` (directly on each service and proxied per relay) asks the supervisor to bring the service back: the process exits shortly after the response is flushed (`202`, `state: restarting`), and the persisted configuration including pending-restart entries is applied on the next startup. Relay restarts are fire-and-forget: the relay acknowledges, exits and drops off the control channel, and registers again once it is back. Managed settings override startup inputs. Schema and validation are provided by `management-schema.json`. The API returns key placeholders, never the private key. It exposes only `id_ed25519.pub` as the public key. diff --git a/src/hbbr.rs b/src/hbbr.rs index 7f6d3e0..ae65848 100644 --- a/src/hbbr.rs +++ b/src/hbbr.rs @@ -1,6 +1,7 @@ use clap::App; mod common; mod management; +mod relay_control; mod relay_server; use flexi_logger::*; use hbb_common::{config::RELAY_PORT, ResultType}; diff --git a/src/management.rs b/src/management.rs index def3bfc..cf54906 100644 --- a/src/management.rs +++ b/src/management.rs @@ -1,10 +1,16 @@ //! Private, opt-in management API. Client protocol ports are never used for HTTP. use axum::{ body::Body, - extract::{Extension, Path, Query}, + extract::{ + ws::{Message, WebSocket, WebSocketUpgrade}, + Extension, Path, Query, + }, http::{Request, StatusCode}, middleware::{self, Next}, - response::{IntoResponse, Response}, + response::{ + sse::{Event as SseEvent, KeepAlive, Sse}, + IntoResponse, Response, + }, routing::get, Json, Router, }; @@ -13,13 +19,18 @@ use once_cell::sync::Lazy; use serde_derive::{Deserialize, Serialize}; use serde_json::{json, Map, Value}; use std::{ - collections::{BTreeMap, HashMap, VecDeque}, + collections::{BTreeMap, HashMap, HashSet, VecDeque}, + convert::Infallible, net::{IpAddr, SocketAddr}, path::PathBuf, - sync::RwLock, - time::Instant, + sync::{ + atomic::{AtomicU64, Ordering}, + RwLock, + }, + time::{Duration, Instant}, }; -use tokio::sync::{oneshot, Mutex}; +use tokio::sync::{broadcast, mpsc, oneshot, Mutex}; +use tokio_stream::{wrappers::ReceiverStream, StreamExt as _}; pub const REDACTED: &str = "__REDACTED__"; static PEERS: Lazy>> = Lazy::new(|| RwLock::new(Vec::new())); @@ -60,6 +71,167 @@ fn apply_dynamic(name: &str, value: &str) -> bool { .unwrap_or(false) } +// Aggregated event bus backing the /v1/events SSE stream. hbbs emits local +// events (peers, bans, relays) and forwards the events reported by connected +// relay servers. +static EVENTS: Lazy> = Lazy::new(|| broadcast::channel(1024).0); + +fn emit_event(kind: &str, data: Value) { + let _ = EVENTS.send(json!({ + "kind": kind, + "timestamp": chrono::Utc::now().to_rfc3339(), + "data": data, + })); +} + +// Registered relay control channels (hbbs only). Each entry describes one +// live reverse connection established by an hbbr instance. +struct RelayConnection { + id: String, + info: Value, + address: Option, + version: String, + generation: u64, + started: Instant, + last_heartbeat: Instant, + sessions: Vec, + tx: mpsc::UnboundedSender, +} + +static RELAYS: Lazy>> = + Lazy::new(|| RwLock::new(HashMap::new())); +static RELAY_GENERATION: AtomicU64 = AtomicU64::new(0); + +// Registered by hbbs to fold discovered relay addresses back into the live +// relay server list (static configuration plus registered relays). +pub type RelaySync = Box) + Send + Sync>; +static RELAY_SYNC: Lazy>> = Lazy::new(|| RwLock::new(None)); + +#[allow(dead_code)] // Only the hbbs binary registers a relay sync callback. +pub fn set_relay_sync(sync: RelaySync) { + *RELAY_SYNC.write().unwrap() = Some(sync); +} + +// Addresses reported by registered relays (used to merge the live relay list). +pub fn registered_relay_addresses() -> Vec { + RELAYS + .read() + .unwrap() + .values() + .filter_map(|relay| relay.address.clone()) + .collect() +} + +fn sync_relay_list() { + let addresses: Vec = RELAYS + .read() + .unwrap() + .values() + .filter_map(|relay| relay.address.clone()) + .collect(); + if let Some(sync) = &*RELAY_SYNC.read().unwrap() { + sync(addresses); + } +} + +fn relay_command(relay_id: &str, payload: Value) -> ResultType<()> { + let relays = RELAYS.read().unwrap(); + let relay = relays + .get(relay_id) + .ok_or_else(|| anyhow!("Unknown relay: {relay_id}"))?; + relay + .tx + .send(Message::Text(payload.to_string())) + .map_err(|_| anyhow!("Relay control channel is closed"))?; + Ok(()) +} + +fn fanout_bans() { + let payload = json!({"type": "bans", "bans": *BANS.read().unwrap()}); + let text = payload.to_string(); + for relay in RELAYS.read().unwrap().values() { + let _ = relay.tx.send(Message::Text(text.clone())); + } +} + +// Pending request-response exchanges with registered relays: hbbs HTTP +// handlers await the completion channel while the relay socket loop resolves +// it from the "response" control message. +struct PendingRelayRequest { + relay: String, + responder: oneshot::Sender, +} + +static RELAY_REQUESTS: Lazy>> = + Lazy::new(|| std::sync::Mutex::new(HashMap::new())); +static RELAY_REQUEST_SEQUENCE: AtomicU64 = AtomicU64::new(1); +const RELAY_REQUEST_TIMEOUT: Duration = Duration::from_secs(10); + +// Sends a command carrying a request id to one relay and awaits the matching +// response, so HTTP handlers can proxy relay-scoped reads and writes. +async fn relay_request(relay_id: &str, mut payload: Value) -> ResultType { + let (responder, waiter) = oneshot::channel(); + let id = RELAY_REQUEST_SEQUENCE.fetch_add(1, Ordering::SeqCst); + if let Some(object) = payload.as_object_mut() { + object.insert("id".to_owned(), json!(id)); + } + RELAY_REQUESTS.lock().unwrap().insert( + id, + PendingRelayRequest { + relay: relay_id.to_owned(), + responder, + }, + ); + if let Err(err) = relay_command(relay_id, payload) { + RELAY_REQUESTS.lock().unwrap().remove(&id); + return Err(err); + } + match tokio::time::timeout(RELAY_REQUEST_TIMEOUT, waiter).await { + Ok(Ok(response)) => Ok(response), + Ok(Err(_)) => Err(anyhow!("Relay control channel closed")), + Err(_) => { + RELAY_REQUESTS.lock().unwrap().remove(&id); + Err(anyhow!("Relay did not answer in time")) + } + } +} + +// Drops the pending exchanges of one relay; awaiting requests fail fast +// instead of waiting for the timeout. +fn cancel_relay_requests(relay_id: &str) { + RELAY_REQUESTS + .lock() + .unwrap() + .retain(|_, pending| pending.relay != relay_id); +} + +// Outbound relay-side events (session lifecycle) consumed by the hbbr control +// client. Only meaningful inside the hbbr binary. Events are enqueued only +// after a consumer has taken the receiver, so a relay without the control +// channel does not accumulate session events forever. +static CONTROL_OUT: Lazy<( + mpsc::UnboundedSender, + std::sync::Mutex>>, +)> = Lazy::new(|| { + let (tx, rx) = mpsc::unbounded_channel(); + (tx, std::sync::Mutex::new(Some(rx))) +}); +static CONTROL_CONSUMED: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false); + +fn control_emit(event: Value) { + if CONTROL_CONSUMED.load(std::sync::atomic::Ordering::Relaxed) { + let _ = CONTROL_OUT.0.send(event); + } +} + +pub fn take_control_receiver() -> Option> { + let receiver = CONTROL_OUT.1.lock().unwrap().take(); + if receiver.is_some() { + CONTROL_CONSUMED.store(true, std::sync::atomic::Ordering::Relaxed); + } + receiver +} + #[derive(Clone, Default, Deserialize, Serialize, PartialEq, Eq)] #[serde(deny_unknown_fields)] pub struct Bans { @@ -128,6 +300,20 @@ pub fn is_banned(id: &str, ip: &str) -> bool { #[allow(dead_code)] pub fn set_peers(peers: Vec) { + fn ids(snapshot: &[Value]) -> HashSet { + snapshot + .iter() + .filter_map(|peer| peer["id"].as_str().map(str::to_owned)) + .collect() + } + let previous = ids(&PEERS.read().unwrap()); + let current = ids(&peers); + for id in current.difference(&previous) { + emit_event("peer_online", json!({"id": id})); + } + for id in previous.difference(¤t) { + emit_event("peer_offline", json!({"id": id})); + } *PEERS.write().unwrap() = peers; } @@ -158,12 +344,14 @@ pub fn open_session( sessions.insert( uuid.to_owned(), Session { - info, + info: info.clone(), ids, ips, cancel: Some(tx), }, ); + control_emit(json!({"type": "event", "event": { + "kind": "session_opened", "session": info, "uuid": uuid}})); Some(rx) } @@ -176,6 +364,34 @@ pub fn update_session(uuid: &str, bytes: usize, speed: usize) { pub fn close_session(uuid: &str) { SESSIONS.write().unwrap().remove(uuid); + control_emit(json!({"type": "event", "event": { + "kind": "session_closed", "uuid": uuid}})); +} + +// Snapshot of the active relay sessions, sorted by uuid. +pub fn sessions_snapshot() -> Vec { + let mut data: Vec = SESSIONS + .read() + .unwrap() + .values() + .map(|s| s.info.clone()) + .collect(); + data.sort_by(|a, b| a["uuid"].as_str().cmp(&b["uuid"].as_str())); + data +} + +// Cancels a session by uuid; returns false when the session is unknown. +pub fn disconnect_session(uuid: &str) -> bool { + let mut sessions = SESSIONS.write().unwrap(); + if let Some(session) = sessions.get_mut(uuid) { + if let Some(tx) = session.cancel.take() { + let _ = tx.send(()); + } + session.info["closing"] = json!(true); + true + } else { + false + } } #[derive(Clone, Default, Deserialize, Serialize)] @@ -256,7 +472,6 @@ struct ApiState { service: &'static str, token: String, started: Instant, - } async fn authorize(request: Request, next: Next) -> Response { @@ -290,7 +505,9 @@ async fn status(Extension(state): Extension) -> Json { ) } -async fn config(Extension(state): Extension) -> Json { +// Live configuration view: schema, effective values and the pending restart +// diff. Shared by the HTTP endpoint and the relay control channel. +pub fn config_view(service: &str) -> Value { let managed = MANAGED.read().unwrap(); let mut pending = Map::new(); for (name, value) in &managed.pending { @@ -302,16 +519,18 @@ async fn config(Extension(state): Extension) -> Json { }), ); } - Json( - json!({"schema":schema(state.service),"values":managed.values.clone(), - "pending_restart":pending}), - ) + json!({"schema":schema(service),"values":managed.values.clone(), + "pending_restart":pending}) +} + +async fn config(Extension(state): Extension) -> Json { + Json(config_view(state.service)) } #[derive(Clone, Default, Deserialize)] #[serde(deny_unknown_fields)] -struct ConfigUpdate { - values: BTreeMap, +pub struct ConfigUpdate { + pub values: BTreeMap, } fn definition_of(definitions: &[Value], name: &str) -> Option { @@ -322,102 +541,118 @@ fn definition_of(definitions: &[Value], name: &str) -> Option { } fn bad_request(message: String) -> Response { - ( - StatusCode::BAD_REQUEST, - Json(json!({"error": message})), - ) - .into_response() + (StatusCode::BAD_REQUEST, Json(json!({"error": message}))).into_response() } -async fn replace_config( - Extension(state): Extension, - Json(update): Json, -) -> Response { - // Validate the submission against the schema before touching any state. - let definitions = schema(state.service); +// Validates a submission against the schema without touching any state. +// Returns the first violation as a message so callers can map it to their +// transport (HTTP bad request or a control channel error response). +pub fn validate_update(service: &str, update: &ConfigUpdate) -> Option { + let definitions = schema(service); for (name, value) in &update.values { let Some(definition) = definition_of(&definitions, name) else { - return bad_request(format!("Unknown managed configuration: {name}")); + return Some(format!("Unknown managed configuration: {name}")); }; if definition["secret"] == true && value == REDACTED { continue; } if let Err(err) = validate_value(&definition, name, value) { - return bad_request(err.to_string()); + return Some(err.to_string()); } } - let service = state.service; - // The task owns the update through persistence and activation even if the - // HTTP caller disconnects while the blocking filesystem work is in flight. - let result = tokio::spawn(async move { - let _update = CONFIG_UPDATES.lock().await; - let (mut values, mut overrides, mut pending) = { - let managed = MANAGED.read().unwrap(); - ( - managed.values.clone(), - managed.overrides.clone(), - managed.pending.clone(), - ) + None +} + +// Validates, persists and activates a managed configuration update. Shared +// by the HTTP handler and the relay control channel (hbbr applies the +// updates forwarded by hbbs). The caller retains the update through +// persistence and activation even if it disconnects while the blocking +// filesystem work is in flight. +pub async fn apply_config_update( + service: &'static str, + update: ConfigUpdate, +) -> ResultType<(Map, Map)> { + let definitions = schema(service); + let _update = CONFIG_UPDATES.lock().await; + let (mut values, mut overrides, mut pending) = { + let managed = MANAGED.read().unwrap(); + ( + managed.values.clone(), + managed.overrides.clone(), + managed.pending.clone(), + ) + }; + for (name, value) in &update.values { + let Some(definition) = definition_of(&definitions, name) else { + continue; }; - for (name, value) in &update.values { - let Some(definition) = definition_of(&definitions, name) else { - continue; - }; - if definition["secret"] == true && value == REDACTED { - continue; - } - overrides.insert(name.clone(), value.clone()); + if definition["secret"] == true && value == REDACTED { + continue; } - // Persist before activation; a write failure keeps the running state. - let persisted = overrides.clone(); - tokio::task::spawn_blocking(move || persist_config(service, &persisted)).await??; - let mut applied = Map::new(); - let mut pending_report = Map::new(); - for (name, value) in &overrides { - let Some(definition) = definition_of(&definitions, name) else { - continue; - }; - let secret = definition["secret"] == true; - // Secrets cannot be compared against the running value because the - // effective view is redacted; only resubmitted secrets are reported. - if secret && !update.values.contains_key(name) { - continue; - } - let display = if secret { - REDACTED.to_owned() - } else { - value.clone() - }; - if !secret && values.get(name) == Some(&display) { - pending.remove(name); - continue; - } - if definition["dynamic"] == true && apply_dynamic(name, value) { - values.insert(name.clone(), display.clone()); - pending.remove(name); - applied.insert(name.clone(), json!(display)); - } else { - let current = values.get(name).cloned().unwrap_or_default(); - pending.insert(name.clone(), display.clone()); - pending_report.insert( - name.clone(), - json!({"current": current, "pending": display}), - ); - } + overrides.insert(name.clone(), value.clone()); + } + // Persist before activation; a write failure keeps the running state. + let persisted = overrides.clone(); + tokio::task::spawn_blocking(move || persist_config(service, &persisted)).await??; + let mut applied = Map::new(); + let mut pending_report = Map::new(); + for (name, value) in &overrides { + let Some(definition) = definition_of(&definitions, name) else { + continue; + }; + let secret = definition["secret"] == true; + // Secrets cannot be compared against the running value because the + // effective view is redacted; only resubmitted secrets are reported. + if secret && !update.values.contains_key(name) { + continue; } - *MANAGED.write().unwrap() = ManagedState { - values, - overrides, - pending, + let display = if secret { + REDACTED.to_owned() + } else { + value.clone() }; - Ok::<_, hbb_common::anyhow::Error>((applied, pending_report)) - }) - .await; + if !secret && values.get(name) == Some(&display) { + pending.remove(name); + continue; + } + if definition["dynamic"] == true && apply_dynamic(name, value) { + values.insert(name.clone(), display.clone()); + pending.remove(name); + applied.insert(name.clone(), json!(display)); + } else { + let current = values.get(name).cloned().unwrap_or_default(); + pending.insert(name.clone(), display.clone()); + pending_report.insert( + name.clone(), + json!({"current": current, "pending": display}), + ); + } + } + *MANAGED.write().unwrap() = ManagedState { + values, + overrides, + pending, + }; + Ok((applied, pending_report)) +} + +async fn replace_config( + Extension(state): Extension, + Json(update): Json, +) -> Response { + if let Some(err) = validate_update(state.service, &update) { + return bad_request(err); + } + let service = state.service; + // The task owns the update through persistence and activation even if + // the HTTP caller disconnects while the blocking filesystem work is in + // flight. + let result = tokio::spawn(apply_config_update(service, update)).await; match result { - Ok(Ok((applied, pending_report))) => Json( - json!({"state":"saved","applied":applied,"pending_restart":pending_report}), - ) - .into_response(), + Ok(Ok((applied, pending_report))) => { + Json(json!({"state":"saved","applied":applied,"pending_restart":pending_report})) + .into_response() + } Ok(Err(err)) => { log::error!("Unable to persist managed configuration: {err}"); StatusCode::INTERNAL_SERVER_ERROR.into_response() @@ -472,12 +707,30 @@ struct LogsQuery { tail: Option, } -async fn logs(Query(query): Query) -> Response { - let tail = query.tail.unwrap_or(200).min(LOG_CAPACITY); +// Recent in-memory log lines. Shared by the HTTP endpoint and the relay +// control channel (hbbs proxies relay log reads over the control channel). +pub fn logs_view(tail: usize) -> Value { + let tail = tail.min(LOG_CAPACITY); let logs = LOGS.read().unwrap(); let skipped = logs.len().saturating_sub(tail); let data: Vec = logs.iter().skip(skipped).cloned().collect(); - Json(json!({"data": data})).into_response() + json!({"data": data}) +} + +async fn logs(Query(query): Query) -> Response { + Json(logs_view(query.tail.unwrap_or(200))).into_response() +} + +// Ask the supervisor (container restart policy or systemd) to bring the +// service back: flush the response, then exit. Dynamic configuration is +// already active; pending entries take effect on the next startup. +async fn restart() -> Response { + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(250)).await; + log::info!("Exiting for restart; the supervisor should bring the service back"); + std::process::exit(0); + }); + (StatusCode::ACCEPTED, Json(json!({"state":"restarting"}))).into_response() } async fn peers(Extension(state): Extension) -> Response { @@ -491,26 +744,14 @@ async fn sessions(Extension(state): Extension) -> Response { if state.service != "hbbr" { return StatusCode::NOT_FOUND.into_response(); } - let mut data: Vec = SESSIONS - .read() - .unwrap() - .values() - .map(|s| s.info.clone()) - .collect(); - data.sort_by(|a, b| a["uuid"].as_str().cmp(&b["uuid"].as_str())); - Json(json!({"data":data})).into_response() + Json(json!({"data":sessions_snapshot()})).into_response() } async fn disconnect(Extension(state): Extension, Path(uuid): Path) -> Response { if state.service != "hbbr" { return StatusCode::NOT_FOUND.into_response(); } - let mut sessions = SESSIONS.write().unwrap(); - if let Some(session) = sessions.get_mut(&uuid) { - if let Some(tx) = session.cancel.take() { - let _ = tx.send(()); - } - session.info["closing"] = json!(true); + if disconnect_session(&uuid) { return (StatusCode::ACCEPTED, Json(json!({"state":"closing"}))).into_response(); } StatusCode::NOT_FOUND.into_response() @@ -520,6 +761,26 @@ async fn bans() -> Json { Json(BANS.read().unwrap().clone()) } +// Validates, persists and activates a ban policy. Shared by the HTTP handler +// and the relay control channel (hbbr applies the policies forwarded by hbbs). +// The update owns persistence through activation even if the caller +// disconnects while the blocking filesystem work is in flight. +pub async fn apply_bans_update(service: &'static str, mut bans: Bans) -> ResultType<()> { + bans.validate()?; + let _update = BAN_UPDATES.lock().await; + if *BANS.read().unwrap() == bans { + return Ok(()); + } + let persisted = bans.clone(); + tokio::task::spawn_blocking(move || persist_bans(service, &persisted)).await??; + // Never hold the protocol admission lock during filesystem I/O. + // Keep the BANS -> SESSIONS order for activation and admission. + let mut current = BANS.write().unwrap(); + *current = bans; + cancel_banned_sessions(¤t); + Ok(()) +} + async fn replace_bans( Extension(state): Extension, Json(mut bans): Json, @@ -531,26 +792,16 @@ async fn replace_bans( ) .into_response(); } - // The task owns the update through persistence and activation even if the - // HTTP caller disconnects while the blocking filesystem work is in flight. - match tokio::spawn(async move { - let _update = BAN_UPDATES.lock().await; - if *BANS.read().unwrap() == bans { - return Ok(()); + let service = state.service; + let result = tokio::spawn(apply_bans_update(service, bans)).await; + match result { + Ok(Ok(())) => { + if service == "hbbs" { + fanout_bans(); + } + emit_event("bans_replaced", json!({"service": service})); + Json(json!({"state":"applied"})).into_response() } - let service = state.service; - let persisted = bans.clone(); - tokio::task::spawn_blocking(move || persist_bans(service, &persisted)).await??; - // Never hold the protocol admission lock during filesystem I/O. - // Keep the BANS -> SESSIONS order for activation and admission. - let mut current = BANS.write().unwrap(); - *current = bans; - cancel_banned_sessions(¤t); - Ok::<(), hbb_common::anyhow::Error>(()) - }) - .await - { - Ok(Ok(())) => Json(json!({"state":"applied"})).into_response(), result => { log::error!("Unable to persist management bans: {result:?}"); StatusCode::INTERNAL_SERVER_ERROR.into_response() @@ -594,17 +845,11 @@ pub fn restore_bans(service: &str) -> ResultType<()> { Ok(()) } -pub fn start(service: &'static str, port: u16) -> ResultType<()> { - let bind = std::env::var("RD_MANAGEMENT_BIND").unwrap_or_default(); - if bind.is_empty() { - return Ok(()); - } - let token = std::env::var("RD_MANAGEMENT_TOKEN").unwrap_or_default(); - if token.len() < 32 { - return Err(anyhow!( - "RD_MANAGEMENT_TOKEN must contain at least 32 characters" - )); - } +// Populates the managed state from the schema, deployment arguments and saved +// overrides. Runs during startup for both services: the HTTP listener is +// optional, but the control channel and its proxied management requests need +// the state regardless. +pub fn init_managed_state(service: &'static str, port: u16) -> ResultType<()> { let overrides = load_config(service)?.values; let mut values = BTreeMap::new(); for definition in schema(service) { @@ -632,6 +877,21 @@ pub fn start(service: &'static str, port: u16) -> ResultType<()> { overrides, pending: BTreeMap::new(), }; + Ok(()) +} + +pub fn start(service: &'static str, port: u16) -> ResultType<()> { + init_managed_state(service, port)?; + let bind = std::env::var("RD_MANAGEMENT_BIND").unwrap_or_default(); + if bind.is_empty() { + return Ok(()); + } + let token = std::env::var("RD_MANAGEMENT_TOKEN").unwrap_or_default(); + if token.len() < 32 { + return Err(anyhow!( + "RD_MANAGEMENT_TOKEN must contain at least 32 characters" + )); + } let state = ApiState { service, token, @@ -650,6 +910,352 @@ pub fn start(service: &'static str, port: u16) -> ResultType<()> { Ok(()) } +// Reverse control channel: hbbr connects here (hbbs only), registers itself, +// reports heartbeats and session events, and receives policy commands. +async fn control(Extension(state): Extension, ws: WebSocketUpgrade) -> Response { + if state.service != "hbbs" { + return StatusCode::NOT_FOUND.into_response(); + } + ws.on_upgrade(move |socket| control_socket(socket)) +} + +fn valid_relay_text(value: &str, max: usize) -> Option { + if !value.is_empty() && value.len() <= max && !value.chars().any(char::is_whitespace) { + Some(value.to_owned()) + } else { + None + } +} + +async fn control_socket(mut socket: WebSocket) { + // Registration is the first frame; the connection is dropped otherwise. + let register = match tokio::time::timeout(Duration::from_secs(10), socket.recv()).await { + Ok(Some(Ok(Message::Text(text)))) => text, + _ => return, + }; + let register: Value = match serde_json::from_str::(®ister) { + Ok(value) if value["type"] == "register" => value, + _ => return, + }; + let relay_id = match register["relay"]["id"] + .as_str() + .and_then(|id| valid_relay_text(id, 100)) + { + Some(id) => id, + None => return, + }; + let address = register["relay"]["address"] + .as_str() + .and_then(|address| valid_relay_text(address, 253)); + let version = register["relay"]["version"] + .as_str() + .unwrap_or_default() + .to_owned(); + let info = register["relay"].clone(); + let generation = RELAY_GENERATION.fetch_add(1, Ordering::SeqCst); + let (tx, mut rx) = mpsc::unbounded_channel::(); + { + // A re-registering relay replaces its previous connection; the stale + // generation prevents the old socket task from clearing the new entry. + let mut relays = RELAYS.write().unwrap(); + if let Some(previous) = relays.remove(&relay_id) { + let _ = previous.tx.send(Message::Close(None)); + } + relays.insert( + relay_id.clone(), + RelayConnection { + id: relay_id.clone(), + info, + address: address.clone(), + version, + generation, + started: Instant::now(), + last_heartbeat: Instant::now(), + sessions: Vec::new(), + tx: tx.clone(), + }, + ); + } + log::info!("Relay {relay_id} registered on the control channel"); + emit_event( + "relay_registered", + json!({"id": relay_id, "address": address, "generation": generation}), + ); + sync_relay_list(); + // Bring the new relay up to date with the active policy and request a + // session snapshot so aggregation is complete right after registration. + let _ = tx.send(Message::Text( + json!({"type": "bans", "bans": *BANS.read().unwrap()}).to_string(), + )); + let _ = tx.send(Message::Text(json!({"type": "get_sessions"}).to_string())); + + loop { + tokio::select! { + inbound = socket.recv() => { + match inbound { + Some(Ok(Message::Text(text))) => { + handle_control_message(&relay_id, generation, &text); + } + Some(Ok(_)) => {} + _ => break, + } + } + outbound = rx.recv() => { + match outbound { + Some(message) => { + if socket.send(message).await.is_err() { + break; + } + } + None => break, + } + } + _ = tokio::time::sleep(Duration::from_secs(30)) => { + let stale = RELAYS + .read() + .unwrap() + .get(&relay_id) + .map(|relay| { + relay.generation == generation + && relay.last_heartbeat.elapsed().as_secs() > 45 + }) + .unwrap_or(true); + if stale { + break; + } + } + } + } + // Drop the write guard before sync_relay_list(): it re-acquires RELAYS + // for reading, which deadlocks while the write lock is held. + let mut left = false; + { + let mut relays = RELAYS.write().unwrap(); + if relays + .get(&relay_id) + .map(|relay| relay.generation == generation) + .unwrap_or(false) + { + relays.remove(&relay_id); + log::info!("Relay {relay_id} left the control channel"); + emit_event("relay_disconnected", json!({"id": relay_id})); + // Pending exchanges die with the connection; waiters fail fast. + cancel_relay_requests(&relay_id); + left = true; + } + } + if left { + sync_relay_list(); + } +} + +fn handle_control_message(relay_id: &str, generation: u64, text: &str) { + let Ok(message) = serde_json::from_str::(text) else { + return; + }; + let mut relays = RELAYS.write().unwrap(); + let Some(relay) = relays.get_mut(relay_id) else { + return; + }; + if relay.generation != generation { + return; + } + match message["type"].as_str() { + Some("heartbeat") => { + relay.last_heartbeat = Instant::now(); + } + Some("sessions") => { + if let Some(data) = message["data"].as_array() { + relay.sessions = data.clone(); + } + } + Some("event") => { + // Relay-reported events join the aggregated event stream. + let mut event = message["event"].clone(); + if let Some(object) = event.as_object_mut() { + object.insert("relay_id".to_owned(), json!(relay_id)); + } + if event["kind"].as_str() == Some("session_opened") { + relay.last_heartbeat = Instant::now(); + } + let kind = event["kind"].as_str().unwrap_or("relay_event").to_owned(); + emit_event(&kind, event); + } + Some("response") => { + // Complete the pending relay_request awaiting this exchange. + if let Some(id) = message["id"].as_u64() { + if let Some(pending) = RELAY_REQUESTS.lock().unwrap().remove(&id) { + let _ = pending.responder.send(message); + } + } + } + _ => {} + } +} + +// Registered relays with health and aggregation metadata (hbbs only). +async fn relays(Extension(state): Extension) -> Response { + if state.service != "hbbs" { + return StatusCode::NOT_FOUND.into_response(); + } + let mut data: Vec = RELAYS + .read() + .unwrap() + .values() + .map(|relay| { + json!({ + "id": relay.id, + "address": relay.address, + "version": relay.version, + "info": relay.info, + "uptime_seconds": relay.started.elapsed().as_secs(), + "last_heartbeat_age_seconds": relay.last_heartbeat.elapsed().as_secs(), + "sessions_count": relay.sessions.len(), + }) + }) + .collect(); + data.sort_by(|a, b| a["id"].as_str().cmp(&b["id"].as_str())); + Json(json!({"data": data})).into_response() +} + +async fn relay_sessions( + Extension(state): Extension, + Path(relay_id): Path, +) -> Response { + if state.service != "hbbs" { + return StatusCode::NOT_FOUND.into_response(); + } + let relays = RELAYS.read().unwrap(); + match relays.get(&relay_id) { + Some(relay) => Json(json!({"data": relay.sessions.clone()})).into_response(), + None => StatusCode::NOT_FOUND.into_response(), + } +} + +async fn relay_disconnect( + Extension(state): Extension, + Path((relay_id, uuid)): Path<(String, String)>, +) -> Response { + if state.service != "hbbs" { + return StatusCode::NOT_FOUND.into_response(); + } + match relay_command(&relay_id, json!({"type": "disconnect", "uuid": uuid})) { + Ok(()) => (StatusCode::ACCEPTED, Json(json!({"state":"closing"}))).into_response(), + Err(_) => StatusCode::NOT_FOUND.into_response(), + } +} + +fn relay_available(relay_id: &str) -> bool { + RELAYS.read().unwrap().contains_key(relay_id) +} + +// The relay answered with an HTTP-equivalent status and body; forward both so +// the proxy endpoints keep the semantics of the direct ones. +fn relay_response(result: ResultType) -> Response { + match result { + Ok(response) => { + let status = response["status"] + .as_u64() + .and_then(|code| StatusCode::from_u16(code as u16).ok()) + .unwrap_or(StatusCode::BAD_GATEWAY); + (status, Json(response["body"].clone())).into_response() + } + Err(err) => ( + StatusCode::BAD_GATEWAY, + Json(json!({"error": err.to_string()})), + ) + .into_response(), + } +} + +// Relay-scoped proxies: management clients talk to hbbs only; hbbs forwards +// configuration and log requests to the target relay over the control channel. +async fn relay_config( + Extension(state): Extension, + Path(relay_id): Path, +) -> Response { + if state.service != "hbbs" || !relay_available(&relay_id) { + return StatusCode::NOT_FOUND.into_response(); + } + relay_response(relay_request(&relay_id, json!({"type": "get_config"})).await) +} + +async fn relay_replace_config( + Extension(state): Extension, + Path(relay_id): Path, + Json(update): Json, +) -> Response { + if state.service != "hbbs" || !relay_available(&relay_id) { + return StatusCode::NOT_FOUND.into_response(); + } + relay_response( + relay_request( + &relay_id, + json!({"type": "set_config", "values": update.values}), + ) + .await, + ) +} + +async fn relay_logs( + Extension(state): Extension, + Path(relay_id): Path, + Query(query): Query, +) -> Response { + if state.service != "hbbs" || !relay_available(&relay_id) { + return StatusCode::NOT_FOUND.into_response(); + } + let tail = query.tail.unwrap_or(200).min(LOG_CAPACITY); + relay_response(relay_request(&relay_id, json!({"type": "get_logs", "tail": tail})).await) +} + +// Restart is fire-and-forget: the relay exits right after acknowledging, so +// there is no stable connection left to answer a correlated request. +async fn relay_restart( + Extension(state): Extension, + Path(relay_id): Path, +) -> Response { + if state.service != "hbbs" { + return StatusCode::NOT_FOUND.into_response(); + } + match relay_command(&relay_id, json!({"type": "restart"})) { + Ok(()) => (StatusCode::ACCEPTED, Json(json!({"state":"restarting"}))).into_response(), + Err(_) => StatusCode::NOT_FOUND.into_response(), + } +} + +// Aggregated server-sent event stream: peer online/offline, relay session +// lifecycle, relay registration and policy changes. +async fn events() -> Sse>> { + let (tx, rx) = mpsc::channel::(64); + tokio::spawn(async move { + let mut subscription = EVENTS.subscribe(); + loop { + match subscription.recv().await { + Ok(event) => { + if tx.send(event).await.is_err() { + break; + } + } + Err(broadcast::error::RecvError::Lagged(_)) => continue, + Err(_) => break, + } + } + }); + let stream = ReceiverStream::new(rx).map(|event| { + Ok::<_, Infallible>( + SseEvent::default() + .event(event["kind"].as_str().unwrap_or("event").to_owned()) + .data(event.to_string()), + ) + }); + Sse::new(stream).keep_alive( + KeepAlive::new() + .interval(Duration::from_secs(15)) + .text("keep-alive"), + ) +} + fn router(state: ApiState) -> Router { Router::new() .route("/v1/status", get(status)) @@ -659,6 +1265,21 @@ fn router(state: ApiState) -> Router { .route("/v1/sessions", get(sessions)) .route("/v1/sessions/:uuid", axum::routing::delete(disconnect)) .route("/v1/bans", get(bans).put(replace_bans)) + .route("/v1/restart", axum::routing::post(restart)) + .route("/v1/events", get(events)) + .route("/v1/control", get(control)) + .route("/v1/relays", get(relays)) + .route("/v1/relays/:id/sessions", get(relay_sessions)) + .route( + "/v1/relays/:id/sessions/:uuid", + axum::routing::delete(relay_disconnect), + ) + .route( + "/v1/relays/:id/config", + get(relay_config).put(relay_replace_config), + ) + .route("/v1/relays/:id/logs", get(relay_logs)) + .route("/v1/relays/:id/restart", axum::routing::post(relay_restart)) .layer(middleware::from_fn(authorize)) .layer(Extension(state)) } @@ -723,7 +1344,6 @@ mod tests { service: "hbbs", token: "test".into(), started: Instant::now(), - }; let a = Bans { device_ids: vec!["policy-a".into()], @@ -791,7 +1411,6 @@ mod tests { service: "hbbs", token: token.into(), started: Instant::now(), - }); let task = tokio::spawn( axum::Server::from_tcp(listener) @@ -846,6 +1465,192 @@ mod tests { "unparseable:header" ); } + fn subscribe_events() -> broadcast::Receiver { + EVENTS.subscribe() + } + + fn register_test_relay( + relay_id: &str, + address: Option<&str>, + ) -> mpsc::UnboundedReceiver { + let (tx, rx) = mpsc::unbounded_channel(); + RELAYS.write().unwrap().insert( + relay_id.to_owned(), + RelayConnection { + id: relay_id.to_owned(), + info: json!({"id": relay_id}), + address: address.map(str::to_owned), + version: "test".to_owned(), + generation: RELAY_GENERATION.fetch_add(1, Ordering::SeqCst), + started: Instant::now(), + last_heartbeat: Instant::now(), + sessions: Vec::new(), + tx, + }, + ); + rx + } + + fn clear_relays() { + RELAYS.write().unwrap().clear(); + } + + #[test] + fn peer_snapshots_emit_online_and_offline_events() { + *PEERS.write().unwrap() = Vec::new(); + let mut events = subscribe_events(); + set_peers(vec![json!({"id": "111111111"})]); + set_peers(vec![json!({"id": "222222222"})]); + let mut kinds = Vec::new(); + while let Ok(event) = events.try_recv() { + kinds.push(( + event["kind"].as_str().unwrap().to_owned(), + event["data"]["id"].as_str().unwrap_or_default().to_owned(), + )); + } + assert!(kinds.contains(&("peer_online".to_owned(), "111111111".into()))); + assert!(kinds.contains(&("peer_offline".to_owned(), "111111111".into()))); + assert!(kinds.contains(&("peer_online".to_owned(), "222222222".into()))); + *PEERS.write().unwrap() = Vec::new(); + } + + #[tokio::test] + async fn apply_bans_update_persists_and_activates_for_relay_use() { + let _test = POLICY_TEST.lock().unwrap(); + let fixture = PolicyFixture::new(); + let bans = Bans { + device_ids: vec!["relay-blocked".into()], + ips: vec![], + }; + apply_bans_update("hbbr", bans).await.unwrap(); + assert!(is_banned("relay-blocked", "")); + let persisted: Bans = + serde_json::from_slice(&std::fs::read(data_path("hbbr", "bans")).unwrap()).unwrap(); + assert!(persisted.device_ids.contains(&"relay-blocked".to_owned())); + drop(fixture); + } + + #[tokio::test] + async fn replace_bans_fans_out_to_registered_relays() { + let _test = POLICY_TEST.lock().unwrap(); + let fixture = PolicyFixture::new(); + clear_relays(); + let mut receiver = register_test_relay("relay-a", Some("relay-a:21117")); + let state = ApiState { + service: "hbbs", + token: "test".into(), + started: Instant::now(), + }; + let response = replace_bans( + Extension(state), + Json(Bans { + device_ids: vec!["fanout-device".into()], + ips: vec![], + }), + ) + .await; + assert_eq!(response.status(), StatusCode::OK); + let message = receiver + .recv() + .await + .expect("the relay must receive the active policy"); + let Message::Text(text) = message else { + panic!("expected a text message"); + }; + let payload: Value = serde_json::from_str(&text).unwrap(); + assert_eq!(payload["type"], "bans"); + assert_eq!(payload["bans"]["device_ids"][0], "fanout-device"); + clear_relays(); + drop(fixture); + } + + #[test] + fn control_events_are_enqueued_only_after_a_consumer_takes_the_receiver() { + let _test = POLICY_TEST.lock().unwrap(); + // Events emitted before any consumer registered are dropped instead + // of piling up in the unbounded channel. + control_emit(json!({"type": "event", "dropped": 1})); + let mut receiver = take_control_receiver().expect("first take returns the receiver"); + assert!( + receiver.try_recv().is_err(), + "events without a consumer must be dropped" + ); + control_emit(json!({"type": "event", "kept": 2})); + assert_eq!(receiver.try_recv().unwrap()["kept"], 2); + assert!( + take_control_receiver().is_none(), + "the receiver is handed out exactly once" + ); + control_emit(json!({"type": "event", "kept": 3})); + assert_eq!(receiver.try_recv().unwrap()["kept"], 3); + } + + #[test] + fn control_messages_update_heartbeat_sessions_and_forward_events() { + let _test = POLICY_TEST.lock().unwrap(); + clear_relays(); + let generation = RELAY_GENERATION.fetch_add(1, Ordering::SeqCst); + let (tx, _rx) = mpsc::unbounded_channel(); + RELAYS.write().unwrap().insert( + "relay-b".to_owned(), + RelayConnection { + id: "relay-b".to_owned(), + info: json!({"id": "relay-b"}), + address: None, + version: "test".to_owned(), + generation, + started: Instant::now(), + last_heartbeat: Instant::now() - Duration::from_secs(120), + sessions: Vec::new(), + tx, + }, + ); + let mut events = subscribe_events(); + handle_control_message("relay-b", generation, r#"{"type":"heartbeat"}"#); + assert!( + RELAYS.read().unwrap()["relay-b"] + .last_heartbeat + .elapsed() + .as_secs() + < 45, + "heartbeat must refresh the liveness timestamp" + ); + handle_control_message( + "relay-b", + generation, + r#"{"type":"sessions","data":[{"uuid":"session-1"}]}"#, + ); + assert_eq!(RELAYS.read().unwrap()["relay-b"].sessions.len(), 1); + // Stale generations are ignored so old sockets cannot corrupt state. + handle_control_message( + "relay-b", + generation + 1, + r#"{"type":"sessions","data":[]}"#, + ); + assert_eq!(RELAYS.read().unwrap()["relay-b"].sessions.len(), 1); + handle_control_message( + "relay-b", + generation, + r#"{"type":"event","event":{"kind":"session_opened","session":{"uuid":"session-1"}}}"#, + ); + // The event bus is shared; wait for our own event instead of + // assuming it is the next one on the channel. + let mut event = None; + for _ in 0..100 { + match events.try_recv() { + Ok(candidate) if candidate["data"]["relay_id"] == "relay-b" => { + event = Some(candidate); + break; + } + Ok(_) => continue, + Err(_) => std::thread::sleep(Duration::from_millis(5)), + } + } + let event = event.expect("relay events join the event bus"); + assert_eq!(event["kind"], "session_opened"); + clear_relays(); + } + #[test] fn bans_cancel_existing_sessions_and_duplicates_are_rejected() { let _test = POLICY_TEST.lock().unwrap(); @@ -953,11 +1758,12 @@ mod tests { assert_eq!(managed.values.get("limit-speed").unwrap(), "64"); assert_eq!(managed.values.get("port").unwrap(), "21117"); assert_eq!(managed.pending.get("port").unwrap(), "21118"); - let persisted: ManagedConfig = serde_json::from_slice( - &std::fs::read(data_path("hbbr", "config")).unwrap(), - ) - .unwrap(); - assert_eq!(persisted.values, managed_values(&[("limit-speed", "64"), ("port", "21118")])); + let persisted: ManagedConfig = + serde_json::from_slice(&std::fs::read(data_path("hbbr", "config")).unwrap()).unwrap(); + assert_eq!( + persisted.values, + managed_values(&[("limit-speed", "64"), ("port", "21118")]) + ); // The live config view reflects the dynamic change and the pending entry. let view = config(Extension(ConfigFixture::state("hbbr"))).await; assert_eq!(view["values"]["limit-speed"], "64"); @@ -1047,14 +1853,9 @@ mod tests { .unwrap(); assert_eq!(body["pending_restart"]["key"]["pending"], REDACTED); assert!(!format!("{body:?}").contains("replacement-secret")); - let persisted: ManagedConfig = serde_json::from_slice( - &std::fs::read(data_path("hbbr", "config")).unwrap(), - ) - .unwrap(); - assert_eq!( - persisted.values.get("key").unwrap(), - "replacement-secret" - ); + let persisted: ManagedConfig = + serde_json::from_slice(&std::fs::read(data_path("hbbr", "config")).unwrap()).unwrap(); + assert_eq!(persisted.values.get("key").unwrap(), "replacement-secret"); // Resubmitting the placeholder keeps the stored secret untouched. let response = replace_config( Extension(ConfigFixture::state("hbbr")), @@ -1064,14 +1865,9 @@ mod tests { ) .await; assert_eq!(response.status(), StatusCode::OK); - let persisted: ManagedConfig = serde_json::from_slice( - &std::fs::read(data_path("hbbr", "config")).unwrap(), - ) - .unwrap(); - assert_eq!( - persisted.values.get("key").unwrap(), - "replacement-secret" - ); + let persisted: ManagedConfig = + serde_json::from_slice(&std::fs::read(data_path("hbbr", "config")).unwrap()).unwrap(); + assert_eq!(persisted.values.get("key").unwrap(), "replacement-secret"); } #[tokio::test] @@ -1097,4 +1893,94 @@ mod tests { // Reset the ring for other tests. LOGS.write().unwrap().clear(); } + #[tokio::test] + async fn relay_requests_proxy_config_logs_and_restart() { + let _test = POLICY_TEST.lock().unwrap(); + let _fixture = ConfigFixture::new(); + clear_relays(); + let state = ConfigFixture::state("hbbs"); + // Unknown relays are reported as missing. + let missing = relay_config(Extension(state.clone()), Path("missing".to_owned())).await; + assert_eq!(missing.status(), StatusCode::NOT_FOUND); + let mut receiver = register_test_relay("relay-c", None); + let generation = RELAYS.read().unwrap()["relay-c"].generation; + let (seen_tx, mut seen_rx) = mpsc::unbounded_channel::(); + // A background task plays the relay: every command is answered with + // an HTTP-equivalent status and body through the control socket loop. + let responder = tokio::spawn(async move { + while let Some(message) = receiver.recv().await { + let Message::Text(text) = message else { + continue; + }; + let payload: Value = serde_json::from_str(&text).unwrap(); + let kind = payload["type"].as_str().unwrap().to_owned(); + let _ = seen_tx.send(kind.clone()); + let id = payload["id"].clone(); + let response = match kind.as_str() { + "get_config" => json!({"type":"response","id":id,"status":200, + "body":{"values":{"limit-speed":"64"}}}), + "get_logs" => json!({"type":"response","id":id,"status":200, + "body":{"data":["relay-log-line"]}}), + "set_config" => json!({"type":"response","id":id,"status":400, + "body":{"error":"invalid value"}}), + "restart" => break, + _ => continue, + }; + handle_control_message("relay-c", generation, &response.to_string()); + } + }); + let response = relay_config(Extension(state.clone()), Path("relay-c".to_owned())).await; + assert_eq!(response.status(), StatusCode::OK); + let body: Value = + serde_json::from_slice(&hyper::body::to_bytes(response.into_body()).await.unwrap()) + .unwrap(); + assert_eq!(body["values"]["limit-speed"], "64"); + let response = relay_logs( + Extension(state.clone()), + Path("relay-c".to_owned()), + Query(LogsQuery { tail: Some(5) }), + ) + .await; + assert_eq!(response.status(), StatusCode::OK); + let body: Value = + serde_json::from_slice(&hyper::body::to_bytes(response.into_body()).await.unwrap()) + .unwrap(); + assert_eq!(body["data"][0], "relay-log-line"); + // Relay-reported statuses flow through (e.g. validation errors). + let response = relay_replace_config( + Extension(state.clone()), + Path("relay-c".to_owned()), + Json(ConfigUpdate { + values: managed_values(&[("limit-speed", "64")]), + }), + ) + .await; + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + let body: Value = + serde_json::from_slice(&hyper::body::to_bytes(response.into_body()).await.unwrap()) + .unwrap(); + assert_eq!(body["error"], "invalid value"); + // Restart is fire-and-forget. + let response = relay_restart(Extension(state.clone()), Path("relay-c".to_owned())).await; + assert_eq!(response.status(), StatusCode::ACCEPTED); + for expected in ["get_config", "get_logs", "set_config", "restart"] { + assert_eq!(seen_rx.recv().await.unwrap(), expected); + } + responder.abort(); + clear_relays(); + } + + #[test] + fn cancelling_relay_requests_releases_waiters() { + let (responder, waiter) = oneshot::channel(); + RELAY_REQUESTS.lock().unwrap().insert( + 987_654, + PendingRelayRequest { + relay: "relay-y".to_owned(), + responder, + }, + ); + cancel_relay_requests("relay-y"); + assert!(waiter.blocking_recv().is_err()); + } } diff --git a/src/relay_control.rs b/src/relay_control.rs new file mode 100644 index 0000000..aeb57f7 --- /dev/null +++ b/src/relay_control.rs @@ -0,0 +1,265 @@ +//! Reverse control channel client. When RD_HBBS_API_URL is set, hbbr connects +//! to the hbbs management API over WebSocket, registers itself, reports +//! heartbeats and session events, receives policy commands (bans, session +//! disconnect, snapshot requests) and answers proxied management requests +//! (configuration reads and writes, log reads, restart). The connection is +//! outbound-only so hbbr never exposes a management listener for aggregation. +use crate::management; +use hbb_common::futures_util::{sink::SinkExt as _, stream::StreamExt as _}; +use hbb_common::tokio::sync::mpsc; +use hbb_common::{anyhow::anyhow, log, tokio, ResultType}; +use serde_json::{json, Value}; +use std::time::{Duration, Instant}; +use tokio_tungstenite::{ + tungstenite::{client::IntoClientRequest, http::HeaderValue, Message}, + MaybeTlsStream, WebSocketStream, +}; + +const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(10); + +type ControlSocket = WebSocketStream>; + +pub fn start(url: &str, token: &str) -> ResultType<()> { + if token.len() < 32 { + return Err(anyhow!( + "RD_MANAGEMENT_TOKEN must contain at least 32 characters" + )); + } + let url = url.to_owned(); + let token = token.to_owned(); + tokio::spawn(async move { + let mut backoff = Duration::from_secs(1); + // The receiver is taken once for the whole task: run() must not take + // it again on reconnects, or every later connection would lose the + // session event feed. + let mut events = management::take_control_receiver(); + loop { + match run(&url, &token, &mut events).await { + Ok(registered) => { + if registered { + // A healthy connection just ended; retry promptly + // instead of keeping the accumulated backoff. + backoff = Duration::from_secs(1); + } + log::info!("hbbs control channel closed; reconnecting") + } + Err(err) => log::warn!("hbbs control channel failed: {err}; reconnecting"), + } + tokio::time::sleep(backoff).await; + backoff = (backoff * 2).min(Duration::from_secs(60)); + } + }); + Ok(()) +} + +// The relay id identifies the node; the advertised address is optional. +// Without RD_RELAY_ADDRESS no address is reported: the id (often a container +// hostname) is not resolvable by clients and must not be broadcast as a relay +// server address. +fn relay_identity() -> (String, Option) { + let id = std::env::var("RD_RELAY_ID") + .ok() + .filter(|value| !value.is_empty()) + .or_else(|| { + std::env::var("HOSTNAME") + .ok() + .filter(|value| !value.is_empty()) + }) + .unwrap_or_else(|| "hbbr".to_owned()); + let address = std::env::var("RD_RELAY_ADDRESS") + .ok() + .and_then(|value| (!value.is_empty()).then_some(value)); + (id, address) +} + +// Returns Ok(true) when the connection had registered successfully before it +// ended, which lets the caller reset its reconnect backoff. +async fn run( + url: &str, + token: &str, + events: &mut Option>, +) -> ResultType { + let mut request = url.into_client_request()?; + request.headers_mut().insert( + "authorization", + HeaderValue::from_str(&format!("Bearer {token}"))?, + ); + let (mut socket, _) = tokio_tungstenite::connect_async(request).await?; + let (id, address) = relay_identity(); + let mut register = json!({ + "type": "register", + "relay": { + "id": id, + "version": crate::version::VERSION, + } + }); + if let Some(address) = &address { + register["relay"]["address"] = json!(address); + } + socket.send(Message::Text(register.to_string())).await?; + log::info!("Registered on the hbbs control channel as {id}"); + // Any outcome from here happened after a successful registration, so the + // caller resets its backoff even when the loop below fails. + run_registered(&mut socket, events).await.map(|_| true) +} + +async fn run_registered( + socket: &mut ControlSocket, + events: &mut Option>, +) -> ResultType<()> { + let started = Instant::now(); + let mut heartbeat = tokio::time::interval(HEARTBEAT_INTERVAL); + heartbeat.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + tokio::select! { + _ = heartbeat.tick() => { + let beat = json!({ + "type": "heartbeat", + "uptime_seconds": started.elapsed().as_secs(), + }); + socket.send(Message::Text(beat.to_string())).await?; + } + event = async { + match events.as_mut() { + Some(receiver) => receiver.recv().await, + None => std::future::pending().await, + } + } => { + let Some(event) = event else { + // The management module dropped its sender; nothing more to forward. + *events = None; + continue; + }; + socket + .send(Message::Text(event.to_string())) + .await?; + // Session lifecycle events are followed by a full snapshot so + // the aggregating side stays consistent without polling. + if event["event"]["kind"].as_str().is_some_and(|kind| { + kind == "session_opened" || kind == "session_closed" + }) { + let snapshot = json!({ + "type": "sessions", + "data": management::sessions_snapshot(), + }); + socket.send(Message::Text(snapshot.to_string())).await?; + } + } + inbound = socket.next() => { + match inbound { + Some(Ok(Message::Text(text))) => { + if let Err(err) = handle_command(&text, socket).await { + log::warn!("hbbs control command failed: {err}"); + } + } + Some(Ok(_)) => {} + _ => return Ok(()), + } + } + } + } +} + +async fn handle_command(text: &str, socket: &mut ControlSocket) -> ResultType<()> { + let message: Value = serde_json::from_str(text)?; + let id = message["id"].clone(); + match message["type"].as_str() { + Some("bans") => { + let bans: management::Bans = serde_json::from_value(message["bans"].clone())?; + // The aggregating hbbs owns the policy; relay-side persistence keeps + // the policy active across relay restarts even if hbbs is down. + management::apply_bans_update("hbbr", bans).await?; + } + Some("get_sessions") => { + let snapshot = json!({ + "type": "sessions", + "data": management::sessions_snapshot(), + }); + socket.send(Message::Text(snapshot.to_string())).await?; + } + Some("disconnect") => { + if let Some(uuid) = message["uuid"].as_str() { + management::disconnect_session(uuid); + } + } + // The remaining commands are proxied management requests: hbbs + // forwards relay-scoped reads and writes and awaits the matching + // response, so each command answers with an HTTP-equivalent status + // and body. + Some("get_config") => { + respond(socket, &id, 200, management::config_view("hbbr")).await?; + } + Some("set_config") => { + let update = match serde_json::from_value::(json!({ + "values": message["values"].clone() + })) { + Ok(update) => update, + Err(err) => { + respond(socket, &id, 400, json!({"error": err.to_string()})).await?; + return Ok(()); + } + }; + if let Some(err) = management::validate_update("hbbr", &update) { + respond(socket, &id, 400, json!({"error": err})).await?; + return Ok(()); + } + match management::apply_config_update("hbbr", update).await { + Ok((applied, pending_restart)) => { + respond( + socket, + &id, + 200, + json!({"state":"saved","applied":applied,"pending_restart":pending_restart}), + ) + .await?; + } + Err(err) => { + log::error!("Unable to persist managed configuration: {err}"); + respond(socket, &id, 500, json!({"error": err.to_string()})).await?; + } + } + } + Some("get_logs") => { + let tail = message["tail"] + .as_u64() + .unwrap_or(200) + .min(management::LOG_CAPACITY as u64) as usize; + respond(socket, &id, 200, management::logs_view(tail)).await?; + } + Some("restart") => { + respond(socket, &id, 202, json!({"state":"restarting"})).await?; + schedule_restart(); + } + _ => {} + } + Ok(()) +} + +// Answers a proxied request on the same control socket; hbbs correlates the +// exchange by request id and forwards the status and body to its client. +async fn respond( + socket: &mut ControlSocket, + id: &Value, + status: u16, + body: Value, +) -> ResultType<()> { + let response = json!({ + "type": "response", + "id": id, + "status": status, + "body": body, + }); + socket.send(Message::Text(response.to_string())).await?; + Ok(()) +} + +// Exits after the acknowledgement is flushed; the supervisor (container +// restart policy or systemd) brings the relay back with the persisted +// configuration and pending-restart entries applied. +fn schedule_restart() { + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(250)).await; + log::info!("Exiting for restart; the supervisor should bring the relay back"); + std::process::exit(0); + }); +} diff --git a/src/relay_server/server.rs b/src/relay_server/server.rs index 702f14a..78f8478 100644 --- a/src/relay_server/server.rs +++ b/src/relay_server/server.rs @@ -50,11 +50,20 @@ pub async fn start_with_bind(bind_addr: Option, port: &str, key: &str) - let websocket = crate::common::listen_tcp(bind_addr, port2).await?; let console = crate::common::listen_console(bind_addr, port).await?; if !management_started { - // Register the dynamic applier before the listener starts so - // early config updates apply immediately instead of being - // reported as pending restart. + // Register the dynamic applier before serving management + // requests so early config updates apply immediately instead + // of being reported as pending restart. crate::management::set_dynamic_apply(Box::new(command::apply_dynamic)); + // The managed state backs the control channel's proxied + // configuration requests even when no local listener is + // bound, so initialize it unconditionally. + crate::management::init_managed_state("hbbr", port)?; crate::management::start("hbbr", port)?; + let control_url = std::env::var("RD_HBBS_API_URL").unwrap_or_default(); + if !control_url.is_empty() { + let token = std::env::var("RD_MANAGEMENT_TOKEN").unwrap_or_default(); + crate::relay_control::start(&control_url, &token)?; + } management_started = true; } connection::io_loop(listener, websocket, console, &key).await; diff --git a/src/rendezvous_server/server.rs b/src/rendezvous_server/server.rs index 5b73fb6..1768df4 100644 --- a/src/rendezvous_server/server.rs +++ b/src/rendezvous_server/server.rs @@ -91,17 +91,47 @@ impl RendezvousServer { log::info!("local-ip: {:?}", rs.inner.local_ip); std::env::set_var("PORT_FOR_API", port.to_string()); rs.parse_relay_servers(&get_arg("relay-servers")); - // Register the hbbs handlers for schema entries marked as dynamic. + // Register the hbbs handlers for schema entries marked as dynamic, and + // fold relay addresses discovered through the control channel back + // into the live relay server list (static configuration plus + // registered relays). { let tx = rs.tx.clone(); - crate::management::set_dynamic_apply(Box::new(move |name, value| match name { - "always-use-relay" => { - let enabled = value.eq_ignore_ascii_case("Y"); - ALWAYS_USE_RELAY.store(enabled, std::sync::atomic::Ordering::SeqCst); - true + let static_relays = std::sync::Arc::new(std::sync::Mutex::new(get_arg( + "relay-servers", + ))); + let dynamic_base = static_relays.clone(); + let dynamic_tx = tx.clone(); + crate::management::set_dynamic_apply(Box::new( + move |name, value| match name { + "always-use-relay" => { + let enabled = value.eq_ignore_ascii_case("Y"); + ALWAYS_USE_RELAY.store(enabled, std::sync::atomic::Ordering::SeqCst); + true + } + "relay-servers" => { + // Dynamic updates replace the static baseline and are + // merged with the registered relay addresses. + *dynamic_base.lock().unwrap() = value.to_owned(); + let mut list = dynamic_base.lock().unwrap().clone(); + for address in crate::management::registered_relay_addresses() { + merge_relay_servers(&mut list, &address); + } + dynamic_tx.send(Data::RelayServers0(list)).is_ok() + } + _ => false, + }, + )); + let sync_tx = tx.clone(); + let sync_base = static_relays; + crate::management::set_relay_sync(Box::new(move |registered: Vec| { + let mut list = sync_base.lock().unwrap().clone(); + for address in registered { + merge_relay_servers(&mut list, &address); + } + if sync_tx.send(Data::RelayServers0(list)).is_err() { + log::warn!("Unable to update the live relay server list"); } - "relay-servers" => tx.send(Data::RelayServers0(value.to_owned())).is_ok(), - _ => false, })); } let mut listener = helper::create_tcp_listener(bind_addr, port).await?; @@ -292,3 +322,19 @@ impl RendezvousServer { } } } + +fn merge_relay_servers(list: &mut String, extra: &str) { + if extra.is_empty() { + return; + } + let mut entries: Vec = list + .split(',') + .map(str::trim) + .filter(|entry| !entry.is_empty()) + .map(str::to_owned) + .collect(); + if !entries.iter().any(|entry| entry == extra) { + entries.push(extra.to_owned()); + } + *list = entries.join(","); +}