Extract reconcile (#85)

sync-server/src/app_state/websocket/broadcasts.rs

@@ -0,0 +1,63 @@
+use std::{collections::HashMap, sync::Arc};
+
+use anyhow::Context;
+use tokio::sync::{Mutex, broadcast};
+
+use super::models::WebSocketServerMessageWithOrigin;
+use crate::{
+    app_state::database::models::VaultId, config::server_config::ServerConfig, errors::server_error,
+};
+
+#[derive(Debug, Clone)]
+pub struct Broadcasts {
+    max_clients_per_vault: usize,
+    tx: Arc<Mutex<HashMap<VaultId, broadcast::Sender<WebSocketServerMessageWithOrigin>>>>,
+}
+
+impl Broadcasts {
+    pub fn new(server_config: &ServerConfig) -> Self {
+        Self {
+            max_clients_per_vault: server_config.max_clients_per_vault,
+            tx: Arc::new(Mutex::new(HashMap::new())),
+        }
+    }
+
+    pub async fn get_receiver(
+        &self,
+        vault: VaultId,
+    ) -> broadcast::Receiver<WebSocketServerMessageWithOrigin> {
+        let tx = self.get_or_create(vault).await;
+
+        tx.subscribe()
+    }
+
+    /// Notify all clients (who are subscribed to the vault) about an update.
+    /// We only log failures.
+    pub async fn send_document_update(
+        &self,
+        vault: VaultId,
+        document: WebSocketServerMessageWithOrigin,
+    ) {
+        let tx = self.get_or_create(vault).await;
+
+        let result = tx
+            .send(document)
+            .context("Cannot broadcast server message to websocket listeners")
+            .map_err(server_error);
+
+        if result.is_err() {
+            log::debug!("Failed to send message: {result:?}");
+        }
+    }
+
+    async fn get_or_create(
+        &self,
+        vault: VaultId,
+    ) -> broadcast::Sender<WebSocketServerMessageWithOrigin> {
+        let mut tx = self.tx.lock().await;
+
+        tx.entry(vault)
+            .or_insert_with(|| broadcast::channel(self.max_clients_per_vault).0.clone())
+            .clone()
+    }
+}

sync-server/src/app_state/websocket/models.rs

@@ -0,0 +1,88 @@
+use std::collections::HashMap;
+
+use serde::{Deserialize, Serialize};
+use ts_rs::TS;
+
+use crate::app_state::database::models::{DeviceId, DocumentVersionWithoutContent, VaultUpdateId};
+
+#[derive(TS, Deserialize, Clone, Debug)]
+#[serde(rename_all = "camelCase")]
+pub struct WebSocketHandshake {
+    pub token: String,
+    pub device_id: DeviceId,
+
+    #[ts(as = "Option<i32>")]
+    pub last_seen_vault_update_id: Option<VaultUpdateId>,
+}
+
+#[derive(TS, Serialize, Deserialize, Clone, Debug)]
+#[serde(rename_all = "camelCase")]
+pub struct CursorSpan {
+    pub start: usize,
+    pub end: usize,
+}
+
+#[derive(TS, Deserialize, Clone, Debug)]
+#[serde(rename_all = "camelCase")]
+pub struct CursorPositionFromClient {
+    pub document_to_cursors: HashMap<String, Vec<CursorSpan>>,
+}
+
+#[derive(TS, Serialize, Clone, Debug)]
+#[serde(rename_all = "camelCase")]
+pub struct ClientCursors {
+    pub user_name: String,
+    pub device_id: DeviceId,
+    pub cursors: HashMap<String, Vec<CursorSpan>>,
+}
+
+#[derive(TS, Serialize, Clone, Debug)]
+#[serde(rename_all = "camelCase")]
+pub struct CursorPositionFromServer {
+    pub clients: Vec<ClientCursors>,
+}
+
+#[derive(TS, Serialize, Clone, Debug)]
+#[serde(rename_all = "camelCase")]
+pub struct WebSocketVaultUpdate {
+    pub documents: Vec<DocumentVersionWithoutContent>,
+    pub is_initial_sync: bool,
+}
+
+#[derive(TS, Deserialize, Clone, Debug)]
+#[serde(rename_all = "camelCase", tag = "type")]
+#[ts(export)]
+pub enum WebSocketClientMessage {
+    Handshake(WebSocketHandshake),
+    CursorPositions(CursorPositionFromClient),
+}
+
+#[derive(TS, Serialize, Clone, Debug)]
+#[serde(rename_all = "camelCase", tag = "type")]
+#[ts(export)]
+pub enum WebSocketServerMessage {
+    VaultUpdate(WebSocketVaultUpdate),
+    CursorPositions(CursorPositionFromServer),
+}
+
+#[derive(Clone, Debug)]
+pub struct WebSocketServerMessageWithOrigin {
+    pub origin_device_id: Option<DeviceId>,
+    pub message: WebSocketServerMessage,
+}
+
+impl WebSocketServerMessageWithOrigin {
+    pub fn new(message: WebSocketServerMessage) -> Self {
+        Self {
+            origin_device_id: None,
+            message,
+        }
+    }
+
+    pub fn with_origin(origin_device_id: DeviceId, message: WebSocketServerMessage) -> Self {
+        Self {
+            origin_device_id: Some(origin_device_id),
+            message,
+        }
+    }
+}

sync-server/src/app_state/websocket/utils.rs

@@ -0,0 +1,80 @@
+use anyhow::Context;
+use axum::extract::ws::{Message, WebSocket};
+use futures::{sink::SinkExt, stream::SplitSink};
+
+use super::models::{WebSocketClientMessage, WebSocketHandshake, WebSocketServerMessage};
+use crate::{
+    app_state::{
+        AppState,
+        database::models::{DocumentVersionWithoutContent, VaultId, VaultUpdateId},
+    },
+    config::user_config::User,
+    errors::{SyncServerError, server_error, unauthenticated_error},
+    server::auth::auth,
+};
+
+pub struct AuthenticatedWebSocketHandshake {
+    pub handshake: WebSocketHandshake,
+    pub user: User,
+}
+
+pub fn get_authenticated_handshake(
+    state: &AppState,
+    vault_id: &VaultId,
+    message: Option<Message>,
+) -> Result<AuthenticatedWebSocketHandshake, SyncServerError> {
+    if let Some(Message::Text(message)) = message {
+        let message: WebSocketClientMessage = serde_json::from_str(&message)
+            .context("Failed to parse message")
+            .map_err(server_error)?;
+
+        match message {
+            WebSocketClientMessage::Handshake(handshake) => {
+                let user = auth(state, handshake.token.trim(), vault_id)?;
+                Ok(AuthenticatedWebSocketHandshake { handshake, user })
+            }
+            WebSocketClientMessage::CursorPositions(_) => Err(unauthenticated_error(
+                anyhow::anyhow!("Expected a handshake message"),
+            )),
+        }
+    } else {
+        Err(unauthenticated_error(anyhow::anyhow!(
+            "Failed to authenticate due to invalid message"
+        )))
+    }
+}
+
+pub async fn get_unseen_documents(
+    state: &AppState,
+    vault_id: &VaultId,
+    last_seen_vault_update_id: Option<VaultUpdateId>,
+) -> Result<Vec<DocumentVersionWithoutContent>, SyncServerError> {
+    if let Some(update_id) = 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)
+    }
+}
+
+pub async fn send_update_over_websocket(
+    update: &WebSocketServerMessage,
+    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)
+}