Add API for propagating cursor locations (#61)
This commit is contained in:
GitHub
parent
f97193e287
commit
e8b9bf40c5
80 changed files with 1930 additions and 2229 deletions
|
|
@ -6,165 +6,176 @@ use axum::{
|
|||
},
|
||||
response::Response,
|
||||
};
|
||||
use futures::{
|
||||
sink::SinkExt,
|
||||
stream::{SplitSink, StreamExt},
|
||||
};
|
||||
use log::{error, info, warn};
|
||||
use schemars::JsonSchema;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use futures::stream::StreamExt;
|
||||
use log::{debug, info};
|
||||
use serde::Deserialize;
|
||||
|
||||
use super::auth::auth;
|
||||
use crate::{
|
||||
app_state::{
|
||||
AppState,
|
||||
database::models::{DeviceId, DocumentVersionWithoutContent, VaultId, VaultUpdateId},
|
||||
database::models::VaultId,
|
||||
websocket::{
|
||||
models::{
|
||||
CursorPositionFromServer, WebSocketClientMessage, WebSocketServerMessage,
|
||||
WebSocketVaultUpdate,
|
||||
},
|
||||
utils::{
|
||||
get_authenticated_handshake, get_unseen_documents, send_update_over_websocket,
|
||||
},
|
||||
},
|
||||
},
|
||||
errors::{SyncServerError, server_error, unauthenticated_error},
|
||||
errors::{SyncServerError, client_error, server_error},
|
||||
utils::normalize::normalize,
|
||||
};
|
||||
|
||||
// This is required for aide to infer the path parameter types and names
|
||||
#[derive(Deserialize, JsonSchema)]
|
||||
pub struct WebsocketPathParams {
|
||||
#[derive(Deserialize)]
|
||||
pub struct WebSocketPathParams {
|
||||
#[serde(deserialize_with = "normalize")]
|
||||
vault_id: VaultId,
|
||||
}
|
||||
|
||||
pub async fn websocket_handler(
|
||||
ws: WebSocketUpgrade,
|
||||
Path(WebsocketPathParams { vault_id }): Path<WebsocketPathParams>,
|
||||
Path(WebSocketPathParams { vault_id }): Path<WebSocketPathParams>,
|
||||
State(state): State<AppState>,
|
||||
) -> Result<Response, SyncServerError> {
|
||||
Ok(ws.on_upgrade(move |socket| websocket_wrapped(state, socket, vault_id)))
|
||||
}
|
||||
|
||||
async fn websocket_wrapped(state: AppState, stream: WebSocket, vault_id: VaultId) {
|
||||
info!("Websocket connection opened on vault '{vault_id}'");
|
||||
info!("WebSocket connection opened on vault '{vault_id}'");
|
||||
|
||||
let result = websocket(state, stream, vault_id.clone()).await;
|
||||
|
||||
if let Err(err) = result {
|
||||
error!("Websocket connection error on vault '{vault_id}': {err}");
|
||||
debug!("WebSocket connection error on vault '{vault_id}': {err}");
|
||||
}
|
||||
|
||||
warn!("Websocket connection closed on vault '{vault_id}'");
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct WebsocketHandshake {
|
||||
pub token: String,
|
||||
pub device_id: DeviceId,
|
||||
pub last_seen_vault_update_id: Option<VaultUpdateId>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct WebsocketVaultUpdate {
|
||||
pub documents: Vec<DocumentVersionWithoutContent>,
|
||||
pub is_initial_sync: bool,
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_lines)]
|
||||
async fn websocket(
|
||||
state: AppState,
|
||||
stream: WebSocket,
|
||||
vault_id: VaultId,
|
||||
) -> Result<(), SyncServerError> {
|
||||
let (mut sender, mut receiver) = stream.split();
|
||||
let (mut sender, mut websocket_receiver) = stream.split();
|
||||
|
||||
let handshake = if let Some(Ok(Message::Text(token))) = receiver.next().await {
|
||||
let handshake: WebsocketHandshake = serde_json::from_str(&token)
|
||||
.context("Failed to parse token")
|
||||
.map_err(server_error)?;
|
||||
|
||||
auth(&state, handshake.token.trim(), &vault_id)?;
|
||||
|
||||
handshake
|
||||
} else {
|
||||
return Err(unauthenticated_error(anyhow::anyhow!(
|
||||
"Failed to authenticate"
|
||||
)));
|
||||
};
|
||||
|
||||
let mut rx = state.broadcasts.get_receiver(vault_id.clone()).await;
|
||||
|
||||
let documents = if let Some(update_id) = handshake.last_seen_vault_update_id {
|
||||
state
|
||||
.database
|
||||
.get_latest_documents_since(&vault_id, update_id, None)
|
||||
let authed_handshake = get_authenticated_handshake(
|
||||
&state,
|
||||
&vault_id,
|
||||
websocket_receiver
|
||||
.next()
|
||||
.await
|
||||
.map_err(server_error)
|
||||
} else {
|
||||
state
|
||||
.database
|
||||
.get_latest_documents(&vault_id, None)
|
||||
.await
|
||||
.map_err(server_error)
|
||||
}?;
|
||||
.transpose()
|
||||
.unwrap_or_default(),
|
||||
)?;
|
||||
|
||||
info!(
|
||||
"WebSocket handshake successful for vault '{vault_id}' for '{}'",
|
||||
authed_handshake.handshake.device_id
|
||||
);
|
||||
|
||||
let mut broadcast_receiver = state.broadcasts.get_receiver(vault_id.clone()).await;
|
||||
|
||||
send_update_over_websocket(
|
||||
&WebsocketVaultUpdate {
|
||||
documents,
|
||||
&WebSocketServerMessage::VaultUpdate(WebSocketVaultUpdate {
|
||||
documents: get_unseen_documents(
|
||||
&state,
|
||||
&vault_id,
|
||||
authed_handshake.handshake.last_seen_vault_update_id,
|
||||
)
|
||||
.await?,
|
||||
is_initial_sync: true,
|
||||
},
|
||||
}),
|
||||
&mut sender,
|
||||
)
|
||||
.await?;
|
||||
|
||||
send_update_over_websocket(
|
||||
&WebSocketServerMessage::CursorPositions(CursorPositionFromServer {
|
||||
clients: state.cursors.get_cursors(&vault_id).await,
|
||||
}),
|
||||
&mut sender,
|
||||
)
|
||||
.await?;
|
||||
|
||||
let device_id = authed_handshake.handshake.device_id.clone();
|
||||
let mut send_task = tokio::spawn(async move {
|
||||
while let Ok(update) = rx.recv().await {
|
||||
if Some(&handshake.device_id) == update.origin_device_id.as_ref() {
|
||||
while let Ok(update) = broadcast_receiver.recv().await {
|
||||
if Some(&device_id) == update.origin_device_id.as_ref() {
|
||||
continue;
|
||||
}
|
||||
|
||||
send_update_over_websocket(
|
||||
&WebsocketVaultUpdate {
|
||||
documents: vec![update.document],
|
||||
is_initial_sync: false,
|
||||
},
|
||||
&mut sender,
|
||||
)
|
||||
.await?;
|
||||
send_update_over_websocket(&update.message, &mut sender).await?;
|
||||
}
|
||||
|
||||
Ok::<(), SyncServerError>(())
|
||||
});
|
||||
|
||||
let mut recv_task =
|
||||
tokio::spawn(
|
||||
async move { while let Some(Ok(Message::Text(_text))) = receiver.next().await {} },
|
||||
);
|
||||
let device_id = authed_handshake.handshake.device_id.clone();
|
||||
let vault_id_clone = vault_id.clone();
|
||||
let cursor_manager = state.cursors.clone();
|
||||
let mut receive_task = tokio::spawn(async move {
|
||||
while let Some(Ok(Message::Text(message))) = websocket_receiver.next().await {
|
||||
let message: WebSocketClientMessage = serde_json::from_str(&message)
|
||||
.context("Failed to parse WebSocket message from client")
|
||||
.map_err(server_error)?;
|
||||
|
||||
match message {
|
||||
WebSocketClientMessage::Handshake(_) => {
|
||||
return Err(client_error(anyhow::anyhow!(
|
||||
"Unexpected handshake message"
|
||||
)));
|
||||
}
|
||||
WebSocketClientMessage::CursorPositions(cursors) => {
|
||||
cursor_manager
|
||||
.update_cursors(
|
||||
vault_id_clone.clone(),
|
||||
authed_handshake.user.name.clone(),
|
||||
&device_id,
|
||||
cursors.document_to_cursors,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok::<(), SyncServerError>(())
|
||||
});
|
||||
|
||||
tokio::select! {
|
||||
_ = &mut send_task => recv_task.abort(),
|
||||
_ = &mut recv_task => send_task.abort(),
|
||||
_ = &mut send_task => receive_task.abort(),
|
||||
_ = &mut receive_task => send_task.abort(),
|
||||
};
|
||||
|
||||
send_task
|
||||
.await
|
||||
.context("Websocket send task failed")
|
||||
.map_err(server_error)??;
|
||||
let result: Result<(), SyncServerError> = (async {
|
||||
send_task
|
||||
.await
|
||||
.context("WebSocket send task failed")
|
||||
.map_err(client_error)
|
||||
.and_then(|err| err)?;
|
||||
|
||||
recv_task
|
||||
.await
|
||||
.context("Websocket receive task failed")
|
||||
.map_err(server_error)?;
|
||||
receive_task
|
||||
.await
|
||||
.context("WebSocket receive task failed")
|
||||
.map_err(client_error)
|
||||
.and_then(|err| err)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn send_update_over_websocket(
|
||||
update: &WebsocketVaultUpdate,
|
||||
sender: &mut SplitSink<WebSocket, Message>,
|
||||
) -> Result<(), SyncServerError> {
|
||||
let serialized_update = serde_json::to_string(update)
|
||||
.context("Failed to serialize update")
|
||||
.map_err(server_error)?;
|
||||
|
||||
sender
|
||||
.send(Message::Text(serialized_update))
|
||||
.await
|
||||
.context("Failed to send message over websocket")
|
||||
.map_err(server_error)
|
||||
Ok(())
|
||||
})
|
||||
.await;
|
||||
|
||||
state
|
||||
.cursors
|
||||
.remove_cursors_of_device(&vault_id, &authed_handshake.handshake.device_id)
|
||||
.await;
|
||||
|
||||
if result.is_err() {
|
||||
info!(
|
||||
"WebSocket disconnected on vault '{vault_id}' for '{}'",
|
||||
authed_handshake.handshake.device_id
|
||||
);
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue