First Version
This commit is contained in:
@@ -0,0 +1,361 @@
|
||||
//! The permanent agent WebSocket toward the relay, speaking **v2 protobuf**
|
||||
//! (docs/relay/relay-protocol.md).
|
||||
//!
|
||||
//! A single WS carries everything: challenge-response auth, the `Authorize` set,
|
||||
//! outbound E2E `Message`s, and inbound `Message` / `ClientPaired` frames. v2
|
||||
//! transport is **binary-only**: every wire frame is a `RelayFrame` protobuf
|
||||
//! message wrapped in `Message::Binary`; WS-level `Ping`/`Pong`/`Close` are
|
||||
//! their own `WsMessage` variants and never appear as protobuf.
|
||||
//!
|
||||
//! Reconnection uses exponential backoff (1,2,4,…,60 s) with jitter, and the
|
||||
//! whole loop is cancellable on stop.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use prost::Message as _;
|
||||
use rand::Rng;
|
||||
use skald_relay_common::crypto;
|
||||
use skald_relay_common::proto::v2::*;
|
||||
use skald_relay_common::proto::v2::relay_frame::Frame;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_tungstenite::tungstenite::Message as WsMessage;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
use crate::state::RelayState;
|
||||
|
||||
/// Run the reconnecting WS loop until `cancel` fires (relay-protocol.md §8).
|
||||
pub(crate) async fn run_loop(
|
||||
state: Arc<RelayState>,
|
||||
mut outbound_rx: mpsc::UnboundedReceiver<Vec<u8>>,
|
||||
cancel: CancellationToken,
|
||||
) {
|
||||
let mut backoff_step: u32 = 0;
|
||||
loop {
|
||||
if cancel.is_cancelled() {
|
||||
return;
|
||||
}
|
||||
|
||||
match connect_once(&state, &mut outbound_rx, &cancel).await {
|
||||
Ok(()) => {
|
||||
// Clean disconnect (cancelled or graceful): reset backoff.
|
||||
backoff_step = 0;
|
||||
}
|
||||
Err(e) => {
|
||||
warn!(crate_name = "skald-relay-client", error = %e, "relay connection ended");
|
||||
}
|
||||
}
|
||||
|
||||
if cancel.is_cancelled() {
|
||||
return;
|
||||
}
|
||||
|
||||
let delay = backoff_delay(backoff_step);
|
||||
backoff_step = backoff_step.saturating_add(1);
|
||||
state.set_connected(false);
|
||||
debug!(crate_name = "skald-relay-client", secs = delay.as_secs_f64(), "reconnect backoff");
|
||||
tokio::select! {
|
||||
_ = cancel.cancelled() => return,
|
||||
_ = tokio::time::sleep(delay) => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Backoff schedule 1,2,4,…,60 s plus up to 50% jitter (relay-protocol.md §8).
|
||||
fn backoff_delay(step: u32) -> Duration {
|
||||
let base = 1u64.checked_shl(step).unwrap_or(60).min(60);
|
||||
let jitter_ms = rand::rng().random_range(0..=(base * 500));
|
||||
Duration::from_millis(base * 1000 + jitter_ms)
|
||||
}
|
||||
|
||||
/// One full connection lifecycle: connect → challenge → auth → authorize → loop.
|
||||
async fn connect_once(
|
||||
state: &Arc<RelayState>,
|
||||
outbound_rx: &mut mpsc::UnboundedReceiver<Vec<u8>>,
|
||||
cancel: &CancellationToken,
|
||||
) -> Result<()> {
|
||||
let url = state.relay_url();
|
||||
info!(crate_name = "skald-relay-client", %url, "connecting to relay");
|
||||
|
||||
let (ws_stream, _resp) = tokio::select! {
|
||||
_ = cancel.cancelled() => return Ok(()),
|
||||
r = tokio_tungstenite::connect_async(&url) => r?,
|
||||
};
|
||||
let (mut sink, mut stream) = ws_stream.split();
|
||||
|
||||
// 1. Wait for the relay's challenge (it speaks first, relay-protocol.md §4).
|
||||
let challenge_nonce = wait_for_challenge(&mut stream).await?;
|
||||
|
||||
// 2. Sign AUTH_DOMAIN ‖ 0x00 ‖ nonce and send the agent Auth frame.
|
||||
let sig = crypto::sign_challenge(&state.identity().signing_key(), &challenge_nonce);
|
||||
let auth = RelayFrame {
|
||||
frame: Some(Frame::Auth(Auth {
|
||||
role: Some(auth::Role::Agent(AuthAgent {
|
||||
agent_ed25519_pub: prost::bytes::Bytes::copy_from_slice(
|
||||
&state.identity().ed25519_pub(),
|
||||
),
|
||||
})),
|
||||
signature: prost::bytes::Bytes::copy_from_slice(&sig),
|
||||
})),
|
||||
};
|
||||
sink.send(WsMessage::Binary(auth.encode_to_vec().into())).await?;
|
||||
|
||||
// 3. Expect AuthOk and verify the namespace_id locally.
|
||||
let ns_raw = wait_for_auth_ok(&mut stream).await?;
|
||||
if ns_raw != state.identity().namespace_id_raw() {
|
||||
return Err(anyhow!(
|
||||
"relay returned mismatched namespace_id (got {}, expected {})",
|
||||
hex::encode(ns_raw),
|
||||
hex::encode(state.identity().namespace_id_raw())
|
||||
));
|
||||
}
|
||||
info!(crate_name = "skald-relay-client", "relay auth ok, namespace verified");
|
||||
state.set_connected(true);
|
||||
|
||||
// 4. Send the current authorize set from the DB (empty on first run).
|
||||
// We push it directly via the sink rather than through `outbound_rx` so it
|
||||
// lands immediately — the queue is only drained inside the main loop below.
|
||||
let authorized = state.authorized_pubkeys_hex().await.unwrap_or_default();
|
||||
let clients: Vec<prost::bytes::Bytes> = authorized
|
||||
.iter()
|
||||
.filter_map(|h| hex::decode(h).ok())
|
||||
.map(prost::bytes::Bytes::from)
|
||||
.collect();
|
||||
let authorize = RelayFrame {
|
||||
frame: Some(Frame::Authorize(Authorize { clients })),
|
||||
};
|
||||
sink.send(WsMessage::Binary(authorize.encode_to_vec().into())).await?;
|
||||
|
||||
// 5. Main dispatch loop: outbound queue, inbound frames, WS-level Ping/Pong.
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = cancel.cancelled() => {
|
||||
let _ = sink.send(WsMessage::Close(None)).await;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Outbound: already-encoded protobuf frames queued by pairing / send
|
||||
// / revoke. The channel carries `Vec<u8>` ready to be shipped as a
|
||||
// binary WS frame.
|
||||
maybe = outbound_rx.recv() => {
|
||||
match maybe {
|
||||
Some(bytes) => sink.send(WsMessage::Binary(bytes.into())).await?,
|
||||
None => return Ok(()), // channel closed → client stopping
|
||||
}
|
||||
}
|
||||
|
||||
// Inbound: relay → agent frames.
|
||||
maybe = stream.next() => {
|
||||
let Some(msg) = maybe else { return Ok(()) }; // stream ended
|
||||
match msg? {
|
||||
WsMessage::Binary(data) => {
|
||||
handle_incoming(state, &data).await;
|
||||
}
|
||||
WsMessage::Ping(p) => sink.send(WsMessage::Pong(p)).await?,
|
||||
WsMessage::Pong(_) => {}
|
||||
WsMessage::Close(_) => return Ok(()),
|
||||
WsMessage::Text(_) | WsMessage::Frame(_) => {
|
||||
// v2 transport is binary-only; ignore text/frame
|
||||
// variants (forward-compat, no protocol-defined reaction).
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Read binary frames until `Challenge` arrives; returns the raw 32-byte nonce.
|
||||
async fn wait_for_challenge<S>(stream: &mut S) -> Result<[u8; 32]>
|
||||
where
|
||||
S: StreamExt<Item = Result<WsMessage, tokio_tungstenite::tungstenite::Error>> + Unpin,
|
||||
{
|
||||
while let Some(msg) = stream.next().await {
|
||||
match msg? {
|
||||
WsMessage::Binary(data) => {
|
||||
let frame = RelayFrame::decode(&data[..])?;
|
||||
if let Some(Frame::Challenge(c)) = frame.frame {
|
||||
if c.nonce.len() != 32 {
|
||||
return Err(anyhow!("challenge nonce is not 32 bytes"));
|
||||
}
|
||||
let mut out = [0u8; 32];
|
||||
out.copy_from_slice(&c.nonce);
|
||||
return Ok(out);
|
||||
}
|
||||
}
|
||||
WsMessage::Close(_) => return Err(anyhow!("closed before challenge")),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Err(anyhow!("connection closed before challenge"))
|
||||
}
|
||||
|
||||
/// Read binary frames until `AuthOk`; returns the raw 32-byte namespace_id.
|
||||
async fn wait_for_auth_ok<S>(stream: &mut S) -> Result<[u8; 32]>
|
||||
where
|
||||
S: StreamExt<Item = Result<WsMessage, tokio_tungstenite::tungstenite::Error>> + Unpin,
|
||||
{
|
||||
while let Some(msg) = stream.next().await {
|
||||
match msg? {
|
||||
WsMessage::Binary(data) => {
|
||||
let frame = RelayFrame::decode(&data[..])?;
|
||||
match frame.frame {
|
||||
Some(Frame::AuthOk(AuthOk { namespace_id })) => {
|
||||
if namespace_id.len() != 32 {
|
||||
return Err(anyhow!("namespace_id is not 32 bytes"));
|
||||
}
|
||||
let mut out = [0u8; 32];
|
||||
out.copy_from_slice(&namespace_id);
|
||||
return Ok(out);
|
||||
}
|
||||
Some(Frame::AuthError(AuthError { code, message })) => {
|
||||
return Err(anyhow!("auth_error from relay: {code} ({message})"));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
WsMessage::Close(_) => return Err(anyhow!("closed before auth_ok")),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Err(anyhow!("connection closed before auth_ok"))
|
||||
}
|
||||
|
||||
/// Dispatch one decoded relay→agent `RelayFrame`. WS-level Ping/Pong are
|
||||
/// handled at the transport layer above; everything that arrives as a binary
|
||||
/// frame is decoded to `RelayFrame` and matched on the `Frame` oneof here.
|
||||
async fn handle_incoming(state: &Arc<RelayState>, data: &[u8]) {
|
||||
let frame = match RelayFrame::decode(data) {
|
||||
Ok(f) => f,
|
||||
Err(e) => {
|
||||
warn!(crate_name = "skald-relay-client", error = %e, "malformed protobuf frame dropped");
|
||||
return;
|
||||
}
|
||||
};
|
||||
let Some(f) = frame.frame else {
|
||||
debug!(crate_name = "skald-relay-client", "empty relay frame dropped");
|
||||
return;
|
||||
};
|
||||
match f {
|
||||
Frame::Message(m) => {
|
||||
// Validate lengths before handing off to the E2E layer.
|
||||
if m.peer.len() != 32 || m.nonce.len() != 12 {
|
||||
warn!(crate_name = "skald-relay-client", "message with wrong peer/nonce length dropped");
|
||||
return;
|
||||
}
|
||||
let mut from = [0u8; 32];
|
||||
from.copy_from_slice(&m.peer);
|
||||
let mut nonce = [0u8; 12];
|
||||
nonce.copy_from_slice(&m.nonce);
|
||||
state.handle_inbound_message(&from, &nonce, &m.ciphertext, m.live).await;
|
||||
}
|
||||
Frame::ClientPaired(cp) => {
|
||||
if cp.client_ed25519_pub.len() != 32 || cp.client_x25519_pub.len() != 32 {
|
||||
warn!(crate_name = "skald-relay-client", "client_paired with wrong pubkey length dropped");
|
||||
return;
|
||||
}
|
||||
let mut ed = [0u8; 32];
|
||||
ed.copy_from_slice(&cp.client_ed25519_pub);
|
||||
let mut x = [0u8; 32];
|
||||
x.copy_from_slice(&cp.client_x25519_pub);
|
||||
// Decode the protobuf `Platform` enum to the lowercase string the DB
|
||||
// expects. The wire value defaults to `0` (`UNSPECIFIED`) — the helper
|
||||
// maps that to `"unknown"`.
|
||||
let platform = platform_i32_to_str(cp.platform);
|
||||
state.handle_client_paired(&ed, &x, platform).await;
|
||||
}
|
||||
Frame::AuthorizeOk(aok) => {
|
||||
debug!(crate_name = "skald-relay-client", authorized = aok.authorized, "authorize_ok");
|
||||
}
|
||||
Frame::PairingReady(_) | Frame::PairingStopOk(_) => {}
|
||||
Frame::PresenceEvent(pe) => {
|
||||
debug!(
|
||||
crate_name = "skald-relay-client",
|
||||
pubkey = %hex::encode(&pe.pubkey),
|
||||
status = pe.status,
|
||||
"presence event"
|
||||
);
|
||||
}
|
||||
Frame::PresenceList(pl) => {
|
||||
debug!(crate_name = "skald-relay-client", online = pl.online.len(), "presence list");
|
||||
}
|
||||
Frame::PeerOffline(po) => {
|
||||
// Expected backstop for route-or-fail live sends (relay-protocol.md
|
||||
// §3): a `live=true` send found the peer gone. A normal protocol
|
||||
// event, not an error.
|
||||
debug!(
|
||||
crate_name = "skald-relay-client",
|
||||
peer = %hex::encode(&po.peer),
|
||||
"peer offline for live send; dropping"
|
||||
);
|
||||
}
|
||||
Frame::Error(e) => {
|
||||
warn!(crate_name = "skald-relay-client", code = %e.code, message = %e.message, "relay error frame");
|
||||
}
|
||||
// Server-to-client or handshake frames the agent never expects inbound.
|
||||
Frame::Challenge(_)
|
||||
| Frame::Auth(_)
|
||||
| Frame::AuthOk(_)
|
||||
| Frame::AuthError(_)
|
||||
| Frame::Authorize(_)
|
||||
| Frame::PairingStart(_)
|
||||
| Frame::PairingStop(_)
|
||||
| Frame::PresenceRequest(_) => {
|
||||
warn!(crate_name = "skald-relay-client", "unexpected relay→agent frame dropped");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Map a protobuf `Platform` enum wire value to the lowercase string the DB
|
||||
/// stores in the `platform` column. Unknown values become `"unknown"`.
|
||||
fn platform_i32_to_str(v: i32) -> &'static str {
|
||||
if v == Platform::Ios as i32 {
|
||||
"ios"
|
||||
} else if v == Platform::Android as i32 {
|
||||
"android"
|
||||
} else {
|
||||
"unknown"
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// `platform_i32_to_str` is total on the wire values the relay emits and
|
||||
/// never panics on bogus inputs (relay-protocol.md §11 forward-compat).
|
||||
#[test]
|
||||
fn platform_conversion() {
|
||||
assert_eq!(platform_i32_to_str(0), "unknown");
|
||||
assert_eq!(platform_i32_to_str(1), "ios");
|
||||
assert_eq!(platform_i32_to_str(2), "android");
|
||||
assert_eq!(platform_i32_to_str(99), "unknown");
|
||||
}
|
||||
|
||||
/// A minimal `Message` frame round-trips through `prost` so the wire
|
||||
/// encoding we emit is the same one the relay will decode.
|
||||
#[test]
|
||||
fn message_frame_round_trip() {
|
||||
let frame = RelayFrame {
|
||||
frame: Some(Frame::Message(Message {
|
||||
ciphertext: vec![0xAA; 64].into(),
|
||||
nonce: vec![0x01; 12].into(),
|
||||
peer: vec![0x02; 32].into(),
|
||||
live: false,
|
||||
})),
|
||||
};
|
||||
let bytes = frame.encode_to_vec();
|
||||
let decoded = RelayFrame::decode(&bytes[..]).expect("decode");
|
||||
match decoded.frame {
|
||||
Some(Frame::Message(m)) => {
|
||||
assert_eq!(m.ciphertext.len(), 64);
|
||||
assert_eq!(m.nonce.len(), 12);
|
||||
assert_eq!(m.peer.len(), 32);
|
||||
assert!(!m.live);
|
||||
}
|
||||
other => panic!("expected Message, got {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user