First Version
This commit is contained in:
@@ -0,0 +1,85 @@
|
||||
[package]
|
||||
name = "skald-relay-server"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
description = "Skald Remote Control relay: zero-trust store-and-forward + push bridge (see data/ios-app/relay.md)"
|
||||
license = "MIT"
|
||||
|
||||
# Main binary (the relay). Deploy entrypoint is /skald-relay-server.
|
||||
# Frame types + crypto live in the shared `skald-relay-common` crate; the
|
||||
# `gen-vectors` reference generator (crypto interop test vectors,
|
||||
# data/ios-app/test-vectors.md §3) moved there too (run with
|
||||
# `cargo run -p skald-relay-common --bin gen-vectors`).
|
||||
|
||||
[dependencies]
|
||||
# --- shared frame types + crypto (frames + verify/namespace subset) ---
|
||||
skald-relay-common = { path = "../skald-relay-common" }
|
||||
|
||||
# --- async runtime + web/WS ---
|
||||
axum = { version = "0.8", features = ["ws"] }
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
tokio-util = { version = "0.7", features = ["rt"] }
|
||||
futures-util = "0.3"
|
||||
|
||||
# --- serde / frames ---
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
async-trait = "0.1"
|
||||
|
||||
# --- persistence ---
|
||||
sqlx = { version = "0.9.0", features = ["runtime-tokio", "sqlite"] }
|
||||
|
||||
# --- encoding used directly in the WS handler ---
|
||||
hex = "0.4"
|
||||
base64 = "0.22"
|
||||
rand = "0.8" # CSPRNG for the 32-byte challenge nonce
|
||||
|
||||
# --- v2 wire format (data/ios-app/v2/relay-protocol.md) ---
|
||||
# `prost`: matches the major used by `skald-relay-common` so the
|
||||
# generated `proto::v2::RelayFrame` types here are the same compile-time
|
||||
# ones the common crate re-exports. `bytes`: `prost` 0.13 generates
|
||||
# `bytes` fields as `prost::bytes::Bytes`; we use the crate directly in
|
||||
# the WS handler to build zero-copy frame payloads.
|
||||
prost = "0.13"
|
||||
bytes = "1"
|
||||
|
||||
# --- observability + misc ---
|
||||
anyhow = "1"
|
||||
tracing = "0.1"
|
||||
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
|
||||
tracing-appender = "0.2"
|
||||
chrono = { version = "0.4", default-features = false, features = ["clock", "std"] }
|
||||
|
||||
# --- push-live senders (optional, off by default) ---
|
||||
# reqwest: HTTP/2 client for APNs (`http2` enables ALPN h2, `rustls-no-provider`
|
||||
# uses rustls without pulling a crypto provider — the binary installs `ring`
|
||||
# as the process default, avoiding any OpenSSL/aws-lc C build; `json` matches
|
||||
# the body type).
|
||||
# jsonwebtoken: ES256 JWT signing for APNs provider auth tokens.
|
||||
# `use_pem` lets us load the .p8 PEM directly; `p256` covers Apple's ES256
|
||||
# algorithm (the only curve APNs accepts) without pulling a separate crypto
|
||||
# crate.
|
||||
# uuid: APNs requires a unique `apns-id` per message so retries are de-duped.
|
||||
reqwest = { version = "0.13", default-features = false, features = ["http2", "json", "rustls-no-provider"], optional = true }
|
||||
jsonwebtoken = { version = "9", features = ["use_pem"], optional = true }
|
||||
uuid = { version = "1", features = ["v4"], optional = true }
|
||||
|
||||
[features]
|
||||
default = []
|
||||
# Live push senders (need real APNs/FCM credentials). Off by default so the
|
||||
# relay builds & runs without external creds; the normative decision/payload
|
||||
# logic is always compiled and unit-tested. See push.rs.
|
||||
push-live = ["dep:reqwest", "dep:jsonwebtoken", "dep:uuid"]
|
||||
|
||||
[dev-dependencies]
|
||||
# WebSocket client used by the protocol integration tests (tests/protocol.rs).
|
||||
tokio-tungstenite = "0.29"
|
||||
# Sign challenges / compute namespace_id from the test harness side.
|
||||
ed25519-dalek = "2"
|
||||
sha2 = "0.10"
|
||||
# The integration test constructs the same protobuf `RelayFrame`s the relay
|
||||
# sends on the wire (data/ios-app/v2/relay-protocol.md) — prost matches the
|
||||
# generator used by `skald-relay-common`, bytes is the `Bytes` type prost 0.13
|
||||
# generates for `bytes` fields.
|
||||
prost = "0.13"
|
||||
bytes = "1"
|
||||
@@ -0,0 +1,12 @@
|
||||
//! The relay's cryptographic operations: verifying the Ed25519 challenge
|
||||
//! signature (crypto.md §8) and deriving the `namespace_id` (crypto.md §7).
|
||||
//!
|
||||
//! The implementation now lives in the shared `skald-relay-common` crate so the
|
||||
//! relay and the mobile-connector plugin can never diverge (see plugin.md §1.1).
|
||||
//! The relay uses only the verify/namespace subset; the full E2E suite is end-to-
|
||||
//! end between agent and client and not touched here. Re-exported so existing
|
||||
//! relay paths (`crate::auth::…`) keep working unchanged.
|
||||
|
||||
pub use skald_relay_common::crypto::{
|
||||
AUTH_DOMAIN, NS_DOMAIN, ct_eq, decode_hex, namespace_id, verify_challenge,
|
||||
};
|
||||
@@ -0,0 +1,167 @@
|
||||
//! Runtime configuration, read from the environment (relay.md §7). Sensible
|
||||
//! defaults so the relay boots with zero config in local dev.
|
||||
//!
|
||||
//! | Env var | Meaning | Default |
|
||||
//! |------------------|------------------------------------------------------|-------------------------------|
|
||||
//! | `RELAY_BIND` | full `ip:port` to listen on | `0.0.0.0:8080` |
|
||||
//! | `PORT` | port only (used if no RELAY_BIND) | — |
|
||||
//! | `RELAY_DB` | SQLite file path | `data/relay.db` |
|
||||
//! | `APNS_KEY_PATH` | (push-live) JSON file with team/key/PEM | `./config/apns-key.json` |
|
||||
//! | `APNS_BUNDLE_ID` | (push-live) iOS bundle id (used as `apns-topic`) | — (required when push-live) |
|
||||
//! | `APNS_SANDBOX` | (push-live) `1`/`true` → api.sandbox.push.apple.com | `0` (production) |
|
||||
|
||||
use std::net::SocketAddr;
|
||||
|
||||
use crate::limits::{
|
||||
PIPE_IDLE_TIMEOUT_SECS, PIPE_MAX_BPS_DEFAULT, PIPE_MAX_FRAME_BYTES, PIPE_MAX_PER_NS,
|
||||
PIPE_PENDING_TTL_SECS,
|
||||
};
|
||||
|
||||
/// Tunables for the `/v1/pipe` data plane (docs/relay/pipe.md §2.3). Defaults in
|
||||
/// [`crate::limits`]; each is overridable via the matching `RELAY_PIPE_*` env var.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PipeConfig {
|
||||
/// Per-connection, per-direction bandwidth cap in bytes/sec. `0` = unlimited.
|
||||
pub max_bps: u64,
|
||||
/// Max concurrent pipes (pending + matched) per namespace.
|
||||
pub max_per_ns: usize,
|
||||
/// Half-open pending TTL (seconds).
|
||||
pub pending_ttl_secs: u64,
|
||||
/// Idle (no-bytes) timeout on a matched pipe (seconds).
|
||||
pub idle_timeout_secs: u64,
|
||||
/// Max data-plane WS frame size (bytes).
|
||||
pub max_frame_bytes: usize,
|
||||
}
|
||||
|
||||
impl Default for PipeConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
max_bps: PIPE_MAX_BPS_DEFAULT,
|
||||
max_per_ns: PIPE_MAX_PER_NS,
|
||||
pending_ttl_secs: PIPE_PENDING_TTL_SECS,
|
||||
idle_timeout_secs: PIPE_IDLE_TIMEOUT_SECS,
|
||||
max_frame_bytes: PIPE_MAX_FRAME_BYTES,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl PipeConfig {
|
||||
fn from_env() -> PipeConfig {
|
||||
fn env_parse<T: std::str::FromStr>(key: &str, default: T) -> T {
|
||||
std::env::var(key).ok().and_then(|s| s.parse().ok()).unwrap_or(default)
|
||||
}
|
||||
let d = PipeConfig::default();
|
||||
PipeConfig {
|
||||
max_bps: env_parse("RELAY_PIPE_MAX_BPS", d.max_bps),
|
||||
max_per_ns: env_parse("RELAY_PIPE_MAX_PER_NS", d.max_per_ns),
|
||||
pending_ttl_secs: env_parse("RELAY_PIPE_PENDING_TTL_SECS", d.pending_ttl_secs),
|
||||
idle_timeout_secs: env_parse("RELAY_PIPE_IDLE_TIMEOUT_SECS", d.idle_timeout_secs),
|
||||
max_frame_bytes: env_parse("RELAY_PIPE_MAX_FRAME_BYTES", d.max_frame_bytes),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// APNs configuration, populated from `config/apns-key.json` and env vars when
|
||||
/// the `push-live` cargo feature is on. The PEM is already newline-decoded by
|
||||
/// `serde_json` so it can be passed straight to `jsonwebtoken`.
|
||||
#[cfg(feature = "push-live")]
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ApnsConfig {
|
||||
pub team_id: String,
|
||||
pub key_id: String,
|
||||
/// PEM-encoded PKCS#8 EC private key (P-256), with real newlines.
|
||||
pub private_key_pem: String,
|
||||
/// Bundle ID for the iOS app (used as `apns-topic`).
|
||||
pub bundle_id: String,
|
||||
/// If true, send to api.sandbox.push.apple.com.
|
||||
pub sandbox: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Config {
|
||||
pub bind: SocketAddr,
|
||||
pub db_path: String,
|
||||
/// Pipe data-plane tunables (docs/relay/pipe.md §2.3).
|
||||
pub pipe: PipeConfig,
|
||||
/// `None` ⇒ the relay falls back to [`LogPusher`] (relay still boots).
|
||||
#[cfg(feature = "push-live")]
|
||||
pub apns: Option<ApnsConfig>,
|
||||
}
|
||||
|
||||
impl Config {
|
||||
pub fn from_env() -> Config {
|
||||
let default_bind: SocketAddr = "0.0.0.0:8080".parse().unwrap();
|
||||
let bind = std::env::var("RELAY_BIND")
|
||||
.ok()
|
||||
.and_then(|s| s.parse().ok())
|
||||
.or_else(|| {
|
||||
std::env::var("PORT")
|
||||
.ok()
|
||||
.and_then(|p| format!("0.0.0.0:{p}").parse().ok())
|
||||
})
|
||||
.unwrap_or(default_bind);
|
||||
let db_path = std::env::var("RELAY_DB").unwrap_or_else(|_| "data/relay.db".into());
|
||||
Config {
|
||||
bind,
|
||||
db_path,
|
||||
pipe: PipeConfig::from_env(),
|
||||
#[cfg(feature = "push-live")]
|
||||
apns: ApnsConfig::load_from_env(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "push-live")]
|
||||
impl ApnsConfig {
|
||||
/// Read `APNS_KEY_PATH` (default `./config/apns-key.json`), `APNS_BUNDLE_ID`,
|
||||
/// and `APNS_SANDBOX`. Returns `None` if anything required is missing — the
|
||||
/// caller in `AppState::build` logs a generic warning and falls back to
|
||||
/// `LogPusher`.
|
||||
pub fn load_from_env() -> Option<ApnsConfig> {
|
||||
let path = std::env::var("APNS_KEY_PATH")
|
||||
.unwrap_or_else(|_| "./config/apns-key.json".into());
|
||||
let raw = match std::fs::read_to_string(&path) {
|
||||
Ok(s) => s,
|
||||
Err(_) => return None, // silent: caller will warn at fallback point
|
||||
};
|
||||
#[derive(serde::Deserialize)]
|
||||
struct KeyFile {
|
||||
team_id: String,
|
||||
key_id: String,
|
||||
private_key: String,
|
||||
}
|
||||
let parsed: KeyFile = match serde_json::from_str(&raw) {
|
||||
Ok(k) => k,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
target: "relay::push",
|
||||
path = %path,
|
||||
error = %e,
|
||||
"apns key file is not valid JSON; APNs disabled"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let bundle_id = match std::env::var("APNS_BUNDLE_ID") {
|
||||
Ok(b) if !b.is_empty() => b,
|
||||
_ => {
|
||||
tracing::warn!(
|
||||
target: "relay::push",
|
||||
"APNS_BUNDLE_ID not set; APNs disabled"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let sandbox = matches!(
|
||||
std::env::var("APNS_SANDBOX").ok().as_deref(),
|
||||
Some("1") | Some("true") | Some("TRUE")
|
||||
);
|
||||
Some(ApnsConfig {
|
||||
team_id: parsed.team_id,
|
||||
key_id: parsed.key_id,
|
||||
private_key_pem: parsed.private_key,
|
||||
bundle_id,
|
||||
sandbox,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,182 @@
|
||||
//! Skald Remote Control relay server library (see data/ios-app/relay.md).
|
||||
//!
|
||||
//! Exposes the building blocks — [`AppState`], [`router`], [`spawn_gc`] — so the
|
||||
//! `main` binary stays thin and integration tests can spin up the real server on
|
||||
//! an ephemeral port.
|
||||
|
||||
pub mod auth;
|
||||
pub mod config;
|
||||
pub mod limits;
|
||||
pub mod pipe;
|
||||
pub mod push;
|
||||
pub mod routing;
|
||||
pub mod store;
|
||||
pub mod types;
|
||||
pub mod ws;
|
||||
|
||||
use std::net::{IpAddr, SocketAddr};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use axum::Router;
|
||||
use axum::extract::{ConnectInfo, State, ws::WebSocketUpgrade};
|
||||
use axum::http::StatusCode;
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use axum::routing::get;
|
||||
|
||||
use config::Config;
|
||||
use limits::{FixedWindow, IP_NEW_CONN_PER_MIN, TRANSPORT_FRAME_CAP, TTL_DAYS};
|
||||
use pipe::PipeRegistry;
|
||||
use push::{LogPusher, Pusher};
|
||||
use routing::Registry;
|
||||
use store::Store;
|
||||
|
||||
/// Shared, cheaply-cloneable application state handed to every connection.
|
||||
#[derive(Clone)]
|
||||
pub struct AppState {
|
||||
pub store: Store,
|
||||
pub registry: Arc<Registry>,
|
||||
/// Stateful proxy registry for `/v1/pipe` (docs/relay/pipe.md §2).
|
||||
pub pipes: Arc<PipeRegistry>,
|
||||
pub ip_limiter: Arc<FixedWindow<IpAddr>>,
|
||||
pub pusher: Arc<dyn Pusher>,
|
||||
conn_seq: Arc<AtomicU64>,
|
||||
pub cfg: Arc<Config>,
|
||||
}
|
||||
|
||||
impl AppState {
|
||||
/// Build the full application state: open the store and wire the default
|
||||
/// (credential-free) push bridge.
|
||||
pub async fn build(cfg: Config) -> anyhow::Result<AppState> {
|
||||
// Ensure the DB directory exists (SQLite creates the file, not the dir).
|
||||
if let Some(parent) = std::path::Path::new(&cfg.db_path).parent()
|
||||
&& !parent.as_os_str().is_empty()
|
||||
{
|
||||
std::fs::create_dir_all(parent)?;
|
||||
}
|
||||
let store = Store::init(&cfg.db_path).await?;
|
||||
// Push bridge: live APNs when `push-live` is enabled and credentials
|
||||
// are present; otherwise the credential-free LogPusher (so the relay
|
||||
// still boots locally — see push.rs).
|
||||
#[cfg(feature = "push-live")]
|
||||
let pusher: Arc<dyn Pusher> = match cfg.apns.as_ref() {
|
||||
Some(apns) => push::build_pusher(apns),
|
||||
None => {
|
||||
tracing::warn!(
|
||||
target: "relay::push",
|
||||
"push-live feature enabled but no APNs config; falling back to LogPusher"
|
||||
);
|
||||
Arc::new(LogPusher)
|
||||
}
|
||||
};
|
||||
#[cfg(not(feature = "push-live"))]
|
||||
let pusher: Arc<dyn Pusher> = Arc::new(LogPusher);
|
||||
Ok(AppState {
|
||||
store,
|
||||
registry: Arc::new(Registry::new()),
|
||||
pipes: Arc::new(PipeRegistry::new()),
|
||||
ip_limiter: Arc::new(FixedWindow::new(
|
||||
Duration::from_secs(60),
|
||||
IP_NEW_CONN_PER_MIN,
|
||||
)),
|
||||
pusher,
|
||||
conn_seq: Arc::new(AtomicU64::new(1)),
|
||||
cfg: Arc::new(cfg),
|
||||
})
|
||||
}
|
||||
|
||||
/// Monotonic per-process connection id (used for safe self-removal).
|
||||
pub fn next_conn_id(&self) -> u64 {
|
||||
self.conn_seq.fetch_add(1, Ordering::Relaxed)
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the axum router: `GET /healthz`, the control WebSocket `GET /v1/ws`,
|
||||
/// and the data-plane pipe WebSocket `GET /v1/pipe` (docs/relay/pipe.md §2).
|
||||
pub fn router(state: AppState) -> Router {
|
||||
Router::new()
|
||||
.route("/healthz", get(healthz))
|
||||
.route("/v1/ws", get(ws_upgrade))
|
||||
.route("/v1/pipe", get(pipe_upgrade))
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
async fn healthz() -> &'static str {
|
||||
"ok"
|
||||
}
|
||||
|
||||
/// `GET /v1/pipe` → data-plane WebSocket upgrade. Same per-IP new-connection
|
||||
/// quota as `/v1/ws`; the per-message cap is the pipe frame size (bulk transfer).
|
||||
async fn pipe_upgrade(
|
||||
ws: WebSocketUpgrade,
|
||||
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
||||
State(state): State<AppState>,
|
||||
) -> Response {
|
||||
let ip = addr.ip();
|
||||
if !state.ip_limiter.allow(&ip) {
|
||||
tracing::warn!(%ip, "rate_limited: too many new pipe connections");
|
||||
return (StatusCode::TOO_MANY_REQUESTS, "rate_limited").into_response();
|
||||
}
|
||||
let max_frame = state.cfg.pipe.max_frame_bytes;
|
||||
ws.max_message_size(max_frame)
|
||||
.on_upgrade(move |socket| pipe::handle_pipe_socket(socket, state, ip))
|
||||
}
|
||||
|
||||
/// `GET /v1/ws` → WebSocket upgrade. Per-IP new-connection quota is enforced
|
||||
/// here (before upgrade) so unauthenticated floods are cheap to reject.
|
||||
async fn ws_upgrade(
|
||||
ws: WebSocketUpgrade,
|
||||
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
||||
State(state): State<AppState>,
|
||||
) -> Response {
|
||||
let ip = addr.ip();
|
||||
if !state.ip_limiter.allow(&ip) {
|
||||
tracing::warn!(%ip, "rate_limited: too many new connections");
|
||||
return (StatusCode::TOO_MANY_REQUESTS, "rate_limited").into_response();
|
||||
}
|
||||
ws.max_message_size(TRANSPORT_FRAME_CAP)
|
||||
.on_upgrade(move |socket| ws::handle_socket(socket, state, ip))
|
||||
}
|
||||
|
||||
/// Periodic garbage collection: drop messages/namespaces past TTL (relay.md §6)
|
||||
/// and prune the IP rate-limiter map.
|
||||
pub fn spawn_gc(state: AppState) {
|
||||
tokio::spawn(async move {
|
||||
let mut tick = tokio::time::interval(Duration::from_secs(3600));
|
||||
loop {
|
||||
tick.tick().await;
|
||||
match state.store.gc(TTL_DAYS).await {
|
||||
Ok((m, n)) if m > 0 || n > 0 => {
|
||||
tracing::info!(messages = m, namespaces = n, "gc removed expired rows");
|
||||
}
|
||||
Ok(_) => {}
|
||||
Err(e) => tracing::error!(error = %e, "gc failed"),
|
||||
}
|
||||
state.ip_limiter.prune();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/// Resolves when SIGINT or SIGTERM arrives (graceful shutdown trigger).
|
||||
pub async fn shutdown_signal() {
|
||||
let ctrl_c = async {
|
||||
let _ = tokio::signal::ctrl_c().await;
|
||||
};
|
||||
#[cfg(unix)]
|
||||
let terminate = async {
|
||||
if let Ok(mut sig) =
|
||||
tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
|
||||
{
|
||||
sig.recv().await;
|
||||
}
|
||||
};
|
||||
#[cfg(not(unix))]
|
||||
let terminate = std::future::pending::<()>();
|
||||
|
||||
tokio::select! {
|
||||
_ = ctrl_c => {},
|
||||
_ = terminate => {},
|
||||
}
|
||||
tracing::info!("shutdown signal received; draining");
|
||||
}
|
||||
@@ -0,0 +1,135 @@
|
||||
//! Normative quotas, timeouts and thresholds (relay-protocol.md §9, relay.md)
|
||||
//! plus a fixed-window rate limiter. The values here are the spec's reasonable
|
||||
//! defaults; the relay may expose them via config later.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::hash::Hash;
|
||||
use std::sync::Mutex;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
/// Maximum size of a WebSocket frame (64 KiB). Above this → `payload_too_large`.
|
||||
/// Applies to all pre-auth frames and to post-auth frames whose `Message` does
|
||||
/// not set the `live` flag (v2 spec §5).
|
||||
pub const MAX_FRAME_BYTES: usize = 64 * 1024;
|
||||
/// v2: max size of a WS binary frame carrying a `Message{live:true}` on an
|
||||
/// authenticated connection (relay-protocol.md §5: `MAX_LIVE_FRAME_BYTES =
|
||||
/// 524288`, i.e. 512 KiB exactly). The relay enforces this manually in `ws.rs`
|
||||
/// based on auth state + `Message.live`.
|
||||
pub const MAX_LIVE_FRAME_BYTES: usize = 512 * 1024;
|
||||
/// axum's per-message cap: must be at least `MAX_LIVE_FRAME_BYTES` so live
|
||||
/// frames can flow. The relay enforces the strict per-frame limits itself
|
||||
/// in `ws.rs` (64 KiB pre-auth, 64 KiB post-auth non-live, 512 KiB
|
||||
/// post-auth-live).
|
||||
pub const TRANSPORT_FRAME_CAP: usize = MAX_LIVE_FRAME_BYTES;
|
||||
|
||||
/// Time allowed to receive `auth` after the `challenge`.
|
||||
pub const CHALLENGE_TIMEOUT_SECS: u64 = 30;
|
||||
/// No traffic for this long → close the connection.
|
||||
pub const IDLE_TIMEOUT_SECS: u64 = 120;
|
||||
/// Keepalive ping interval.
|
||||
pub const PING_INTERVAL_SECS: u64 = 30;
|
||||
|
||||
/// Max queued messages per recipient; above this → `queue_full`.
|
||||
pub const QUEUE_MAX_PER_DEST: i64 = 200;
|
||||
/// Threshold (bytes of base64(ciphertext)) at/below which we do content-in-push.
|
||||
pub const CONTENT_PUSH_MAX_B64: usize = 3500;
|
||||
|
||||
/// TTL for the store-and-forward queue and for idle namespaces.
|
||||
pub const TTL_DAYS: i64 = 7;
|
||||
|
||||
/// Pairing window TTL.
|
||||
pub const PAIRING_TTL_DEFAULT: u64 = 300;
|
||||
pub const PAIRING_TTL_MAX: u64 = 600;
|
||||
|
||||
/// Anti-flood quotas on the public endpoint.
|
||||
pub const IP_NEW_CONN_PER_MIN: u32 = 30;
|
||||
pub const CONN_MSG_PER_MIN: u32 = 60;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Pipe data plane (docs/relay/pipe.md §2.3). The relay becomes a stateful
|
||||
// connection proxy for `/v1/pipe`; these bound its resource use. All are
|
||||
// overridable via `RELAY_PIPE_*` env vars (see config.rs).
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// First side dialed, second never showed → reap the half-open pending.
|
||||
pub const PIPE_PENDING_TTL_SECS: u64 = 30;
|
||||
/// No bytes for this long on a matched pipe → close (reclaim dead pipes).
|
||||
pub const PIPE_IDLE_TIMEOUT_SECS: u64 = 120;
|
||||
/// Max concurrent matched/pending pipes per namespace.
|
||||
pub const PIPE_MAX_PER_NS: usize = 8;
|
||||
/// Max size of one data-plane WS binary frame (bulk transfer; separate from the
|
||||
/// message-channel caps).
|
||||
pub const PIPE_MAX_FRAME_BYTES: usize = 1024 * 1024;
|
||||
/// Per-connection bandwidth cap in bytes/sec, **per direction**. `0` = unlimited.
|
||||
pub const PIPE_MAX_BPS_DEFAULT: u64 = 0;
|
||||
|
||||
/// Thread-safe fixed-window rate limiter, generic over the key.
|
||||
///
|
||||
/// One `allow()` per event: returns `false` when the current window's quota is
|
||||
/// exceeded. The window resets automatically once it elapses.
|
||||
pub struct FixedWindow<K: Eq + Hash + Clone> {
|
||||
window: Duration,
|
||||
max: u32,
|
||||
map: Mutex<HashMap<K, (Instant, u32)>>,
|
||||
}
|
||||
|
||||
impl<K: Eq + Hash + Clone> FixedWindow<K> {
|
||||
pub fn new(window: Duration, max: u32) -> Self {
|
||||
Self {
|
||||
window,
|
||||
max,
|
||||
map: Mutex::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Record an event for `key`. `true` = allowed, `false` = quota exceeded.
|
||||
pub fn allow(&self, key: &K) -> bool {
|
||||
let mut map = self.map.lock().unwrap();
|
||||
let now = Instant::now();
|
||||
let entry = map.entry(key.clone()).or_insert((now, 0));
|
||||
if now.duration_since(entry.0) >= self.window {
|
||||
*entry = (now, 0);
|
||||
}
|
||||
entry.1 += 1;
|
||||
entry.1 <= self.max
|
||||
}
|
||||
|
||||
/// Opportunistic pruning of expired windows (called by the GC task).
|
||||
pub fn prune(&self) {
|
||||
let mut map = self.map.lock().unwrap();
|
||||
let now = Instant::now();
|
||||
map.retain(|_, (start, _)| now.duration_since(*start) < self.window);
|
||||
}
|
||||
}
|
||||
|
||||
/// Per-connection (non-shared) rate counter: messages per minute.
|
||||
pub struct ConnRate {
|
||||
window_start: Instant,
|
||||
count: u32,
|
||||
}
|
||||
|
||||
impl ConnRate {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
window_start: Instant::now(),
|
||||
count: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// `true` if under quota, `false` if the connection exceeded `CONN_MSG_PER_MIN`.
|
||||
pub fn allow_message(&mut self) -> bool {
|
||||
let now = Instant::now();
|
||||
if now.duration_since(self.window_start) >= Duration::from_secs(60) {
|
||||
self.window_start = now;
|
||||
self.count = 0;
|
||||
}
|
||||
self.count += 1;
|
||||
self.count <= CONN_MSG_PER_MIN
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ConnRate {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
//! Skald Remote Control relay server — binary entrypoint (see data/ios-app/relay.md).
|
||||
//!
|
||||
//! Thin wrapper: load config → build [`AppState`] → start the GC task → serve
|
||||
//! the axum router until a shutdown signal arrives. All logic lives in the lib.
|
||||
|
||||
use std::net::SocketAddr;
|
||||
|
||||
use skald_relay_server::config::Config;
|
||||
use skald_relay_server::{AppState, router, shutdown_signal, spawn_gc};
|
||||
use tracing_subscriber::layer::SubscriberExt;
|
||||
use tracing_subscriber::util::SubscriberInitExt;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
// Persist logs to `logs/skald-relay.log` (rolling daily), mirroring the main
|
||||
// app, and also mirror to stdout for terminal development. Raise verbosity
|
||||
// with RUST_LOG, e.g. `RUST_LOG=skald_relay_server=debug` (or `=trace` for
|
||||
// full frame-level tracing). The `_log_guard` must live for the whole
|
||||
// program so the non-blocking writer flushes.
|
||||
std::fs::create_dir_all("logs")?;
|
||||
let file_appender = tracing_appender::rolling::daily("logs", "skald-relay.log");
|
||||
let (non_blocking, _log_guard) = tracing_appender::non_blocking(file_appender);
|
||||
|
||||
let filter = tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| "skald_relay_server=info,info".into());
|
||||
|
||||
tracing_subscriber::registry()
|
||||
.with(filter)
|
||||
.with(
|
||||
tracing_subscriber::fmt::layer()
|
||||
.with_writer(non_blocking)
|
||||
.with_ansi(false),
|
||||
)
|
||||
.with(tracing_subscriber::fmt::layer())
|
||||
.init();
|
||||
|
||||
let cfg = Config::from_env();
|
||||
let bind = cfg.bind;
|
||||
let db_path = cfg.db_path.clone();
|
||||
|
||||
let state = AppState::build(cfg).await?;
|
||||
tracing::info!(db = %db_path, "store ready");
|
||||
|
||||
spawn_gc(state.clone());
|
||||
|
||||
let listener = tokio::net::TcpListener::bind(bind).await?;
|
||||
tracing::info!(%bind, "relay listening on /v1/ws + /v1/pipe");
|
||||
|
||||
axum::serve(
|
||||
listener,
|
||||
router(state).into_make_service_with_connect_info::<SocketAddr>(),
|
||||
)
|
||||
.with_graceful_shutdown(shutdown_signal())
|
||||
.await?;
|
||||
|
||||
tracing::info!("relay stopped");
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,432 @@
|
||||
//! `/v1/pipe` data plane: the relay as a **stateful connection proxy**
|
||||
//! (docs/relay/pipe.md §2). The relay never reads pipe payloads — it
|
||||
//! authenticates each side (signature + namespace membership + cross-dest),
|
||||
//! matches the two sides by `connection_id`, then splices opaque ciphertext
|
||||
//! frames bidirectionally with a per-direction bandwidth cap.
|
||||
//!
|
||||
//! State machine per socket: `challenge → pipe_auth → pending → matched →
|
||||
//! streaming → teardown`. The **first** side to authenticate parks in the
|
||||
//! registry until the second arrives (within `pending_ttl`); the **second**
|
||||
//! hands its socket halves to the first, which then owns the bidirectional
|
||||
//! splice for the pipe's lifetime. Either side closing/erroring tears down both
|
||||
//! (no orphans).
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::net::IpAddr;
|
||||
use std::sync::Mutex;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use axum::extract::ws::{Message as WsMsg, WebSocket};
|
||||
use futures_util::stream::{SplitSink, SplitStream, StreamExt};
|
||||
use futures_util::SinkExt;
|
||||
use rand::RngCore;
|
||||
use skald_relay_common::crypto;
|
||||
use skald_relay_common::pipe::{self, PipeAuth, PipeChallenge};
|
||||
use tokio::sync::oneshot;
|
||||
|
||||
use crate::AppState;
|
||||
use crate::config::PipeConfig;
|
||||
use crate::limits::CHALLENGE_TIMEOUT_SECS;
|
||||
|
||||
type WsSink = SplitSink<WebSocket, WsMsg>;
|
||||
type WsStream = SplitStream<WebSocket>;
|
||||
|
||||
/// The authenticated identity of one pipe side.
|
||||
#[derive(Clone, Copy)]
|
||||
struct PeerMeta {
|
||||
/// This side's ed25519 pubkey.
|
||||
pubkey: [u8; 32],
|
||||
/// `SHA256(intended counterparty pubkey)`.
|
||||
dest: [u8; 32],
|
||||
}
|
||||
|
||||
/// The second side's socket halves, handed to the first side once the relay has
|
||||
/// verified the cross-dest match. Identity was checked before the handoff, so
|
||||
/// only the halves travel.
|
||||
struct PeerArrival {
|
||||
sink: WsSink,
|
||||
stream: WsStream,
|
||||
}
|
||||
|
||||
/// A half-open pipe: the first side authenticated, waiting for the second.
|
||||
struct PendingPipe {
|
||||
ns: String,
|
||||
meta: PeerMeta,
|
||||
/// The first side awaits this; the second side sends its halves through it.
|
||||
peer_tx: oneshot::Sender<PeerArrival>,
|
||||
}
|
||||
|
||||
/// Why an insert was refused.
|
||||
#[derive(Debug)]
|
||||
enum InsertError {
|
||||
/// `connection_id` already has a pending side.
|
||||
Duplicate,
|
||||
/// The namespace is at its concurrent-pipe cap.
|
||||
TooMany,
|
||||
}
|
||||
|
||||
/// In-memory pipe registry shared across all `/v1/pipe` connection tasks.
|
||||
#[derive(Default)]
|
||||
pub struct PipeRegistry {
|
||||
inner: Mutex<Inner>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct Inner {
|
||||
/// keyed by `connection_id` hex.
|
||||
pending: HashMap<String, PendingPipe>,
|
||||
/// namespace_id hex → number of active pipes (pending + matched). Each pipe
|
||||
/// is counted once (by its first side) for its whole lifetime.
|
||||
counts: HashMap<String, usize>,
|
||||
}
|
||||
|
||||
impl PipeRegistry {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Register the first side. Increments the namespace pipe count on success;
|
||||
/// the caller MUST call [`release`](Self::release) exactly once when done.
|
||||
fn try_insert(
|
||||
&self,
|
||||
cid_hex: &str,
|
||||
pending: PendingPipe,
|
||||
max_per_ns: usize,
|
||||
) -> Result<(), InsertError> {
|
||||
let mut g = self.inner.lock().unwrap();
|
||||
if g.pending.contains_key(cid_hex) {
|
||||
return Err(InsertError::Duplicate);
|
||||
}
|
||||
let count = g.counts.get(&pending.ns).copied().unwrap_or(0);
|
||||
if count >= max_per_ns {
|
||||
return Err(InsertError::TooMany);
|
||||
}
|
||||
*g.counts.entry(pending.ns.clone()).or_insert(0) += 1;
|
||||
g.pending.insert(cid_hex.to_string(), pending);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Take a pending side (the second arrival claims it). Does NOT touch the
|
||||
/// count — that is released by the first side.
|
||||
fn take(&self, cid_hex: &str) -> Option<PendingPipe> {
|
||||
self.inner.lock().unwrap().pending.remove(cid_hex)
|
||||
}
|
||||
|
||||
/// Release the first side: drop any lingering pending entry for `cid_hex`
|
||||
/// and decrement the namespace count. Call exactly once per [`try_insert`].
|
||||
fn release(&self, cid_hex: &str, ns: &str) {
|
||||
let mut g = self.inner.lock().unwrap();
|
||||
g.pending.remove(cid_hex);
|
||||
if let Some(c) = g.counts.get_mut(ns) {
|
||||
*c = c.saturating_sub(1);
|
||||
if *c == 0 {
|
||||
g.counts.remove(ns);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Drive one accepted `/v1/pipe` WebSocket to completion (called from `lib.rs`
|
||||
/// after the axum upgrade).
|
||||
pub async fn handle_pipe_socket(socket: WebSocket, state: AppState, peer_ip: IpAddr) {
|
||||
let (mut sink, mut stream) = socket.split();
|
||||
let cfg = state.cfg.pipe.clone();
|
||||
|
||||
// 1. Challenge: the relay speaks first (mirrors the main WS, pipe.md §2.1).
|
||||
let mut nonce = [0u8; 32];
|
||||
rand::rngs::OsRng.fill_bytes(&mut nonce);
|
||||
let chal = PipeChallenge { nonce: nonce.to_vec() };
|
||||
if sink.send(WsMsg::Binary(pipe::encode(&chal).into())).await.is_err() {
|
||||
return;
|
||||
}
|
||||
|
||||
// 2. Read the auth frame within the challenge timeout.
|
||||
let Some(auth) =
|
||||
read_pipe_auth(&mut stream, Duration::from_secs(CHALLENGE_TIMEOUT_SECS), cfg.max_frame_bytes)
|
||||
.await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
|
||||
// 3. Validate field lengths.
|
||||
let (Some(cid), Some(pubkey), Some(dest), Some(ns_raw), Some(sig)) = (
|
||||
pipe::to_array::<32>(&auth.connection_id),
|
||||
pipe::to_array::<32>(&auth.pubkey),
|
||||
pipe::to_array::<32>(&auth.dest),
|
||||
pipe::to_array::<32>(&auth.namespace_id),
|
||||
pipe::to_array::<64>(&auth.signature),
|
||||
) else {
|
||||
return close(&mut sink, "bad_request").await;
|
||||
};
|
||||
|
||||
// 3a. Signature proves control of `pubkey` and binds nonce + connection_id.
|
||||
if !crypto::verify_pipe_auth(&pubkey, &nonce, &cid, &sig) {
|
||||
return close(&mut sink, "invalid_signature").await;
|
||||
}
|
||||
// 3b. Namespace membership: the agent, or an authorized client.
|
||||
let ns = hex::encode(ns_raw);
|
||||
if !is_member(&state, &ns, &pubkey).await {
|
||||
return close(&mut sink, "unauthorized").await;
|
||||
}
|
||||
|
||||
let cid_hex = hex::encode(cid);
|
||||
let meta = PeerMeta { pubkey, dest };
|
||||
|
||||
// 4. Rendezvous by connection_id.
|
||||
if let Some(pending) = state.pipes.take(&cid_hex) {
|
||||
// We are the SECOND side: verify cross-refs, then hand our halves over.
|
||||
let cross_ok = pending.ns == ns
|
||||
&& crypto::sha256(&pending.meta.pubkey) == dest
|
||||
&& crypto::sha256(&pubkey) == pending.meta.dest;
|
||||
if !cross_ok {
|
||||
// Dropping `pending.peer_tx` also unblocks + tears down the first side.
|
||||
tracing::debug!(target: "relay::pipe", ns = %short(&ns), %peer_ip, "cross-dest mismatch");
|
||||
return close(&mut sink, "not_found").await;
|
||||
}
|
||||
let arrival = PeerArrival { sink, stream };
|
||||
if pending.peer_tx.send(arrival).is_err() {
|
||||
tracing::debug!(target: "relay::pipe", ns = %short(&ns), "first side gone before match");
|
||||
}
|
||||
// The first side now owns the splice; our halves moved into it.
|
||||
return;
|
||||
}
|
||||
|
||||
// We are the FIRST side: register and park until the second arrives.
|
||||
let (peer_tx, peer_rx) = oneshot::channel::<PeerArrival>();
|
||||
let pending = PendingPipe { ns: ns.clone(), meta, peer_tx };
|
||||
match state.pipes.try_insert(&cid_hex, pending, cfg.max_per_ns) {
|
||||
Ok(()) => {}
|
||||
Err(InsertError::Duplicate) => return close(&mut sink, "duplicate_connection").await,
|
||||
Err(InsertError::TooMany) => return close(&mut sink, "too_many_pipes").await,
|
||||
}
|
||||
|
||||
let ttl = Duration::from_secs(cfg.pending_ttl_secs);
|
||||
match tokio::time::timeout(ttl, peer_rx).await {
|
||||
Ok(Ok(arrival)) => {
|
||||
tracing::info!(target: "relay::pipe", ns = %short(&ns), %peer_ip, "pipe matched; streaming");
|
||||
splice(sink, stream, arrival.sink, arrival.stream, &cfg).await;
|
||||
}
|
||||
Ok(Err(_)) => {
|
||||
// Second side dropped its sender (closed / cross-dest mismatch).
|
||||
let _ = close(&mut sink, "peer_aborted").await;
|
||||
}
|
||||
Err(_) => {
|
||||
tracing::debug!(target: "relay::pipe", ns = %short(&ns), "pending TTL expired");
|
||||
let _ = close(&mut sink, "timeout").await;
|
||||
}
|
||||
}
|
||||
state.pipes.release(&cid_hex, &ns);
|
||||
}
|
||||
|
||||
/// `true` if `pubkey` is the agent of `ns` or an authorized client.
|
||||
async fn is_member(state: &AppState, ns: &str, pubkey: &[u8; 32]) -> bool {
|
||||
if matches!(state.store.agent_pub(ns).await, Ok(Some(a)) if &a == pubkey) {
|
||||
return true;
|
||||
}
|
||||
state.store.is_authorized_client(ns, pubkey).await.unwrap_or(false)
|
||||
}
|
||||
|
||||
/// Read binary frames until the first one decodes as [`PipeAuth`]; `None` on
|
||||
/// timeout, oversize, non-binary, malformed, or early close.
|
||||
async fn read_pipe_auth(
|
||||
stream: &mut WsStream,
|
||||
within: Duration,
|
||||
max_frame: usize,
|
||||
) -> Option<PipeAuth> {
|
||||
let deadline = tokio::time::sleep(within);
|
||||
tokio::pin!(deadline);
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = &mut deadline => return None,
|
||||
msg = stream.next() => match msg {
|
||||
Some(Ok(WsMsg::Binary(data))) => {
|
||||
if data.len() > max_frame {
|
||||
return None;
|
||||
}
|
||||
return pipe::decode::<PipeAuth>(&data).ok();
|
||||
}
|
||||
Some(Ok(WsMsg::Ping(_))) | Some(Ok(WsMsg::Pong(_))) => continue,
|
||||
_ => return None, // close, error, or text (pipe is binary-only)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// What the splice loop should do after handling one frame.
|
||||
enum Flow {
|
||||
Continue,
|
||||
Close,
|
||||
}
|
||||
|
||||
/// Bidirectionally forward binary frames between the two sides until either
|
||||
/// closes/errors or the pipe goes idle. WS-level Ping is answered on the
|
||||
/// originating socket; data is rate-limited per direction. On exit both sides
|
||||
/// are closed (no half-close in v1).
|
||||
async fn splice(
|
||||
mut a_sink: WsSink,
|
||||
mut a_stream: WsStream,
|
||||
mut b_sink: WsSink,
|
||||
mut b_stream: WsStream,
|
||||
cfg: &PipeConfig,
|
||||
) {
|
||||
let idle = Duration::from_secs(cfg.idle_timeout_secs);
|
||||
let mut bucket_ab = TokenBucket::new(cfg.max_bps, cfg.max_frame_bytes);
|
||||
let mut bucket_ba = TokenBucket::new(cfg.max_bps, cfg.max_frame_bytes);
|
||||
|
||||
loop {
|
||||
let timeout = tokio::time::sleep(idle);
|
||||
tokio::pin!(timeout);
|
||||
tokio::select! {
|
||||
_ = &mut timeout => break,
|
||||
ma = a_stream.next() => {
|
||||
if let Flow::Close =
|
||||
forward(ma, &mut b_sink, &mut a_sink, &mut bucket_ab, cfg.max_frame_bytes).await
|
||||
{
|
||||
break;
|
||||
}
|
||||
}
|
||||
mb = b_stream.next() => {
|
||||
if let Flow::Close =
|
||||
forward(mb, &mut a_sink, &mut b_sink, &mut bucket_ba, cfg.max_frame_bytes).await
|
||||
{
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
let _ = a_sink.send(WsMsg::Close(None)).await;
|
||||
let _ = b_sink.send(WsMsg::Close(None)).await;
|
||||
}
|
||||
|
||||
/// Handle one inbound frame from a side: forward `Binary` to `to_sink` (rate
|
||||
/// limited), answer `Ping` on `same_sink`, ignore `Pong`, and treat
|
||||
/// close/error/text/oversize as end-of-pipe.
|
||||
async fn forward(
|
||||
msg: Option<Result<WsMsg, axum::Error>>,
|
||||
to_sink: &mut WsSink,
|
||||
same_sink: &mut WsSink,
|
||||
bucket: &mut Option<TokenBucket>,
|
||||
max_frame: usize,
|
||||
) -> Flow {
|
||||
let Some(Ok(m)) = msg else { return Flow::Close };
|
||||
match m {
|
||||
WsMsg::Binary(data) => {
|
||||
if data.len() > max_frame {
|
||||
return Flow::Close;
|
||||
}
|
||||
if let Some(b) = bucket {
|
||||
b.consume(data.len()).await;
|
||||
}
|
||||
if to_sink.send(WsMsg::Binary(data)).await.is_err() {
|
||||
return Flow::Close;
|
||||
}
|
||||
Flow::Continue
|
||||
}
|
||||
WsMsg::Ping(p) => {
|
||||
let _ = same_sink.send(WsMsg::Pong(p)).await;
|
||||
Flow::Continue
|
||||
}
|
||||
WsMsg::Pong(_) => Flow::Continue,
|
||||
// Close, or a text frame (pipe is binary-only) → tear down.
|
||||
_ => Flow::Close,
|
||||
}
|
||||
}
|
||||
|
||||
/// Send a debug-logged close. The data plane has no error frame; the client
|
||||
/// reads a close during/after the handshake as a rejection.
|
||||
async fn close(sink: &mut WsSink, reason: &str) {
|
||||
tracing::debug!(target: "relay::pipe", reason, "closing pipe socket");
|
||||
let _ = sink.send(WsMsg::Close(None)).await;
|
||||
}
|
||||
|
||||
/// Truncate a namespace_id for logging.
|
||||
fn short(s: &str) -> String {
|
||||
let n = s.len().min(8);
|
||||
format!("{}…", &s[..n])
|
||||
}
|
||||
|
||||
/// Simple token bucket for the per-direction byte-rate cap. `None` (via
|
||||
/// [`TokenBucket::new`] with `max_bps == 0`) means unlimited. Burst is bounded
|
||||
/// to `max(rate, frame_cap)` so a single max-size frame can always pass.
|
||||
struct TokenBucket {
|
||||
rate: f64,
|
||||
burst: f64,
|
||||
allowance: f64,
|
||||
last: Instant,
|
||||
}
|
||||
|
||||
impl TokenBucket {
|
||||
fn new(max_bps: u64, frame_cap: usize) -> Option<TokenBucket> {
|
||||
if max_bps == 0 {
|
||||
return None;
|
||||
}
|
||||
let rate = max_bps as f64;
|
||||
let burst = rate.max(frame_cap as f64);
|
||||
Some(TokenBucket { rate, burst, allowance: burst, last: Instant::now() })
|
||||
}
|
||||
|
||||
/// Block until `bytes` worth of tokens are available, then deduct them.
|
||||
async fn consume(&mut self, bytes: usize) {
|
||||
let now = Instant::now();
|
||||
let elapsed = now.duration_since(self.last).as_secs_f64();
|
||||
self.last = now;
|
||||
self.allowance = (self.allowance + elapsed * self.rate).min(self.burst);
|
||||
let need = bytes as f64;
|
||||
if self.allowance < need {
|
||||
let wait = (need - self.allowance) / self.rate;
|
||||
tokio::time::sleep(Duration::from_secs_f64(wait)).await;
|
||||
self.allowance = 0.0;
|
||||
} else {
|
||||
self.allowance -= need;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn registry_enforces_per_ns_cap_and_releases() {
|
||||
let reg = PipeRegistry::new();
|
||||
let meta = PeerMeta { pubkey: [1; 32], dest: [2; 32] };
|
||||
let mk = |ns: &str| {
|
||||
let (tx, _rx) = oneshot::channel();
|
||||
PendingPipe { ns: ns.into(), meta, peer_tx: tx }
|
||||
};
|
||||
assert!(reg.try_insert("a", mk("ns"), 2).is_ok());
|
||||
assert!(reg.try_insert("b", mk("ns"), 2).is_ok());
|
||||
// Third in the same ns is over the cap.
|
||||
assert!(matches!(reg.try_insert("c", mk("ns"), 2), Err(InsertError::TooMany)));
|
||||
// Duplicate connection_id is refused regardless of cap.
|
||||
assert!(matches!(reg.try_insert("a", mk("ns"), 9), Err(InsertError::Duplicate)));
|
||||
// Releasing one frees a slot.
|
||||
reg.release("a", "ns");
|
||||
assert!(reg.try_insert("c", mk("ns"), 2).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn take_returns_pending_once() {
|
||||
let reg = PipeRegistry::new();
|
||||
let (tx, _rx) = oneshot::channel();
|
||||
let meta = PeerMeta { pubkey: [1; 32], dest: [2; 32] };
|
||||
reg.try_insert("x", PendingPipe { ns: "ns".into(), meta, peer_tx: tx }, 4).unwrap();
|
||||
assert!(reg.take("x").is_some());
|
||||
assert!(reg.take("x").is_none());
|
||||
reg.release("x", "ns");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn token_bucket_unlimited_is_noop() {
|
||||
assert!(TokenBucket::new(0, 1024).is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn token_bucket_passes_large_frame() {
|
||||
// A frame bigger than the per-second rate must still go through (burst
|
||||
// is bounded to the frame cap), just after a bounded wait.
|
||||
let mut b = TokenBucket::new(1000, 4096).unwrap();
|
||||
b.consume(4096).await; // initial burst
|
||||
b.consume(4096).await; // forces a wait, then succeeds
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,401 @@
|
||||
//! Push bridge (APNs / FCM). The **normative**, testable part always lives here:
|
||||
//! the content-in-push vs wake-only decision (relay.md §5, 3500-byte base64
|
||||
//! threshold) and the JSON payload construction. The actual send to Apple/Google
|
||||
//! sits behind the [`Pusher`] trait: the default [`LogPusher`] needs no
|
||||
//! credentials (it logs a redacted decision), so the relay also boots locally.
|
||||
//! Live senders sit behind the `push-live` feature.
|
||||
|
||||
use crate::limits::CONTENT_PUSH_MAX_B64;
|
||||
use async_trait::async_trait;
|
||||
use base64::{Engine, engine::general_purpose::STANDARD as B64};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
/// Device platform (relay-protocol.md): selects APNs vs FCM.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum Platform {
|
||||
Ios,
|
||||
Android,
|
||||
}
|
||||
|
||||
impl Platform {
|
||||
pub fn parse(s: &str) -> Option<Platform> {
|
||||
match s {
|
||||
"ios" => Some(Platform::Ios),
|
||||
"android" => Some(Platform::Android),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Platform::Ios => "ios",
|
||||
Platform::Android => "android",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Result of the push-mode decision (relay.md §5).
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum PushKind {
|
||||
/// The encrypted blob fits the limit: include it (NSE/app decrypts E2E).
|
||||
Content,
|
||||
/// Blob too large: wake only; the device opens a WS and drains the queue.
|
||||
Wake,
|
||||
}
|
||||
|
||||
/// Everything needed to build a push.
|
||||
///
|
||||
/// `ciphertext` carries the **raw** bytes — exactly what v2 transports on the
|
||||
/// WebSocket (protobuf `Message.ciphertext`). The relay base64-encodes them
|
||||
/// when building the APNs/FCM JSON envelope (`d.c` field): that field is still
|
||||
/// base64 because that's what the APNs/FCM wire protocols expect, unchanged
|
||||
/// from v1. The base64 step therefore lives **inside** `apns_payload()` /
|
||||
/// `fcm_payload()` — callers never pre-encode.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PushItem {
|
||||
pub namespace_id: String,
|
||||
pub from_hex: String,
|
||||
pub nonce_hex: String,
|
||||
/// Raw ciphertext bytes (v2 `Message.ciphertext` as it travels on the
|
||||
/// WebSocket). The relay base64-encodes them when building APNs/FCM JSON.
|
||||
pub ciphertext: Vec<u8>,
|
||||
}
|
||||
|
||||
impl PushItem {
|
||||
/// Normative selection rule: content-in-push if
|
||||
/// `len(base64(ciphertext)) <= CONTENT_PUSH_MAX_B64`, otherwise wake-only.
|
||||
/// We measure the base64 length on demand (cheap, no allocation kept).
|
||||
pub fn kind(&self) -> PushKind {
|
||||
if B64.encode(&self.ciphertext).len() <= CONTENT_PUSH_MAX_B64 {
|
||||
PushKind::Content
|
||||
} else {
|
||||
PushKind::Wake
|
||||
}
|
||||
}
|
||||
|
||||
/// APNs payload (relay.md §5.1/5.2). `aps.alert` is a generic fallback:
|
||||
/// **never** sensitive content.
|
||||
pub fn apns_payload(&self) -> Value {
|
||||
match self.kind() {
|
||||
PushKind::Content => json!({
|
||||
"aps": {
|
||||
"alert": { "title": "Skald", "body": "Azione richiesta" },
|
||||
"badge": 1,
|
||||
"sound": "default",
|
||||
"mutable-content": 1,
|
||||
"category": "skald_inbox"
|
||||
},
|
||||
"d": {
|
||||
"ns": self.namespace_id,
|
||||
"from": self.from_hex,
|
||||
"n": self.nonce_hex,
|
||||
"c": B64.encode(&self.ciphertext)
|
||||
}
|
||||
}),
|
||||
PushKind::Wake => json!({
|
||||
"aps": {
|
||||
"alert": { "title": "Skald", "body": "Azione richiesta" },
|
||||
"badge": 1,
|
||||
"sound": "default",
|
||||
"content-available": 1
|
||||
},
|
||||
"d": { "ns": self.namespace_id, "wake": true }
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// FCM HTTP v1 payload (relay.md §5.3): **data-only**, high priority, so the
|
||||
/// app always handles decryption even in the background.
|
||||
pub fn fcm_payload(&self, device_token: &str) -> Value {
|
||||
let mut data = serde_json::Map::new();
|
||||
data.insert("ns".into(), json!(self.namespace_id));
|
||||
match self.kind() {
|
||||
PushKind::Content => {
|
||||
data.insert("from".into(), json!(self.from_hex));
|
||||
data.insert("n".into(), json!(self.nonce_hex));
|
||||
data.insert("c".into(), json!(B64.encode(&self.ciphertext)));
|
||||
}
|
||||
PushKind::Wake => {
|
||||
data.insert("wake".into(), json!("true"));
|
||||
}
|
||||
}
|
||||
json!({
|
||||
"message": {
|
||||
"token": device_token,
|
||||
"android": { "priority": "high" },
|
||||
"data": Value::Object(data)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Push-send abstraction. Implemented by [`LogPusher`] (default) and, behind the
|
||||
/// `push-live` feature, by the real APNs/FCM senders.
|
||||
#[async_trait]
|
||||
pub trait Pusher: Send + Sync {
|
||||
async fn notify(&self, device_token: &str, platform: Platform, item: &PushItem);
|
||||
}
|
||||
|
||||
/// Default pusher: sends nothing, only logs a redacted decision. Lets
|
||||
/// store-and-forward work locally without Apple/Google credentials.
|
||||
pub struct LogPusher;
|
||||
|
||||
#[async_trait]
|
||||
impl Pusher for LogPusher {
|
||||
async fn notify(&self, device_token: &str, platform: Platform, item: &PushItem) {
|
||||
let kind = item.kind();
|
||||
// Never log the content: only metadata and truncated identifiers.
|
||||
tracing::info!(
|
||||
target: "relay::push",
|
||||
platform = platform.as_str(),
|
||||
kind = ?kind,
|
||||
ns = %short(&item.namespace_id),
|
||||
token = %short(device_token),
|
||||
ct_b64_len = B64.encode(&item.ciphertext).len(),
|
||||
"would deliver push (no push credentials configured: LogPusher)"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Truncate an identifier for logging (never log full sensitive strings).
|
||||
fn short(s: &str) -> String {
|
||||
let n = s.len().min(8);
|
||||
format!("{}…", &s[..n])
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Live push senders (feature `push-live`). The normative decision/payload
|
||||
// logic above stays feature-free and is what the unit tests cover; the real
|
||||
// network calls to Apple/Google live behind the gate and need no test
|
||||
// fixtures (they need real credentials).
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[cfg(feature = "push-live")]
|
||||
mod live {
|
||||
use super::*;
|
||||
use crate::config::ApnsConfig;
|
||||
use jsonwebtoken::{Algorithm, EncodingKey, Header, encode};
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
use uuid::Uuid;
|
||||
|
||||
/// APNs HTTP/2 sender (provider-auth via ES256 JWT, per Apple docs). Caches
|
||||
/// the JWT in memory and refreshes it at 30 min (token is valid 60 min).
|
||||
pub struct ApnsPusher {
|
||||
config: ApnsConfig,
|
||||
jwt: tokio::sync::RwLock<JwtState>,
|
||||
client: reqwest::Client,
|
||||
}
|
||||
|
||||
/// Cached provider-auth JWT. Re-signed lazily when the remaining TTL
|
||||
/// drops below [`REFRESH_AFTER`].
|
||||
struct JwtState {
|
||||
token: String,
|
||||
expires_at: Instant,
|
||||
}
|
||||
|
||||
/// Refresh threshold (Apple allows up to 60 min; we renew at the halfway
|
||||
/// point so a clock-skew rejection is unlikely).
|
||||
const REFRESH_AFTER: Duration = Duration::from_secs(30 * 60);
|
||||
/// TTL Apple assigns to a freshly issued provider JWT.
|
||||
const JWT_TTL: Duration = Duration::from_secs(60 * 60);
|
||||
|
||||
impl ApnsPusher {
|
||||
pub fn new(config: ApnsConfig) -> Self {
|
||||
let client = reqwest::Client::new();
|
||||
let jwt = tokio::sync::RwLock::new(JwtState {
|
||||
token: String::new(),
|
||||
// Start expired so the first `notify()` triggers a sign.
|
||||
expires_at: Instant::now(),
|
||||
});
|
||||
Self {
|
||||
config,
|
||||
jwt,
|
||||
client,
|
||||
}
|
||||
}
|
||||
|
||||
/// Return a valid provider JWT, signing a fresh one if the cached one
|
||||
/// is within [`REFRESH_AFTER`] of its TTL.
|
||||
async fn jwt(&self) -> anyhow::Result<String> {
|
||||
// Fast path: cached token is still good.
|
||||
{
|
||||
let state = self.jwt.read().await;
|
||||
if state.expires_at.saturating_duration_since(Instant::now()) > REFRESH_AFTER {
|
||||
return Ok(state.token.clone());
|
||||
}
|
||||
}
|
||||
// Slow path: take the write lock, double-check (another task may
|
||||
// have refreshed while we were waiting), then sign.
|
||||
let mut state = self.jwt.write().await;
|
||||
if state.expires_at.saturating_duration_since(Instant::now()) > REFRESH_AFTER {
|
||||
return Ok(state.token.clone());
|
||||
}
|
||||
let token = self.generate_jwt()?;
|
||||
state.token = token.clone();
|
||||
state.expires_at = Instant::now() + JWT_TTL;
|
||||
Ok(token)
|
||||
}
|
||||
|
||||
/// Sign a fresh provider JWT (ES256 over the team's P-256 key).
|
||||
fn generate_jwt(&self) -> anyhow::Result<String> {
|
||||
let mut header = Header::new(Algorithm::ES256);
|
||||
header.kid = Some(self.config.key_id.clone());
|
||||
|
||||
let iat = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)?
|
||||
.as_secs() as i64;
|
||||
let claims = serde_json::json!({
|
||||
"iss": self.config.team_id,
|
||||
"iat": iat,
|
||||
});
|
||||
|
||||
let key = EncodingKey::from_ec_pem(self.config.private_key_pem.as_bytes())?;
|
||||
Ok(encode(&header, &claims, &key)?)
|
||||
}
|
||||
|
||||
/// POST the APNs payload over HTTP/2 (negotiated via ALPN by reqwest).
|
||||
async fn send_apns(&self, device_token: &str, item: &PushItem) -> anyhow::Result<()> {
|
||||
let token = self.jwt().await?;
|
||||
let host = if self.config.sandbox {
|
||||
"https://api.sandbox.push.apple.com"
|
||||
} else {
|
||||
"https://api.push.apple.com"
|
||||
};
|
||||
let url = format!("{host}/3/device/{device_token}");
|
||||
let push_type = match item.kind() {
|
||||
PushKind::Content => "alert",
|
||||
PushKind::Wake => "background",
|
||||
};
|
||||
let body = item.apns_payload();
|
||||
let apns_id = Uuid::new_v4().to_string();
|
||||
|
||||
let resp = self
|
||||
.client
|
||||
.post(&url)
|
||||
.header("apns-topic", &self.config.bundle_id)
|
||||
.header("apns-push-type", push_type)
|
||||
.header("apns-id", &apns_id)
|
||||
.header("authorization", format!("bearer {token}"))
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let status = resp.status();
|
||||
if !status.is_success() {
|
||||
// Apple returns a JSON `{"reason": "..."}` body on errors; safe
|
||||
// to log (it never echoes our payload content).
|
||||
let reason = resp.text().await.unwrap_or_default();
|
||||
tracing::warn!(
|
||||
target: "relay::push",
|
||||
status = %status,
|
||||
apns_id = %apns_id,
|
||||
reason = %reason,
|
||||
"APNs request failed"
|
||||
);
|
||||
} else {
|
||||
tracing::info!(
|
||||
target: "relay::push",
|
||||
apns_id = %apns_id,
|
||||
"APNs request accepted"
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Pusher for ApnsPusher {
|
||||
async fn notify(&self, device_token: &str, platform: Platform, item: &PushItem) {
|
||||
// Defense in depth: an empty token would build `/3/device/` and get
|
||||
// a `MissingDeviceToken` 400 from Apple. Callers already filter
|
||||
// these out, but never spend a request on a token we know is empty.
|
||||
if device_token.is_empty() {
|
||||
tracing::debug!(
|
||||
target: "relay::push",
|
||||
"skipping APNs send: empty device token"
|
||||
);
|
||||
return;
|
||||
}
|
||||
// FcmPusher is not implemented yet: this sender only knows APNs.
|
||||
if platform != Platform::Ios {
|
||||
tracing::debug!(
|
||||
target: "relay::push",
|
||||
platform = platform.as_str(),
|
||||
"ApnsPusher ignoring non-iOS notification (no FcmPusher yet)"
|
||||
);
|
||||
return;
|
||||
}
|
||||
if let Err(e) = self.send_apns(device_token, item).await {
|
||||
// Never echo device_token or content — only the truncated
|
||||
// identifier and a generic error class.
|
||||
tracing::warn!(
|
||||
target: "relay::push",
|
||||
device_token = %short(device_token),
|
||||
error = %e,
|
||||
"APNs send failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the live APNs pusher. Caller falls back to [`LogPusher`] if
|
||||
/// `cfg.apns` is `None` (key file missing, bundle id unset, …).
|
||||
pub fn build_pusher(cfg: &ApnsConfig) -> Arc<dyn Pusher> {
|
||||
Arc::new(ApnsPusher::new(cfg.clone()))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "push-live")]
|
||||
pub use live::build_pusher;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn item(ct_len: usize) -> PushItem {
|
||||
PushItem {
|
||||
namespace_id: "a".repeat(64),
|
||||
from_hex: "b".repeat(64),
|
||||
nonce_hex: "c".repeat(24),
|
||||
ciphertext: vec![0xAA; ct_len],
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn threshold_is_inclusive_3500() {
|
||||
// 2625 raw bytes → 3500 base64 bytes (2625 % 3 == 0 → no padding needed).
|
||||
let at_boundary = item((CONTENT_PUSH_MAX_B64 / 4) * 3);
|
||||
assert_eq!(at_boundary.kind(), PushKind::Content);
|
||||
// 2626 raw bytes → 3504 base64 bytes (2626 % 3 == 1 → 2 padding chars) → Wake.
|
||||
assert_eq!(
|
||||
item((CONTENT_PUSH_MAX_B64 / 4) * 3 + 1).kind(),
|
||||
PushKind::Wake
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn apns_content_has_blob_and_mutable() {
|
||||
let p = item(100).apns_payload();
|
||||
assert_eq!(p["aps"]["mutable-content"], 1);
|
||||
assert_eq!(p["d"]["c"], B64.encode(&vec![0xAA; 100]));
|
||||
assert_eq!(p["d"]["n"], "c".repeat(24));
|
||||
assert!(p["d"].get("wake").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn apns_wake_has_no_content() {
|
||||
let p = item(CONTENT_PUSH_MAX_B64 + 50).apns_payload();
|
||||
assert_eq!(p["aps"]["content-available"], 1);
|
||||
assert_eq!(p["d"]["wake"], true);
|
||||
assert!(p["d"].get("c").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn fcm_is_data_only_high_priority() {
|
||||
let p = item(100).fcm_payload("tok123");
|
||||
assert_eq!(p["message"]["token"], "tok123");
|
||||
assert_eq!(p["message"]["android"]["priority"], "high");
|
||||
assert_eq!(p["message"]["data"]["c"], B64.encode(&vec![0xAA; 100]));
|
||||
assert!(p["message"].get("notification").is_none());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,326 @@
|
||||
//! In-memory registry of live connections (relay.md §4). Maps each namespace to
|
||||
//! its single agent connection and its set of client connections. Used to
|
||||
//! forward messages live; when the recipient is absent the caller falls back to
|
||||
//! store-and-forward + push.
|
||||
//!
|
||||
//! Concurrency: a plain `std::sync::Mutex` guards the map. We never hold the
|
||||
//! lock across an `.await`: lookups clone the cheap `mpsc::Sender` and release
|
||||
//! the lock before sending. Stale-connection eviction uses a per-connection
|
||||
//! `CancellationToken` plus a unique id so a connection only ever removes its
|
||||
//! own entry.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Mutex;
|
||||
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use crate::types::proto::RelayFrame;
|
||||
|
||||
/// Items sent to a connection's writer task (the task that owns the WS sink).
|
||||
///
|
||||
/// v2 transport: every control/data frame is a protobuf [`RelayFrame`] carried
|
||||
/// inside a WebSocket **binary** message. WS-level Ping/Pong are used for
|
||||
/// keepalive (relay-protocol.md §5) and are therefore their own variants so
|
||||
/// the writer does not have to encode them as protobuf.
|
||||
pub enum WsOut {
|
||||
/// A protobuf `RelayFrame` to be encoded and sent as a WebSocket **binary** frame.
|
||||
Frame(RelayFrame),
|
||||
/// A WS-level Pong (reply to an inbound WS Ping).
|
||||
Pong(Vec<u8>),
|
||||
/// A WS-level Ping (keepalive). Payload is opaque; a 0-byte payload is fine.
|
||||
Ping(Vec<u8>),
|
||||
/// Ask the writer to close the socket (eviction / fatal error).
|
||||
Close,
|
||||
}
|
||||
|
||||
/// A handle to one live WebSocket's writer task.
|
||||
#[derive(Clone)]
|
||||
pub struct ConnHandle {
|
||||
/// Unique id of the connection (identity check on self-removal).
|
||||
pub id: u64,
|
||||
/// Sender into the connection's writer task.
|
||||
pub tx: mpsc::Sender<WsOut>,
|
||||
/// Cancels the connection (used to evict a replaced/revoked peer).
|
||||
pub cancel: CancellationToken,
|
||||
/// ed25519 pubkey of the peer authenticated on this connection. Agents and
|
||||
/// clients both have one; used to build `PresenceList.online[]` and to
|
||||
/// populate `PresenceEvent.pubkey` (v2 spec §4).
|
||||
pub pubkey: [u8; 32],
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct NamespaceConns {
|
||||
/// The single agent connection for this namespace, if any. The agent's
|
||||
/// pubkey lives on the [`ConnHandle`].
|
||||
agent: Option<ConnHandle>,
|
||||
/// keyed by client ed25519 pubkey, hex.
|
||||
clients: HashMap<String, ConnHandle>,
|
||||
}
|
||||
|
||||
/// Thread-safe registry shared across all connection tasks.
|
||||
#[derive(Default)]
|
||||
pub struct Registry {
|
||||
inner: Mutex<HashMap<String, NamespaceConns>>,
|
||||
}
|
||||
|
||||
impl Registry {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Register the agent connection for `ns`, returning the previous one (if
|
||||
/// any) so the caller can cancel it (one agent per namespace).
|
||||
pub fn register_agent(&self, ns: &str, handle: ConnHandle) -> Option<ConnHandle> {
|
||||
let mut map = self.inner.lock().unwrap();
|
||||
let entry = map.entry(ns.to_string()).or_default();
|
||||
entry.agent.replace(handle)
|
||||
}
|
||||
|
||||
/// Register a client connection, returning the previous one for the same
|
||||
/// pubkey (if any) so the caller can cancel it (one connection per device).
|
||||
pub fn register_client(
|
||||
&self,
|
||||
ns: &str,
|
||||
pubkey_hex: &str,
|
||||
handle: ConnHandle,
|
||||
) -> Option<ConnHandle> {
|
||||
let mut map = self.inner.lock().unwrap();
|
||||
let entry = map.entry(ns.to_string()).or_default();
|
||||
entry.clients.insert(pubkey_hex.to_string(), handle)
|
||||
}
|
||||
|
||||
/// Live sender of the namespace's agent, if connected.
|
||||
pub fn agent_tx(&self, ns: &str) -> Option<mpsc::Sender<WsOut>> {
|
||||
let map = self.inner.lock().unwrap();
|
||||
map.get(ns)
|
||||
.and_then(|c| c.agent.as_ref())
|
||||
.map(|h| h.tx.clone())
|
||||
}
|
||||
|
||||
/// Live sender of a client, if connected.
|
||||
pub fn client_tx(&self, ns: &str, pubkey_hex: &str) -> Option<mpsc::Sender<WsOut>> {
|
||||
let map = self.inner.lock().unwrap();
|
||||
map.get(ns)
|
||||
.and_then(|c| c.clients.get(pubkey_hex))
|
||||
.map(|h| h.tx.clone())
|
||||
}
|
||||
|
||||
/// Remove the agent entry, but only if it is still the connection with `id`.
|
||||
pub fn remove_agent(&self, ns: &str, id: u64) {
|
||||
let mut map = self.inner.lock().unwrap();
|
||||
if let Some(conns) = map.get_mut(ns) {
|
||||
if conns.agent.as_ref().is_some_and(|h| h.id == id) {
|
||||
conns.agent = None;
|
||||
}
|
||||
Self::gc_empty(&mut map, ns);
|
||||
}
|
||||
}
|
||||
|
||||
/// Remove a client entry, but only if it is still the connection with `id`.
|
||||
pub fn remove_client(&self, ns: &str, pubkey_hex: &str, id: u64) {
|
||||
let mut map = self.inner.lock().unwrap();
|
||||
if let Some(conns) = map.get_mut(ns) {
|
||||
if conns.clients.get(pubkey_hex).is_some_and(|h| h.id == id) {
|
||||
conns.clients.remove(pubkey_hex);
|
||||
}
|
||||
Self::gc_empty(&mut map, ns);
|
||||
}
|
||||
}
|
||||
|
||||
/// Evict a client by pubkey regardless of id (revocation). Returns the
|
||||
/// handle so the caller can cancel it.
|
||||
pub fn evict_client(&self, ns: &str, pubkey_hex: &str) -> Option<ConnHandle> {
|
||||
let mut map = self.inner.lock().unwrap();
|
||||
let handle = map.get_mut(ns).and_then(|c| c.clients.remove(pubkey_hex));
|
||||
Self::gc_empty(&mut map, ns);
|
||||
handle
|
||||
}
|
||||
|
||||
/// All pubkeys currently connected in `ns`: the agent (if connected)
|
||||
/// followed by every connected client. Used to build
|
||||
/// `PresenceList.online[]` in response to `PresenceRequest` (v2 spec §4).
|
||||
pub fn list_online(&self, ns: &str) -> Vec<[u8; 32]> {
|
||||
let map = self.inner.lock().unwrap();
|
||||
let Some(conns) = map.get(ns) else {
|
||||
return Vec::new();
|
||||
};
|
||||
let mut out: Vec<[u8; 32]> = Vec::with_capacity(1 + conns.clients.len());
|
||||
if let Some(h) = &conns.agent {
|
||||
out.push(h.pubkey);
|
||||
}
|
||||
for (_, h) in &conns.clients {
|
||||
out.push(h.pubkey);
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Broadcast `frame` to every connection in `ns`, optionally skipping the
|
||||
/// connection with `id == skip_id`. Used for `PresenceEvent` (skip the
|
||||
/// source so it doesn't see its own presence change).
|
||||
///
|
||||
/// Errors are silently dropped: a slow/blocked peer must not stall the
|
||||
/// sender while we hold the registry mutex. If the channel is full the
|
||||
/// frame is dropped for that peer — acceptable for presence (the peer
|
||||
/// will see the next periodic refresh or a later event).
|
||||
///
|
||||
/// Returns the number of targets the frame was **offered** to (i.e.
|
||||
/// `try_send` did not fail because the channel was closed). Returns 0 if
|
||||
/// the namespace is unknown.
|
||||
pub fn broadcast_ns(&self, ns: &str, frame: RelayFrame, skip_id: Option<u64>) -> usize {
|
||||
let map = self.inner.lock().unwrap();
|
||||
let Some(conns) = map.get(ns) else {
|
||||
return 0;
|
||||
};
|
||||
let mut n = 0usize;
|
||||
if let Some(h) = &conns.agent
|
||||
&& skip_id != Some(h.id)
|
||||
{
|
||||
if h.tx.try_send(WsOut::Frame(frame.clone())).is_ok() {
|
||||
n += 1;
|
||||
}
|
||||
}
|
||||
for (_, h) in &conns.clients {
|
||||
if skip_id == Some(h.id) {
|
||||
continue;
|
||||
}
|
||||
if h.tx.try_send(WsOut::Frame(frame.clone())).is_ok() {
|
||||
n += 1;
|
||||
}
|
||||
}
|
||||
n
|
||||
}
|
||||
|
||||
fn gc_empty(map: &mut HashMap<String, NamespaceConns>, ns: &str) {
|
||||
if let Some(conns) = map.get(ns)
|
||||
&& conns.agent.is_none()
|
||||
&& conns.clients.is_empty()
|
||||
{
|
||||
map.remove(ns);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn handle(id: u64, pubkey: [u8; 32]) -> (ConnHandle, mpsc::Receiver<WsOut>) {
|
||||
let (tx, rx) = mpsc::channel(4);
|
||||
(
|
||||
ConnHandle {
|
||||
id,
|
||||
tx,
|
||||
cancel: CancellationToken::new(),
|
||||
pubkey,
|
||||
},
|
||||
rx,
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn agent_replacement_returns_old() {
|
||||
let reg = Registry::new();
|
||||
let (h1, _r1) = handle(1, [0xAA; 32]);
|
||||
let (h2, _r2) = handle(2, [0xBB; 32]);
|
||||
assert!(reg.register_agent("ns", h1).is_none());
|
||||
let old = reg.register_agent("ns", h2).expect("old agent");
|
||||
assert_eq!(old.id, 1);
|
||||
assert!(reg.agent_tx("ns").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn self_removal_respects_identity() {
|
||||
let reg = Registry::new();
|
||||
let (h1, _r1) = handle(1, [0xAA; 32]);
|
||||
let (h2, _r2) = handle(2, [0xBB; 32]);
|
||||
reg.register_agent("ns", h1);
|
||||
// A newer connection replaced id=1 with id=2.
|
||||
reg.register_agent("ns", h2);
|
||||
// The old connection (id=1) cleaning up must NOT drop the new one.
|
||||
reg.remove_agent("ns", 1);
|
||||
assert!(reg.agent_tx("ns").is_some());
|
||||
// The current connection (id=2) removes itself → gone.
|
||||
reg.remove_agent("ns", 2);
|
||||
assert!(reg.agent_tx("ns").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn evict_client_returns_handle() {
|
||||
let reg = Registry::new();
|
||||
let (h, _r) = handle(7, [0xCC; 32]);
|
||||
reg.register_client("ns", "ab", h);
|
||||
assert!(reg.client_tx("ns", "ab").is_some());
|
||||
let evicted = reg.evict_client("ns", "ab").expect("handle");
|
||||
assert_eq!(evicted.id, 7);
|
||||
assert!(reg.client_tx("ns", "ab").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_online_returns_agent_and_clients() {
|
||||
let reg = Registry::new();
|
||||
let agent_pub = [0xAAu8; 32];
|
||||
let client_pub = [0xBBu8; 32];
|
||||
let (h1, _r1) = handle(1, agent_pub);
|
||||
let (h2, _r2) = handle(2, client_pub);
|
||||
reg.register_agent("ns", h1);
|
||||
reg.register_client("ns", &hex::encode(client_pub), h2);
|
||||
let online = reg.list_online("ns");
|
||||
assert_eq!(online.len(), 2);
|
||||
assert!(online.contains(&agent_pub));
|
||||
assert!(online.contains(&client_pub));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_online_empty_when_namespace_unknown() {
|
||||
let reg = Registry::new();
|
||||
assert!(reg.list_online("nope").is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_online_agent_only_when_no_clients() {
|
||||
let reg = Registry::new();
|
||||
let agent_pub = [0xAAu8; 32];
|
||||
let (h, _r) = handle(1, agent_pub);
|
||||
reg.register_agent("ns", h);
|
||||
let online = reg.list_online("ns");
|
||||
assert_eq!(online, vec![agent_pub]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn broadcast_ns_skips_source() {
|
||||
let reg = Registry::new();
|
||||
let (h1, mut r1) = handle(1, [0xAA; 32]);
|
||||
let (h2, mut r2) = handle(2, [0xBB; 32]);
|
||||
reg.register_agent("ns", h1);
|
||||
reg.register_client("ns", &hex::encode([0xBB; 32]), h2);
|
||||
let frame = RelayFrame { frame: None };
|
||||
let n = reg.broadcast_ns("ns", frame, Some(1)); // skip id=1 (agent)
|
||||
assert_eq!(n, 1);
|
||||
// Agent (id=1) should NOT see the frame.
|
||||
assert!(r1.try_recv().is_err());
|
||||
// Client (id=2) should see it.
|
||||
assert!(r2.try_recv().is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn broadcast_ns_with_no_skip_targets_all() {
|
||||
let reg = Registry::new();
|
||||
let (h1, mut r1) = handle(1, [0xAA; 32]);
|
||||
let (h2, mut r2) = handle(2, [0xBB; 32]);
|
||||
reg.register_agent("ns", h1);
|
||||
reg.register_client("ns", &hex::encode([0xBB; 32]), h2);
|
||||
let frame = RelayFrame { frame: None };
|
||||
let n = reg.broadcast_ns("ns", frame, None);
|
||||
assert_eq!(n, 2);
|
||||
assert!(r1.try_recv().is_ok());
|
||||
assert!(r2.try_recv().is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn broadcast_ns_unknown_namespace_returns_zero() {
|
||||
let reg = Registry::new();
|
||||
let frame = RelayFrame { frame: None };
|
||||
assert_eq!(reg.broadcast_ns("nope", frame, None), 0);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,558 @@
|
||||
//! SQLite persistence (relay.md §3). No sensitive data in the clear: the
|
||||
//! `ciphertext` blobs are E2E, the pubkeys are public identifiers. The API is
|
||||
//! designed to be swappable for Postgres+Redis post-v1 (relay.md §3 "scale
|
||||
//! path"); for now there is a single writer (the SQLite-on-EFS constraint).
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::str::FromStr;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::Result;
|
||||
use sqlx::Row;
|
||||
use sqlx::sqlite::{SqliteConnectOptions, SqlitePool, SqlitePoolOptions};
|
||||
|
||||
use crate::auth::ct_eq;
|
||||
|
||||
/// Current unix milliseconds (application timestamp encoding, index.md §5).
|
||||
pub fn now_ms() -> i64 {
|
||||
chrono::Utc::now().timestamp_millis()
|
||||
}
|
||||
|
||||
fn to_arr<const N: usize>(v: &[u8]) -> Option<[u8; N]> {
|
||||
if v.len() != N {
|
||||
return None;
|
||||
}
|
||||
let mut out = [0u8; N];
|
||||
out.copy_from_slice(v);
|
||||
Some(out)
|
||||
}
|
||||
|
||||
/// Persisted client state.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ClientRow {
|
||||
pub x25519_pub: [u8; 32],
|
||||
pub device_token: Option<String>,
|
||||
pub platform: String,
|
||||
pub state: String, // 'pending' | 'authorized'
|
||||
}
|
||||
|
||||
/// A store-and-forward queued message.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct QueuedMsg {
|
||||
pub id: i64,
|
||||
pub from_pub: [u8; 32],
|
||||
pub nonce: [u8; 12],
|
||||
pub ciphertext: Vec<u8>,
|
||||
pub created_at: i64,
|
||||
}
|
||||
|
||||
/// A client in `pending` state (for re-sending `client_paired` to the agent).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PendingClient {
|
||||
pub ed25519_pub: [u8; 32],
|
||||
pub x25519_pub: [u8; 32],
|
||||
pub platform: String,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Store {
|
||||
pool: SqlitePool,
|
||||
}
|
||||
|
||||
impl Store {
|
||||
/// Open/create the DB and apply the schema (idempotent). No WAL: the deploy
|
||||
/// EFS/NFS does not support it; `busy_timeout` serializes the single writer.
|
||||
pub async fn init(path: &str) -> Result<Store> {
|
||||
let opts = SqliteConnectOptions::from_str(&format!("sqlite://{path}"))?
|
||||
.create_if_missing(true)
|
||||
.busy_timeout(Duration::from_secs(5))
|
||||
.foreign_keys(true);
|
||||
let pool = SqlitePoolOptions::new()
|
||||
.max_connections(4)
|
||||
.connect_with(opts)
|
||||
.await?;
|
||||
|
||||
for stmt in SCHEMA {
|
||||
// SCHEMA entries are 'static string literals (audited, no user data).
|
||||
sqlx::query(*stmt).execute(&pool).await?;
|
||||
}
|
||||
Ok(Store { pool })
|
||||
}
|
||||
|
||||
// ----- namespaces ---------------------------------------------------------
|
||||
|
||||
/// Create the namespace if absent (binding it immutably to the pubkey) and
|
||||
/// bump `last_active`. Idempotent.
|
||||
pub async fn upsert_namespace(&self, ns: &str, agent_pub: &[u8; 32]) -> Result<()> {
|
||||
let now = now_ms();
|
||||
sqlx::query(
|
||||
"INSERT INTO namespaces (namespace_id, agent_ed25519_pub, created_at, last_active)
|
||||
VALUES (?1, ?2, ?3, ?3)
|
||||
ON CONFLICT(namespace_id) DO UPDATE SET last_active = ?3",
|
||||
)
|
||||
.bind(ns)
|
||||
.bind(&agent_pub[..])
|
||||
.bind(now)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// `true` if the namespace exists.
|
||||
pub async fn namespace_exists(&self, ns: &str) -> Result<bool> {
|
||||
let row = sqlx::query("SELECT 1 FROM namespaces WHERE namespace_id = ?1")
|
||||
.bind(ns)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
Ok(row.is_some())
|
||||
}
|
||||
|
||||
/// The namespace agent's ed25519 pubkey (None if it does not exist).
|
||||
pub async fn agent_pub(&self, ns: &str) -> Result<Option<[u8; 32]>> {
|
||||
let row = sqlx::query("SELECT agent_ed25519_pub FROM namespaces WHERE namespace_id = ?1")
|
||||
.bind(ns)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
Ok(row.and_then(|r| {
|
||||
let b: Vec<u8> = r.get(0);
|
||||
to_arr::<32>(&b)
|
||||
}))
|
||||
}
|
||||
|
||||
pub async fn touch_namespace(&self, ns: &str) -> Result<()> {
|
||||
sqlx::query("UPDATE namespaces SET last_active = ?2 WHERE namespace_id = ?1")
|
||||
.bind(ns)
|
||||
.bind(now_ms())
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ----- pairing ------------------------------------------------------------
|
||||
|
||||
/// Open/replace the pairing window. `expiry_ms` already computed.
|
||||
pub async fn pairing_start(&self, ns: &str, token: &[u8; 32], expiry_ms: i64) -> Result<()> {
|
||||
sqlx::query(
|
||||
"UPDATE namespaces
|
||||
SET pairing_token = ?2, pairing_expiry = ?3, pairing_consumed = 0
|
||||
WHERE namespace_id = ?1",
|
||||
)
|
||||
.bind(ns)
|
||||
.bind(&token[..])
|
||||
.bind(expiry_ms)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn pairing_stop(&self, ns: &str) -> Result<()> {
|
||||
sqlx::query(
|
||||
"UPDATE namespaces
|
||||
SET pairing_token = NULL, pairing_expiry = NULL, pairing_consumed = 0
|
||||
WHERE namespace_id = ?1",
|
||||
)
|
||||
.bind(ns)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Try to consume the pairing token (single-use, constant-time, not
|
||||
/// expired). Returns `true` if pairing is allowed to proceed.
|
||||
pub async fn consume_pairing_token(&self, ns: &str, token: &[u8; 32]) -> Result<bool> {
|
||||
let row = sqlx::query(
|
||||
"SELECT pairing_token, pairing_expiry, pairing_consumed
|
||||
FROM namespaces WHERE namespace_id = ?1",
|
||||
)
|
||||
.bind(ns)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
|
||||
let Some(row) = row else { return Ok(false) };
|
||||
let stored: Option<Vec<u8>> = row.get(0);
|
||||
let expiry: Option<i64> = row.get(1);
|
||||
let consumed: i64 = row.get(2);
|
||||
|
||||
let (Some(stored), Some(expiry)) = (stored, expiry) else {
|
||||
return Ok(false); // no open window
|
||||
};
|
||||
if consumed != 0 || expiry <= now_ms() || !ct_eq(&stored, &token[..]) {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
// Atomic guard against concurrent double-consume.
|
||||
let res = sqlx::query(
|
||||
"UPDATE namespaces SET pairing_consumed = 1
|
||||
WHERE namespace_id = ?1 AND pairing_consumed = 0",
|
||||
)
|
||||
.bind(ns)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(res.rows_affected() == 1)
|
||||
}
|
||||
|
||||
// ----- clients ------------------------------------------------------------
|
||||
|
||||
/// Register/update a client as `pending` (after a successful pairing).
|
||||
pub async fn upsert_pending_client(
|
||||
&self,
|
||||
ns: &str,
|
||||
ed_pub: &[u8; 32],
|
||||
x_pub: &[u8; 32],
|
||||
device_token: &str,
|
||||
platform: &str,
|
||||
) -> Result<()> {
|
||||
sqlx::query(
|
||||
"INSERT INTO clients
|
||||
(namespace_id, client_ed25519_pub, client_x25519_pub, device_token, platform, state, last_seen)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, 'pending', ?6)
|
||||
ON CONFLICT(namespace_id, client_ed25519_pub) DO UPDATE SET
|
||||
client_x25519_pub = ?3,
|
||||
device_token = CASE WHEN ?4 = '' THEN clients.device_token ELSE ?4 END,
|
||||
platform = ?5, state = 'pending', last_seen = ?6",
|
||||
)
|
||||
.bind(ns)
|
||||
.bind(&ed_pub[..])
|
||||
.bind(&x_pub[..])
|
||||
.bind(device_token)
|
||||
.bind(platform)
|
||||
.bind(now_ms())
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn get_client(&self, ns: &str, ed_pub: &[u8; 32]) -> Result<Option<ClientRow>> {
|
||||
let row = sqlx::query(
|
||||
"SELECT client_x25519_pub, device_token, platform, state
|
||||
FROM clients WHERE namespace_id = ?1 AND client_ed25519_pub = ?2",
|
||||
)
|
||||
.bind(ns)
|
||||
.bind(&ed_pub[..])
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
Ok(row.and_then(|r| {
|
||||
let x: Vec<u8> = r.get(0);
|
||||
Some(ClientRow {
|
||||
x25519_pub: to_arr::<32>(&x)?,
|
||||
device_token: r.get::<Option<String>, _>(1),
|
||||
platform: r.get(2),
|
||||
state: r.get(3),
|
||||
})
|
||||
}))
|
||||
}
|
||||
|
||||
pub async fn is_authorized_client(&self, ns: &str, ed_pub: &[u8; 32]) -> Result<bool> {
|
||||
let row = sqlx::query(
|
||||
"SELECT 1 FROM clients
|
||||
WHERE namespace_id = ?1 AND client_ed25519_pub = ?2 AND state = 'authorized'",
|
||||
)
|
||||
.bind(ns)
|
||||
.bind(&ed_pub[..])
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
Ok(row.is_some())
|
||||
}
|
||||
|
||||
/// Update the client's push token (APNs/FCM rotate it) + last_seen.
|
||||
///
|
||||
/// An **empty** `device_token` is treated as "no token available right now"
|
||||
/// (e.g. the device connected before its APNs registration completed) and
|
||||
/// must NOT clobber a previously stored, valid token — otherwise every push
|
||||
/// to that client fails with `MissingDeviceToken`. In that case we still
|
||||
/// bump `last_seen` but keep the existing token.
|
||||
pub async fn update_client_device_token(
|
||||
&self,
|
||||
ns: &str,
|
||||
ed_pub: &[u8; 32],
|
||||
device_token: &str,
|
||||
) -> Result<()> {
|
||||
sqlx::query(
|
||||
"UPDATE clients
|
||||
SET device_token = CASE WHEN ?3 = '' THEN device_token ELSE ?3 END,
|
||||
last_seen = ?4
|
||||
WHERE namespace_id = ?1 AND client_ed25519_pub = ?2",
|
||||
)
|
||||
.bind(ns)
|
||||
.bind(&ed_pub[..])
|
||||
.bind(device_token)
|
||||
.bind(now_ms())
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn list_pending_clients(&self, ns: &str) -> Result<Vec<PendingClient>> {
|
||||
let rows = sqlx::query(
|
||||
"SELECT client_ed25519_pub, client_x25519_pub, platform
|
||||
FROM clients WHERE namespace_id = ?1 AND state = 'pending'",
|
||||
)
|
||||
.bind(ns)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
let mut out = Vec::new();
|
||||
for r in rows {
|
||||
let ed: Vec<u8> = r.get(0);
|
||||
let x: Vec<u8> = r.get(1);
|
||||
if let (Some(ed), Some(x)) = (to_arr::<32>(&ed), to_arr::<32>(&x)) {
|
||||
out.push(PendingClient {
|
||||
ed25519_pub: ed,
|
||||
x25519_pub: x,
|
||||
platform: r.get(2),
|
||||
});
|
||||
}
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Apply `authorize` (replace semantics, relay-protocol.md §6). Returns
|
||||
/// `(authorized count, revoked pubkeys)`. Revoked clients must then be
|
||||
/// disconnected; their queue has already been purged here.
|
||||
pub async fn apply_authorize(
|
||||
&self,
|
||||
ns: &str,
|
||||
new_list: &[[u8; 32]],
|
||||
) -> Result<(i64, Vec<[u8; 32]>)> {
|
||||
let new_set: HashSet<Vec<u8>> = new_list.iter().map(|k| k.to_vec()).collect();
|
||||
|
||||
let existing =
|
||||
sqlx::query("SELECT client_ed25519_pub FROM clients WHERE namespace_id = ?1")
|
||||
.bind(ns)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
let mut revoked = Vec::new();
|
||||
for r in existing {
|
||||
let pub_bytes: Vec<u8> = r.get(0);
|
||||
if new_set.contains(&pub_bytes) {
|
||||
// Present in the new list → authorized (leaves pending).
|
||||
sqlx::query(
|
||||
"UPDATE clients SET state = 'authorized'
|
||||
WHERE namespace_id = ?1 AND client_ed25519_pub = ?2",
|
||||
)
|
||||
.bind(ns)
|
||||
.bind(&pub_bytes)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
} else {
|
||||
// Absent → revoked: purge queue, forget device_token, remove.
|
||||
self.purge_queue_for_bytes(ns, &pub_bytes).await?;
|
||||
sqlx::query(
|
||||
"DELETE FROM clients WHERE namespace_id = ?1 AND client_ed25519_pub = ?2",
|
||||
)
|
||||
.bind(ns)
|
||||
.bind(&pub_bytes)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
if let Some(k) = to_arr::<32>(&pub_bytes) {
|
||||
revoked.push(k);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let count: i64 = sqlx::query(
|
||||
"SELECT COUNT(*) FROM clients WHERE namespace_id = ?1 AND state = 'authorized'",
|
||||
)
|
||||
.bind(ns)
|
||||
.fetch_one(&self.pool)
|
||||
.await?
|
||||
.get(0);
|
||||
|
||||
Ok((count, revoked))
|
||||
}
|
||||
|
||||
// ----- queue (store-and-forward) -----------------------------------------
|
||||
|
||||
pub async fn queue_count(&self, ns: &str, to_pub: &[u8; 32]) -> Result<i64> {
|
||||
let n: i64 =
|
||||
sqlx::query("SELECT COUNT(*) FROM queue WHERE namespace_id = ?1 AND to_pub = ?2")
|
||||
.bind(ns)
|
||||
.bind(&to_pub[..])
|
||||
.fetch_one(&self.pool)
|
||||
.await?
|
||||
.get(0);
|
||||
Ok(n)
|
||||
}
|
||||
|
||||
/// Enqueue a message. `Ok(false)` if the recipient's queue is full.
|
||||
pub async fn enqueue(
|
||||
&self,
|
||||
ns: &str,
|
||||
to_pub: &[u8; 32],
|
||||
from_pub: &[u8; 32],
|
||||
nonce: &[u8; 12],
|
||||
ciphertext: &[u8],
|
||||
max_per_dest: i64,
|
||||
) -> Result<bool> {
|
||||
if self.queue_count(ns, to_pub).await? >= max_per_dest {
|
||||
return Ok(false);
|
||||
}
|
||||
sqlx::query(
|
||||
"INSERT INTO queue (namespace_id, to_pub, from_pub, nonce, ciphertext, created_at)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
|
||||
)
|
||||
.bind(ns)
|
||||
.bind(&to_pub[..])
|
||||
.bind(&from_pub[..])
|
||||
.bind(&nonce[..])
|
||||
.bind(ciphertext)
|
||||
.bind(now_ms())
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
pub async fn fetch_pending(&self, ns: &str, to_pub: &[u8; 32]) -> Result<Vec<QueuedMsg>> {
|
||||
let rows = sqlx::query(
|
||||
"SELECT id, from_pub, nonce, ciphertext, created_at
|
||||
FROM queue WHERE namespace_id = ?1 AND to_pub = ?2 ORDER BY id ASC",
|
||||
)
|
||||
.bind(ns)
|
||||
.bind(&to_pub[..])
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
let mut out = Vec::new();
|
||||
for r in rows {
|
||||
let from: Vec<u8> = r.get(1);
|
||||
let nonce: Vec<u8> = r.get(2);
|
||||
if let (Some(from), Some(nonce)) = (to_arr::<32>(&from), to_arr::<12>(&nonce)) {
|
||||
out.push(QueuedMsg {
|
||||
id: r.get(0),
|
||||
from_pub: from,
|
||||
nonce,
|
||||
ciphertext: r.get(3),
|
||||
created_at: r.get(4),
|
||||
});
|
||||
}
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
pub async fn delete_pending(&self, id: i64) -> Result<()> {
|
||||
sqlx::query("DELETE FROM queue WHERE id = ?1")
|
||||
.bind(id)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn purge_queue_for_bytes(&self, ns: &str, to_pub: &[u8]) -> Result<()> {
|
||||
sqlx::query("DELETE FROM queue WHERE namespace_id = ?1 AND to_pub = ?2")
|
||||
.bind(ns)
|
||||
.bind(to_pub)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ----- garbage collection -------------------------------------------------
|
||||
|
||||
/// Delete messages older than `ttl_days` and namespaces idle for `ttl_days`
|
||||
/// (cascade to clients + queue). Returns `(messages, namespaces)` removed.
|
||||
pub async fn gc(&self, ttl_days: i64) -> Result<(u64, u64)> {
|
||||
let cutoff = now_ms() - ttl_days * 24 * 60 * 60 * 1000;
|
||||
let msgs = sqlx::query("DELETE FROM queue WHERE created_at < ?1")
|
||||
.bind(cutoff)
|
||||
.execute(&self.pool)
|
||||
.await?
|
||||
.rows_affected();
|
||||
let namespaces = sqlx::query("DELETE FROM namespaces WHERE last_active < ?1")
|
||||
.bind(cutoff)
|
||||
.execute(&self.pool)
|
||||
.await?
|
||||
.rows_affected();
|
||||
Ok((msgs, namespaces))
|
||||
}
|
||||
}
|
||||
|
||||
const SCHEMA: &[&str] = &[
|
||||
"CREATE TABLE IF NOT EXISTS namespaces (
|
||||
namespace_id TEXT PRIMARY KEY,
|
||||
agent_ed25519_pub BLOB NOT NULL UNIQUE,
|
||||
created_at INTEGER NOT NULL,
|
||||
last_active INTEGER NOT NULL,
|
||||
pairing_token BLOB,
|
||||
pairing_expiry INTEGER,
|
||||
pairing_consumed INTEGER NOT NULL DEFAULT 0
|
||||
)",
|
||||
"CREATE TABLE IF NOT EXISTS clients (
|
||||
namespace_id TEXT NOT NULL REFERENCES namespaces(namespace_id) ON DELETE CASCADE,
|
||||
client_ed25519_pub BLOB NOT NULL,
|
||||
client_x25519_pub BLOB NOT NULL,
|
||||
device_token TEXT,
|
||||
platform TEXT NOT NULL,
|
||||
state TEXT NOT NULL,
|
||||
last_seen INTEGER,
|
||||
PRIMARY KEY (namespace_id, client_ed25519_pub)
|
||||
)",
|
||||
"CREATE TABLE IF NOT EXISTS queue (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
namespace_id TEXT NOT NULL REFERENCES namespaces(namespace_id) ON DELETE CASCADE,
|
||||
to_pub BLOB NOT NULL,
|
||||
from_pub BLOB NOT NULL,
|
||||
nonce BLOB NOT NULL,
|
||||
ciphertext BLOB NOT NULL,
|
||||
created_at INTEGER NOT NULL
|
||||
)",
|
||||
"CREATE INDEX IF NOT EXISTS idx_queue_dest ON queue(namespace_id, to_pub, id)",
|
||||
];
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
|
||||
async fn temp_store() -> Store {
|
||||
static SEQ: AtomicU64 = AtomicU64::new(0);
|
||||
let nanos = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_nanos();
|
||||
let seq = SEQ.fetch_add(1, Ordering::Relaxed);
|
||||
let path = std::env::temp_dir().join(format!(
|
||||
"relay-store-ut-{nanos}-{}-{seq}.db",
|
||||
std::process::id()
|
||||
));
|
||||
Store::init(&path.to_string_lossy()).await.expect("init store")
|
||||
}
|
||||
|
||||
/// An empty `device_token` (device connected before APNs registration
|
||||
/// finished) must NOT wipe a previously stored, valid token — otherwise
|
||||
/// every push fails with `MissingDeviceToken`.
|
||||
#[tokio::test]
|
||||
async fn empty_device_token_does_not_clobber() {
|
||||
let s = temp_store().await;
|
||||
let ns = "a".repeat(64);
|
||||
let ed = [1u8; 32];
|
||||
let x = [2u8; 32];
|
||||
s.upsert_namespace(&ns, &[9u8; 32]).await.unwrap();
|
||||
|
||||
// Pair with a real token.
|
||||
s.upsert_pending_client(&ns, &ed, &x, "realtoken", "ios").await.unwrap();
|
||||
assert_eq!(
|
||||
s.get_client(&ns, &ed).await.unwrap().unwrap().device_token.as_deref(),
|
||||
Some("realtoken")
|
||||
);
|
||||
|
||||
// A later connect with an empty token keeps the existing one.
|
||||
s.update_client_device_token(&ns, &ed, "").await.unwrap();
|
||||
assert_eq!(
|
||||
s.get_client(&ns, &ed).await.unwrap().unwrap().device_token.as_deref(),
|
||||
Some("realtoken")
|
||||
);
|
||||
|
||||
// A non-empty token still updates.
|
||||
s.update_client_device_token(&ns, &ed, "rotated").await.unwrap();
|
||||
assert_eq!(
|
||||
s.get_client(&ns, &ed).await.unwrap().unwrap().device_token.as_deref(),
|
||||
Some("rotated")
|
||||
);
|
||||
|
||||
// Re-pairing with an empty token also preserves the stored one.
|
||||
s.upsert_pending_client(&ns, &ed, &x, "", "ios").await.unwrap();
|
||||
assert_eq!(
|
||||
s.get_client(&ns, &ed).await.unwrap().unwrap().device_token.as_deref(),
|
||||
Some("rotated")
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
//! Wire-protocol types re-exported from the shared `skald-relay-common` crate
|
||||
//! (see plugin.md §1.1). The relay and the mobile-connector plugin use the
|
||||
//! same byte-level frames so they can never diverge.
|
||||
//!
|
||||
//! - **v2 (current)**: [`proto`] — protobuf types for the binary WebSocket
|
||||
//! transport (data/ios-app/v2/relay-protocol.md). Every wire frame is a
|
||||
//! `RelayFrame` carrying one of the sub-messages (Challenge, Auth, Message,
|
||||
//! PresenceEvent, …). This is the only transport the relay speaks now.
|
||||
|
||||
/// v2 protobuf frames — namespaced. The WS layer reads/writes these.
|
||||
pub mod proto {
|
||||
pub use skald_relay_common::proto::v2::*;
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
#[test]
|
||||
fn v2_proto_types_exposed() {
|
||||
let _v2_frame: skald_relay_common::proto::v2::RelayFrame =
|
||||
skald_relay_common::proto::v2::RelayFrame { frame: None };
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,203 @@
|
||||
//! End-to-end tests for the `/v1/pipe` data plane (docs/relay/pipe.md §2).
|
||||
//!
|
||||
//! Two raw WebSocket peers authenticate to the relay (MsgPack `pipe_auth`,
|
||||
//! Ed25519 signature), get matched by `connection_id`, and stream opaque bytes
|
||||
//! the relay never reads. Covers the happy path plus the auth/cross-dest
|
||||
//! rejections.
|
||||
|
||||
use std::net::SocketAddr;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use ed25519_dalek::SigningKey;
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use skald_relay_common::crypto;
|
||||
use skald_relay_common::pipe::{self, PipeAuth, PipeChallenge};
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
|
||||
use skald_relay_server::config::{Config, PipeConfig};
|
||||
use skald_relay_server::{AppState, router};
|
||||
|
||||
type Ws =
|
||||
tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
|
||||
|
||||
/// Boot a relay with a throwaway DB, returning its addr and the shared state so
|
||||
/// tests can seed the namespace / authorized clients directly.
|
||||
async fn spawn_relay() -> (SocketAddr, AppState) {
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
static COUNTER: AtomicU64 = AtomicU64::new(0);
|
||||
let nanos = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_nanos();
|
||||
let seq = COUNTER.fetch_add(1, Ordering::Relaxed);
|
||||
let db = std::env::temp_dir().join(format!("relay-pipe-it-{nanos}-{}-{seq}.db", std::process::id()));
|
||||
let cfg = Config {
|
||||
bind: "127.0.0.1:0".parse().unwrap(),
|
||||
db_path: db.to_string_lossy().into(),
|
||||
pipe: PipeConfig::default(),
|
||||
};
|
||||
let state = AppState::build(cfg).await.expect("build state");
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let serve_state = state.clone();
|
||||
tokio::spawn(async move {
|
||||
axum::serve(
|
||||
listener,
|
||||
router(serve_state).into_make_service_with_connect_info::<SocketAddr>(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
});
|
||||
(addr, state)
|
||||
}
|
||||
|
||||
/// An identity = its Ed25519 signing key + derived pubkeys.
|
||||
struct Id {
|
||||
sk: SigningKey,
|
||||
ed_pub: [u8; 32],
|
||||
}
|
||||
|
||||
fn id_from_seed(seed: u8) -> Id {
|
||||
let dk = crypto::derive_keys(&[seed; 32]);
|
||||
Id { sk: SigningKey::from_bytes(&dk.ed25519_priv), ed_pub: dk.ed25519_pub }
|
||||
}
|
||||
|
||||
/// Seed a namespace owned by `agent` and authorize `client` in it.
|
||||
async fn seed_namespace(state: &AppState, agent: &Id, client: &Id) -> [u8; 32] {
|
||||
let (ns_raw, ns_hex) = crypto::namespace_id(&agent.ed_pub);
|
||||
state.store.upsert_namespace(&ns_hex, &agent.ed_pub).await.unwrap();
|
||||
let client_x = crypto::derive_keys(&[0xC1; 32]).x25519_pub; // any 32B is fine for membership
|
||||
state
|
||||
.store
|
||||
.upsert_pending_client(&ns_hex, &client.ed_pub, &client_x, "", "ios")
|
||||
.await
|
||||
.unwrap();
|
||||
state.store.apply_authorize(&ns_hex, &[client.ed_pub]).await.unwrap();
|
||||
ns_raw
|
||||
}
|
||||
|
||||
/// Connect to `/v1/pipe`, complete the challenge→auth handshake for `me`
|
||||
/// targeting `peer_ed`, and return the live socket. `dest_override` lets a test
|
||||
/// declare the wrong counterparty (cross-dest rejection).
|
||||
async fn dial_and_auth(
|
||||
addr: SocketAddr,
|
||||
me: &Id,
|
||||
peer_ed: &[u8; 32],
|
||||
ns_raw: &[u8; 32],
|
||||
connection_id: &[u8; 32],
|
||||
corrupt_sig: bool,
|
||||
dest_override: Option<[u8; 32]>,
|
||||
) -> Ws {
|
||||
let url = format!("ws://{addr}/v1/pipe");
|
||||
let (mut ws, _) = tokio_tungstenite::connect_async(&url).await.expect("connect");
|
||||
|
||||
// Relay speaks first: PipeChallenge.
|
||||
let nonce = loop {
|
||||
match ws.next().await.expect("frame").expect("ws ok") {
|
||||
Message::Binary(data) => {
|
||||
let c: PipeChallenge = pipe::decode(&data).expect("challenge");
|
||||
break pipe::to_array::<32>(&c.nonce).expect("32B nonce");
|
||||
}
|
||||
Message::Ping(_) | Message::Pong(_) => continue,
|
||||
other => panic!("expected challenge, got {other:?}"),
|
||||
}
|
||||
};
|
||||
|
||||
let mut sig = crypto::sign_pipe_auth(&me.sk, &nonce, connection_id);
|
||||
if corrupt_sig {
|
||||
sig[0] ^= 0x01;
|
||||
}
|
||||
let dest = dest_override.unwrap_or_else(|| crypto::sha256(peer_ed));
|
||||
let auth = PipeAuth {
|
||||
connection_id: connection_id.to_vec(),
|
||||
pubkey: me.ed_pub.to_vec(),
|
||||
dest: dest.to_vec(),
|
||||
namespace_id: ns_raw.to_vec(),
|
||||
signature: sig.to_vec(),
|
||||
};
|
||||
ws.send(Message::Binary(pipe::encode(&auth).into())).await.expect("send auth");
|
||||
ws
|
||||
}
|
||||
|
||||
/// Read the next binary frame, or `None` if the socket closed/ended.
|
||||
async fn next_binary(ws: &mut Ws) -> Option<Vec<u8>> {
|
||||
loop {
|
||||
match ws.next().await {
|
||||
Some(Ok(Message::Binary(d))) => return Some(d.to_vec()),
|
||||
Some(Ok(Message::Ping(_))) | Some(Ok(Message::Pong(_))) => continue,
|
||||
Some(Ok(Message::Close(_))) | None => return None,
|
||||
Some(Ok(_)) => continue,
|
||||
Some(Err(_)) => return None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pipe_matches_and_splices_bytes_both_ways() {
|
||||
let (addr, state) = spawn_relay().await;
|
||||
let agent = id_from_seed(1);
|
||||
let client = id_from_seed(2);
|
||||
let ns_raw = seed_namespace(&state, &agent, &client).await;
|
||||
let cid = [0x7Au8; 32];
|
||||
|
||||
// Agent dials first (becomes pending), client second (matches).
|
||||
let mut a = dial_and_auth(addr, &agent, &client.ed_pub, &ns_raw, &cid, false, None).await;
|
||||
// Small delay so A is registered pending before B arrives.
|
||||
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
|
||||
let mut b = dial_and_auth(addr, &client, &agent.ed_pub, &ns_raw, &cid, false, None).await;
|
||||
|
||||
// A → B
|
||||
a.send(Message::Binary(b"hello-from-a".to_vec().into())).await.unwrap();
|
||||
assert_eq!(next_binary(&mut b).await.as_deref(), Some(&b"hello-from-a"[..]));
|
||||
// B → A
|
||||
b.send(Message::Binary(b"hello-from-b".to_vec().into())).await.unwrap();
|
||||
assert_eq!(next_binary(&mut a).await.as_deref(), Some(&b"hello-from-b"[..]));
|
||||
|
||||
// Closing one tears down the other (no orphans).
|
||||
a.close(None).await.unwrap();
|
||||
assert_eq!(next_binary(&mut b).await, None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pipe_rejects_bad_signature() {
|
||||
let (addr, state) = spawn_relay().await;
|
||||
let agent = id_from_seed(3);
|
||||
let client = id_from_seed(4);
|
||||
let ns_raw = seed_namespace(&state, &agent, &client).await;
|
||||
let cid = [0x01u8; 32];
|
||||
|
||||
// Corrupt signature → relay closes without registering a pending pipe.
|
||||
let mut a = dial_and_auth(addr, &agent, &client.ed_pub, &ns_raw, &cid, true, None).await;
|
||||
assert_eq!(next_binary(&mut a).await, None, "relay must close on bad signature");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pipe_rejects_cross_dest_mismatch() {
|
||||
let (addr, state) = spawn_relay().await;
|
||||
let agent = id_from_seed(5);
|
||||
let client = id_from_seed(6);
|
||||
let ns_raw = seed_namespace(&state, &agent, &client).await;
|
||||
let cid = [0x02u8; 32];
|
||||
|
||||
// A targets the client correctly; B (the client) declares the wrong dest
|
||||
// (points at a stranger, not the agent) → cross-ref fails, both torn down.
|
||||
let stranger = crypto::sha256(&[0xEE; 32]);
|
||||
let mut a = dial_and_auth(addr, &agent, &client.ed_pub, &ns_raw, &cid, false, None).await;
|
||||
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
|
||||
let mut b =
|
||||
dial_and_auth(addr, &client, &agent.ed_pub, &ns_raw, &cid, false, Some(stranger)).await;
|
||||
|
||||
assert_eq!(next_binary(&mut b).await, None, "mismatched second side is closed");
|
||||
assert_eq!(next_binary(&mut a).await, None, "first side is torn down too");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pipe_rejects_non_member() {
|
||||
let (addr, state) = spawn_relay().await;
|
||||
let agent = id_from_seed(7);
|
||||
let client = id_from_seed(8);
|
||||
let _ns_raw = seed_namespace(&state, &agent, &client).await;
|
||||
let (ns_raw, _) = crypto::namespace_id(&agent.ed_pub);
|
||||
let outsider = id_from_seed(9); // never authorized in this namespace
|
||||
let cid = [0x03u8; 32];
|
||||
|
||||
let mut a = dial_and_auth(addr, &outsider, &agent.ed_pub, &ns_raw, &cid, false, None).await;
|
||||
assert_eq!(next_binary(&mut a).await, None, "non-member must be rejected");
|
||||
}
|
||||
@@ -0,0 +1,697 @@
|
||||
//! End-to-end protocol tests for the v2 relay transport
|
||||
//! (data/iOS-app/v2/relay-protocol.md). Speaks protobuf binary frames over
|
||||
//! WebSocket against a real axum server bound to an ephemeral port.
|
||||
//!
|
||||
//! Every post-upgrade WS frame is a binary frame (opcode `0x2`) that carries
|
||||
//! exactly one `RelayFrame` protobuf message. The relay speaks first
|
||||
//! (`Challenge`), then the client authenticates with an Ed25519 signature over
|
||||
//! `AUTH_DOMAIN ‖ 0x00 ‖ challenge_nonce_raw(32B)`; see
|
||||
//! `skald_relay_common::crypto::challenge_message`.
|
||||
|
||||
use std::net::SocketAddr;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use bytes::Bytes;
|
||||
use ed25519_dalek::{Signer, SigningKey};
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use prost::Message as _;
|
||||
use sha2::{Digest, Sha256};
|
||||
use skald_relay_common::proto::v2::{
|
||||
self, Auth, AuthAgent, AuthClient, AuthError, AuthOk, AuthPairing, Authorize, AuthorizeOk,
|
||||
ClientPaired, Message as ProtoMessage, PairingReady, PairingStart, PeerOffline, PresenceEvent,
|
||||
PresenceList, PresenceRequest, RelayFrame,
|
||||
};
|
||||
use skald_relay_common::proto::v2::auth::Role as AuthRole;
|
||||
use skald_relay_common::proto::v2::relay_frame::Frame;
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
|
||||
use skald_relay_server::config::Config;
|
||||
use skald_relay_server::{AppState, router};
|
||||
|
||||
type Ws =
|
||||
tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test harness
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Boot a relay on a random port with a throwaway SQLite file. Returns its addr.
|
||||
///
|
||||
/// Each test gets its own DB file. We use `std::process::id()` + a per-call
|
||||
/// counter (incremented atomically across the whole process) so two tests
|
||||
/// calling `spawn_relay()` in parallel — even on the same nanosecond — never
|
||||
/// collide on the file path. A `spawn-relay-tests` counter is also fine, but
|
||||
/// `AtomicU64` is independent of any test framework / test name.
|
||||
async fn spawn_relay() -> SocketAddr {
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
static COUNTER: AtomicU64 = AtomicU64::new(0);
|
||||
let nanos = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.expect("clock")
|
||||
.as_nanos();
|
||||
let seq = COUNTER.fetch_add(1, Ordering::Relaxed);
|
||||
let db = std::env::temp_dir().join(format!(
|
||||
"relay-it-{nanos}-{}-{seq}.db",
|
||||
std::process::id()
|
||||
));
|
||||
let cfg = Config {
|
||||
bind: "127.0.0.1:0".parse().expect("bind addr"),
|
||||
db_path: db.to_string_lossy().into(),
|
||||
pipe: skald_relay_server::config::PipeConfig::default(),
|
||||
};
|
||||
let state = AppState::build(cfg).await.expect("build state");
|
||||
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("bind");
|
||||
let addr = listener.local_addr().expect("local_addr");
|
||||
tokio::spawn(async move {
|
||||
axum::serve(
|
||||
listener,
|
||||
router(state).into_make_service_with_connect_info::<SocketAddr>(),
|
||||
)
|
||||
.await
|
||||
.expect("serve");
|
||||
});
|
||||
addr
|
||||
}
|
||||
|
||||
async fn connect(addr: SocketAddr) -> Ws {
|
||||
let url = format!("ws://{addr}/v1/ws");
|
||||
let (ws, _) = tokio_tungstenite::connect_async(url)
|
||||
.await
|
||||
.expect("connect");
|
||||
ws
|
||||
}
|
||||
|
||||
/// Send a protobuf `RelayFrame` as a WebSocket **binary** frame (v2 transport,
|
||||
/// relay-protocol.md §1).
|
||||
async fn send(ws: &mut Ws, frame: &RelayFrame) {
|
||||
let bytes = frame.encode_to_vec();
|
||||
ws.send(Message::Binary(bytes.into()))
|
||||
.await
|
||||
.expect("send binary");
|
||||
}
|
||||
|
||||
/// Read the next `RelayFrame`. WS-level Ping/Pong are silently consumed
|
||||
/// (axum/tokio-tungstenite handle the actual pong). A `Text` frame or a WS
|
||||
/// `Close` is a protocol violation under v2 — we panic with a clear message.
|
||||
async fn recv(ws: &mut Ws) -> RelayFrame {
|
||||
loop {
|
||||
let m = ws.next().await.expect("stream open").expect("ws frame");
|
||||
match m {
|
||||
Message::Binary(b) => {
|
||||
return RelayFrame::decode(b.as_ref()).expect("decode protobuf");
|
||||
}
|
||||
Message::Ping(_) | Message::Pong(_) => continue,
|
||||
Message::Close(f) => panic!("unexpected ws close: {f:?}"),
|
||||
Message::Text(t) => panic!("unexpected text frame in v2 transport: {t}"),
|
||||
other => panic!("unexpected ws frame: {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Frame builders + crypto helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Sign the v2 challenge message: `AUTH_DOMAIN ‖ 0x00 ‖ nonce(32B)`.
|
||||
/// Mirrors `skald_relay_common::crypto::challenge_message` exactly.
|
||||
fn sign_challenge(sk: &SigningKey, nonce: &[u8; 32]) -> [u8; 64] {
|
||||
let mut msg = Vec::with_capacity(b"skald-relay-auth-v1".len() + 1 + 32);
|
||||
msg.extend_from_slice(b"skald-relay-auth-v1");
|
||||
msg.push(0);
|
||||
msg.extend_from_slice(nonce);
|
||||
sk.sign(&msg).to_bytes()
|
||||
}
|
||||
|
||||
/// `namespace_id` = `hex(SHA256(NS_DOMAIN ‖ 0x00 ‖ agent_ed25519_pub))`
|
||||
/// (crypto.md §7). Returns the raw 32-byte value and the lowercase hex string.
|
||||
fn namespace_id(pubkey: &[u8; 32]) -> ([u8; 32], String) {
|
||||
let mut h = Sha256::new();
|
||||
h.update(b"skald-namespace-v1");
|
||||
h.update([0u8]);
|
||||
h.update(pubkey);
|
||||
let raw = h.finalize();
|
||||
let mut out = [0u8; 32];
|
||||
out.copy_from_slice(&raw);
|
||||
(out, hex::encode(raw))
|
||||
}
|
||||
|
||||
/// Read the relay's first frame — must be `RelayFrame::Challenge{nonce}`.
|
||||
async fn read_challenge(ws: &mut Ws) -> [u8; 32] {
|
||||
let frame = recv(ws).await;
|
||||
match frame.frame {
|
||||
Some(Frame::Challenge(c)) => c.nonce.as_ref().try_into().expect("32B challenge"),
|
||||
other => panic!("expected Challenge, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
/// `Auth{role=Agent(...), signature}` — agent handshake.
|
||||
fn auth_agent_frame(sk: &SigningKey, challenge: &[u8; 32]) -> RelayFrame {
|
||||
let sig = sign_challenge(sk, challenge);
|
||||
let pubkey = sk.verifying_key().to_bytes();
|
||||
RelayFrame {
|
||||
frame: Some(Frame::Auth(Auth {
|
||||
signature: Bytes::copy_from_slice(&sig),
|
||||
role: Some(AuthRole::Agent(AuthAgent {
|
||||
agent_ed25519_pub: Bytes::copy_from_slice(&pubkey),
|
||||
})),
|
||||
})),
|
||||
}
|
||||
}
|
||||
|
||||
/// `Auth{role=Client(...), signature}` — client handshake.
|
||||
fn auth_client_frame(sk: &SigningKey, challenge: &[u8; 32], ns_hex: &str) -> RelayFrame {
|
||||
let sig = sign_challenge(sk, challenge);
|
||||
let pubkey = sk.verifying_key().to_bytes();
|
||||
let ns_raw: [u8; 32] = hex::decode(ns_hex).expect("ns hex").try_into().expect("32B ns");
|
||||
RelayFrame {
|
||||
frame: Some(Frame::Auth(Auth {
|
||||
signature: Bytes::copy_from_slice(&sig),
|
||||
role: Some(AuthRole::Client(AuthClient {
|
||||
namespace_id: Bytes::copy_from_slice(&ns_raw),
|
||||
client_ed25519_pub: Bytes::copy_from_slice(&pubkey),
|
||||
device_token: "devtok".into(),
|
||||
platform: v2::Platform::Ios as i32,
|
||||
})),
|
||||
})),
|
||||
}
|
||||
}
|
||||
|
||||
/// `Auth{role=Pairing(...), signature}` — short-lived pairing connection.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn auth_pairing_frame(
|
||||
sk: &SigningKey,
|
||||
challenge: &[u8; 32],
|
||||
ns_hex: &str,
|
||||
token: &[u8; 32],
|
||||
x25519_pub: &[u8; 32],
|
||||
) -> RelayFrame {
|
||||
let sig = sign_challenge(sk, challenge);
|
||||
let pubkey = sk.verifying_key().to_bytes();
|
||||
let ns_raw: [u8; 32] = hex::decode(ns_hex).expect("ns hex").try_into().expect("32B ns");
|
||||
RelayFrame {
|
||||
frame: Some(Frame::Auth(Auth {
|
||||
signature: Bytes::copy_from_slice(&sig),
|
||||
role: Some(AuthRole::Pairing(AuthPairing {
|
||||
namespace_id: Bytes::copy_from_slice(&ns_raw),
|
||||
client_ed25519_pub: Bytes::copy_from_slice(&pubkey),
|
||||
client_x25519_pub: Bytes::copy_from_slice(x25519_pub),
|
||||
pairing_token: Bytes::copy_from_slice(token),
|
||||
device_token: "devtok".into(),
|
||||
platform: v2::Platform::Ios as i32,
|
||||
})),
|
||||
})),
|
||||
}
|
||||
}
|
||||
|
||||
/// Authenticate as `agent`; returns the live connection and the namespace hex.
|
||||
async fn auth_agent(addr: SocketAddr, sk: &SigningKey) -> (Ws, String) {
|
||||
let pubkey = sk.verifying_key().to_bytes();
|
||||
let mut ws = connect(addr).await;
|
||||
let challenge = read_challenge(&mut ws).await;
|
||||
send(&mut ws, &auth_agent_frame(sk, &challenge)).await;
|
||||
let frame = recv(&mut ws).await;
|
||||
let AuthOk { namespace_id: ns_bytes } = match frame.frame {
|
||||
Some(Frame::AuthOk(ok)) => ok,
|
||||
other => panic!("expected AuthOk, got {other:?}"),
|
||||
};
|
||||
let ns_hex = hex::encode(&ns_bytes);
|
||||
let (want_raw, want_hex) = namespace_id(&pubkey);
|
||||
assert_eq!(
|
||||
ns_hex, want_hex,
|
||||
"AuthOk.namespace_id must match SHA256(NS_DOMAIN‖0x00‖pubkey)"
|
||||
);
|
||||
// The wire carries the raw 32B value; compare bytes too.
|
||||
assert_eq!(ns_bytes.as_ref(), want_raw.as_ref());
|
||||
(ws, ns_hex)
|
||||
}
|
||||
|
||||
/// Authenticate as `client`; returns the live connection. Caller is
|
||||
/// responsible for draining the agent-side `PresenceEvent{ONLINE}` that the
|
||||
/// relay broadcasts on auth_ok.
|
||||
async fn auth_client(addr: SocketAddr, sk: &SigningKey, ns_hex: &str) -> Ws {
|
||||
let mut ws = connect(addr).await;
|
||||
let challenge = read_challenge(&mut ws).await;
|
||||
send(&mut ws, &auth_client_frame(sk, &challenge, ns_hex)).await;
|
||||
let frame = recv(&mut ws).await;
|
||||
match frame.frame {
|
||||
Some(Frame::AuthOk(_)) => {}
|
||||
other => panic!("expected AuthOk, got {other:?}"),
|
||||
}
|
||||
ws
|
||||
}
|
||||
|
||||
/// `PairingStart{pairing_token, ttl}` — open a pairing window on the agent.
|
||||
async fn send_pairing_start(ws: &mut Ws, token: &[u8; 32], ttl: u32) {
|
||||
let frame = RelayFrame {
|
||||
frame: Some(Frame::PairingStart(PairingStart {
|
||||
pairing_token: Bytes::copy_from_slice(token),
|
||||
ttl,
|
||||
})),
|
||||
};
|
||||
send(ws, &frame).await;
|
||||
}
|
||||
|
||||
/// `Authorize{clients[]}` — replace-semantics on the authorized set.
|
||||
async fn send_authorize(ws: &mut Ws, clients: &[[u8; 32]]) {
|
||||
let frame = RelayFrame {
|
||||
frame: Some(Frame::Authorize(Authorize {
|
||||
clients: clients
|
||||
.iter()
|
||||
.map(|c| Bytes::copy_from_slice(c))
|
||||
.collect(),
|
||||
})),
|
||||
};
|
||||
send(ws, &frame).await;
|
||||
}
|
||||
|
||||
/// End-to-end pairing flow on a side connection: `challenge → auth(pairing)
|
||||
/// → AuthOk → close`. Returns the freshly-paired `client_pub` and the
|
||||
/// `x25519_pub` we lied about — the relay never inspects X25519 material.
|
||||
async fn pair_client(
|
||||
addr: SocketAddr,
|
||||
client_sk: &SigningKey,
|
||||
ns_hex: &str,
|
||||
token: &[u8; 32],
|
||||
x25519_pub: [u8; 32],
|
||||
) -> [u8; 32] {
|
||||
let mut pairing = connect(addr).await;
|
||||
let c = read_challenge(&mut pairing).await;
|
||||
send(
|
||||
&mut pairing,
|
||||
&auth_pairing_frame(client_sk, &c, ns_hex, token, &x25519_pub),
|
||||
)
|
||||
.await;
|
||||
let ok = recv(&mut pairing).await;
|
||||
match ok.frame {
|
||||
Some(Frame::AuthOk(_)) => {}
|
||||
other => panic!("pairing expected AuthOk, got {other:?}"),
|
||||
};
|
||||
// The relay sends a Close after AuthOk on a pairing connection — draining
|
||||
// the next frame is optional; let it drop here.
|
||||
drop(pairing);
|
||||
client_sk.verifying_key().to_bytes()
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Agent-side happy path: the relay speaks first, returns `AuthOk` with the
|
||||
/// correct 32-byte `namespace_id`, and accepts no further peer (just registers
|
||||
/// the agent in the registry). Ported from the v1 test, but speaks binary
|
||||
/// protobuf.
|
||||
#[tokio::test]
|
||||
async fn agent_handshake_creates_namespace() {
|
||||
let addr = spawn_relay().await;
|
||||
let sk = SigningKey::from_bytes(&[1u8; 32]);
|
||||
let (_agent, ns) = auth_agent(addr, &sk).await;
|
||||
let (want_raw, want_hex) = namespace_id(&sk.verifying_key().to_bytes());
|
||||
assert_eq!(hex::encode(want_raw), ns);
|
||||
assert_eq!(ns, want_hex);
|
||||
}
|
||||
|
||||
/// A signature that doesn't cover the real challenge must be rejected with
|
||||
/// `AuthError{code = "invalid_signature"}`. The relay closes the socket right
|
||||
/// after; we drain the Close so the test doesn't panic on it.
|
||||
#[tokio::test]
|
||||
async fn bad_signature_is_rejected() {
|
||||
let addr = spawn_relay().await;
|
||||
let sk = SigningKey::from_bytes(&[1u8; 32]);
|
||||
let pubkey = sk.verifying_key().to_bytes();
|
||||
let mut ws = connect(addr).await;
|
||||
let _challenge = read_challenge(&mut ws).await;
|
||||
// Sign a different message — the signature won't verify against the
|
||||
// real challenge nonce.
|
||||
let bogus = sk.sign(b"not the challenge").to_bytes();
|
||||
send(
|
||||
&mut ws,
|
||||
&RelayFrame {
|
||||
frame: Some(Frame::Auth(Auth {
|
||||
signature: Bytes::copy_from_slice(&bogus),
|
||||
role: Some(AuthRole::Agent(AuthAgent {
|
||||
agent_ed25519_pub: Bytes::copy_from_slice(&pubkey),
|
||||
})),
|
||||
})),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
let err = recv(&mut ws).await;
|
||||
let AuthError { code, message: _ } = match err.frame {
|
||||
Some(Frame::AuthError(e)) => e,
|
||||
other => panic!("expected AuthError, got {other:?}"),
|
||||
};
|
||||
assert_eq!(code, "invalid_signature");
|
||||
// The relay follows AuthError with a Close — drain it so we exit cleanly.
|
||||
let close = ws.next().await.expect("stream").expect("ws frame");
|
||||
assert!(matches!(close, Message::Close(_)));
|
||||
}
|
||||
|
||||
/// A client that never paired/was authorized cannot connect as `client`. The
|
||||
/// relay must answer `AuthError{code = "unauthorized"}` and close.
|
||||
#[tokio::test]
|
||||
async fn unauthorized_client_is_rejected() {
|
||||
let addr = spawn_relay().await;
|
||||
let agent_sk = SigningKey::from_bytes(&[1u8; 32]);
|
||||
let (_agent, ns) = auth_agent(addr, &agent_sk).await;
|
||||
|
||||
let client_sk = SigningKey::from_bytes(&[2u8; 32]);
|
||||
let mut ws = connect(addr).await;
|
||||
let challenge = read_challenge(&mut ws).await;
|
||||
send(&mut ws, &auth_client_frame(&client_sk, &challenge, &ns)).await;
|
||||
let err = recv(&mut ws).await;
|
||||
let AuthError { code, .. } = match err.frame {
|
||||
Some(Frame::AuthError(e)) => e,
|
||||
other => panic!("expected AuthError, got {other:?}"),
|
||||
};
|
||||
assert_eq!(code, "unauthorized");
|
||||
let close = ws.next().await.expect("stream").expect("ws frame");
|
||||
assert!(matches!(close, Message::Close(_)));
|
||||
}
|
||||
|
||||
/// End-to-end pairing → `Authorize` → E2E `Message` flow. The relay must:
|
||||
/// 1. Accept a `PairingStart` from the agent and respond with `PairingReady`.
|
||||
/// 2. Accept a short-lived `auth(pairing)` connection, close it, and forward
|
||||
/// `ClientPaired` to the agent.
|
||||
/// 3. Accept an `Authorize` and reply with `AuthorizeOk{authorized: 1}`.
|
||||
/// 4. Accept the `auth(client)` connection, send `AuthOk`, and broadcast
|
||||
/// `PresenceEvent{ONLINE}` to the agent.
|
||||
/// 5. Forward `Message{live:false}` agent→client, rewriting `peer = from` and
|
||||
/// passing `ciphertext`/`nonce` byte-for-byte; same for client→agent.
|
||||
#[tokio::test]
|
||||
async fn pairing_authorize_and_live_message() {
|
||||
let addr = spawn_relay().await;
|
||||
let agent_sk = SigningKey::from_bytes(&[1u8; 32]);
|
||||
let agent_pub = agent_sk.verifying_key().to_bytes();
|
||||
let (mut agent, ns) = auth_agent(addr, &agent_sk).await;
|
||||
|
||||
// 1) Agent opens a pairing window.
|
||||
let token = [0x11u8; 32];
|
||||
send_pairing_start(&mut agent, &token, 300).await;
|
||||
let ready = recv(&mut agent).await;
|
||||
let PairingReady { ttl } = match ready.frame {
|
||||
Some(Frame::PairingReady(p)) => p,
|
||||
other => panic!("expected PairingReady, got {other:?}"),
|
||||
};
|
||||
assert_eq!(ttl, 300);
|
||||
|
||||
// 2) Client pairs on a side connection.
|
||||
let client_sk = SigningKey::from_bytes(&[2u8; 32]);
|
||||
let client_x = [0x33u8; 32]; // opaque X25519 pubkey; relay never inspects
|
||||
let client_pub = pair_client(addr, &client_sk, &ns, &token, client_x).await;
|
||||
assert_eq!(client_pub, client_sk.verifying_key().to_bytes());
|
||||
|
||||
// 3) Agent is told a device paired.
|
||||
let paired = recv(&mut agent).await;
|
||||
let ClientPaired {
|
||||
client_ed25519_pub,
|
||||
client_x25519_pub,
|
||||
platform,
|
||||
} = match paired.frame {
|
||||
Some(Frame::ClientPaired(p)) => p,
|
||||
other => panic!("expected ClientPaired, got {other:?}"),
|
||||
};
|
||||
assert_eq!(client_ed25519_pub.as_ref(), &client_pub[..]);
|
||||
assert_eq!(client_x25519_pub.as_ref(), &client_x[..]);
|
||||
assert_eq!(platform, v2::Platform::Ios as i32);
|
||||
|
||||
// 4) Agent authorizes the client.
|
||||
send_authorize(&mut agent, &[client_pub]).await;
|
||||
let authorized = recv(&mut agent).await;
|
||||
let AuthorizeOk { authorized } = match authorized.frame {
|
||||
Some(Frame::AuthorizeOk(a)) => a,
|
||||
other => panic!("expected AuthorizeOk, got {other:?}"),
|
||||
};
|
||||
assert_eq!(authorized, 1);
|
||||
|
||||
// 5) Client connects as the authorized role.
|
||||
let mut client = auth_client(addr, &client_sk, &ns).await;
|
||||
|
||||
// 5a) Drain the agent-side PresenceEvent{ONLINE} for the new client.
|
||||
let pe = recv(&mut agent).await;
|
||||
let PresenceEvent { pubkey, status } = match pe.frame {
|
||||
Some(Frame::PresenceEvent(p)) => p,
|
||||
other => panic!("expected PresenceEvent, got {other:?}"),
|
||||
};
|
||||
assert_eq!(pubkey.as_ref(), &client_pub[..]);
|
||||
assert_eq!(status, v2::Status::Online as i32);
|
||||
|
||||
// 6) Agent → client Message{live:false}. The relay stamps `peer = from`
|
||||
// (the agent's pubkey) and forwards `ciphertext`/`nonce` byte-for-byte.
|
||||
let nonce = [0u8; 12];
|
||||
let ciphertext = b"hello world";
|
||||
send(
|
||||
&mut agent,
|
||||
&RelayFrame {
|
||||
frame: Some(Frame::Message(ProtoMessage {
|
||||
ciphertext: Bytes::copy_from_slice(ciphertext),
|
||||
nonce: Bytes::copy_from_slice(&nonce),
|
||||
peer: Bytes::copy_from_slice(&client_pub),
|
||||
live: false,
|
||||
})),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
let msg = recv(&mut client).await;
|
||||
let ProtoMessage {
|
||||
ciphertext: ct,
|
||||
nonce: n,
|
||||
peer: from,
|
||||
live,
|
||||
} = match msg.frame {
|
||||
Some(Frame::Message(m)) => m,
|
||||
other => panic!("expected Message, got {other:?}"),
|
||||
};
|
||||
assert_eq!(ct.as_ref(), ciphertext);
|
||||
assert_eq!(n.as_ref(), &nonce[..]);
|
||||
assert_eq!(from.as_ref(), &agent_pub[..]);
|
||||
assert!(!live, "relay must rewrite live=false on delivery");
|
||||
|
||||
// 7) Client → agent reply routes back.
|
||||
let reply_ct = b"reply";
|
||||
let reply_nonce: [u8; 12] = [0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 1];
|
||||
send(
|
||||
&mut client,
|
||||
&RelayFrame {
|
||||
frame: Some(Frame::Message(ProtoMessage {
|
||||
ciphertext: Bytes::copy_from_slice(reply_ct),
|
||||
nonce: Bytes::copy_from_slice(&reply_nonce),
|
||||
peer: Bytes::copy_from_slice(&agent_pub),
|
||||
live: false,
|
||||
})),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
let back = recv(&mut agent).await;
|
||||
let ProtoMessage {
|
||||
ciphertext: ct,
|
||||
nonce: n,
|
||||
peer: from,
|
||||
..
|
||||
} = match back.frame {
|
||||
Some(Frame::Message(m)) => m,
|
||||
other => panic!("expected Message back, got {other:?}"),
|
||||
};
|
||||
assert_eq!(ct.as_ref(), reply_ct);
|
||||
assert_eq!(n.as_ref(), &reply_nonce[..]);
|
||||
assert_eq!(from.as_ref(), &client_pub[..]);
|
||||
}
|
||||
|
||||
/// v2 live channel: `Message{live:true}` is route-or-fail. If the destination
|
||||
/// isn't connected, the relay answers the sender with `PeerOffline{peer}` —
|
||||
/// no enqueue, no push.
|
||||
///
|
||||
/// To exercise this we register an authorized client that never connects as
|
||||
/// `client` (so its `client_tx` is None), then have the agent live-send to
|
||||
/// it. The relay must return `PeerOffline`.
|
||||
#[tokio::test]
|
||||
async fn live_message_to_offline_peer_returns_peer_offline() {
|
||||
let addr = spawn_relay().await;
|
||||
let agent_sk = SigningKey::from_bytes(&[1u8; 32]);
|
||||
let (mut agent, ns) = auth_agent(addr, &agent_sk).await;
|
||||
|
||||
let client_sk = SigningKey::from_bytes(&[2u8; 32]);
|
||||
let client_pub = client_sk.verifying_key().to_bytes();
|
||||
// Pair + authorize (mirrors the full flow) — the client never connects.
|
||||
let token = [0x11u8; 32];
|
||||
send_pairing_start(&mut agent, &token, 300).await;
|
||||
let _ = recv(&mut agent).await; // PairingReady
|
||||
let _paired = pair_client(addr, &client_sk, &ns, &token, [0x33; 32]).await;
|
||||
let cp = recv(&mut agent).await; // ClientPaired
|
||||
assert!(matches!(cp.frame, Some(Frame::ClientPaired(_))));
|
||||
send_authorize(&mut agent, &[client_pub]).await;
|
||||
let _ = recv(&mut agent).await; // AuthorizeOk
|
||||
|
||||
// The client is registered as authorized but never connects as `client`,
|
||||
// so its `client_tx` is None. The relay must return PeerOffline.
|
||||
let nonce = [0u8; 12];
|
||||
let ct = vec![0u8; 32];
|
||||
send(
|
||||
&mut agent,
|
||||
&RelayFrame {
|
||||
frame: Some(Frame::Message(ProtoMessage {
|
||||
ciphertext: Bytes::copy_from_slice(&ct),
|
||||
nonce: Bytes::copy_from_slice(&nonce),
|
||||
peer: Bytes::copy_from_slice(&client_pub),
|
||||
live: true,
|
||||
})),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
let resp = recv(&mut agent).await;
|
||||
let PeerOffline { peer } = match resp.frame {
|
||||
Some(Frame::PeerOffline(p)) => p,
|
||||
other => panic!("expected PeerOffline, got {other:?}"),
|
||||
};
|
||||
assert_eq!(peer.as_ref(), &client_pub[..]);
|
||||
}
|
||||
|
||||
/// `PresenceRequest` → `PresenceList{online[]}` snapshot, scoped to the
|
||||
/// requester's namespace, includes every connected peer (agent + clients).
|
||||
#[tokio::test]
|
||||
async fn presence_list_returns_online_peers() {
|
||||
let addr = spawn_relay().await;
|
||||
let agent_sk = SigningKey::from_bytes(&[1u8; 32]);
|
||||
let agent_pub = agent_sk.verifying_key().to_bytes();
|
||||
let (mut agent, ns) = auth_agent(addr, &agent_sk).await;
|
||||
|
||||
// Pair + authorize a client, then connect it.
|
||||
let client_sk = SigningKey::from_bytes(&[2u8; 32]);
|
||||
let client_pub = client_sk.verifying_key().to_bytes();
|
||||
let token = [0x11u8; 32];
|
||||
send_pairing_start(&mut agent, &token, 300).await;
|
||||
let _ = recv(&mut agent).await; // PairingReady
|
||||
let _paired = pair_client(addr, &client_sk, &ns, &token, [0x33; 32]).await;
|
||||
let _ = recv(&mut agent).await; // ClientPaired
|
||||
send_authorize(&mut agent, &[client_pub]).await;
|
||||
let _ = recv(&mut agent).await; // AuthorizeOk
|
||||
|
||||
let _client = auth_client(addr, &client_sk, &ns).await;
|
||||
|
||||
// Drain the ONLINE presence event from the agent.
|
||||
let pe = recv(&mut agent).await;
|
||||
let PresenceEvent { pubkey, status } = match pe.frame {
|
||||
Some(Frame::PresenceEvent(p)) => p,
|
||||
other => panic!("expected PresenceEvent, got {other:?}"),
|
||||
};
|
||||
assert_eq!(pubkey.as_ref(), &client_pub[..]);
|
||||
assert_eq!(status, v2::Status::Online as i32);
|
||||
|
||||
// Now ask for the namespace's presence snapshot.
|
||||
send(
|
||||
&mut agent,
|
||||
&RelayFrame {
|
||||
frame: Some(Frame::PresenceRequest(PresenceRequest {})),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
let list = recv(&mut agent).await;
|
||||
let PresenceList { online } = match list.frame {
|
||||
Some(Frame::PresenceList(p)) => p,
|
||||
other => panic!("expected PresenceList, got {other:?}"),
|
||||
};
|
||||
let mut got: Vec<[u8; 32]> = online
|
||||
.iter()
|
||||
.map(|b| b.as_ref().try_into().expect("32B pubkey"))
|
||||
.collect();
|
||||
got.sort();
|
||||
let mut want = vec![agent_pub, client_pub];
|
||||
want.sort();
|
||||
assert_eq!(got, want, "PresenceList must contain agent + client pubkeys");
|
||||
}
|
||||
|
||||
/// `PresenceEvent{ONLINE}` is broadcast at the peer's `auth_ok`. When the
|
||||
/// client disconnects, `PresenceEvent{OFFLINE}` is broadcast to the other
|
||||
/// members of the namespace.
|
||||
#[tokio::test]
|
||||
async fn presence_event_on_auth_ok_and_disconnect() {
|
||||
let addr = spawn_relay().await;
|
||||
let agent_sk = SigningKey::from_bytes(&[1u8; 32]);
|
||||
let (mut _agent, ns) = auth_agent(addr, &agent_sk).await;
|
||||
|
||||
let client_sk = SigningKey::from_bytes(&[2u8; 32]);
|
||||
let client_pub = client_sk.verifying_key().to_bytes();
|
||||
let token = [0x11u8; 32];
|
||||
send_pairing_start(&mut _agent, &token, 300).await;
|
||||
let _ = recv(&mut _agent).await; // PairingReady
|
||||
let _paired = pair_client(addr, &client_sk, &ns, &token, [0x33; 32]).await;
|
||||
let _ = recv(&mut _agent).await; // ClientPaired
|
||||
send_authorize(&mut _agent, &[client_pub]).await;
|
||||
let _ = recv(&mut _agent).await; // AuthorizeOk
|
||||
|
||||
// Connect the client. The agent must see PresenceEvent{ONLINE} for it.
|
||||
let mut client = auth_client(addr, &client_sk, &ns).await;
|
||||
let pe_on = recv(&mut _agent).await;
|
||||
let PresenceEvent { pubkey, status } = match pe_on.frame {
|
||||
Some(Frame::PresenceEvent(p)) => p,
|
||||
other => panic!("expected PresenceEvent (online), got {other:?}"),
|
||||
};
|
||||
assert_eq!(pubkey.as_ref(), &client_pub[..]);
|
||||
assert_eq!(status, v2::Status::Online as i32);
|
||||
|
||||
// Drop the client. The relay must broadcast PresenceEvent{OFFLINE} to the
|
||||
// agent. Dropping a tungstenite stream sends a WS Close; the agent's
|
||||
// reader task observes the end-of-stream and runs the disconnect
|
||||
// cleanup.
|
||||
drop(client);
|
||||
// Give the agent's reader task time to detect the close and broadcast
|
||||
// OFFLINE. 100ms is plenty on a fast loopback; bump if flaky on CI.
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
let pe_off = recv(&mut _agent).await;
|
||||
let PresenceEvent { pubkey, status } = match pe_off.frame {
|
||||
Some(Frame::PresenceEvent(p)) => p,
|
||||
other => panic!("expected PresenceEvent (offline), got {other:?}"),
|
||||
};
|
||||
assert_eq!(pubkey.as_ref(), &client_pub[..]);
|
||||
assert_eq!(status, v2::Status::Offline as i32);
|
||||
}
|
||||
|
||||
/// `Message{live:false}` to an offline peer: the relay enqueues the message
|
||||
/// and never returns `PeerOffline`. (The `live=true` counterpart is covered
|
||||
/// by `live_message_to_offline_peer_returns_peer_offline`.) We assert the
|
||||
/// negative invariant: no `PeerOffline` arrives at the sender.
|
||||
#[tokio::test]
|
||||
async fn store_and_forward_when_peer_offline() {
|
||||
let addr = spawn_relay().await;
|
||||
let agent_sk = SigningKey::from_bytes(&[1u8; 32]);
|
||||
let (mut agent, _ns) = auth_agent(addr, &agent_sk).await;
|
||||
|
||||
let client_sk = SigningKey::from_bytes(&[2u8; 32]);
|
||||
let client_pub = client_sk.verifying_key().to_bytes();
|
||||
let token = [0x11u8; 32];
|
||||
send_pairing_start(&mut agent, &token, 300).await;
|
||||
let _ = recv(&mut agent).await; // PairingReady
|
||||
let _paired = pair_client(addr, &client_sk, &_ns, &token, [0x33; 32]).await;
|
||||
let _ = recv(&mut agent).await; // ClientPaired
|
||||
send_authorize(&mut agent, &[client_pub]).await;
|
||||
let _ = recv(&mut agent).await; // AuthorizeOk
|
||||
|
||||
// Send `live=false` to the offline client. We give the relay a moment
|
||||
// to enqueue and (best-effort) push, then assert that the next frame
|
||||
// the agent reads is NOT a PeerOffline.
|
||||
let nonce = [0u8; 12];
|
||||
let ct = vec![0u8; 32];
|
||||
send(
|
||||
&mut agent,
|
||||
&RelayFrame {
|
||||
frame: Some(Frame::Message(ProtoMessage {
|
||||
ciphertext: Bytes::copy_from_slice(&ct),
|
||||
nonce: Bytes::copy_from_slice(&nonce),
|
||||
peer: Bytes::copy_from_slice(&client_pub),
|
||||
live: false,
|
||||
})),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
// No frame should be coming back. Race against a short timeout.
|
||||
let r = tokio::time::timeout(std::time::Duration::from_millis(150), recv(&mut agent)).await;
|
||||
match r {
|
||||
Err(_) => { /* expected: no response on live=false */ }
|
||||
Ok(RelayFrame {
|
||||
frame: Some(Frame::PeerOffline(_)),
|
||||
}) => panic!("live=false must NOT trigger PeerOffline"),
|
||||
Ok(other) => panic!("unexpected frame on live=false: {other:?}"),
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user