Federated relay tier: peer flooding, per-member isolation, registry admission and relay discovery

src/relay.rs

@@ -3,7 +3,7 @@ use std::net::SocketAddr;
 use std::sync::{Arc, Mutex};
 use std::time::{Duration, Instant};
 
-use anyhow::Result;
+use anyhow::{Result, anyhow};
 use axum::extract::{Query, State};
 use axum::http::{StatusCode, header};
 use axum::response::{IntoResponse, Response};
@@ -15,13 +15,40 @@ use tokio::net::TcpListener;
 
 use crate::crypto::{now_ts, poll_signing_bytes, random_nonce, verify_signature};
 use crate::message::{Envelope, TYPE_AGGREGATE, TYPE_QUERY, TYPE_RESPONSE, timestamp_is_fresh};
+use crate::registry::Watcher;
 
 pub const DEFAULT_CAPACITY: usize = 256;
 const CHALLENGE_TTL: Duration = Duration::from_secs(120);
+const SEEN_TTL: Duration = Duration::from_secs(30);
+
+#[derive(Clone, Debug)]
+pub struct RelayOptions {
+    pub capacity: usize,
+    pub peers: Vec<String>,
+    pub url: Option<String>,
+    pub registry: Option<String>,
+    pub ma_key: Option<String>,
+    pub max_hops: usize,
+}
+
+impl Default for RelayOptions {
+    fn default() -> Self {
+        Self {
+            capacity: DEFAULT_CAPACITY,
+            peers: Vec::new(),
+            url: None,
+            registry: None,
+            ma_key: None,
+            max_hops: 3,
+        }
+    }
+}
 
 #[derive(Default)]
 struct MemberQueue {
     items: VecDeque<(u64, Envelope)>,
+    lagging: bool,
+    missed: u64,
 }
 
 struct Challenge {
@@ -32,57 +59,136 @@ struct Challenge {
 struct Inner {
     members: HashMap<String, MemberQueue>,
     challenges: HashMap<String, Challenge>,
+    seen: HashMap<String, Instant>,
     seq: u64,
 }
 
 pub struct Relay {
     inner: Mutex<Inner>,
-    capacity: usize,
+    options: RelayOptions,
+    registry: Option<Arc<Watcher>>,
+    client: reqwest::Client,
 }
 
 impl Relay {
-    pub fn new(capacity: usize) -> Arc<Self> {
-        Arc::new(Self {
+    pub fn new(options: RelayOptions) -> Result<Arc<Self>> {
+        let registry = match (options.registry.as_deref(), options.ma_key.as_deref()) {
+            (Some(source), Some(ma_key)) => Some(Watcher::new(source, ma_key, None)),
+            (Some(_), None) => return Err(anyhow!("registry configured without ma_key")),
+            _ => None,
+        };
+        if !options.peers.is_empty() && options.url.is_none() {
+            return Err(anyhow!("--url is required when --peer is set"));
+        }
+        let relay = Arc::new(Self {
             inner: Mutex::new(Inner {
                 members: HashMap::new(),
                 challenges: HashMap::new(),
+                seen: HashMap::new(),
                 seq: 0,
             }),
-            capacity,
-        })
+            options,
+            registry,
+            client: reqwest::Client::builder()
+                .timeout(Duration::from_secs(5))
+                .build()?,
+        });
+        if let Some(watcher) = &relay.registry {
+            watcher.load_initial();
+        }
+        Ok(relay)
     }
 
-    fn push(&self, targets: Option<&str>, envelope: Envelope) -> Result<usize, StatusCode> {
-        let mut inner = self.inner.lock().expect("relay lock");
-        if inner
-            .members
-            .values()
-            .any(|q| q.items.len() >= self.capacity)
-        {
-            return Err(StatusCode::TOO_MANY_REQUESTS);
+    fn admit(&self, envelope: &Envelope) -> bool {
+        match &self.registry {
+            Some(watcher) => {
+                watcher.refresh_if_changed();
+                watcher.authorized(&envelope.key, now_ts()).is_some()
+            }
+            None => true,
         }
+    }
+
+    fn mark_seen(&self, signature: &str) -> bool {
+        let mut inner = self.inner.lock().expect("relay lock");
+        let now = Instant::now();
+        inner
+            .seen
+            .retain(|_, at| now.duration_since(*at) < SEEN_TTL);
+        if inner.seen.contains_key(signature) {
+            return false;
+        }
+        inner.seen.insert(signature.to_string(), now);
+        true
+    }
+
+    fn push_local(&self, target: Option<&str>, envelope: Envelope) -> Result<usize, StatusCode> {
+        let mut inner = self.inner.lock().expect("relay lock");
         inner.seq += 1;
         let seq = inner.seq;
-        match targets {
+        match target {
             Some(member) => {
                 let Some(queue) = inner.members.get_mut(member) else {
                     return Err(StatusCode::NOT_FOUND);
                 };
+                if queue.items.len() >= self.options.capacity {
+                    queue.lagging = true;
+                    queue.missed += 1;
+                    return Err(StatusCode::TOO_MANY_REQUESTS);
+                }
                 queue.items.push_back((seq, envelope));
                 Ok(1)
             }
             None => {
-                let from = envelope.key.clone();
-                inner.members.entry(from).or_default();
+                let publisher = envelope.key.clone();
+                inner.members.entry(publisher).or_default();
                 let mut delivered = 0;
                 for queue in inner.members.values_mut() {
-                    queue.items.push_back((seq, envelope.clone()));
-                    delivered += 1;
+                    if queue.items.len() >= self.options.capacity {
+                        queue.lagging = true;
+                        queue.missed += 1;
+                    } else {
+                        queue.items.push_back((seq, envelope.clone()));
+                        delivered += 1;
+                    }
                 }
                 Ok(delivered)
             }
         }
     }
+
+    fn forward(&self, envelope: Envelope, origin: Option<&str>, hops: usize) {
+        if self.options.peers.is_empty() {
+            return;
+        }
+        let Some(url) = self.options.url.clone() else {
+            return;
+        };
+        let peers: Vec<String> = self
+            .options
+            .peers
+            .iter()
+            .filter(|peer| Some(peer.as_str()) != origin)
+            .cloned()
+            .collect();
+        if peers.is_empty() {
+            return;
+        }
+        let client = self.client.clone();
+        tokio::spawn(async move {
+            for peer in peers {
+                let target = format!(
+                    "{}/v1/federation?origin={}&hops={}",
+                    peer.trim_end_matches('/'),
+                    url,
+                    hops
+                );
+                if let Err(error) = client.post(&target).json(&envelope).send().await {
+                    eprintln!("federation forward to {peer} failed: {error}");
+                }
+            }
+        });
+    }
 }
 
 pub fn router(relay: Arc<Relay>) -> Router {
@@ -90,13 +196,30 @@ pub fn router(relay: Arc<Relay>) -> Router {
         .route("/health", get(health))
         .route("/v1/challenge", get(challenge))
         .route("/v1/publish", post(publish))
+        .route("/v1/federation", post(federation))
         .route("/v1/unicast", post(unicast))
         .route("/v1/poll", get(poll))
         .with_state(relay)
 }
 
-pub async fn run(listener: TcpListener, capacity: usize) -> Result<()> {
-    let relay = Relay::new(capacity);
+pub async fn run(listener: TcpListener, options: RelayOptions) -> Result<()> {
+    let relay = Relay::new(options)?;
+    if let Some(watcher) = relay.registry.clone() {
+        if watcher.is_url() {
+            let client = relay.client.clone();
+            if let Err(error) = watcher.fetch(&client).await {
+                eprintln!("initial registry fetch failed: {error}");
+            }
+            tokio::spawn(async move {
+                loop {
+                    tokio::time::sleep(Duration::from_secs(60)).await;
+                    if let Err(error) = watcher.fetch(&client).await {
+                        eprintln!("registry fetch failed: {error}");
+                    }
+                }
+            });
+        }
+    }
     axum::serve(listener, router(relay)).await?;
     Ok(())
 }
@@ -152,14 +275,65 @@ async fn publish(State(relay): State<Arc<Relay>>, Json(envelope): Json<Envelope>
     if envelope.msg_type != TYPE_QUERY {
         return bad_request("relay carries broadcast queries only");
     }
-    match relay.push(None, envelope) {
-        Ok(delivered) => (
-            StatusCode::ACCEPTED,
-            Json(json!({ "delivered": delivered })),
-        )
-            .into_response(),
-        Err(status) => backpressure(status),
+    if !relay.admit(&envelope) {
+        return bad_request("sender not admitted");
     }
+    relay.mark_seen(&envelope.sig);
+    let delivered = relay.push_local(None, envelope.clone()).unwrap_or(0);
+    relay.forward(envelope, None, 1);
+    (
+        StatusCode::ACCEPTED,
+        Json(json!({ "delivered": delivered })),
+    )
+        .into_response()
+}
+
+#[derive(Deserialize)]
+struct FederationParams {
+    origin: Option<String>,
+    hops: Option<usize>,
+}
+
+async fn federation(
+    State(relay): State<Arc<Relay>>,
+    Query(params): Query<FederationParams>,
+    Json(envelope): Json<Envelope>,
+) -> Response {
+    let Some(origin) = params.origin.clone() else {
+        return bad_request("origin required");
+    };
+    if !relay.options.peers.iter().any(|peer| peer == &origin) {
+        return bad_request("unknown peer origin");
+    }
+    if envelope.verify().is_err() {
+        return bad_request("invalid signature");
+    }
+    if !timestamp_is_fresh(envelope.ts, now_ts()) {
+        return bad_request("stale timestamp");
+    }
+    if envelope.msg_type != TYPE_QUERY {
+        return bad_request("relay carries broadcast queries only");
+    }
+    if !relay.admit(&envelope) {
+        return bad_request("sender not admitted");
+    }
+    if !relay.mark_seen(&envelope.sig) {
+        return (
+            StatusCode::ACCEPTED,
+            Json(json!({ "delivered": 0, "duplicate": true })),
+        )
+            .into_response();
+    }
+    let delivered = relay.push_local(None, envelope.clone()).unwrap_or(0);
+    let hops = params.hops.unwrap_or(1);
+    if hops < relay.options.max_hops {
+        relay.forward(envelope, Some(&origin), hops + 1);
+    }
+    (
+        StatusCode::ACCEPTED,
+        Json(json!({ "delivered": delivered })),
+    )
+        .into_response()
 }
 
 #[derive(Deserialize)]
@@ -181,14 +355,17 @@ async fn unicast(
     if envelope.msg_type != TYPE_RESPONSE && envelope.msg_type != TYPE_AGGREGATE {
         return bad_request("unicast carries responses and aggregates only");
     }
-    match relay.push(Some(&params.to), envelope) {
+    if !relay.admit(&envelope) {
+        return bad_request("sender not admitted");
+    }
+    match relay.push_local(Some(&params.to), envelope) {
         Ok(_) => (StatusCode::OK, Json(json!({ "delivered": true }))).into_response(),
         Err(StatusCode::NOT_FOUND) => (
             StatusCode::NOT_FOUND,
             Json(json!({ "error": "member not connected" })),
         )
             .into_response(),
-        Err(status) => backpressure(status),
+        Err(status) => backpressure(status, 0),
     }
 }
 
@@ -225,15 +402,29 @@ async fn poll(State(relay): State<Arc<Relay>>, Query(params): Query<PollParams>)
     let timeout = Duration::from_millis(params.timeout_ms.unwrap_or(25_000).min(60_000));
     let deadline = Instant::now() + timeout;
     loop {
-        let batch: Vec<Envelope> = {
+        let (batch, lagged): (Vec<Envelope>, Option<u64>) = {
             let mut inner = relay.inner.lock().expect("relay lock");
             let queue = inner.members.entry(params.member.clone()).or_default();
-            queue
-                .items
-                .drain(..)
-                .map(|(_, envelope)| envelope)
-                .collect()
+            if queue.lagging {
+                let missed = queue.missed;
+                queue.lagging = false;
+                queue.missed = 0;
+                queue.items.clear();
+                (Vec::new(), Some(missed))
+            } else {
+                (
+                    queue
+                        .items
+                        .drain(..)
+                        .map(|(_, envelope)| envelope)
+                        .collect(),
+                    None,
+                )
+            }
         };
+        if let Some(missed) = lagged {
+            return backpressure(StatusCode::TOO_MANY_REQUESTS, missed);
+        }
         if !batch.is_empty() {
             return (StatusCode::OK, Json(json!({ "messages": batch }))).into_response();
         }
@@ -252,11 +443,11 @@ fn unauthorized(message: &str) -> Response {
     (StatusCode::UNAUTHORIZED, Json(json!({ "error": message }))).into_response()
 }
 
-fn backpressure(status: StatusCode) -> Response {
+fn backpressure(status: StatusCode, missed: u64) -> Response {
     (
         status,
         [(header::RETRY_AFTER, "1")],
-        Json(json!({ "error": "transport backpressure" })),
+        Json(json!({ "error": "lagging", "missed": missed })),
     )
         .into_response()
 }