use super::{AppState, Call, CallKind, Failure, cookie, same, valid_token};
use crate::action::Error;
use crate::ctx::{BoxFuture, Ctx, Service};
use axum::Router;
use axum::extract::rejection::QueryRejection;
use axum::extract::ws::{CloseFrame, Message, WebSocket, WebSocketUpgrade};
use axum::extract::{Path, Query, State};
use axum::http::{HeaderMap, StatusCode, header};
use axum::response::{IntoResponse, Response};
use axum::routing::get;
use futures_util::{SinkExt, StreamExt};
use serde::Deserialize;
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use std::time::{Duration, Instant};
use tokio::sync::mpsc;
use yrs::updates::decoder::Decode;
use yrs::{Doc, ReadTxn, StateVector, Transact, Update};
const JS: &str = include_str!("../../assets/live.js");
const LOADER: &str = "/_zeyn/live.js";
pub(super) const MARKER: &str = "data-zeyn-live=\"";
pub const FRAME_UPDATE: u8 = 1;
pub const CLOSE_SLOW: u16 = 1013;
pub const CLOSE_POLICY: u16 = 1008;
pub const CLOSE_TOO_BIG: u16 = 1009;
#[derive(Clone, Copy, Debug)]
pub struct Limits {
pub per_user: usize,
pub per_room: usize,
pub rooms: usize,
pub queue: usize,
pub message_bytes: usize,
pub doc_bytes: usize,
pub idle: Duration,
pub heartbeat: Duration,
pub persist_every: Duration,
}
impl Default for Limits {
fn default() -> Self {
Limits {
per_user: 8,
per_room: 32,
rooms: 256,
queue: 32,
message_bytes: 64 * 1024,
doc_bytes: 1024 * 1024,
idle: Duration::from_secs(30),
heartbeat: Duration::from_secs(20),
persist_every: Duration::from_secs(2),
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct RoomKey {
pub tenant: String,
pub kind: String,
pub doc: String,
}
pub trait DocStore: Send + Sync + 'static {
fn load(&self, key: RoomKey) -> BoxFuture<crate::Result<Option<Vec<u8>>>>;
fn save(&self, key: RoomKey, state: Vec<u8>) -> BoxFuture<crate::Result<()>>;
}
#[derive(Clone, Default)]
pub struct MemoryDocStore(Arc<Mutex<HashMap<RoomKey, Vec<u8>>>>);
impl MemoryDocStore {
pub fn get(&self, key: &RoomKey) -> Option<Vec<u8>> {
self.0.lock().ok()?.get(key).cloned()
}
}
impl DocStore for MemoryDocStore {
fn load(&self, key: RoomKey) -> BoxFuture<crate::Result<Option<Vec<u8>>>> {
let state = self.get(&key);
Box::pin(async move { Ok(state) })
}
fn save(&self, key: RoomKey, state: Vec<u8>) -> BoxFuture<crate::Result<()>> {
if let Ok(mut m) = self.0.lock() {
m.insert(key, state);
}
Box::pin(async { Ok(()) })
}
}
struct DocKind {
read: &'static str,
write: &'static str,
store: Arc<dyn DocStore>,
}
struct Config {
kinds: HashMap<&'static str, DocKind>,
topics: HashMap<&'static str, &'static str>,
limits: Limits,
origins: Vec<String>,
}
pub struct Realtime {
hub: Hub,
config: Config,
}
impl Default for Realtime {
fn default() -> Self {
Self::new()
}
}
impl Realtime {
pub fn new() -> Self {
Realtime {
hub: Hub::default(),
config: Config {
kinds: HashMap::new(),
topics: HashMap::new(),
limits: Limits::default(),
origins: Vec::new(),
},
}
}
pub fn documents(
mut self,
kind: &'static str,
read: &'static str,
write: &'static str,
store: impl DocStore,
) -> Self {
self.config.kinds.insert(kind, DocKind { read, write, store: Arc::new(store) });
self
}
pub fn topic(mut self, name: &'static str, permission: &'static str) -> Self {
self.config.topics.insert(name, permission);
self
}
pub fn limits(mut self, limits: Limits) -> Self {
self.config.limits = limits;
self
}
pub fn allow_origin(mut self, origin: impl Into<String>) -> Self {
self.config.origins.push(origin.into().trim_end_matches('/').to_string());
self
}
pub fn hub(&self) -> Hub {
self.hub.clone()
}
pub(crate) fn install(self) -> Hub {
let _ = self.hub.0.config.set(self.config);
self.hub
}
}
type Outbox = mpsc::Sender<Message>;
struct Room {
doc: Doc,
clients: Mutex<Vec<(u64, Outbox)>>,
idle_since: Mutex<Option<Instant>>,
dirty: AtomicBool,
bytes: AtomicUsize,
store: Arc<dyn DocStore>,
}
#[derive(Default)]
struct Inner {
config: OnceLock<Config>,
rooms: Mutex<HashMap<RoomKey, Arc<Room>>>,
users: Mutex<HashMap<String, usize>>,
topics: Mutex<HashMap<(String, String), Vec<(u64, Outbox)>>>,
next: AtomicU64,
sweeping: AtomicBool,
}
#[derive(Clone, Default)]
pub struct Hub(Arc<Inner>);
impl Service for Hub {}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct HubStats {
pub rooms: usize,
pub clients: usize,
pub subscribers: usize,
pub doc_bytes: usize,
}
impl Hub {
pub fn notify(&self, ctx: &Ctx, topic: &str) {
self.notify_tenant(ctx.tenant().map(|t| t.id()), topic);
}
pub fn notify_tenant(&self, tenant: Option<&str>, topic: &str) {
let key = (tenant.unwrap_or("").to_string(), topic.to_string());
let text = serde_json::json!({ "topic": topic }).to_string();
if let Ok(mut topics) = self.0.topics.lock() {
if let Some(subs) = topics.get_mut(&key) {
subs.retain(|(_, tx)| {
!matches!(
tx.try_send(Message::Text(text.clone().into())),
Err(mpsc::error::TrySendError::Closed(_))
)
});
}
}
}
pub fn stats(&self) -> HubStats {
let rooms = self.0.rooms.lock().map(|r| r.values().cloned().collect::<Vec<_>>()).unwrap_or_default();
HubStats {
rooms: rooms.len(),
clients: rooms.iter().map(|r| r.clients.lock().map(|c| c.len()).unwrap_or(0)).sum(),
subscribers: self.0.topics.lock().map(|t| t.values().map(Vec::len).sum()).unwrap_or(0),
doc_bytes: rooms.iter().map(|r| r.bytes.load(Ordering::Relaxed)).sum(),
}
}
pub async fn flush(&self) {
let rooms: Vec<(RoomKey, Arc<Room>)> =
self.0.rooms.lock().map(|r| r.iter().map(|(k, v)| (k.clone(), v.clone())).collect()).unwrap_or_default();
for (key, room) in rooms {
persist(&key, &room).await;
}
}
fn config(&self) -> Option<&Config> {
self.0.config.get()
}
fn take_user(&self, user: &str, max: usize) -> bool {
let Ok(mut users) = self.0.users.lock() else { return false };
let n = users.entry(user.to_string()).or_default();
if *n >= max {
return false;
}
*n += 1;
true
}
fn release_user(&self, user: &str) {
if let Ok(mut users) = self.0.users.lock() {
if let Some(n) = users.get_mut(user) {
*n -= 1;
if *n == 0 {
users.remove(user);
}
}
}
}
async fn open(&self, key: &RoomKey, kind: &DocKind, limits: &Limits) -> Option<Arc<Room>> {
if let Some(room) = self.0.rooms.lock().ok()?.get(key) {
return Some(room.clone());
}
if self.0.rooms.lock().ok()?.len() >= limits.rooms {
return None;
}
let state = kind.store.load(key.clone()).await.ok()?;
let doc = Doc::new();
let mut bytes = 0;
if let Some(state) = state {
bytes = state.len();
let update = Update::decode_v1(&state).ok()?;
doc.transact_mut().apply_update(update).ok()?;
}
let room = Arc::new(Room {
doc,
clients: Mutex::new(Vec::new()),
idle_since: Mutex::new(None),
dirty: AtomicBool::new(false),
bytes: AtomicUsize::new(bytes),
store: kind.store.clone(),
});
let mut rooms = self.0.rooms.lock().ok()?;
if rooms.len() >= limits.rooms && !rooms.contains_key(key) {
return None;
}
Some(rooms.entry(key.clone()).or_insert(room).clone())
}
fn start_sweeper(&self, limits: Limits) {
if self.0.sweeping.swap(true, Ordering::AcqRel) {
return;
}
let hub = Arc::downgrade(&self.0);
tokio::spawn(async move {
let mut tick = tokio::time::interval(limits.persist_every);
loop {
tick.tick().await;
let Some(inner) = hub.upgrade() else { return };
Hub(inner).sweep(limits.idle).await;
}
});
}
async fn sweep(&self, idle: Duration) {
let rooms: Vec<(RoomKey, Arc<Room>)> =
self.0.rooms.lock().map(|r| r.iter().map(|(k, v)| (k.clone(), v.clone())).collect()).unwrap_or_default();
for (key, room) in rooms {
persist(&key, &room).await;
let expired = room.idle_since.lock().ok().and_then(|s| *s).is_some_and(|since| since.elapsed() >= idle);
if expired {
if let Ok(mut map) = self.0.rooms.lock() {
let empty = room.clients.lock().map(|c| c.is_empty()).unwrap_or(true);
if empty && !room.dirty.load(Ordering::Acquire) {
map.remove(&key);
}
}
}
}
}
}
async fn persist(key: &RoomKey, room: &Room) {
if !room.dirty.swap(false, Ordering::AcqRel) {
return;
}
let state = room.doc.transact().encode_state_as_update_v1(&StateVector::default());
room.bytes.store(state.len(), Ordering::Relaxed);
if let Err(e) = room.store.save(key.clone(), state).await {
room.dirty.store(true, Ordering::Release);
tracing::warn!("realtime: saving {}/{} failed: {e}", key.kind, key.doc);
}
}
pub(super) fn routes(router: Router<Arc<AppState>>) -> Router<Arc<AppState>> {
router
.route(LOADER, get(|| async { super::asset(JS, "text/javascript; charset=utf-8") }))
.route("/_zeyn/live/doc/{kind}/{doc}", get(doc_socket))
.route("/_zeyn/live/topics", get(topic_socket))
}
pub(super) fn is_asset(path: &str) -> bool {
path == LOADER
}
pub(super) fn head(st: &AppState, screen: &str, page: &mut String) {
if st.realtime.is_some() && screen.contains(MARKER) {
let tag =
format!(r#"<script type="module" src="{}{LOADER}"></script></head>"#, crate::render::escape(&st.prefix));
*page = page.replacen("</head>", &tag, 1);
}
}
pub(super) fn holds(st: &AppState, screen: &str) -> bool {
st.realtime.is_some() && screen.contains(MARKER)
}
#[derive(Deserialize)]
struct LiveQuery {
#[serde(default)]
csrf: String,
#[serde(default)]
t: String,
}
fn valid_doc(id: &str) -> bool {
(1..=64).contains(&id.len()) && id.bytes().all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_' | b'.'))
}
fn origin_allowed(headers: &HeaderMap, allowed: &[String]) -> bool {
let Some(origin) = headers.get(header::ORIGIN).and_then(|v| v.to_str().ok()) else { return false };
if allowed.iter().any(|a| a == origin) {
return true;
}
let host = headers.get(header::HOST).and_then(|v| v.to_str().ok()).unwrap_or("");
let bare = origin.split_once("://").map(|(_, rest)| rest).unwrap_or("");
!host.is_empty() && bare.eq_ignore_ascii_case(host)
}
async fn admit(st: &AppState, headers: &HeaderMap, csrf: &str) -> Result<(Ctx, Hub, String), Response> {
let ctx = st.licensed(st.ctx(headers)).await;
let Some(hub) = st.realtime.clone() else {
return Err(st.error_response(&ctx, Failure::App(Error::NotFound)));
};
let Some(config) = hub.config() else {
return Err(st.error_response(&ctx, Failure::App(Error::NotFound)));
};
let token_ok = cookie(headers, "zeyn_csrf").is_some_and(|c| valid_token(c) && same(c, csrf));
if !origin_allowed(headers, &config.origins) || !token_ok {
return Err(st.error_response(&ctx, Failure::App(Error::Denied(crate::i18n::Text::key("zeyn.error.csrf")))));
}
let Some(user) = ctx.identity().user.clone() else {
return Err(st.error_response(&ctx, Failure::App(Error::Forbidden)));
};
let user = format!("{}/{user}", ctx.tenant().map(|t| t.id()).unwrap_or(""));
if !hub.take_user(&user, config.limits.per_user) {
let mut r =
st.error_response(&ctx, Failure::App(Error::Denied(crate::i18n::Text::key("zeyn.error.rate_limited"))));
*r.status_mut() = StatusCode::TOO_MANY_REQUESTS;
return Err(r);
}
Ok((ctx, hub, user))
}
fn busy() -> Response {
(StatusCode::SERVICE_UNAVAILABLE, [(header::RETRY_AFTER, "5")], "busy").into_response()
}
async fn doc_socket(
State(st): State<Arc<AppState>>,
Path((kind, doc)): Path<(String, String)>,
query: Result<Query<LiveQuery>, QueryRejection>,
headers: HeaderMap,
ws: Result<WebSocketUpgrade, axum::extract::ws::rejection::WebSocketUpgradeRejection>,
) -> Response {
let csrf = query.map(|q| q.0.csrf).unwrap_or_default();
let (ctx, hub, user) = match admit(&st, &headers, &csrf).await {
Ok(a) => a,
Err(r) => return r,
};
let config = hub.config().expect("checked in admit");
let refuse = |hub: &Hub, failure: Failure| {
hub.release_user(&user);
st.error_response(&ctx, failure)
};
let Some((&kind_name, kind_def)) = config.kinds.get_key_value(kind.as_str()).filter(|_| valid_doc(&doc)) else {
return refuse(&hub, Failure::App(Error::NotFound));
};
if !ctx.identity().can(&crate::action::Access::Permission(kind_def.read)) {
return refuse(&hub, Failure::App(Error::Forbidden));
}
let call = Call { kind: CallKind::Live, screen: kind_name, action: None, via_agent: ctx.via_agent() };
if let Err(f) = st.admit_call(&ctx, &call) {
return refuse(&hub, f);
}
let ws = match ws {
Ok(ws) => ws,
Err(e) => {
hub.release_user(&user);
return e.into_response();
}
};
let limits = config.limits;
let key = RoomKey { tenant: ctx.tenant().map(|t| t.id().to_string()).unwrap_or_default(), kind, doc };
let Some(room) = hub.open(&key, kind_def, &limits).await else {
hub.release_user(&user);
return busy();
};
let (tx, rx) = mpsc::channel(limits.queue);
let id = hub.0.next.fetch_add(1, Ordering::Relaxed);
{
let rooms = hub.0.rooms.lock();
let current = rooms.as_ref().ok().and_then(|r| r.get(&key)).is_some_and(|r| Arc::ptr_eq(r, &room));
let clients = room.clients.lock();
let (true, Ok(mut clients)) = (current, clients) else {
hub.release_user(&user);
return busy();
};
if clients.len() >= limits.per_room {
hub.release_user(&user);
return busy();
}
clients.push((id, tx));
if let Ok(mut idle) = room.idle_since.lock() {
*idle = None;
}
}
hub.start_sweeper(limits);
let can_write = ctx.identity().can(&crate::action::Access::Permission(kind_def.write));
ws.max_message_size(limits.message_bytes).on_upgrade(move |socket| async move {
run_doc(socket, &room, id, rx, can_write, &limits).await;
if let Ok(mut clients) = room.clients.lock() {
clients.retain(|(c, _)| *c != id);
if clients.is_empty() {
if let Ok(mut idle) = room.idle_since.lock() {
*idle = Some(Instant::now());
}
}
}
hub.release_user(&user);
})
}
fn frame(update: &[u8]) -> Message {
let mut out = Vec::with_capacity(update.len() + 1);
out.push(FRAME_UPDATE);
out.extend_from_slice(update);
Message::Binary(out.into())
}
fn close(code: u16, reason: &'static str) -> Message {
Message::Close(Some(CloseFrame { code, reason: reason.into() }))
}
async fn run_doc(
socket: WebSocket,
room: &Room,
id: u64,
mut rx: mpsc::Receiver<Message>,
can_write: bool,
limits: &Limits,
) {
let (mut sink, mut stream) = socket.split();
let state = room.doc.transact().encode_state_as_update_v1(&StateVector::default());
if sink.send(frame(&state)).await.is_err() {
return;
}
drop(state);
let mut beat = tokio::time::interval(limits.heartbeat);
beat.tick().await;
let mut heard = Instant::now();
loop {
tokio::select! {
out = rx.recv() => match out {
Some(m) => if sink.send(m).await.is_err() { return },
None => {
let _ = sink.send(close(CLOSE_SLOW, "too slow")).await;
return;
}
},
incoming = stream.next() => {
heard = Instant::now();
match incoming {
Some(Ok(Message::Binary(data))) => {
if let Err(code) = apply(room, id, &data, can_write, limits) {
let _ = sink.send(close(code, "refused")).await;
return;
}
}
Some(Ok(Message::Ping(_) | Message::Pong(_))) => {}
Some(Ok(Message::Text(_))) => {
let _ = sink.send(close(CLOSE_POLICY, "binary only")).await;
return;
}
Some(Ok(Message::Close(_))) | Some(Err(_)) | None => return,
}
}
_ = beat.tick() => {
if heard.elapsed() > limits.heartbeat * 3 || sink.send(Message::Ping(Default::default())).await.is_err() {
return;
}
}
}
}
}
fn apply(room: &Room, from: u64, data: &[u8], can_write: bool, limits: &Limits) -> Result<(), u16> {
let (Some(&FRAME_UPDATE), true) = (data.first(), can_write) else { return Err(CLOSE_POLICY) };
let body = &data[1..];
if room.bytes.load(Ordering::Relaxed) + body.len() > limits.doc_bytes {
return Err(CLOSE_TOO_BIG);
}
let update = Update::decode_v1(body).map_err(|_| CLOSE_POLICY)?;
room.doc.transact_mut().apply_update(update).map_err(|_| CLOSE_POLICY)?;
room.bytes.fetch_add(body.len(), Ordering::Relaxed);
room.dirty.store(true, Ordering::Release);
let message = Message::Binary(axum::body::Bytes::copy_from_slice(data));
if let Ok(mut clients) = room.clients.lock() {
clients.retain(|(id, tx)| *id == from || tx.try_send(message.clone()).is_ok());
}
Ok(())
}
async fn topic_socket(
State(st): State<Arc<AppState>>,
query: Result<Query<LiveQuery>, QueryRejection>,
headers: HeaderMap,
ws: Result<WebSocketUpgrade, axum::extract::ws::rejection::WebSocketUpgradeRejection>,
) -> Response {
let q = query.map(|q| q.0).unwrap_or(LiveQuery { csrf: String::new(), t: String::new() });
let (ctx, hub, user) = match admit(&st, &headers, &q.csrf).await {
Ok(a) => a,
Err(r) => return r,
};
let config = hub.config().expect("checked in admit");
let mut topics: Vec<&'static str> = Vec::new();
for name in q.t.split(',').filter(|t| !t.is_empty()).take(16) {
let allowed = config
.topics
.get_key_value(name)
.filter(|(_, p)| ctx.identity().can(&crate::action::Access::Permission(p)));
match allowed {
Some((n, _)) if !topics.contains(n) => topics.push(n),
Some(_) => {}
None => {
hub.release_user(&user);
return st.error_response(&ctx, Failure::App(Error::Forbidden));
}
}
}
if !topics.is_empty() {
let call = Call { kind: CallKind::Live, screen: "", action: None, via_agent: ctx.via_agent() };
if let Err(f) = st.admit_call(&ctx, &call) {
hub.release_user(&user);
return st.error_response(&ctx, f);
}
}
let ws = match ws {
Ok(ws) if !topics.is_empty() => ws,
Ok(_) => {
hub.release_user(&user);
return st.error_response(&ctx, Failure::App(Error::NotFound));
}
Err(e) => {
hub.release_user(&user);
return e.into_response();
}
};
let tenant = ctx.tenant().map(|t| t.id().to_string()).unwrap_or_default();
let limits = config.limits;
let (tx, mut rx) = mpsc::channel::<Message>(limits.queue);
let id = hub.0.next.fetch_add(1, Ordering::Relaxed);
if let Ok(mut map) = hub.0.topics.lock() {
for t in &topics {
map.entry((tenant.clone(), t.to_string())).or_default().push((id, tx.clone()));
}
}
drop(tx);
ws.max_message_size(1024).on_upgrade(move |socket| async move {
let (mut sink, mut stream) = socket.split();
let mut beat = tokio::time::interval(limits.heartbeat);
beat.tick().await;
let mut heard = Instant::now();
loop {
tokio::select! {
out = rx.recv() => match out {
Some(m) => if sink.send(m).await.is_err() { break },
None => break,
},
incoming = stream.next() => match incoming {
Some(Ok(Message::Close(_))) | Some(Err(_)) | None => break,
Some(Ok(_)) => heard = Instant::now(),
},
_ = beat.tick() => {
if heard.elapsed() > limits.heartbeat * 3 || sink.send(Message::Ping(Default::default())).await.is_err() {
break;
}
}
}
}
if let Ok(mut map) = hub.0.topics.lock() {
for t in &topics {
let key = (tenant.clone(), t.to_string());
if let Some(subs) = map.get_mut(&key) {
subs.retain(|(c, _)| *c != id);
if subs.is_empty() {
map.remove(&key);
}
}
}
}
hub.release_user(&user);
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn origins_must_match_the_host_or_the_allow_list() {
let mut h = HeaderMap::new();
h.insert(header::HOST, "app.test:3000".parse().unwrap());
assert!(!origin_allowed(&h, &[]));
h.insert(header::ORIGIN, "http://app.test:3000".parse().unwrap());
assert!(origin_allowed(&h, &[]));
h.insert(header::ORIGIN, "https://evil.test".parse().unwrap());
assert!(!origin_allowed(&h, &[]));
assert!(origin_allowed(&h, &["https://evil.test".into()]));
}
#[test]
fn a_client_with_a_full_queue_is_dropped_from_the_room() {
let room = Room {
doc: Doc::new(),
clients: Mutex::new(Vec::new()),
idle_since: Mutex::new(None),
dirty: AtomicBool::new(false),
bytes: AtomicUsize::new(0),
store: Arc::new(MemoryDocStore::default()),
};
let (slow, _slow_rx) = mpsc::channel(1);
let (fast, mut fast_rx) = mpsc::channel(8);
slow.try_send(Message::Ping(Default::default())).unwrap();
room.clients.lock().unwrap().extend([(1, slow), (2, fast)]);
let other = Doc::new();
let body = other.get_or_insert_text("t");
yrs::Text::insert(&body, &mut other.transact_mut(), 0, "hi");
let update = other.transact().encode_state_as_update_v1(&StateVector::default());
let frame = [vec![FRAME_UPDATE], update].concat();
assert!(apply(&room, 3, &frame, true, &Limits::default()).is_ok());
assert_eq!(room.clients.lock().unwrap().iter().map(|(id, _)| *id).collect::<Vec<_>>(), [2]);
assert!(fast_rx.try_recv().is_ok());
assert!(room.dirty.load(Ordering::Acquire));
assert_eq!(apply(&room, 3, &frame, false, &Limits::default()), Err(CLOSE_POLICY));
assert_eq!(apply(&room, 3, &[FRAME_UPDATE, 0xff, 0xff], true, &Limits::default()), Err(CLOSE_POLICY));
}
#[test]
fn document_ids_are_plain() {
assert!(valid_doc("main") && valid_doc("order-7.v2"));
assert!(!valid_doc("") && !valid_doc("a/b") && !valid_doc(&"x".repeat(65)));
}
}