Extract reconcile (#85)

sync-server/src/app_state/cursors.rs

@@ -0,0 +1,128 @@
+use core::time::Duration;
+use std::{collections::HashMap, sync::Arc};
+
+use tokio::sync::Mutex;
+
+use super::{
+    database::models::{DeviceId, VaultId},
+    websocket::{
+        broadcasts::Broadcasts,
+        models::{
+            ClientCursors, CursorPositionFromServer, CursorSpan, WebSocketServerMessage,
+            WebSocketServerMessageWithOrigin,
+        },
+    },
+};
+use crate::config::database_config::DatabaseConfig;
+
+#[derive(Clone, Debug)]
+pub struct Cursors {
+    config: DatabaseConfig,
+    broadcasts: Broadcasts,
+    vault_to_cursors: Arc<Mutex<HashMap<VaultId, Vec<ClientCursorsWithTimeToLive>>>>,
+}
+
+impl Cursors {
+    pub fn new(config: &DatabaseConfig, broadcasts: &Broadcasts) -> Self {
+        Self {
+            config: config.clone(),
+            broadcasts: broadcasts.clone(),
+            vault_to_cursors: Arc::new(Mutex::new(HashMap::new())),
+        }
+    }
+
+    pub async fn update_cursors(
+        &self,
+        vault_id: VaultId,
+        user_name: String,
+        device_id: &DeviceId,
+        document_to_cursors: HashMap<String, Vec<CursorSpan>>,
+    ) {
+        let mut vault_to_cursors = self.vault_to_cursors.lock().await;
+
+        let all_device_cursors = vault_to_cursors.entry(vault_id).or_insert_with(Vec::new);
+
+        all_device_cursors.retain(|c| &c.client_cursors.device_id != device_id);
+        all_device_cursors.push(ClientCursorsWithTimeToLive::new(ClientCursors {
+            user_name,
+            device_id: device_id.to_string(),
+            cursors: document_to_cursors,
+        }));
+
+        drop(vault_to_cursors); // Explicitly drop the lock before broadcasting to avoid deadlock
+        self.broadcast_cursors().await;
+    }
+
+    pub async fn get_cursors(&self, vault_id: &VaultId) -> Vec<ClientCursors> {
+        let vault_to_cursors = self.vault_to_cursors.lock().await;
+        vault_to_cursors
+            .get(vault_id)
+            .map(|cursors| {
+                cursors
+                    .iter()
+                    .cloned()
+                    .map(|with_ttl| with_ttl.client_cursors)
+                    .collect::<Vec<_>>()
+            })
+            .unwrap_or_default()
+    }
+
+    pub fn start_background_task(self) {
+        tokio::spawn(async move {
+            loop {
+                self.remove_expired_cursors().await;
+                tokio::time::sleep(Duration::from_secs(1)).await;
+            }
+        });
+    }
+
+    async fn remove_expired_cursors(&self) {
+        let mut vault_to_cursors = self.vault_to_cursors.lock().await;
+
+        for (_vault_id, cursors) in vault_to_cursors.iter_mut() {
+            cursors.retain(|cursor| !cursor.is_expired(self.config.cursor_timeout));
+        }
+    }
+
+    async fn broadcast_cursors(&self) {
+        let vault_to_cursors = self.vault_to_cursors.lock().await;
+
+        for (vault_id, cursors) in vault_to_cursors.iter() {
+            self.broadcasts
+                .send_document_update(
+                    vault_id.clone(),
+                    WebSocketServerMessageWithOrigin::new(WebSocketServerMessage::CursorPositions(
+                        CursorPositionFromServer {
+                            clients: cursors.iter().map(|c| c.client_cursors.clone()).collect(),
+                        },
+                    )),
+                )
+                .await;
+        }
+    }
+
+    pub async fn remove_cursors_of_device(&self, vault_id: &str, device_id: &str) {
+        let mut vault_to_cursors = self.vault_to_cursors.lock().await;
+
+        if let Some(cursors) = vault_to_cursors.get_mut(vault_id) {
+            cursors.retain(|c| c.client_cursors.device_id != device_id);
+        }
+    }
+}
+
+#[derive(Clone, Debug)]
+struct ClientCursorsWithTimeToLive {
+    client_cursors: ClientCursors,
+    last_updated: std::time::Instant,
+}
+
+impl ClientCursorsWithTimeToLive {
+    fn new(client_cursors: ClientCursors) -> Self {
+        Self {
+            client_cursors,
+            last_updated: std::time::Instant::now(),
+        }
+    }
+
+    pub fn is_expired(&self, ttl: Duration) -> bool { self.last_updated.elapsed() > ttl }
+}

sync-server/src/app_state/database.rs

@@ -0,0 +1,425 @@
+use core::time::Duration;
+use std::{collections::HashMap, sync::Arc};
+
+use anyhow::{Context as _, Result};
+use models::{
+    DocumentId, DocumentVersionWithoutContent, StoredDocumentVersion, VaultId, VaultUpdateId,
+};
+use sqlx::{sqlite::SqliteConnectOptions, types::chrono::Utc};
+
+pub mod models;
+use sqlx::{Pool, Sqlite, sqlite::SqlitePoolOptions};
+use tokio::sync::Mutex;
+use uuid::fmt::Hyphenated;
+
+use super::websocket::{
+    broadcasts::Broadcasts,
+    models::{WebSocketServerMessage, WebSocketServerMessageWithOrigin, WebSocketVaultUpdate},
+};
+use crate::config::database_config::DatabaseConfig;
+
+#[derive(Clone, Debug)]
+pub struct Database {
+    config: DatabaseConfig,
+    broadcasts: Broadcasts,
+    connection_pools: Arc<Mutex<HashMap<VaultId, Pool<Sqlite>>>>,
+}
+
+pub type Transaction<'a> = sqlx::Transaction<'a, Sqlite>;
+
+impl Database {
+    pub async fn try_new(config: &DatabaseConfig, broadcasts: &Broadcasts) -> Result<Self> {
+        tokio::fs::create_dir_all(&config.databases_directory_path)
+            .await
+            .with_context(|| {
+                format!(
+                    "Failed to create databases directory: {}",
+                    config.databases_directory_path.to_string_lossy()
+                )
+            })?;
+
+        let mut connection_pools = std::collections::HashMap::new();
+
+        let mut entries = tokio::fs::read_dir(&config.databases_directory_path).await?;
+        while let Some(entry) = entries.next_entry().await? {
+            if !entry.file_name().to_string_lossy().ends_with(".sqlite") {
+                continue;
+            }
+
+            let vault: VaultId = entry
+                .file_name()
+                .to_string_lossy()
+                .trim_end_matches(".sqlite")
+                .to_owned();
+
+            connection_pools.insert(
+                vault.clone(),
+                Self::create_vault_database(config, &vault).await?,
+            );
+        }
+
+        Ok(Self {
+            config: config.clone(),
+            connection_pools: Arc::new(Mutex::new(connection_pools)),
+            broadcasts: broadcasts.clone(),
+        })
+    }
+
+    async fn create_vault_database(
+        config: &DatabaseConfig,
+        vault: &VaultId,
+    ) -> Result<Pool<Sqlite>> {
+        let file_name = config
+            .databases_directory_path
+            .join(format!("{vault}.sqlite"));
+
+        let connection_options = SqliteConnectOptions::new()
+            .filename(file_name.clone())
+            .create_if_missing(true)
+            .busy_timeout(Duration::from_secs(3600))
+            .journal_mode(sqlx::sqlite::SqliteJournalMode::Wal);
+
+        let pool = SqlitePoolOptions::new()
+            .max_connections(config.max_connections_per_vault)
+            .test_before_acquire(true)
+            .connect_with(connection_options)
+            .await
+            .with_context(|| format!("Cannot open database at {}", file_name.display()))?;
+
+        Self::run_migrations(&pool).await?;
+
+        Ok(pool)
+    }
+
+    async fn run_migrations(pool: &Pool<Sqlite>) -> Result<()> {
+        sqlx::migrate!("src/app_state/database/migrations")
+            .run(pool)
+            .await
+            .context("Cannot check for pending migrations")
+    }
+
+    async fn get_connection_pool(&self, vault: &VaultId) -> Result<Pool<Sqlite>> {
+        let mut pools = self.connection_pools.lock().await;
+        if !pools.contains_key(vault) {
+            let pool = Self::create_vault_database(&self.config, vault).await?;
+            pools.insert(vault.clone(), pool);
+        }
+
+        let pool = pools
+            .get(vault)
+            .expect("Pool was just inserted or already exists");
+
+        Ok(pool.clone())
+    }
+
+    /// Attempting to write from this transaction might result in a
+    /// database locked error. Use this transaction for read-only operations.
+    pub async fn create_readonly_transaction(
+        &self,
+        vault: &VaultId,
+    ) -> Result<Transaction<'static>> {
+        self.get_connection_pool(vault)
+            .await?
+            .begin()
+            .await
+            .context("Cannot create transaction")
+    }
+
+    pub async fn create_write_transaction(&self, vault: &VaultId) -> Result<Transaction<'static>> {
+        let mut transaction = self.create_readonly_transaction(vault).await?;
+
+        // sqlx doesn't support immediate transactions for sqlite: https://github.com/launchbadge/sqlx/issues/481
+        sqlx::query!("END; BEGIN IMMEDIATE;")
+            .execute(&mut *transaction)
+            .await?;
+
+        Ok(transaction)
+    }
+
+    /// Return the latest state of all documents in the vault
+    pub async fn get_latest_documents(
+        &self,
+        vault: &VaultId,
+        transaction: Option<&mut Transaction<'_>>,
+    ) -> Result<Vec<DocumentVersionWithoutContent>> {
+        let query = sqlx::query!(
+            r#"
+            select
+                vault_update_id,
+                document_id as "document_id: Hyphenated",
+                relative_path,
+                updated_date as "updated_date: chrono::DateTime<Utc>",
+                is_deleted,
+                user_id,
+                device_id,
+                length(content) as "content_size: u64"
+            from latest_document_versions
+            order by vault_update_id
+            "#,
+        );
+
+        if let Some(transaction) = transaction {
+            query.fetch_all(&mut **transaction).await
+        } else {
+            query
+                .fetch_all(&self.get_connection_pool(vault).await?)
+                .await
+        }
+        .context("Cannot fetch latest documents")
+        .map(|rows| {
+            rows.into_iter()
+                .map(|row| DocumentVersionWithoutContent {
+                    vault_update_id: row.vault_update_id,
+                    document_id: row.document_id.into(),
+                    relative_path: row.relative_path,
+                    updated_date: row.updated_date,
+                    is_deleted: row.is_deleted,
+                    user_id: row.user_id,
+                    device_id: row.device_id,
+                    content_size: row
+                        .content_size
+                        .expect("Content size can't be null but sqlx can't infer it"),
+                })
+                .collect()
+        })
+    }
+
+    /// Return the latest state of all documents (including deleted) in the
+    /// vault which have changed since the given update id
+    pub async fn get_latest_documents_since(
+        &self,
+        vault: &VaultId,
+        vault_update_id: VaultUpdateId,
+        transaction: Option<&mut Transaction<'_>>,
+    ) -> Result<Vec<DocumentVersionWithoutContent>> {
+        let query = sqlx::query!(
+            r#"
+            select
+                vault_update_id,
+                document_id as "document_id: Hyphenated",
+                relative_path,
+                updated_date as "updated_date: chrono::DateTime<Utc>",
+                is_deleted,
+                user_id,
+                device_id,
+                length(content) as "content_size: u64"
+            from latest_document_versions
+            where vault_update_id > ?
+            order by vault_update_id
+            "#,
+            vault_update_id
+        );
+
+        if let Some(transaction) = transaction {
+            query.fetch_all(&mut **transaction).await
+        } else {
+            query
+                .fetch_all(&self.get_connection_pool(vault).await?)
+                .await
+        }
+        .with_context(|| {
+            format!("Cannot fetch latest documents since vault_update_id {vault_update_id}")
+        })
+        .map(|rows| {
+            rows.into_iter()
+                .map(|row| DocumentVersionWithoutContent {
+                    vault_update_id: row.vault_update_id,
+                    document_id: row.document_id.into(),
+                    relative_path: row.relative_path,
+                    updated_date: row.updated_date,
+                    is_deleted: row.is_deleted,
+                    user_id: row.user_id,
+                    device_id: row.device_id,
+                    content_size: row
+                        .content_size
+                        .expect("Content size can't be null but sqlx can't infer it"),
+                })
+                .collect()
+        })
+    }
+
+    pub async fn get_max_update_id_in_vault(
+        &self,
+        vault: &VaultId,
+        transaction: Option<&mut Transaction<'_>>,
+    ) -> Result<i64> {
+        let query = sqlx::query!(
+            r#"
+            select coalesce(max(vault_update_id), 0) as max_vault_update_id
+            from documents
+            "#,
+        );
+
+        if let Some(transaction) = transaction {
+            query.fetch_one(&mut **transaction).await
+        } else {
+            query
+                .fetch_one(&self.get_connection_pool(vault).await?)
+                .await
+        }
+        .map(|row| row.max_vault_update_id)
+        .context("Cannot fetch max update id in vault")
+    }
+
+    pub async fn get_latest_document_by_path(
+        &self,
+        vault: &VaultId,
+        relative_path: &str,
+        transaction: Option<&mut Transaction<'_>>,
+    ) -> Result<Option<StoredDocumentVersion>> {
+        let query = sqlx::query_as!(
+            StoredDocumentVersion,
+            r#"
+            select
+                vault_update_id,
+                document_id as "document_id: Hyphenated",
+                relative_path,
+                updated_date as "updated_date: chrono::DateTime<Utc>",
+                content,
+                is_deleted,
+                user_id,
+                device_id
+            from latest_document_versions
+            where relative_path = ?
+            order by vault_update_id desc  -- `latest_document_versions` only contains a single latest version of each document, however,
+                                           -- multiple documents can have the same `relative_path`, if they have been deleted. That's
+                                           -- why we only care about the latest version of the document with the given relative path.
+            limit 1
+            "#,
+            relative_path
+        );
+
+        if let Some(transaction) = transaction {
+            query.fetch_optional(&mut **transaction).await
+        } else {
+            query
+                .fetch_optional(&self.get_connection_pool(vault).await?)
+                .await
+        }
+        .context("Cannot fetch latest document version")
+    }
+
+    pub async fn get_latest_document(
+        &self,
+        vault: &VaultId,
+        document_id: &DocumentId,
+        transaction: Option<&mut Transaction<'_>>,
+    ) -> Result<Option<StoredDocumentVersion>> {
+        let document_id = document_id.as_hyphenated();
+        let query = sqlx::query_as!(
+            StoredDocumentVersion,
+            r#"
+            select
+                vault_update_id,
+                document_id as "document_id: Hyphenated",
+                relative_path,
+                updated_date as "updated_date: chrono::DateTime<Utc>",
+                content,
+                is_deleted,
+                user_id,
+                device_id
+            from latest_document_versions
+            where document_id = ?
+            "#,
+            document_id
+        );
+
+        if let Some(transaction) = transaction {
+            query.fetch_optional(&mut **transaction).await
+        } else {
+            query
+                .fetch_optional(&self.get_connection_pool(vault).await?)
+                .await
+        }
+        .context("Cannot fetch latest document version")
+    }
+
+    pub async fn get_document_version(
+        &self,
+        vault: &VaultId,
+        vault_update_id: VaultUpdateId,
+        transaction: Option<&mut Transaction<'_>>,
+    ) -> Result<Option<StoredDocumentVersion>> {
+        let query = sqlx::query_as!(
+            StoredDocumentVersion,
+            r#"
+            select
+                vault_update_id,
+                document_id as "document_id: Hyphenated",
+                relative_path,
+                updated_date as "updated_date: chrono::DateTime<Utc>",
+                content,
+                is_deleted,
+                user_id,
+                device_id
+            from documents
+            where vault_update_id = ?"#,
+            vault_update_id
+        );
+
+        if let Some(transaction) = transaction {
+            query.fetch_optional(&mut **transaction).await
+        } else {
+            query
+                .fetch_optional(&self.get_connection_pool(vault).await?)
+                .await
+        }
+        .context("Cannot fetch document version")
+    }
+
+    pub async fn insert_document_version(
+        &self,
+        vault_id: &VaultId,
+        version: &StoredDocumentVersion,
+        transaction: Option<&mut Transaction<'_>>,
+    ) -> Result<()> {
+        let document_id = version.document_id.as_hyphenated();
+        let query = sqlx::query!(
+            r#"
+            insert into documents (
+                vault_update_id,
+                document_id,
+                relative_path,
+                updated_date,
+                content,
+                is_deleted,
+                user_id,
+                device_id
+            )
+            values (?, ?, ?, ?, ?, ?, ?, ?)
+            "#,
+            version.vault_update_id,
+            document_id,
+            version.relative_path,
+            version.updated_date,
+            version.content,
+            version.is_deleted,
+            version.user_id,
+            version.device_id
+        );
+
+        if let Some(transaction) = transaction {
+            query.execute(&mut **transaction).await
+        } else {
+            query
+                .execute(&self.get_connection_pool(vault_id).await?)
+                .await
+        }
+        .context("Cannot insert document version")?;
+
+        self.broadcasts
+            .send_document_update(
+                vault_id.clone(),
+                WebSocketServerMessageWithOrigin::with_origin(
+                    version.device_id.clone(),
+                    WebSocketServerMessage::VaultUpdate(WebSocketVaultUpdate {
+                        documents: vec![version.clone().into()],
+                        is_initial_sync: false,
+                    }),
+                ),
+            )
+            .await;
+
+        Ok(())
+    }
+}

sync-server/src/app_state/database/migrations/20241207143519_bootstrap.sql

@@ -0,0 +1,21 @@
+CREATE TABLE IF NOT EXISTS documents (
+    vault_update_id INTEGER NOT NULL PRIMARY KEY,
+    document_id TEXT NOT NULL,
+    relative_path TEXT NOT NULL,
+    updated_date TIMESTAMP NOT NULL,
+    content BLOB NOT NULL,
+    is_deleted BOOLEAN NOT NULL
+);
+
+CREATE VIEW IF NOT EXISTS latest_document_versions AS
+SELECT d.*
+FROM documents d
+INNER JOIN (
+    SELECT MAX(vault_update_id) AS max_version_id
+    FROM documents
+    GROUP BY document_id
+) max_versions
+ON d.vault_update_id = max_versions.max_version_id;
+
+CREATE INDEX IF NOT EXISTS idx_documents_vault_id_relative_path
+ON documents (relative_path);

sync-server/src/app_state/database/migrations/20250522192949_add_provenance_columns.sql

@@ -0,0 +1,2 @@
+ALTER TABLE documents ADD COLUMN user_id TEXT NOT NULL DEFAULT "";
+ALTER TABLE documents ADD COLUMN device_id TEXT NOT NULL DEFAULT "";

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

@@ -0,0 +1,89 @@
+use base64::{Engine as _, engine::general_purpose::STANDARD};
+use chrono::{DateTime, Utc};
+use serde::Serialize;
+use ts_rs::TS;
+
+pub type VaultId = String;
+pub type VaultUpdateId = i64;
+
+pub type DocumentId = uuid::Uuid;
+pub type UserId = String;
+pub type DeviceId = String;
+
+#[derive(Debug, Clone)]
+pub struct StoredDocumentVersion {
+    pub vault_update_id: VaultUpdateId,
+    pub document_id: DocumentId,
+    pub relative_path: String,
+    pub updated_date: DateTime<Utc>,
+    pub content: Vec<u8>,
+    pub is_deleted: bool,
+    pub user_id: UserId,
+    pub device_id: DeviceId,
+}
+
+impl PartialEq<Self> for StoredDocumentVersion {
+    fn eq(&self, other: &Self) -> bool { self.vault_update_id == other.vault_update_id }
+}
+
+#[derive(TS, Debug, Clone, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct DocumentVersionWithoutContent {
+    #[ts(as = "i32")]
+    pub vault_update_id: VaultUpdateId,
+
+    pub document_id: DocumentId,
+    pub relative_path: String,
+    pub updated_date: DateTime<Utc>,
+    pub is_deleted: bool,
+    pub user_id: UserId,
+    pub device_id: DeviceId,
+
+    #[ts(as = "i32")]
+    pub content_size: u64,
+}
+
+impl From<StoredDocumentVersion> for DocumentVersionWithoutContent {
+    fn from(value: StoredDocumentVersion) -> Self {
+        Self {
+            vault_update_id: value.vault_update_id,
+            document_id: value.document_id,
+            relative_path: value.relative_path,
+            updated_date: value.updated_date,
+            is_deleted: value.is_deleted,
+            user_id: value.user_id,
+            device_id: value.device_id,
+            content_size: value.content.len() as u64,
+        }
+    }
+}
+
+#[derive(TS, Debug, Clone, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct DocumentVersion {
+    #[ts(as = "i32")]
+    pub vault_update_id: VaultUpdateId,
+
+    pub document_id: DocumentId,
+    pub relative_path: String,
+    pub updated_date: DateTime<Utc>,
+    pub content_base64: String,
+    pub is_deleted: bool,
+    pub user_id: UserId,
+    pub device_id: DeviceId,
+}
+
+impl From<StoredDocumentVersion> for DocumentVersion {
+    fn from(value: StoredDocumentVersion) -> Self {
+        Self {
+            vault_update_id: value.vault_update_id,
+            document_id: value.document_id,
+            relative_path: value.relative_path,
+            updated_date: value.updated_date,
+            content_base64: STANDARD.encode(&value.content),
+            is_deleted: value.is_deleted,
+            user_id: value.user_id,
+            device_id: value.device_id,
+        }
+    }
+}

sync-server/src/app_state/websocket.rs

@@ -0,0 +1,3 @@
+pub mod broadcasts;
+pub mod models;
+pub mod utils;

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)
+}