use std::collections::{BTreeSet, HashMap, VecDeque};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, LazyLock, Mutex};
use std::time::Duration;
use axum::body::{Body, Bytes};
use axum::http::{HeaderValue, StatusCode, header};
use axum::response::{IntoResponse, Response};
use base64::Engine as _;
use chrono::{DateTime, Utc};
use hmac::{Hmac, KeyInit, Mac};
use serde_json::Value;
use sha2::{Digest, Sha256};
use tokio::sync::{broadcast, mpsc};
use trust_tasks_rs::TrustTask;
use trust_tasks_rs::specs::vtc::admin::events::{
event::v0_1 as event, event::v0_2 as event_v0_2, subscribe::v0_1 as subscribe,
subscribe::v0_2 as subscribe_v0_2,
};
use vti_common::error::AppError;
use crate::server::AppState;
pub mod codes {
use trust_tasks_rs::specs::vtc::admin::events::subscribe::v0_1::error_codes as e;
pub const NOT_ADMINISTRATOR: &str = e::NOT_ADMINISTRATOR.code;
pub const STREAM_UNAVAILABLE: &str = e::STREAM_UNAVAILABLE.code;
pub const TOO_MANY_STREAMS: &str = e::TOO_MANY_STREAMS.code;
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum Version {
V0_1,
V0_2,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum Topic {
Actions,
Acknowledgements,
JoinRequests,
Members,
SingleAdminMode,
Config,
Rooms,
}
impl Topic {
pub const ALL: [Topic; 7] = [
Topic::Actions,
Topic::Acknowledgements,
Topic::JoinRequests,
Topic::Members,
Topic::SingleAdminMode,
Topic::Config,
Topic::Rooms,
];
pub fn wire(self) -> &'static str {
match self {
Topic::Actions => "actions",
Topic::Acknowledgements => "acknowledgements",
Topic::JoinRequests => "joinRequests",
Topic::Members => "members",
Topic::SingleAdminMode => "singleAdminMode",
Topic::Config => "config",
Topic::Rooms => "rooms",
}
}
pub(crate) fn from_wire(name: &str) -> Option<Topic> {
Topic::ALL.into_iter().find(|t| t.wire() == name)
}
pub(crate) fn in_version(self, version: Version) -> bool {
version == Version::V0_2 || self != Topic::Rooms
}
fn is_fingerprinted(self) -> bool {
self.is_count() || self == Topic::Rooms
}
pub fn is_count(self) -> bool {
matches!(
self,
Topic::Actions | Topic::Acknowledgements | Topic::JoinRequests
)
}
pub(crate) fn from_subscribe(t: &subscribe::Topic) -> Option<Topic> {
Topic::from_wire(&t.to_string()).filter(|t| t.in_version(Version::V0_1))
}
pub(crate) fn from_subscribe_v0_2(t: &subscribe_v0_2::Topic) -> Option<Topic> {
Topic::from_wire(&t.to_string())
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum Signal {
Changed(Topic),
Authority,
}
#[derive(Clone, Copy, Debug)]
pub(crate) struct Change {
pub signal: Signal,
pub at: DateTime<Utc>,
}
pub const HISTORY_SECS: i64 = 10 * 60;
pub const HISTORY_MAX: usize = 4096;
const BUS_CAPACITY: usize = 1024;
struct History {
position: u64,
dropped_through: u64,
entries: VecDeque<(u64, Topic, DateTime<Utc>)>,
}
struct Bus {
tx: broadcast::Sender<Change>,
history: Mutex<History>,
}
static BUS: LazyLock<Bus> = LazyLock::new(|| Bus {
tx: broadcast::channel(BUS_CAPACITY).0,
history: Mutex::new(History {
position: 0,
dropped_through: 0,
entries: VecDeque::new(),
}),
});
fn lock_history() -> std::sync::MutexGuard<'static, History> {
BUS.history.lock().unwrap_or_else(|p| p.into_inner())
}
pub fn notify(topic: Topic) {
let at = Utc::now();
{
let mut h = lock_history();
h.position += 1;
let position = h.position;
h.entries.retain(|e| e.1 != topic);
h.entries.push_back((position, topic, at));
let horizon = at - chrono::TimeDelta::seconds(HISTORY_SECS);
while h.entries.len() > HISTORY_MAX || h.entries.front().is_some_and(|e| e.2 < horizon) {
if let Some((p, _, _)) = h.entries.pop_front() {
h.dropped_through = p;
}
}
}
let _ = BUS.tx.send(Change {
signal: Signal::Changed(topic),
at,
});
}
pub fn notify_rooms() {
notify(Topic::Rooms);
}
pub(crate) fn notify_actions() {
notify(Topic::Actions);
notify(Topic::Acknowledgements);
}
pub(crate) fn notify_authority() {
let _ = BUS.tx.send(Change {
signal: Signal::Authority,
at: Utc::now(),
});
notify(Topic::Members);
notify_actions();
}
fn position() -> u64 {
lock_history().position
}
fn changed_since(since: u64) -> Option<Vec<(Topic, DateTime<Utc>)>> {
let h = lock_history();
if since > h.position || since < h.dropped_through {
return None;
}
let mut latest: Vec<(Topic, DateTime<Utc>)> = Vec::new();
for (p, topic, at) in h.entries.iter() {
if *p <= since {
continue;
}
match latest.iter_mut().find(|(t, _)| t == topic) {
Some(slot) => slot.1 = *at,
None => latest.push((*topic, *at)),
}
}
Some(latest)
}
static TOKEN_KEY: LazyLock<[u8; 32]> = LazyLock::new(|| {
let mut k = [0u8; 32];
rand::fill(&mut k);
k
});
const TOKEN_PREFIX: &str = "e1";
const TOKEN_MAC_LEN: usize = 12;
fn keyed(label: &[u8], caller: &str, extra: &[u8]) -> [u8; 32] {
let mut mac = <Hmac<Sha256>>::new_from_slice(&*TOKEN_KEY).expect("HMAC accepts any key size");
mac.update(label);
mac.update(&(caller.len() as u64).to_be_bytes());
mac.update(caller.as_bytes());
mac.update(extra);
mac.finalize().into_bytes().into()
}
fn mint_token(caller: &str, position: u64, n: u32) -> String {
let mut clear = [0u8; 12];
clear[..8].copy_from_slice(&position.to_be_bytes());
clear[8..].copy_from_slice(&n.to_be_bytes());
let mask = keyed(b"vtc-admin-events/mask\0", caller, &[]);
let mut masked = clear;
for (b, m) in masked.iter_mut().zip(mask.iter()) {
*b ^= m;
}
let tag = keyed(b"vtc-admin-events/mac\0", caller, &clear);
let b64 = base64::engine::general_purpose::URL_SAFE_NO_PAD;
format!(
"{TOKEN_PREFIX}.{}.{}",
b64.encode(masked),
b64.encode(&tag[..TOKEN_MAC_LEN])
)
}
fn redeem_token(caller: &str, token: &str) -> Option<u64> {
let mut parts = token.split('.');
let (Some(TOKEN_PREFIX), Some(masked), Some(tag), None) =
(parts.next(), parts.next(), parts.next(), parts.next())
else {
return None;
};
let b64 = base64::engine::general_purpose::URL_SAFE_NO_PAD;
let masked: [u8; 12] = b64.decode(masked).ok()?.try_into().ok()?;
let tag = b64.decode(tag).ok()?;
if tag.len() != TOKEN_MAC_LEN {
return None;
}
let mask = keyed(b"vtc-admin-events/mask\0", caller, &[]);
let mut clear = masked;
for (b, m) in clear.iter_mut().zip(mask.iter()) {
*b ^= m;
}
let mut mac = <Hmac<Sha256>>::new_from_slice(&*TOKEN_KEY).expect("HMAC accepts any key size");
mac.update(b"vtc-admin-events/mac\0");
mac.update(&(caller.len() as u64).to_be_bytes());
mac.update(caller.as_bytes());
mac.update(&clear);
mac.verify_truncated_left(&tag).ok()?;
Some(u64::from_be_bytes(clear[..8].try_into().ok()?))
}
pub const MAX_STREAMS_PER_SUBJECT: usize = 5;
pub const MAX_STREAMS_TOTAL: usize = 256;
static OPEN: LazyLock<Mutex<(usize, HashMap<String, usize>)>> =
LazyLock::new(|| Mutex::new((0, HashMap::new())));
pub(crate) struct StreamPermit {
subject: String,
}
fn acquire(subject: &str) -> Option<StreamPermit> {
let mut open = OPEN.lock().unwrap_or_else(|p| p.into_inner());
let mine = open.1.get(subject).copied().unwrap_or(0);
if mine >= MAX_STREAMS_PER_SUBJECT || open.0 >= MAX_STREAMS_TOTAL {
return None;
}
open.0 += 1;
open.1.insert(subject.to_string(), mine + 1);
Some(StreamPermit {
subject: subject.to_string(),
})
}
impl Drop for StreamPermit {
fn drop(&mut self) {
let mut open = OPEN.lock().unwrap_or_else(|p| p.into_inner());
open.0 = open.0.saturating_sub(1);
if let Some(n) = open.1.get_mut(&self.subject) {
*n = n.saturating_sub(1);
if *n == 0 {
open.1.remove(&self.subject);
}
}
}
}
pub const DEFAULT_HEARTBEAT_SECS: u64 = 25;
static HEARTBEAT_SECS: AtomicU64 = AtomicU64::new(DEFAULT_HEARTBEAT_SECS);
#[doc(hidden)]
pub fn set_heartbeat_seconds(secs: u64) {
HEARTBEAT_SECS.store(secs.clamp(5, 60), Ordering::Relaxed);
}
fn heartbeat() -> Duration {
Duration::from_secs(HEARTBEAT_SECS.load(Ordering::Relaxed))
}
pub const STREAM_LIFETIME: Duration = Duration::from_secs(60 * 60);
const GATHER: Duration = Duration::from_millis(150);
const MIN_HINT_INTERVAL: Duration = Duration::from_secs(1);
fn write_stall() -> Duration {
heartbeat() * 2
}
pub(crate) fn read_capability(topic: Topic) -> Option<crate::acl::Capability> {
match topic {
Topic::Config => Some(crate::acl::Capability::ConfigAdmin),
Topic::Rooms => Some(crate::acl::Capability::RoomsAdmin),
Topic::Actions
| Topic::Acknowledgements
| Topic::SingleAdminMode
| Topic::Members
| Topic::JoinRequests => None,
}
}
#[derive(Debug, Clone)]
pub(crate) struct Standing {
pub principal: String,
pub topics: BTreeSet<Topic>,
pub unrestricted: bool,
pub expires_at: Option<DateTime<Utc>>,
}
pub(crate) async fn standing(state: &AppState, signer: &str) -> Result<Option<Standing>, AppError> {
let own = crate::acl::get_acl_entry(&state.acl_ks, signer).await?;
let (principal, entry, delegated_until) = match own {
Some(entry) => (signer.to_string(), entry, None),
None => {
let Some(d) =
crate::acl::console_key::resolve_delegated_admin(&state.console_keys_ks, signer)
.await?
else {
return Ok(None);
};
match crate::acl::get_acl_entry(&state.acl_ks, &d.admin_did).await? {
Some(entry) => (d.admin_did.clone(), entry, Some(d.expires_at)),
None => return Ok(None),
}
}
};
if !entry.is_administrator() {
return Ok(None);
}
let topics: BTreeSet<Topic> = if entry.is_suspended() {
[Topic::Actions, Topic::SingleAdminMode]
.into_iter()
.collect()
} else {
Topic::ALL
.into_iter()
.filter(|t| read_capability(*t).is_none_or(|cap| entry.can(cap, None)))
.collect()
};
let entry_until = entry
.expires_at
.and_then(|e| DateTime::<Utc>::from_timestamp(e as i64, 0));
let expires_at = match (entry_until, delegated_until) {
(Some(a), Some(b)) => Some(a.min(b)),
(a, b) => a.or(b),
};
Ok(Some(Standing {
principal,
topics,
unrestricted: entry.can(crate::acl::Capability::AuditRead, None),
expires_at,
}))
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct Observed {
count: u64,
digest: [u8; 32],
}
#[derive(Debug, Default)]
struct Observation {
topics: HashMap<Topic, Observed>,
rooms_wake: Option<DateTime<Utc>>,
}
async fn observe(
state: &AppState,
standing: &Standing,
topics: &BTreeSet<Topic>,
) -> Result<Observation, AppError> {
let mut out = HashMap::new();
let mut rooms_wake = None;
if topics.contains(&Topic::Rooms) {
let (digest, wake) = rooms_fingerprint(state).await?;
out.insert(Topic::Rooms, Observed { count: 0, digest });
rooms_wake = wake;
}
if topics.contains(&Topic::Actions) || topics.contains(&Topic::Acknowledgements) {
let page = crate::admin_actions::list(
state,
&standing.principal,
standing.unrestricted,
crate::admin_actions::View::WaitingForMe,
None,
0,
0,
crate::admin_actions::WireVersion::V0_2,
)
.await?;
if topics.contains(&Topic::Actions) {
let mut visible: Vec<Value> = crate::admin_actions::all(state)
.await?
.iter()
.filter(|r| {
crate::admin_actions::visible_to(r, &standing.principal, standing.unrestricted)
})
.filter_map(|r| serde_json::to_value(r).ok())
.collect();
visible.sort_by(|a, b| a["id"].as_str().cmp(&b["id"].as_str()));
out.insert(
Topic::Actions,
Observed {
count: page.waiting_for_me,
digest: digest(&(page.waiting_for_me, &visible)),
},
);
}
if topics.contains(&Topic::Acknowledgements) {
let owed = page.ext["operatorWritesUnacknowledged"]
.as_array()
.cloned()
.unwrap_or_default();
out.insert(
Topic::Acknowledgements,
Observed {
count: owed.len() as u64,
digest: digest(&owed),
},
);
}
}
if topics.contains(&Topic::JoinRequests) {
let pending = crate::routes::join_requests::read::list_join_requests_inner(
state,
crate::routes::join_requests::read::ListJoinRequestsQuery {
status: Some(crate::join::JoinStatus::Pending),
cursor: None,
limit: Some(1),
},
)
.await?
.total_estimate
.unwrap_or(0);
let requests = crate::join::list_join_requests(&state.join_requests_ks).await?;
let mut shape: Vec<Value> = requests
.iter()
.filter_map(|r| serde_json::to_value(r).ok())
.collect();
shape.sort_by(|a, b| a["id"].as_str().cmp(&b["id"].as_str()));
out.insert(
Topic::JoinRequests,
Observed {
count: pending,
digest: digest(&shape),
},
);
}
Ok(Observation {
topics: out,
rooms_wake,
})
}
fn band(used: u64, limit: u64) -> u8 {
if limit == 0 {
return 100;
}
let (used, limit) = (u128::from(used) * 100, u128::from(limit));
[100u8, 95, 80]
.into_iter()
.find(|pct| used >= limit * u128::from(*pct))
.unwrap_or(0)
}
async fn rooms_fingerprint(
state: &AppState,
) -> Result<([u8; 32], Option<DateTime<Utc>>), AppError> {
let now = Utc::now().timestamp().max(0) as u64;
let files = &state.room_files;
let with_files: std::collections::HashSet<String> =
files.rooms_with_files().await?.into_iter().collect();
let mut rooms = vti_rooms::storage::list_rooms(&state.rooms_ks).await?;
rooms.sort_by(|a, b| a.room_id.cmp(&b.room_id));
let mut shape: Vec<Value> = Vec::with_capacity(rooms.len());
let mut wake: Option<u64> = None;
for room in &rooms {
let room_band = if with_files.contains(&room.room_id) {
let usage = files.room_usage(&room.room_id).await?;
let limits = files.effective_limits(room).await?.room;
let bytes = limits
.max_bytes
.map_or(0, |l| band(usage.bytes + usage.reserved_bytes, l));
let count = limits.max_files.map_or(0, |l| band(usage.files, l));
bytes.max(count)
} else {
0
};
shape.push(serde_json::json!([
room.room_id,
room.lifecycle(now).as_str(),
room_band
]));
if let Some(next) = room.next_lifecycle_change(now) {
wake = Some(wake.map_or(next, |w| w.min(next)));
}
}
let holdings = files.config_holdings().await?;
let mut configs = state.room_storage.configs().await?;
configs.sort_by(|a, b| a["id"].as_str().cmp(&b["id"].as_str()));
for config in configs {
let id = config["id"].as_str().unwrap_or_default();
let held = holdings.get(id).map_or(0, |h| h.bytes);
let config_band = config["capacityBytes"]
.as_u64()
.map_or(0, |cap| band(held, cap));
shape.push(serde_json::json!([config, config_band]));
}
let wake = wake.and_then(|w| DateTime::<Utc>::from_timestamp(w as i64, 0));
Ok((digest(&shape), wake))
}
fn digest<T: serde::Serialize>(v: &T) -> [u8; 32] {
Sha256::digest(serde_json::to_vec(v).unwrap_or_default()).into()
}
#[derive(Clone)]
pub(crate) struct StreamSlot {
inner: Arc<SlotInner>,
}
struct SlotInner {
accepts_stream: bool,
last_event_id: Option<String>,
grant: Mutex<Option<Grant>>,
}
tokio::task_local! {
static STREAM_SLOT: StreamSlot;
}
impl StreamSlot {
pub(crate) fn new(accepts_stream: bool, last_event_id: Option<String>) -> Self {
Self {
inner: Arc::new(SlotInner {
accepts_stream,
last_event_id,
grant: Mutex::new(None),
}),
}
}
pub(crate) fn accepts_stream(&self) -> bool {
self.inner.accepts_stream
}
pub(crate) fn last_event_id(&self) -> Option<&str> {
self.inner.last_event_id.as_deref()
}
pub(crate) fn grant(&self, grant: Grant) {
*self.inner.grant.lock().unwrap_or_else(|p| p.into_inner()) = Some(grant);
}
pub(crate) fn take(&self) -> Option<Grant> {
self.inner
.grant
.lock()
.unwrap_or_else(|p| p.into_inner())
.take()
}
pub(crate) fn scope<'a, T>(
&self,
fut: std::pin::Pin<Box<dyn std::future::Future<Output = T> + Send + 'a>>,
) -> tokio::task::futures::TaskLocalFuture<
StreamSlot,
std::pin::Pin<Box<dyn std::future::Future<Output = T> + Send + 'a>>,
> {
STREAM_SLOT.scope(self.clone(), fut)
}
}
pub(crate) fn current_slot() -> Option<StreamSlot> {
STREAM_SLOT.try_with(|s| s.clone()).ok()
}
pub(crate) fn accepts_event_stream(accept: Option<&str>) -> bool {
let Some(accept) = accept else { return false };
accept.split(',').any(|range| {
let mut parts = range.split(';').map(str::trim);
let media = parts.next().unwrap_or_default();
if !media.eq_ignore_ascii_case("text/event-stream") {
return false;
}
parts
.filter_map(|p| p.split_once('='))
.find(|(k, _)| k.trim().eq_ignore_ascii_case("q"))
.is_none_or(|(_, q)| q.trim().parse::<f32>().is_ok_and(|q| q > 0.0))
})
}
pub(crate) struct Grant {
rx: broadcast::Receiver<Change>,
_permit: StreamPermit,
signer: String,
vtc_did: String,
parent_thread: String,
topics: BTreeSet<Topic>,
version: Version,
backlog: Vec<(Topic, DateTime<Utc>)>,
baseline: HashMap<Topic, Observed>,
rooms_wake: Option<DateTime<Utc>>,
deadline: DateTime<Utc>,
first_id: String,
position: u64,
}
pub(crate) struct Opened {
pub response: Value,
pub grant: Grant,
}
pub(crate) enum Refusal {
NoTopics,
TooMany,
Internal(AppError),
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn open(
state: &AppState,
signer: &str,
standing: &Standing,
version: Version,
requested: &BTreeSet<Topic>,
since: Option<&str>,
parent_thread: String,
document_expires: Option<DateTime<Utc>>,
) -> Result<Opened, Refusal> {
let topics: BTreeSet<Topic> = requested
.intersection(&standing.topics)
.copied()
.filter(|t| t.in_version(version))
.collect();
if topics.is_empty() {
return Err(Refusal::NoTopics);
}
let permit = acquire(&standing.principal).ok_or(Refusal::TooMany)?;
let rx = BUS.tx.subscribe();
let position = position();
let (resumed, backlog) = match since
.and_then(|t| redeem_token(signer, t))
.and_then(changed_since)
{
Some(changed) => (
true,
changed
.into_iter()
.filter(|(t, _)| topics.contains(t))
.collect(),
),
None => (false, Vec::new()),
};
let observation = observe(state, standing, &topics)
.await
.map_err(Refusal::Internal)?;
let mut baseline = observation.topics;
for (topic, _) in &backlog {
baseline.remove(topic);
}
let now = Utc::now();
let lifetime =
chrono::TimeDelta::from_std(STREAM_LIFETIME).unwrap_or(chrono::TimeDelta::hours(1));
let deadline = [Some(now + lifetime), document_expires, standing.expires_at]
.into_iter()
.flatten()
.min()
.unwrap_or(now + lifetime);
let first_id = mint_token(signer, position, 0);
let response =
subscribe_response(version, &topics, resumed, &first_id).map_err(Refusal::Internal)?;
let vtc_did = state
.config
.read()
.await
.vtc_did
.clone()
.unwrap_or_default();
Ok(Opened {
response,
grant: Grant {
rx,
_permit: permit,
signer: signer.to_string(),
vtc_did,
parent_thread,
topics,
version,
backlog,
baseline,
rooms_wake: observation.rooms_wake,
deadline,
first_id,
position,
},
})
}
fn subscribe_response(
version: Version,
topics: &BTreeSet<Topic>,
resumed: bool,
token: &str,
) -> Result<Value, AppError> {
let internal = |e: &dyn std::fmt::Display| AppError::Internal(e.to_string());
macro_rules! build {
($m:ident) => {{
let topics = topics
.iter()
.map(|t| t.wire().parse::<$m::Topic>().map_err(|e| internal(&e)))
.collect::<Result<Vec<_>, _>>()?;
let response = $m::Response::builder()
.topics(topics)
.heartbeat_seconds(heartbeat().as_secs() as i64)
.resumed(resumed)
.resume_token($m::ResumeToken::try_from(token).map_err(|e| internal(&e))?);
let response = $m::Response::try_from(response).map_err(|e| internal(&e))?;
serde_json::to_value(response).map_err(|e| internal(&e))
}};
}
match version {
Version::V0_1 => build!(subscribe),
Version::V0_2 => build!(subscribe_v0_2),
}
}
fn hint_document(
grant: &Grant,
topic: Topic,
at: DateTime<Utc>,
count: Option<u64>,
token: &str,
) -> Option<Value> {
macro_rules! build {
($m:ident) => {{
let mut builder = $m::Payload::builder()
.topic(topic.wire().parse::<$m::Topic>().ok()?)
.at(at)
.resume_token($m::ResumeToken::try_from(token).ok()?);
if let Some(c) = count {
builder = builder.count(Some(c));
}
let payload = $m::Payload::try_from(builder).ok()?;
let mut doc =
TrustTask::for_payload(format!("urn:uuid:{}", uuid::Uuid::new_v4()), payload);
doc.parent_thread_id = Some(grant.parent_thread.clone());
doc.issuer = Some(grant.vtc_did.clone());
doc.recipient = Some(grant.signer.clone());
doc.issued_at = DateTime::<Utc>::from_timestamp(Utc::now().timestamp(), 0);
serde_json::to_value(&doc).ok()
}};
}
match grant.version {
Version::V0_1 => build!(event),
Version::V0_2 => build!(event_v0_2),
}
}
fn frame(id: &str, document: &[u8]) -> Option<Bytes> {
if document.contains(&b'\n') || document.contains(&b'\r') || id.contains(['\n', '\r', '\0']) {
return None;
}
let mut out = Vec::with_capacity(document.len() + id.len() + 16);
out.extend_from_slice(b"id: ");
out.extend_from_slice(id.as_bytes());
out.extend_from_slice(b"\ndata: ");
out.extend_from_slice(document);
out.extend_from_slice(b"\n\n");
Some(Bytes::from(out))
}
const HEARTBEAT_FRAME: &[u8] = b": heartbeat\n\n";
pub(crate) fn respond(state: AppState, grant: Grant, response_document: Vec<u8>) -> Response {
let (tx, mut rx) = mpsc::channel::<Bytes>(16);
let Some(first) = frame(&grant.first_id, &response_document) else {
tracing::error!("subscribe #response does not frame as one SSE event");
return (
StatusCode::INTERNAL_SERVER_ERROR,
[(header::CONTENT_TYPE, "application/json")],
response_document,
)
.into_response();
};
let _ = tx.try_send(first);
tokio::spawn(run(state, grant, tx));
let body = Body::from_stream(futures_util::stream::poll_fn(move |cx| {
rx.poll_recv(cx)
.map(|b| b.map(Ok::<_, std::convert::Infallible>))
}));
let mut res = Response::new(body);
*res.status_mut() = StatusCode::OK;
let h = res.headers_mut();
h.insert(
header::CONTENT_TYPE,
HeaderValue::from_static("text/event-stream"),
);
h.insert(header::CACHE_CONTROL, HeaderValue::from_static("no-store"));
h.insert("x-accel-buffering", HeaderValue::from_static("no"));
res
}
#[derive(Debug)]
enum Ended {
ClientGone,
Shutdown,
Deadline,
Stalled,
StandingShrank,
BusClosed,
}
async fn write(tx: &mpsc::Sender<Bytes>, bytes: Bytes) -> Result<(), Ended> {
match tokio::time::timeout(write_stall(), tx.send(bytes)).await {
Ok(Ok(())) => Ok(()),
Ok(Err(_)) => Err(Ended::ClientGone),
Err(_) => Err(Ended::Stalled),
}
}
struct Stream {
state: AppState,
grant: Grant,
pending: HashMap<Topic, DateTime<Utc>>,
last_sent: HashMap<Topic, tokio::time::Instant>,
recheck: bool,
rooms_wake: Option<tokio::time::Instant>,
n: u32,
}
async fn run(state: AppState, mut grant: Grant, tx: mpsc::Sender<Bytes>) {
let mut shutdown = state.shutdown_tx.subscribe();
let backlog = std::mem::take(&mut grant.backlog);
let rooms_wake = grant.rooms_wake.map(wake_instant);
let mut stream = Stream {
state,
grant,
pending: backlog.into_iter().collect(),
last_sent: HashMap::new(),
recheck: false,
rooms_wake,
n: 0,
};
let ended = stream.drive(&tx, &mut shutdown).await;
tracing::debug!(?ended, "admin event stream ended");
}
fn wake_instant(at: DateTime<Utc>) -> tokio::time::Instant {
let left = (at - Utc::now()).to_std().unwrap_or(Duration::ZERO);
tokio::time::Instant::now() + left.min(STREAM_LIFETIME) + Duration::from_secs(1)
}
impl Stream {
fn deadline_instant(&self) -> tokio::time::Instant {
let left = (self.grant.deadline - Utc::now())
.to_std()
.unwrap_or(Duration::ZERO);
tokio::time::Instant::now() + left
}
fn next_flush(&self) -> Option<tokio::time::Instant> {
if self.pending.is_empty() && !self.recheck {
return None;
}
let now = tokio::time::Instant::now();
let earliest = self
.pending
.keys()
.map(|t| {
self.last_sent
.get(t)
.map_or(now, |s| (*s + MIN_HINT_INTERVAL).max(now))
})
.min()
.unwrap_or(now);
Some(earliest.max(now) + GATHER)
}
async fn drive(
&mut self,
tx: &mpsc::Sender<Bytes>,
shutdown: &mut tokio::sync::watch::Receiver<bool>,
) -> Ended {
if *shutdown.borrow() {
return Ended::Shutdown;
}
let mut beat = tokio::time::Instant::now() + heartbeat();
loop {
let deadline = self.deadline_instant();
let flush_at = self.next_flush();
let rooms_wake = self.rooms_wake;
tokio::select! {
biased;
_ = tx.closed() => return Ended::ClientGone,
_ = shutdown.changed() => return Ended::Shutdown,
_ = tokio::time::sleep_until(deadline) => return Ended::Deadline,
msg = self.grant.rx.recv() => match msg {
Ok(change) => self.absorb(change),
Err(broadcast::error::RecvError::Lagged(_)) => {
let now = Utc::now();
for t in self.grant.topics.clone() {
self.pending.insert(t, now);
}
self.recheck = true;
}
Err(broadcast::error::RecvError::Closed) => return Ended::BusClosed,
},
_ = async { tokio::time::sleep_until(flush_at.expect("guarded")).await },
if flush_at.is_some() =>
{
match self.flush(tx).await {
Ok(true) => beat = tokio::time::Instant::now() + heartbeat(),
Ok(false) => {}
Err(ended) => return ended,
}
}
_ = async { tokio::time::sleep_until(rooms_wake.expect("guarded")).await },
if rooms_wake.is_some() =>
{
self.rooms_wake = None;
let now = Utc::now();
let at = self.pending.entry(Topic::Rooms).or_insert(now);
*at = (*at).max(now);
}
_ = tokio::time::sleep_until(beat) => {
if let Err(ended) = self.check_standing().await {
return ended;
}
if let Err(ended) = write(tx, Bytes::from_static(HEARTBEAT_FRAME)).await {
return ended;
}
beat = tokio::time::Instant::now() + heartbeat();
}
}
}
}
fn absorb(&mut self, change: Change) {
match change.signal {
Signal::Changed(topic) if self.grant.topics.contains(&topic) => {
let at = self.pending.entry(topic).or_insert(change.at);
*at = (*at).max(change.at);
}
Signal::Changed(_) => {}
Signal::Authority => self.recheck = true,
}
}
async fn check_standing(&mut self) -> Result<Standing, Ended> {
self.recheck = false;
let standing = match standing(&self.state, &self.grant.signer).await {
Ok(Some(s)) => s,
Ok(None) => return Err(Ended::StandingShrank),
Err(e) => {
tracing::warn!(error = %e, "admin event stream: standing unreadable; ending");
return Err(Ended::StandingShrank);
}
};
if !self.grant.topics.is_subset(&standing.topics) {
return Err(Ended::StandingShrank);
}
if let Some(until) = standing.expires_at {
if until <= Utc::now() {
return Err(Ended::Deadline);
}
self.grant.deadline = self.grant.deadline.min(until);
}
Ok(standing)
}
async fn flush(&mut self, tx: &mpsc::Sender<Bytes>) -> Result<bool, Ended> {
let standing = self.check_standing().await?;
let now = tokio::time::Instant::now();
let due: BTreeSet<Topic> = self
.pending
.keys()
.filter(|t| {
self.last_sent
.get(*t)
.is_none_or(|s| now.duration_since(*s) >= MIN_HINT_INTERVAL)
})
.copied()
.collect();
if due.is_empty() {
return Ok(false);
}
let fingerprinted: BTreeSet<Topic> = due
.iter()
.copied()
.filter(|t| t.is_fingerprinted())
.collect();
let observed = if fingerprinted.is_empty() {
HashMap::new()
} else {
match observe(&self.state, &standing, &fingerprinted).await {
Ok(o) => {
if fingerprinted.contains(&Topic::Rooms) {
self.rooms_wake = o.rooms_wake.map(wake_instant);
}
o.topics
}
Err(e) => {
tracing::warn!(error = %e, "admin event stream: topic unreadable; ending");
return Err(Ended::StandingShrank);
}
}
};
let mut wrote = false;
for topic in due {
let Some(at) = self.pending.remove(&topic) else {
continue;
};
let count = if topic.is_fingerprinted() {
let Some(seen) = observed.get(&topic).copied() else {
continue;
};
if self.grant.baseline.get(&topic) == Some(&seen) {
continue;
}
self.grant.baseline.insert(topic, seen);
topic.is_count().then_some(seen.count)
} else {
None
};
self.n = self.n.saturating_add(1);
let token = mint_token(
&self.grant.signer,
position().max(self.grant.position),
self.n,
);
let Some(bytes) = self.hint(topic, at, count, &token).await else {
continue;
};
write(tx, bytes).await?;
self.last_sent.insert(topic, tokio::time::Instant::now());
wrote = true;
}
Ok(wrote)
}
async fn hint(
&self,
topic: Topic,
at: DateTime<Utc>,
count: Option<u64>,
token: &str,
) -> Option<Bytes> {
let mut value = hint_document(&self.grant, topic, at, count, token)?;
if let Some(signer) = self.state.credential_signer.as_ref()
&& let Err(e) = signer.sign_operational_doc(&mut value).await
{
tracing::warn!(error = %e, "admin event hint sent unsigned");
}
let bytes = serde_json::to_vec(&value).ok()?;
frame(token, &bytes)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_token_round_trips_for_its_caller_only() {
let t = mint_token("did:key:zA", 42, 3);
assert!(
t.chars()
.all(|c| c.is_ascii_alphanumeric() || "._~-".contains(c))
);
assert!(subscribe::ResumeToken::try_from(t.as_str()).is_ok());
assert_eq!(redeem_token("did:key:zA", &t), Some(42));
assert_eq!(redeem_token("did:key:zB", &t), None);
assert_eq!(redeem_token("did:key:zA", "evt.000141"), None);
assert_eq!(redeem_token("did:key:zA", ""), None);
}
#[test]
fn two_callers_tokens_for_one_position_share_nothing() {
let a = mint_token("did:key:zA", 7, 0);
let b = mint_token("did:key:zB", 7, 0);
assert_ne!(a.split('.').nth(1), b.split('.').nth(1));
}
#[test]
fn tokens_on_one_stream_are_distinct_positions() {
assert_ne!(
mint_token("did:key:zA", 7, 1),
mint_token("did:key:zA", 7, 2)
);
}
#[test]
fn accept_negotiation() {
assert!(accepts_event_stream(Some(
"text/event-stream, application/json;q=0.5"
)));
assert!(accepts_event_stream(Some("TEXT/EVENT-STREAM")));
assert!(!accepts_event_stream(Some("application/json")));
assert!(!accepts_event_stream(Some("text/event-stream;q=0")));
assert!(!accepts_event_stream(None));
}
#[test]
fn a_frame_is_one_id_and_one_data_line() {
let f = frame("e1.a.b", br#"{"a":"x\ny"}"#).unwrap();
assert_eq!(
&f[..],
b"id: e1.a.b\ndata: {\"a\":\"x\\ny\"}\n\n".as_slice()
);
assert!(frame("e1", b"{\n}").is_none());
}
#[test]
fn caps_hold_per_subject() {
let subject = format!("did:key:zCap{}", uuid::Uuid::new_v4());
let permits: Vec<_> = (0..MAX_STREAMS_PER_SUBJECT)
.map(|_| acquire(&subject).expect("under the cap"))
.collect();
assert!(acquire(&subject).is_none());
drop(permits);
assert!(acquire(&subject).is_some());
}
#[test]
fn bands_are_the_consoles() {
assert_eq!(band(0, 100), 0);
assert_eq!(band(79, 100), 0);
assert_eq!(band(80, 100), 80);
assert_eq!(band(94, 100), 80);
assert_eq!(band(95, 100), 95);
assert_eq!(band(100, 100), 100);
assert_eq!(band(250, 100), 100);
assert_eq!(band(0, 0), 100, "a limit of 0 refuses everything: full");
assert_eq!(band(u64::MAX, u64::MAX), 100);
}
#[test]
fn the_topic_set_is_closed_per_version() {
for t in Topic::ALL {
assert_eq!(Topic::from_wire(t.wire()), Some(t));
assert!(t.wire().parse::<subscribe_v0_2::Topic>().is_ok());
assert!(t.wire().parse::<event_v0_2::Topic>().is_ok());
assert_eq!(
t.wire().parse::<subscribe::Topic>().is_ok(),
t.in_version(Version::V0_1)
);
assert_eq!(
t.wire().parse::<event::Topic>().is_ok(),
t.in_version(Version::V0_1)
);
}
assert!(!Topic::Rooms.in_version(Version::V0_1));
assert!(!Topic::Rooms.is_count(), "rooms carries no count");
}
#[test]
fn rooms_is_read_with_the_rooms_capability() {
assert_eq!(
read_capability(Topic::Rooms),
Some(crate::acl::Capability::RoomsAdmin)
);
}
#[test]
fn the_history_keeps_one_entry_per_topic_and_still_resumes() {
let before = position();
for _ in 0..50 {
notify(Topic::Rooms);
}
let h = lock_history();
assert!(
h.entries.iter().filter(|e| e.1 == Topic::Rooms).count() <= 1,
"a busy seam does not crowd the history"
);
drop(h);
let changed = changed_since(before).expect("still covered");
assert!(changed.iter().any(|(t, _)| *t == Topic::Rooms));
}
#[test]
fn history_answers_what_changed_after_a_position() {
let before = position();
notify(Topic::Config);
let changed = changed_since(before).expect("covered");
assert!(changed.iter().any(|(t, _)| *t == Topic::Config));
assert!(changed_since(position() + 1).is_none());
}
}