125 lines
4.7 KiB
Rust
125 lines
4.7 KiB
Rust
use crate::{
|
|
app_state::{
|
|
AppState,
|
|
database::models::VaultId,
|
|
websocket::{
|
|
broadcasts::Notification,
|
|
models::{CursorPositionFromServer, WebSocketClientMessage, WebSocketServerMessage},
|
|
utils::{get_authenticated_handshake, send_update_over_websocket},
|
|
},
|
|
},
|
|
consts::WEBSOCKET_HANDSHAKE_TIMEOUT,
|
|
errors::{SyncServerError, client_error, server_error},
|
|
};
|
|
use axum::{
|
|
extract::{
|
|
Path, State,
|
|
ws::{Message, WebSocket, WebSocketUpgrade},
|
|
},
|
|
http::StatusCode,
|
|
response::{IntoResponse, Response},
|
|
};
|
|
use futures::StreamExt;
|
|
use log::{debug, info};
|
|
use std::time::Duration;
|
|
use tokio::sync::broadcast::error::RecvError;
|
|
|
|
use super::VaultPath;
|
|
|
|
#[axum::debug_handler]
|
|
pub async fn websocket_handler(
|
|
ws: WebSocketUpgrade,
|
|
Path(VaultPath(vault_id)): Path<VaultPath>,
|
|
State(state): State<AppState>,
|
|
) -> Result<Response, SyncServerError> {
|
|
debug!("Upgrading WebSocket connection for vault `{vault_id}`");
|
|
|
|
let Some(permit) = state.broadcasts.try_admit(&vault_id).await else {
|
|
return Ok((
|
|
StatusCode::TOO_MANY_REQUESTS,
|
|
"Vault connection limit reached",
|
|
)
|
|
.into_response());
|
|
};
|
|
|
|
Ok(ws.on_upgrade(move |socket| async move {
|
|
let _permit = permit;
|
|
if let Err(error) = websocket(state, socket, vault_id).await {
|
|
debug!("WebSocket disconnected: {error}");
|
|
}
|
|
}))
|
|
}
|
|
|
|
async fn websocket(
|
|
state: AppState,
|
|
socket: WebSocket,
|
|
vault: VaultId,
|
|
) -> Result<(), SyncServerError> {
|
|
let (mut sender, mut receiver) = socket.split();
|
|
let authenticated = get_authenticated_handshake(
|
|
&state,
|
|
&vault,
|
|
tokio::time::timeout(WEBSOCKET_HANDSHAKE_TIMEOUT, receiver.next())
|
|
.await
|
|
.map_err(|_| client_error(anyhow::anyhow!("WebSocket handshake deadline expired")))?
|
|
.transpose()
|
|
.unwrap_or_default(),
|
|
)?;
|
|
|
|
let device = authenticated.handshake.device_id;
|
|
let session = state.cursors.register_connection(&vault, &device).await;
|
|
let mut last_checkpoint = None;
|
|
let mut notifications = state.broadcasts.get_receiver(vault.clone()).await;
|
|
let mut timer = tokio::time::interval(Duration::from_secs(2));
|
|
timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
|
|
|
info!("WebSocket connected to vault {vault}");
|
|
|
|
let result = async {
|
|
loop {
|
|
let checkpoint = state.database.history_checkpoint(&vault).await.map_err(server_error)?;
|
|
if last_checkpoint.as_ref() != Some(&checkpoint) {
|
|
send_update_over_websocket(&WebSocketServerMessage::VaultChanged, &mut sender).await?;
|
|
last_checkpoint = Some(checkpoint);
|
|
}
|
|
|
|
tokio::select! {
|
|
_ = timer.tick() => {},
|
|
notification = notifications.recv() => match notification {
|
|
Ok(update) => if let Notification::Cursors(mut positions) = update {
|
|
positions.clients.retain(|c| c.device_id != device);
|
|
send_update_over_websocket(&WebSocketServerMessage::CursorPositions(positions), &mut sender).await?;
|
|
},
|
|
|
|
Err(RecvError::Lagged(_)) => {}, // Check the durable head on the next iteration.
|
|
Err(RecvError::Closed) => break,
|
|
},
|
|
|
|
incoming = receiver.next() => match incoming {
|
|
Some(Ok(Message::Text(text))) => {
|
|
let message: WebSocketClientMessage = serde_json::from_str(&text).map_err(|error| client_error(error.into()))?;
|
|
match message {
|
|
WebSocketClientMessage::Handshake(_) => return Err(client_error(anyhow::anyhow!("Unexpected handshake"))),
|
|
WebSocketClientMessage::CursorPositions(positions) => {
|
|
state.cursors.update_cursors(vault.clone(), authenticated.user.name.clone(), &device, session, positions.documents_with_cursors).await;
|
|
let clients = state.cursors.get_cursors(&vault).await.into_iter().filter(|c| c.device_id != device).collect();
|
|
send_update_over_websocket(&WebSocketServerMessage::CursorPositions(CursorPositionFromServer { clients }), &mut sender).await?;
|
|
}
|
|
}
|
|
},
|
|
Some(Ok(Message::Ping(_) | Message::Pong(_))) => {},
|
|
_ => break,
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}.await;
|
|
|
|
state
|
|
.cursors
|
|
.remove_cursors_of_device(&vault, &device, session)
|
|
.await;
|
|
|
|
result
|
|
}
|