First Version
This commit is contained in:
@@ -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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user