Implement cursor broadcasting backend

This commit is contained in:
Andras Schmelczer 2025-06-01 09:50:52 +01:00
parent 483e03e2de
commit eb1cc61042
No known key found for this signature in database
GPG key ID: FC8F2C3D3D1A718C
19 changed files with 488 additions and 191 deletions

View file

@ -17,9 +17,8 @@ use super::{
use crate::{
app_state::{
AppState,
broadcasts::VaultUpdate,
database::models::{
DeviceId, DocumentId, DocumentVersionWithoutContent, StoredDocumentVersion, VaultId,
DocumentId, DocumentVersionWithoutContent, StoredDocumentVersion, VaultId,
},
},
config::user_config::User,
@ -41,7 +40,7 @@ pub struct CreateDocumentPathParams {
pub async fn create_document_multipart(
Path(CreateDocumentPathParams { vault_id }): Path<CreateDocumentPathParams>,
Extension(user): Extension<User>,
TypedHeader(user_agent): TypedHeader<DeviceIdHeader>,
TypedHeader(device_id): TypedHeader<DeviceIdHeader>,
State(state): State<AppState>,
TypedMultipart(axum_typed_multipart::TypedMultipart(request)): TypedMultipart<
CreateDocumentVersionMultipart,
@ -49,12 +48,11 @@ pub async fn create_document_multipart(
) -> Result<Json<DocumentVersionWithoutContent>, SyncServerError> {
internal_create_document(
user,
user_agent,
device_id,
state,
vault_id,
request.document_id,
request.relative_path,
request.device_id,
request.content.contents.to_vec(),
)
.await
@ -67,7 +65,7 @@ pub async fn create_document_multipart(
pub async fn create_document_json(
Path(CreateDocumentPathParams { vault_id }): Path<CreateDocumentPathParams>,
Extension(user): Extension<User>,
TypedHeader(user_agent): TypedHeader<DeviceIdHeader>,
TypedHeader(device_id): TypedHeader<DeviceIdHeader>,
State(state): State<AppState>,
Json(request): Json<CreateDocumentVersion>,
) -> Result<Json<DocumentVersionWithoutContent>, SyncServerError> {
@ -77,12 +75,11 @@ pub async fn create_document_json(
internal_create_document(
user,
user_agent,
device_id,
state,
vault_id,
request.document_id,
request.relative_path,
request.device_id,
content_bytes,
)
.await
@ -91,12 +88,11 @@ pub async fn create_document_json(
#[allow(clippy::too_many_arguments)]
async fn internal_create_document(
user: User,
user_agent: DeviceIdHeader,
device_id: DeviceIdHeader,
state: AppState,
vault_id: VaultId,
document_id: Option<DocumentId>,
relative_path: String,
device_id: Option<DeviceId>,
content: Vec<u8>,
) -> Result<Json<DocumentVersionWithoutContent>, SyncServerError> {
let mut transaction = state
@ -140,7 +136,7 @@ async fn internal_create_document(
updated_date: chrono::Utc::now(),
is_deleted: false,
user_id: user.name,
device_id: user_agent.0,
device_id: device_id.0,
};
state
@ -155,16 +151,5 @@ async fn internal_create_document(
.context("Failed to commit successful transaction")
.map_err(server_error)?;
state
.broadcasts
.send(
vault_id,
VaultUpdate {
origin_device_id: device_id,
document: new_version.clone().into(),
},
)
.await;
Ok(Json(new_version.into()))
}

View file

@ -12,7 +12,6 @@ use super::{device_id_header::DeviceIdHeader, requests::DeleteDocumentVersion};
use crate::{
app_state::{
AppState,
broadcasts::VaultUpdate,
database::models::{
DocumentId, DocumentVersionWithoutContent, StoredDocumentVersion, VaultId,
},
@ -38,7 +37,7 @@ pub async fn delete_document(
document_id,
}): Path<DeleteDocumentPathParams>,
Extension(user): Extension<User>,
TypedHeader(user_agent): TypedHeader<DeviceIdHeader>,
TypedHeader(device_id): TypedHeader<DeviceIdHeader>,
State(state): State<AppState>,
Json(request): Json<DeleteDocumentVersion>,
) -> Result<Json<DocumentVersionWithoutContent>, SyncServerError> {
@ -69,7 +68,7 @@ pub async fn delete_document(
updated_date: chrono::Utc::now(),
is_deleted: true,
user_id: user.name,
device_id: user_agent.0,
device_id: device_id.0,
};
state
@ -84,16 +83,5 @@ pub async fn delete_document(
.context("Failed to commit successful transaction")
.map_err(server_error)?;
state
.broadcasts
.send(
vault_id,
VaultUpdate {
origin_device_id: request.device_id,
document: new_version.clone().into(),
},
)
.await;
Ok(Json(new_version.into()))
}

View file

@ -4,7 +4,7 @@ use axum_typed_multipart::TryFromMultipart;
use schemars::JsonSchema;
use serde::{self, Deserialize};
use crate::app_state::database::models::{DeviceId, DocumentId, VaultUpdateId};
use crate::app_state::database::models::{DocumentId, VaultUpdateId};
#[derive(Debug, Deserialize, JsonSchema)]
#[serde(rename_all = "camelCase")]
@ -16,7 +16,6 @@ pub struct CreateDocumentVersion {
pub document_id: Option<DocumentId>,
pub relative_path: String,
pub content_base64: String,
pub device_id: Option<DeviceId>,
}
#[derive(Debug, TryFromMultipart, JsonSchema)]
@ -25,7 +24,6 @@ pub struct CreateDocumentVersionMultipart {
pub relative_path: String,
#[form_data(limit = "unlimited")]
pub content: FieldData<Bytes>,
pub device_id: Option<DeviceId>,
}
#[derive(Debug, Deserialize, JsonSchema)]
@ -34,7 +32,6 @@ pub struct UpdateDocumentVersion {
pub parent_version_id: VaultUpdateId,
pub relative_path: String,
pub content_base64: String,
pub device_id: Option<DeviceId>,
}
#[derive(Debug, TryFromMultipart, JsonSchema)]
@ -44,12 +41,10 @@ pub struct UpdateDocumentVersionMultipart {
pub relative_path: String,
#[form_data(limit = "unlimited")]
pub content: FieldData<Bytes>,
pub device_id: Option<DeviceId>,
}
#[derive(Debug, Deserialize, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct DeleteDocumentVersion {
pub relative_path: String,
pub device_id: Option<DeviceId>,
}

View file

@ -19,8 +19,7 @@ use super::{
use crate::{
app_state::{
AppState,
broadcasts::VaultUpdate,
database::models::{DeviceId, DocumentId, StoredDocumentVersion, VaultId, VaultUpdateId},
database::models::{DocumentId, StoredDocumentVersion, VaultId, VaultUpdateId},
},
config::user_config::User,
errors::{SyncServerError, client_error, not_found_error, server_error},
@ -43,7 +42,7 @@ pub async fn update_document_multipart(
document_id,
}): Path<UpdateDocumentPathParams>,
Extension(user): Extension<User>,
TypedHeader(user_agent): TypedHeader<DeviceIdHeader>,
TypedHeader(device_id): TypedHeader<DeviceIdHeader>,
State(state): State<AppState>,
TypedMultipart(axum_typed_multipart::TypedMultipart(request)): TypedMultipart<
UpdateDocumentVersionMultipart,
@ -51,13 +50,12 @@ pub async fn update_document_multipart(
) -> Result<Json<DocumentUpdateResponse>, SyncServerError> {
internal_update_document(
user,
user_agent,
device_id,
state,
vault_id,
document_id,
request.parent_version_id,
request.relative_path,
request.device_id,
request.content.contents.to_vec(),
)
.await
@ -70,7 +68,7 @@ pub async fn update_document_json(
document_id,
}): Path<UpdateDocumentPathParams>,
Extension(user): Extension<User>,
TypedHeader(user_agent): TypedHeader<DeviceIdHeader>,
TypedHeader(device_id): TypedHeader<DeviceIdHeader>,
State(state): State<AppState>,
Json(request): Json<UpdateDocumentVersion>,
) -> Result<Json<DocumentUpdateResponse>, SyncServerError> {
@ -80,13 +78,12 @@ pub async fn update_document_json(
internal_update_document(
user,
user_agent,
device_id,
state,
vault_id,
document_id,
request.parent_version_id,
request.relative_path,
request.device_id,
content_bytes,
)
.await
@ -95,13 +92,12 @@ pub async fn update_document_json(
#[allow(clippy::too_many_arguments, clippy::too_many_lines)]
async fn internal_update_document(
user: User,
user_agent: DeviceIdHeader,
device_id: DeviceIdHeader,
state: AppState,
vault_id: VaultId,
document_id: DocumentId,
parent_version_id: VaultUpdateId,
relative_path: String,
device_id: Option<DeviceId>,
content: Vec<u8>,
) -> Result<Json<DocumentUpdateResponse>, SyncServerError> {
// No need for a transaction as document versions are immutable
@ -215,7 +211,7 @@ async fn internal_update_document(
updated_date: chrono::Utc::now(),
is_deleted: false,
user_id: user.name,
device_id: user_agent.0,
device_id: device_id.0,
};
state
@ -230,17 +226,6 @@ async fn internal_update_document(
.context("Failed to commit successful transaction")
.map_err(server_error)?;
state
.broadcasts
.send(
vault_id,
VaultUpdate {
origin_device_id: device_id,
document: new_version.clone().into(),
},
)
.await;
Ok(Json(if is_different_from_request_content {
DocumentUpdateResponse::MergingUpdate(new_version.into())
} else {

View file

@ -6,64 +6,52 @@ use axum::{
},
response::Response,
};
use futures::{
sink::SinkExt,
stream::{SplitSink, StreamExt},
};
use futures::stream::StreamExt;
use log::{error, info, warn};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
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_handshake, get_unseen_documents, send_update_over_websocket},
},
},
errors::{SyncServerError, server_error, unauthenticated_error},
errors::{SyncServerError, client_error, server_error, unauthenticated_error},
utils::normalize::normalize,
};
// This is required for aide to infer the path parameter types and names
#[derive(Deserialize, JsonSchema)]
pub struct WebsocketPathParams {
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}");
error!("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,
warn!("WebSocket connection closed on vault '{vault_id}'");
}
async fn websocket(
@ -73,68 +61,71 @@ async fn websocket(
) -> Result<(), SyncServerError> {
let (mut sender, mut 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
let handshake = if let Some(Ok(message)) = receiver.next().await {
get_handshake(&state, &vault_id, message)?
} else {
return Err(unauthenticated_error(anyhow::anyhow!(
"Failed to authenticate"
"Failed to authenticate due to invalid message"
)));
};
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)
.await
.map_err(server_error)
} else {
state
.database
.get_latest_documents(&vault_id, None)
.await
.map_err(server_error)
}?;
send_update_over_websocket(
&WebsocketVaultUpdate {
documents,
&WebSocketServerMessage::VaultUpdate(WebSocketVaultUpdate {
documents: get_unseen_documents(&state, &vault_id, 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 = 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() {
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 = handshake.device_id.clone();
let mut recv_task = tokio::spawn(async move {
while let Some(Ok(Message::Text(message))) = receiver.next().await {
let message: WebSocketClientMessage = serde_json::from_str(&message)
.context("Failed to parse message")
.map_err(server_error)?;
match message {
WebSocketClientMessage::Handshake(_) => {
return Err(client_error(anyhow::anyhow!(
"Unexpected handshake message"
)));
}
WebSocketClientMessage::CursorPositions(cursors) => {
state
.cursors
.update_cursors(vault_id.clone(), &device_id, cursors.document_to_cursors)
.await;
}
}
}
Ok::<(), SyncServerError>(())
});
tokio::select! {
_ = &mut send_task => recv_task.abort(),
@ -143,28 +134,13 @@ async fn websocket(
send_task
.await
.context("Websocket send task failed")
.context("WebSocket send task failed")
.map_err(server_error)??;
recv_task
.await
.context("Websocket receive task failed")
.map_err(server_error)?;
.context("WebSocket receive task failed")
.map_err(server_error)??;
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)
}