Add WebSocket support (#12)

backend/Cargo.lock

@@ -459,6 +459,16 @@ dependencies = [
  "clap_derive",
 ]
 
+[[package]]
+name = "clap-verbosity-flag"
+version = "3.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2678fade3b77aa3a8ff3aae87e9c008d3fb00473a41c71fbf74e91c8c7b37e84"
+dependencies = [
+ "clap",
+ "log",
+]
+
 [[package]]
 name = "clap_builder"
 version = "4.5.32"
@@ -2069,9 +2079,9 @@ dependencies = [
 
 [[package]]
 name = "serde_json"
-version = "1.0.133"
+version = "1.0.140"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c7fceb2473b9166b2294ef05efcb65a3db80803f0b03ef86a5fc88a2b85ee377"
+checksum = "20068b6e96dc6c9bd23e01df8827e6c7e1f2fddd43c21810382803c136b99373"
 dependencies = [
  "itoa",
  "memchr",
@@ -2524,6 +2534,8 @@ dependencies = [
  "axum_typed_multipart",
  "chrono",
  "clap",
+ "clap-verbosity-flag",
+ "futures",
  "log",
  "rand",
  "reconcile",
@@ -2531,6 +2543,7 @@ dependencies = [
  "sanitize-filename",
  "schemars",
  "serde",
+ "serde_json",
  "serde_yaml",
  "sqlx",
  "sync_lib",

backend/Cargo.toml

@@ -22,6 +22,7 @@ thiserror = { version = "1.0.66", default-features = false  }
 codegen-units = 1
 lto = true
 opt-level = 3
+strip="debuginfo"  # Keep some info for better panics
 
 [workspace.lints.rust]
 unsafe_code = "forbid"

backend/Dockerfile

@@ -8,7 +8,7 @@ RUN cargo install sqlx-cli
 COPY . .
 
 RUN sqlx database create --database-url sqlite://db.sqlite3
-RUN sqlx migrate run --source sync_server/src/database/migrations --database-url sqlite://db.sqlite3
+RUN sqlx migrate run --source sync_server/src/app_state/database/migrations --database-url sqlite://db.sqlite3
 
 RUN cargo build --package sync_server --release --target x86_64-unknown-linux-musl
 

backend/config-e2e.yml

@@ -1,10 +1,13 @@
 database:
   databases_directory_path: databases
   max_connections: 12
+
 server:
   host: 0.0.0.0
   port: 3000
   max_body_size_mb: 512
+  max_clients_per_vault: 256
+
 users:
   user_tokens:
   - name: admin

backend/sync_server/Cargo.toml

@@ -34,6 +34,9 @@ sanitize-filename = "0.6.0"
 axum-jsonschema = { version = "0.8.0", features = ["aide"] }
 regex = "1.11.1"
 clap = { version = "4.5.32", features = ["derive"] }
+futures = "0.3.31"
+serde_json = "1.0.140"
+clap-verbosity-flag = "3.0.2"
 
 [lints]
 workspace = true

backend/sync_server/README.md

@@ -1,4 +1,4 @@
 
 cargo install sqlx-cli
 rm db.sqlite3; sqlx database create --database-url sqlite://db.sqlite3
-sqlx migrate run --source sync_server/src/database/migrations --database-url sqlite://db.sqlite3
+sqlx migrate run --source sync_server/src/app_state/database/migrations --database-url sqlite://db.sqlite3

backend/sync_server/src/app_state.rs

@@ -0,0 +1,34 @@
+pub mod broadcasts;
+pub mod database;
+
+use std::ffi::OsString;
+
+use anyhow::Result;
+use broadcasts::Broadcasts;
+use database::Database;
+
+use crate::{config::Config, consts::DEFAULT_CONFIG_PATH};
+
+#[derive(Clone, Debug)]
+pub struct AppState {
+    pub config: Config,
+    pub database: Database,
+    pub broadcasts: Broadcasts,
+}
+
+impl AppState {
+    pub async fn try_new(config_path: Option<OsString>) -> Result<Self> {
+        let config_path = config_path.unwrap_or_else(|| OsString::from(DEFAULT_CONFIG_PATH));
+        let path = std::path::PathBuf::from(config_path);
+
+        let config = Config::read_or_create(&path).await?;
+        let database = Database::try_new(&config.database).await?;
+        let broadcasts = Broadcasts::new(&config.server);
+
+        Ok(Self {
+            config,
+            database,
+            broadcasts,
+        })
+    }
+}

backend/sync_server/src/app_state/broadcasts.rs

@@ -0,0 +1,57 @@
+use std::{collections::HashMap, sync::Arc};
+
+use anyhow::Context;
+use tokio::sync::{Mutex, broadcast};
+
+use super::database::models::{DocumentVersionWithoutContent, VaultId};
+use crate::{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<DocumentVersionWithoutContent>>>>,
+}
+
+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<DocumentVersionWithoutContent> {
+        let tx = self.get_or_create(vault).await;
+
+        tx.subscribe()
+    }
+
+    /// Sent a document update to all clients subscribed to the vault.
+    /// We ignore & log failures.
+    pub async fn send(&self, vault: VaultId, document: DocumentVersionWithoutContent) {
+        let tx = self.get_or_create(vault).await;
+
+        let result = tx
+            .send(document)
+            .context("Cannot broadcast update 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<DocumentVersionWithoutContent> {
+        let mut tx = self.tx.lock().await;
+
+        tx.entry(vault)
+            .or_insert_with(|| broadcast::channel(self.max_clients_per_vault).0.clone())
+            .clone()
+    }
+}

backend/sync_server/src/app_state/database.rs

@@ -0,0 +1,357 @@
+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 crate::config::database_config::DatabaseConfig;
+
+#[derive(Clone, Debug)]
+pub struct Database {
+    config: DatabaseConfig,
+    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) -> 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)),
+        })
+    }
+
+    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)
+            .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_as!(
+            DocumentVersionWithoutContent,
+            r#"
+            select 
+                vault_update_id,
+                document_id as "document_id: Hyphenated", 
+                relative_path,
+                updated_date as "updated_date: chrono::DateTime<Utc>",
+                is_deleted
+            from latest_document_versions
+            order by vault_update_id desc
+            "#,
+        );
+
+        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")
+    }
+
+    /// 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_as!(
+            DocumentVersionWithoutContent,
+            r#"
+            select
+                vault_update_id,
+                document_id as "document_id: Hyphenated",
+                relative_path,
+                updated_date as "updated_date: chrono::DateTime<Utc>",
+                is_deleted
+            from latest_document_versions
+            where vault_update_id > ?
+            order by vault_update_id desc
+            "#,
+            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}")
+        })
+    }
+
+    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
+            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
+            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
+            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: &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
+            )
+            values (?, ?, ?, ?, ?, ?)
+            "#,
+            version.vault_update_id,
+            document_id,
+            version.relative_path,
+            version.updated_date,
+            version.content,
+            version.is_deleted
+        );
+
+        if let Some(transaction) = transaction {
+            query.execute(&mut **transaction).await
+        } else {
+            query.execute(&self.get_connection_pool(vault).await?).await
+        }
+        .context("Cannot insert document version")?;
+
+        Ok(())
+    }
+}

backend/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);

backend/sync_server/src/app_state/database/models.rs

@@ -0,0 +1,68 @@
+use chrono::{DateTime, Utc};
+use schemars::JsonSchema;
+use serde::Serialize;
+use sync_lib::bytes_to_base64;
+
+pub type VaultId = String;
+pub type VaultUpdateId = i64;
+pub type DocumentId = uuid::Uuid;
+
+#[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,
+}
+
+impl PartialEq<Self> for StoredDocumentVersion {
+    fn eq(&self, other: &Self) -> bool { self.vault_update_id == other.vault_update_id }
+}
+
+#[derive(Debug, Clone, Serialize, JsonSchema)]
+#[serde(rename_all = "camelCase")]
+pub struct DocumentVersionWithoutContent {
+    pub vault_update_id: VaultUpdateId,
+    pub document_id: DocumentId,
+    pub relative_path: String,
+    pub updated_date: DateTime<Utc>,
+    pub is_deleted: bool,
+}
+
+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,
+        }
+    }
+}
+
+#[derive(Debug, Clone, Serialize, JsonSchema)]
+#[serde(rename_all = "camelCase")]
+pub struct DocumentVersion {
+    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,
+}
+
+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: bytes_to_base64(&value.content),
+            is_deleted: value.is_deleted,
+        }
+    }
+}

backend/sync_server/src/cli.rs

@@ -1 +1,2 @@
 pub mod args;
+pub mod color_when;

backend/sync_server/src/cli/args.rs

@@ -1,38 +1,26 @@
 use std::ffi::OsString;
 
-use clap::{Parser, ValueEnum};
+use clap::Parser;
+use clap_verbosity_flag::{InfoLevel, Verbosity};
 
-/// Server for backing the VaultLink plugin
+use crate::cli::color_when::ColorWhen;
+
+/// Server for backing the `VaultLink` plugin
 #[derive(Parser, Debug)]
 #[command(version, about, long_about = None)]
 pub struct Args {
     #[arg(index = 1)]
     pub config_path: Option<OsString>,
 
+    #[command(flatten)]
+    pub verbose: Verbosity<InfoLevel>,
+
     #[arg(
             long,
-            require_equals = true,
             value_name = "WHEN",
-            num_args = 0..=1,
             default_value_t = ColorWhen::Auto,
             default_missing_value = "always",
             value_enum
         )]
     pub color: ColorWhen,
 }
-
-#[derive(ValueEnum, Copy, Clone, Debug, PartialEq, Eq)]
-pub enum ColorWhen {
-    Always,
-    Auto,
-    Never,
-}
-
-impl std::fmt::Display for ColorWhen {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        self.to_possible_value()
-            .expect("no values are skipped")
-            .get_name()
-            .fmt(f)
-    }
-}

backend/sync_server/src/cli/color_when.rs

@@ -0,0 +1,31 @@
+use std::io::IsTerminal;
+
+use clap::ValueEnum;
+
+#[derive(ValueEnum, Copy, Clone, Debug, PartialEq, Eq)]
+pub enum ColorWhen {
+    Always,
+    Auto,
+    Never,
+}
+
+impl ColorWhen {
+    pub fn use_colors(self) -> bool {
+        match self {
+            ColorWhen::Always => true,
+            ColorWhen::Auto => {
+                std::env::var_os("NO_COLOR").is_none() && std::io::stderr().is_terminal()
+            }
+            ColorWhen::Never => false,
+        }
+    }
+}
+
+impl std::fmt::Display for ColorWhen {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.to_possible_value()
+            .expect("no values are skipped")
+            .get_name()
+            .fmt(f)
+    }
+}

backend/sync_server/src/config/server_config.rs

@@ -1,7 +1,10 @@
 use log::debug;
 use serde::{Deserialize, Serialize};
 
-use crate::consts::{DEFAULT_HOST, DEFAULT_MAX_BODY_SIZE_MB, DEFAULT_PORT};
+use crate::consts::{
+    DEFAULT_HOST, DEFAULT_MAX_BODY_SIZE_MB, DEFAULT_MAX_CLIENTS_PER_VAULT, DEFAULT_PORT,
+};
+
 #[derive(Debug, Deserialize, Serialize, Clone)]
 pub struct ServerConfig {
     #[serde(default = "default_host")]
@@ -12,6 +15,9 @@ pub struct ServerConfig {
 
     #[serde(default = "default_max_body_size_mb")]
     pub max_body_size_mb: usize,
+
+    #[serde(default = "default_max_clients_per_vault")]
+    pub max_clients_per_vault: usize,
 }
 
 fn default_host() -> String {
@@ -29,12 +35,18 @@ fn default_max_body_size_mb() -> usize {
     DEFAULT_MAX_BODY_SIZE_MB
 }
 
+fn default_max_clients_per_vault() -> usize {
+    debug!("Using default max clients per vault: {DEFAULT_MAX_CLIENTS_PER_VAULT}");
+    DEFAULT_MAX_CLIENTS_PER_VAULT
+}
+
 impl Default for ServerConfig {
     fn default() -> Self {
         Self {
             host: default_host(),
             port: default_port(),
             max_body_size_mb: default_max_body_size_mb(),
+            max_clients_per_vault: default_max_clients_per_vault(),
         }
     }
 }

backend/sync_server/src/consts.rs

@@ -4,3 +4,4 @@ pub const DEFAULT_HOST: &str = "127.0.0.1";
 pub const DEFAULT_PORT: u16 = 3000;
 pub const DEFAULT_MAX_CONNECTIONS: u32 = 12;
 pub const DEFAULT_MAX_BODY_SIZE_MB: usize = 4096;
+pub const DEFAULT_MAX_CLIENTS_PER_VAULT: usize = 256;

backend/sync_server/src/main.rs

@@ -1,38 +1,79 @@
+mod app_state;
 mod cli;
 mod config;
 mod consts;
-mod database;
 mod errors;
 mod server;
 mod utils;
 
+use std::process::ExitCode;
+
 use anyhow::{Context as _, Result};
 use clap::Parser;
 use cli::args::Args;
 use errors::{SyncServerError, init_error};
 use log::info;
 use server::create_server;
-use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
+use tracing_subscriber::{EnvFilter, fmt::format, util::SubscriberInitExt};
 
 #[tokio::main]
-async fn main() -> Result<(), SyncServerError> {
+async fn main() -> ExitCode {
     let args = Args::parse();
 
-    tracing_subscriber::registry()
-        .with(
-            tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| {
-                format!(
-                    "{}=debug,tower_http=debug,axum::rejection=trace",
-                    env!("CARGO_CRATE_NAME")
-                )
-                .into()
-            }),
+    let mut result = set_up_logging(&args);
+
+    if result.is_ok() {
+        result = start_server(args).await;
+    }
+
+    match result {
+        Ok(()) => ExitCode::SUCCESS,
+        Err(e) => {
+            eprintln!("Failed to set up logging: {e}");
+            ExitCode::FAILURE
+        }
+    }
+}
+
+fn set_up_logging(args: &Args) -> Result<(), SyncServerError> {
+    let level_filter = match args.verbose.log_level_filter() {
+        // We don't want to allow disabling all logging
+        log::LevelFilter::Off | log::LevelFilter::Error => tracing::Level::ERROR,
+        log::LevelFilter::Warn => tracing::Level::WARN,
+        log::LevelFilter::Info => tracing::Level::INFO,
+        log::LevelFilter::Debug => tracing::Level::DEBUG,
+        log::LevelFilter::Trace => tracing::Level::TRACE,
+    };
+
+    let env_filter = EnvFilter::builder()
+        .with_default_directive(level_filter.into())
+        .from_env()
+        .context("Failed to create logging env filter")
+        .map_err(init_error)?;
+
+    let use_colors = args.color.use_colors();
+
+    let is_debug_mode = args.verbose.log_level_filter() >= log::LevelFilter::Debug;
+
+    tracing_subscriber::fmt()
+        .with_ansi(use_colors)
+        .with_env_filter(env_filter)
+        .event_format(
+            format()
+                .without_time()
+                .with_target(is_debug_mode)
+                .with_line_number(is_debug_mode)
+                .compact(),
         )
-        .with(tracing_subscriber::fmt::layer())
+        .finish()
         .try_init()
         .context("Failed to initialise tracing")
         .map_err(init_error)?;
 
+    Ok(())
+}
+
+async fn start_server(args: Args) -> Result<(), SyncServerError> {
     info!(
         "Starting VaultLink server version {}",
         env!("CARGO_PKG_VERSION")

backend/sync_server/src/server.rs

@@ -1,3 +1,16 @@
+mod auth;
+mod create_document;
+mod delete_document;
+mod fetch_document_version;
+mod fetch_document_version_content;
+mod fetch_latest_document_version;
+mod fetch_latest_documents;
+mod ping;
+mod requests;
+mod responses;
+mod update_document;
+mod websocket;
+
 use std::{ffi::OsString, sync::Arc};
 
 use aide::{
@@ -10,7 +23,6 @@ use aide::{
     transform::TransformOpenApi,
 };
 use anyhow::{Context as _, Result, anyhow};
-use app_state::AppState;
 use axum::{
     Extension, Json,
     extract::{DefaultBodyLimit, Request},
@@ -32,21 +44,10 @@ use tower_http::{
 use tracing::{Level, info_span};
 
 use crate::{
+    app_state::AppState,
     config::server_config::ServerConfig,
     errors::{SerializedError, not_found_error},
 };
-mod app_state;
-mod auth;
-mod create_document;
-mod delete_document;
-mod fetch_document_version;
-mod fetch_document_version_content;
-mod fetch_latest_document_version;
-mod fetch_latest_documents;
-mod ping;
-mod requests;
-mod responses;
-mod update_document;
 
 pub async fn create_server(config_path: Option<OsString>) -> Result<()> {
     aide::r#gen::on_error(|err| error!("{err}"));
@@ -65,6 +66,7 @@ pub async fn create_server(config_path: Option<OsString>) -> Result<()> {
             "/vaults/:vault_id/documents",
             get(fetch_latest_documents::fetch_latest_documents),
         )
+        .route("/vaults/:vault_id/ws", get(websocket::websocket_handler))
         .api_route(
             "/vaults/:vault_id/documents",
             post(create_document::create_document_multipart),

backend/sync_server/src/server/auth.rs

@@ -1,9 +1,10 @@
-use super::app_state::AppState;
 use crate::{
+    app_state::AppState,
     config::user_config::User,
     errors::{SyncServerError, unauthorized_error},
 };
 
+// TODO: turn this into a middleware
 pub fn auth(app_state: &AppState, token: &str) -> Result<User, SyncServerError> {
     app_state
         .config

backend/sync_server/src/server/create_document.rs

@@ -11,12 +11,16 @@ use serde::Deserialize;
 use sync_lib::base64_to_bytes;
 
 use super::{
-    app_state::AppState,
     auth::auth,
     requests::{CreateDocumentVersion, CreateDocumentVersionMultipart},
 };
 use crate::{
-    database::models::{DocumentId, DocumentVersionWithoutContent, StoredDocumentVersion, VaultId},
+    app_state::{
+        AppState,
+        database::models::{
+            DocumentId, DocumentVersionWithoutContent, StoredDocumentVersion, VaultId,
+        },
+    },
     errors::{SyncServerError, client_error, server_error},
     utils::sanitize_path,
 };
@@ -77,7 +81,7 @@ pub async fn create_document_json(
 
 async fn internal_create_document(
     auth_header: Authorization<Bearer>,
-    mut state: AppState,
+    state: AppState,
     vault_id: VaultId,
     document_id: Option<DocumentId>,
     relative_path: String,
@@ -139,5 +143,10 @@ async fn internal_create_document(
         .context("Failed to commit successful transaction")
         .map_err(server_error)?;
 
+    state
+        .broadcasts
+        .send(vault_id, new_version.clone().into())
+        .await;
+
     Ok(Json(new_version.into()))
 }