use anyhow::{anyhow, Result};
use axum::{
body::Bytes,
extract::{
ws::{Message, WebSocket},
Path, Query, Request, State, WebSocketUpgrade,
},
http::StatusCode,
middleware::{self, Next},
response::{IntoResponse, Response},
routing::{get, post},
Json, Router,
};
use axum_extra::typed_header::TypedHeader;
use dashmap::{mapref::one::MappedRef, DashMap};
use futures::{SinkExt, StreamExt};
use serde::Deserialize;
use serde_json::{json, Value};
use std::{
net::SocketAddr,
sync::{Arc, RwLock},
time::Duration,
};
use tokio::{
net::TcpListener,
sync::mpsc::{channel, Receiver},
};
use tokio_util::{sync::CancellationToken, task::TaskTracker};
use tracing::{span, Instrument, Level};
use url::Url;
use y_sweet_core::{
api_types::{
validate_doc_name, AuthDocRequest, ClientToken, DocCreationRequest, NewDocResponse,
},
auth::Authenticator,
doc_connection::DocConnection,
doc_sync::DocWithSyncKv,
store::Store,
sync::awareness::Awareness,
sync_kv::SyncKv,
};
fn current_time_epoch_millis() -> u64 {
let now = std::time::SystemTime::now();
let duration_since_epoch = now.duration_since(std::time::UNIX_EPOCH).unwrap();
duration_since_epoch.as_millis() as u64
}
#[derive(Debug)]
pub struct AppError(StatusCode, anyhow::Error);
impl std::error::Error for AppError {}
impl IntoResponse for AppError {
fn into_response(self) -> Response {
(self.0, format!("Something went wrong: {}", self.1)).into_response()
}
}
impl<E> From<(StatusCode, E)> for AppError
where
E: Into<anyhow::Error>,
{
fn from((status_code, err): (StatusCode, E)) -> Self {
Self(status_code, err.into())
}
}
impl std::fmt::Display for AppError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "Status code: {} {}", self.0, self.1)?;
Ok(())
}
}
pub struct Server {
docs: Arc<DashMap<String, DocWithSyncKv>>,
doc_worker_tracker: TaskTracker,
store: Option<Arc<Box<dyn Store>>>,
checkpoint_freq: Duration,
authenticator: Option<Authenticator>,
url_prefix: Option<Url>,
cancellation_token: CancellationToken,
doc_gc: bool,
}
impl Server {
pub async fn new(
store: Option<Box<dyn Store>>,
checkpoint_freq: Duration,
authenticator: Option<Authenticator>,
url_prefix: Option<Url>,
cancellation_token: CancellationToken,
doc_gc: bool,
) -> Result<Self> {
Ok(Self {
docs: Arc::new(DashMap::new()),
doc_worker_tracker: TaskTracker::new(),
store: store.map(Arc::new),
checkpoint_freq,
authenticator,
url_prefix,
cancellation_token,
doc_gc,
})
}
pub async fn doc_exists(&self, doc_id: &str) -> bool {
if self.docs.contains_key(doc_id) {
return true;
}
if let Some(store) = &self.store {
store
.exists(&format!("{}/data.ysweet", doc_id))
.await
.unwrap_or_default()
} else {
false
}
}
pub async fn create_doc(&self) -> Result<String> {
let doc_id = nanoid::nanoid!();
self.load_doc(&doc_id).await?;
tracing::info!(doc_id=?doc_id, "Created doc");
Ok(doc_id)
}
pub async fn load_doc(&self, doc_id: &str) -> Result<()> {
let (send, recv) = channel(1024);
let dwskv = DocWithSyncKv::new(doc_id, self.store.clone(), move || {
send.try_send(()).unwrap();
})
.await?;
dwskv
.sync_kv()
.persist()
.await
.map_err(|e| anyhow!("Error persisting: {:?}", e))?;
{
let sync_kv = dwskv.sync_kv();
let checkpoint_freq = self.checkpoint_freq;
let doc_id = doc_id.to_string();
let persistence_cancellation_token = CancellationToken::new();
self.doc_worker_tracker.spawn(
Self::doc_persistence_worker(
recv,
sync_kv,
checkpoint_freq,
doc_id.clone(),
persistence_cancellation_token.clone(),
)
.instrument(span!(Level::INFO, "save_loop", doc_id=?doc_id)),
);
if self.doc_gc {
self.doc_worker_tracker.spawn(
Self::doc_gc_worker(
self.docs.clone(),
doc_id.clone(),
checkpoint_freq,
persistence_cancellation_token,
self.cancellation_token.clone(),
)
.instrument(span!(Level::INFO, "gc_loop", doc_id=?doc_id)),
);
}
}
self.docs.insert(doc_id.to_string(), dwskv);
Ok(())
}
async fn doc_gc_worker(
docs: Arc<DashMap<String, DocWithSyncKv>>,
doc_id: String,
checkpoint_freq: Duration,
persistence_cancellation_token: CancellationToken,
cancellation_token: CancellationToken,
) {
let mut checkpoints_without_refs = 0;
loop {
tokio::select! {
_ = tokio::time::sleep(checkpoint_freq) => {
if let Some(doc) = docs.get(&doc_id) {
let awareness = Arc::downgrade(&doc.awareness());
if awareness.strong_count() > 1 {
checkpoints_without_refs = 0;
tracing::debug!("doc is still alive - it has {} references", awareness.strong_count());
} else {
checkpoints_without_refs += 1;
tracing::info!("doc has only one reference, candidate for GC. checkpoints_without_refs: {}", checkpoints_without_refs);
}
} else {
break;
}
if checkpoints_without_refs >= 2 {
tracing::info!("GCing doc");
docs.remove(&doc_id);
break;
}
}
_ = cancellation_token.cancelled() => {
break;
}
};
}
persistence_cancellation_token.cancel();
tracing::info!("Exiting gc_loop");
}
async fn doc_persistence_worker(
mut recv: Receiver<()>,
sync_kv: Arc<SyncKv>,
checkpoint_freq: Duration,
doc_id: String,
cancellation_token: CancellationToken,
) {
let mut last_save = std::time::Instant::now();
loop {
let is_done = tokio::select! {
v = recv.recv() => v.is_none(),
_ = cancellation_token.cancelled() => true,
};
tracing::info!("Received signal. done: {}", is_done);
let now = std::time::Instant::now();
if !is_done && now - last_save < checkpoint_freq {
let sleep = tokio::time::sleep(checkpoint_freq - (now - last_save));
tokio::pin!(sleep);
tracing::info!("Throttling.");
loop {
tokio::select! {
_ = &mut sleep => {
break;
}
v = recv.recv() => {
tracing::info!("Received dirty while throttling.");
if v.is_none() {
break;
}
}
_ = cancellation_token.cancelled() => {
tracing::info!("Received cancellation while throttling.");
break;
}
}
tracing::info!("Done throttling.");
}
}
tracing::info!("Persisting.");
if let Err(e) = sync_kv.persist().await {
tracing::error!(?e, "Error persisting.");
} else {
tracing::info!("Done persisting.");
}
last_save = std::time::Instant::now();
if is_done {
break;
}
}
tracing::info!("Terminating loop for {}", doc_id);
}
pub async fn get_or_create_doc(
&self,
doc_id: &str,
) -> Result<MappedRef<String, DocWithSyncKv, DocWithSyncKv>> {
if !self.docs.contains_key(doc_id) {
tracing::info!(doc_id=?doc_id, "Loading doc");
self.load_doc(doc_id).await?;
}
Ok(self
.docs
.get(doc_id)
.ok_or_else(|| anyhow!("Failed to get-or-create doc"))?
.map(|d| d))
}
pub fn check_auth(
&self,
header: Option<TypedHeader<headers::Authorization<headers::authorization::Bearer>>>,
) -> Result<(), AppError> {
if let Some(auth) = &self.authenticator {
if let Some(TypedHeader(headers::Authorization(bearer))) = header {
if let Ok(()) =
auth.verify_server_token(bearer.token(), current_time_epoch_millis())
{
return Ok(());
}
}
Err((StatusCode::UNAUTHORIZED, anyhow!("Unauthorized.")))?
} else {
Ok(())
}
}
pub fn check_doc_auth(
&self,
header: Option<TypedHeader<headers::Authorization<headers::authorization::Bearer>>>,
doc_id: &str,
) -> Result<(), AppError> {
if let Some(auth) = &self.authenticator {
if let Some(TypedHeader(headers::Authorization(bearer))) = header {
if let Ok(()) =
auth.verify_doc_token(bearer.token(), doc_id, current_time_epoch_millis())
{
return Ok(());
}
}
Err((StatusCode::UNAUTHORIZED, anyhow!("Unauthorized.")))?
} else {
Ok(())
}
}
pub async fn redact_error_middleware(req: Request, next: Next) -> impl IntoResponse {
let resp = next.run(req).await;
if resp.status().is_server_error() || resp.status().is_client_error() {
return resp.status().into_response();
}
resp
}
pub fn routes(self: &Arc<Self>) -> Router {
Router::new()
.route("/check_store", get(check_store))
.route("/doc/ws/:doc_id", get(handle_socket_upgrade))
.route("/doc/new", post(new_doc))
.route("/doc/:doc_id/auth", post(auth_doc))
.route("/doc/:doc_id/as-update", get(get_doc_as_update))
.route("/doc/:doc_id/update", post(update_doc))
.route("/d/:doc_id/as-update", get(get_doc_as_update))
.route("/d/:doc_id/update", post(update_doc))
.with_state(self.clone())
}
pub fn single_doc_routes(self: &Arc<Self>) -> Router {
Router::new()
.route("/ws/:doc_id", get(handle_socket_upgrade_single))
.route("/as-update", get(get_doc_as_update_single))
.route("/update", post(update_doc_single))
.with_state(self.clone())
}
async fn serve_internal(
self: Arc<Self>,
addr: &SocketAddr,
redact_errors: bool,
routes: Router,
) -> Result<()> {
let listener = TcpListener::bind(addr).await?;
let token = self.cancellation_token.clone();
let app = if redact_errors {
routes
} else {
routes.layer(middleware::from_fn(Self::redact_error_middleware))
};
axum::serve(listener, app.into_make_service())
.with_graceful_shutdown(async move { token.cancelled().await })
.await?;
self.doc_worker_tracker.close();
self.doc_worker_tracker.wait().await;
Ok(())
}
pub async fn serve(self, addr: &SocketAddr, redact_errors: bool) -> Result<()> {
let s = Arc::new(self);
let routes = s.routes();
s.serve_internal(addr, redact_errors, routes).await
}
pub async fn serve_doc(self, addr: &SocketAddr, redact_errors: bool) -> Result<()> {
let s = Arc::new(self);
let routes = s.single_doc_routes();
s.serve_internal(addr, redact_errors, routes).await
}
fn verify_doc_token(&self, token: Option<&str>, doc: &str) -> Result<(), AppError> {
if let Some(authenticator) = &self.authenticator {
if let Some(token) = token {
authenticator
.verify_doc_token(token, doc, current_time_epoch_millis())
.map_err(|e| (StatusCode::FORBIDDEN, e))?;
} else {
Err((StatusCode::UNAUTHORIZED, anyhow!("No token provided.")))?
}
}
Ok(())
}
fn get_single_doc_id(&self) -> Result<String, AppError> {
self.docs
.iter()
.next()
.map(|entry| entry.key().clone())
.ok_or_else(|| AppError(StatusCode::NOT_FOUND, anyhow!("No document found")))
}
}
#[derive(Deserialize)]
struct HandlerParams {
token: Option<String>,
}
async fn get_doc_as_update(
State(server_state): State<Arc<Server>>,
Path(doc_id): Path<String>,
authorization: Option<TypedHeader<headers::Authorization<headers::authorization::Bearer>>>,
) -> Result<Response, AppError> {
server_state.check_doc_auth(authorization, &doc_id)?;
let dwskv = server_state
.get_or_create_doc(&doc_id)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
let update = dwskv.as_update();
tracing::debug!("update: {:?}", update);
Ok(update.into_response())
}
async fn get_doc_as_update_single(
State(server_state): State<Arc<Server>>,
authorization: Option<TypedHeader<headers::Authorization<headers::authorization::Bearer>>>,
) -> Result<Response, AppError> {
let doc_id = server_state.get_single_doc_id()?;
get_doc_as_update(State(server_state), Path(doc_id), authorization).await
}
async fn update_doc(
Path(doc_id): Path<String>,
State(server_state): State<Arc<Server>>,
authorization: Option<TypedHeader<headers::Authorization<headers::authorization::Bearer>>>,
body: Bytes,
) -> Result<Response, AppError> {
server_state.check_doc_auth(authorization, &doc_id)?;
let dwskv = server_state
.get_or_create_doc(&doc_id)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
if let Err(err) = dwskv.apply_update(&body) {
tracing::error!(?err, "Failed to apply update");
return Err(AppError(StatusCode::INTERNAL_SERVER_ERROR, err));
}
Ok(StatusCode::OK.into_response())
}
async fn update_doc_single(
State(server_state): State<Arc<Server>>,
authorization: Option<TypedHeader<headers::Authorization<headers::authorization::Bearer>>>,
body: Bytes,
) -> Result<Response, AppError> {
let doc_id = server_state.get_single_doc_id()?;
update_doc(Path(doc_id), State(server_state), authorization, body).await
}
async fn handle_socket_upgrade(
ws: WebSocketUpgrade,
Path(doc_id): Path<String>,
Query(params): Query<HandlerParams>,
State(server_state): State<Arc<Server>>,
) -> Result<Response, AppError> {
server_state.verify_doc_token(params.token.as_deref(), &doc_id)?;
let dwskv = server_state
.get_or_create_doc(&doc_id)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
let awareness = dwskv.awareness();
let cancellation_token = server_state.cancellation_token.clone();
Ok(ws.on_upgrade(move |socket| handle_socket(socket, awareness, cancellation_token)))
}
async fn handle_socket_upgrade_single(
ws: WebSocketUpgrade,
Path(doc_id): Path<String>,
Query(params): Query<HandlerParams>,
State(server_state): State<Arc<Server>>,
) -> Result<Response, AppError> {
let single_doc_id = server_state.get_single_doc_id()?;
if doc_id != single_doc_id {
return Err(AppError(
StatusCode::NOT_FOUND,
anyhow!("Document not found"),
));
}
handle_socket_upgrade(ws, Path(single_doc_id), Query(params), State(server_state)).await
}
async fn handle_socket(
socket: WebSocket,
awareness: Arc<RwLock<Awareness>>,
cancellation_token: CancellationToken,
) {
let (mut sink, mut stream) = socket.split();
let (send, mut recv) = channel(1024);
tokio::spawn(async move {
while let Some(msg) = recv.recv().await {
let _ = sink.send(Message::Binary(msg)).await;
}
});
let connection = DocConnection::new(awareness, move |bytes| {
if let Err(e) = send.try_send(bytes.to_vec()) {
tracing::warn!(?e, "Error sending message");
}
});
loop {
tokio::select! {
Some(msg) = stream.next() => {
let msg = match msg {
Ok(Message::Binary(bytes)) => bytes,
Ok(Message::Close(_)) => break,
Err(_e) => {
continue;
}
msg => {
tracing::warn!(?msg, "Received non-binary message");
continue;
}
};
if let Err(e) = connection.send(&msg).await {
tracing::warn!(?e, "Error handling message");
}
}
_ = cancellation_token.cancelled() => {
tracing::debug!("Closing doc connection due to server cancel...");
break;
}
}
}
}
async fn check_store(
authorization: Option<TypedHeader<headers::Authorization<headers::authorization::Bearer>>>,
State(server_state): State<Arc<Server>>,
) -> Result<Json<Value>, AppError> {
server_state.check_auth(authorization)?;
if server_state.store.is_none() {
return Ok(Json(json!({"ok": false, "error": "No store set."})));
};
Ok(Json(json!({"ok": true})))
}
async fn new_doc(
authorization: Option<TypedHeader<headers::Authorization<headers::authorization::Bearer>>>,
State(server_state): State<Arc<Server>>,
Json(body): Json<DocCreationRequest>,
) -> Result<Json<NewDocResponse>, AppError> {
server_state.check_auth(authorization)?;
let doc_id = if let Some(doc_id) = body.doc_id {
if !validate_doc_name(doc_id.as_str()) {
Err((StatusCode::BAD_REQUEST, anyhow!("Invalid document name")))?
}
server_state
.get_or_create_doc(doc_id.as_str())
.await
.map_err(|e| {
tracing::error!(?e, "Failed to create doc");
(StatusCode::INTERNAL_SERVER_ERROR, e)
})?;
doc_id
} else {
server_state.create_doc().await.map_err(|d| {
tracing::error!(?d, "Failed to create doc");
(StatusCode::INTERNAL_SERVER_ERROR, d)
})?
};
Ok(Json(NewDocResponse { doc_id }))
}
async fn auth_doc(
authorization: Option<TypedHeader<headers::Authorization<headers::authorization::Bearer>>>,
TypedHeader(host): TypedHeader<headers::Host>,
State(server_state): State<Arc<Server>>,
Path(doc_id): Path<String>,
Json(_body): Json<AuthDocRequest>,
) -> Result<Json<ClientToken>, AppError> {
server_state.check_auth(authorization)?;
if !server_state.doc_exists(&doc_id).await {
Err((StatusCode::NOT_FOUND, anyhow!("Doc {} not found", doc_id)))?;
}
let token = if let Some(auth) = &server_state.authenticator {
let token = auth.gen_doc_token(&doc_id, current_time_epoch_millis());
Some(token)
} else {
None
};
let url = if let Some(url_prefix) = &server_state.url_prefix {
let mut url = url_prefix.clone();
let scheme = if url.scheme() == "https" { "wss" } else { "ws" };
url.set_scheme(scheme).unwrap();
url = url.join("/doc/ws").unwrap();
url.to_string()
} else {
format!("ws://{host}/doc/ws")
};
let base_url = if let Some(url_prefix) = &server_state.url_prefix {
format!("{url_prefix}/d/{doc_id}")
} else {
format!("http://{host}/d/{doc_id}")
};
Ok(Json(ClientToken {
url,
base_url: Some(base_url),
doc_id,
token,
}))
}
#[cfg(test)]
mod test {
use super::*;
use std::collections::HashMap;
use y_sweet_core::api_types::Authorization;
#[tokio::test]
async fn test_auth_doc() {
let server_state = Server::new(
None,
Duration::from_secs(60),
None,
None,
CancellationToken::new(),
true,
)
.await
.unwrap();
let doc_id = server_state.create_doc().await.unwrap();
let token = auth_doc(
None,
TypedHeader(headers::Host::from(http::uri::Authority::from_static(
"localhost",
))),
State(Arc::new(server_state)),
Path(doc_id.clone()),
Json(AuthDocRequest {
authorization: Authorization::Full,
user_id: None,
metadata: HashMap::new(),
}),
)
.await
.unwrap();
assert_eq!(token.url, "ws://localhost/doc/ws");
assert_eq!(token.doc_id, doc_id);
assert!(token.token.is_none());
}
#[tokio::test]
async fn test_auth_doc_with_prefix() {
let prefix: Url = "https://foo.bar".parse().unwrap();
let server_state = Server::new(
None,
Duration::from_secs(60),
None,
Some(prefix),
CancellationToken::new(),
true,
)
.await
.unwrap();
let doc_id = server_state.create_doc().await.unwrap();
let token = auth_doc(
None,
TypedHeader(headers::Host::from(http::uri::Authority::from_static(
"localhost",
))),
State(Arc::new(server_state)),
Path(doc_id.clone()),
Json(AuthDocRequest {
authorization: Authorization::Full,
user_id: None,
metadata: HashMap::new(),
}),
)
.await
.unwrap();
assert_eq!(token.url, "wss://foo.bar/doc/ws");
assert_eq!(token.doc_id, doc_id);
assert!(token.token.is_none());
}
}