rustpad/src/state.rs

use crate::database::Database;
use serde::Serialize;
use std::{
    collections::HashMap,
    sync::{
        Arc,
        atomic::{AtomicU64, Ordering},
    },
};
use tokio::sync::{RwLock, broadcast};

const CHANNEL_CAPACITY: usize = 256;

#[derive(Debug, Clone)]
pub struct SmtpConfig {
    pub host: String,
    pub port: u16,
    pub username: String,
    pub password: String,
    pub from: String,
    pub public_url: String,
}

#[derive(Debug, Clone)]
pub struct NoteUpdate {
    pub content: String,
    pub revision_id: i64,
    pub updated_at: String,
    pub author: Option<String>,
    pub owner_map: String,
}

#[derive(Debug, Clone, Serialize)]
pub struct PresenceUser {
    pub name: String,
    pub color: Option<String>,
}

#[derive(Debug, Clone)]
struct PresenceConnection {
    identity: String,
    user: PresenceUser,
}

#[derive(Debug, Clone)]
pub enum RoomEvent {
    Document(NoteUpdate),
    Presence(Vec<PresenceUser>),
    Chat { sender: String, text: String },
}

#[derive(Debug)]
pub struct AppState {
    pub db: Database,
    pub asset_version: String,
    pub storage: crate::storage::Storage,
    pub upload_max_size_bytes: usize,
    pub file_cache_max_age_seconds: u64,
    pub smtp: Option<SmtpConfig>,
    pub registration_enabled: bool,
    pub account_confirmation_required: bool,
    pub share_confirmation_required: bool,
    pub frontend_log_level: String,
    pub anonymous_access_token_ttl_days: i64,
    pub user_session_ttl_days: i64,
    channels: RwLock<HashMap<String, broadcast::Sender<RoomEvent>>>,
    presence: RwLock<HashMap<String, HashMap<u64, PresenceConnection>>>,
    next_connection_id: AtomicU64,
}

impl AppState {
    pub fn new(
        db: Database,
        asset_version: String,
        storage: crate::storage::Storage,
        upload_max_size_bytes: usize,
        file_cache_max_age_seconds: u64,
        smtp: Option<SmtpConfig>,
        registration_enabled: bool,
        account_confirmation_required: bool,
        share_confirmation_required: bool,
        frontend_log_level: String,
        anonymous_access_token_ttl_days: i64,
        user_session_ttl_days: i64,
    ) -> Self {
        Self {
            db,
            asset_version,
            storage,
            upload_max_size_bytes,
            file_cache_max_age_seconds,
            smtp,
            registration_enabled,
            account_confirmation_required,
            share_confirmation_required,
            frontend_log_level,
            anonymous_access_token_ttl_days,
            user_session_ttl_days,
            channels: RwLock::new(HashMap::new()),
            presence: RwLock::new(HashMap::new()),
            next_connection_id: AtomicU64::new(1),
        }
    }
    async fn channel_for_key(&self, key: String) -> broadcast::Sender<RoomEvent> {
        if let Some(sender) = self.channels.read().await.get(&key) {
            return sender.clone();
        }
        let mut channels = self.channels.write().await;
        channels
            .entry(key)
            .or_insert_with(|| broadcast::channel(CHANNEL_CAPACITY).0)
            .clone()
    }
    pub fn note_room_key(workspace_slug: &str, note_slug: &str) -> String {
        format!("workspace:{workspace_slug}/{note_slug}")
    }
    pub fn pad_room_key(slug: &str) -> String {
        format!("pad:{slug}")
    }
    pub async fn note_channel(
        &self,
        workspace_slug: &str,
        note_slug: &str,
    ) -> broadcast::Sender<RoomEvent> {
        self.channel_for_key(Self::note_room_key(workspace_slug, note_slug))
            .await
    }
    pub async fn pad_channel(&self, slug: &str) -> broadcast::Sender<RoomEvent> {
        self.channel_for_key(Self::pad_room_key(slug)).await
    }
    pub async fn join_room(
        &self,
        key: &str,
        nickname: String,
        color: Option<String>,
        identity: Option<String>,
    ) -> (u64, Vec<PresenceUser>) {
        let id = self.next_connection_id.fetch_add(1, Ordering::Relaxed);
        let identity = identity.unwrap_or_else(|| format!("connection:{id}"));
        let mut presence = self.presence.write().await;
        let room = presence.entry(key.to_owned()).or_default();
        room.insert(
            id,
            PresenceConnection {
                identity,
                user: PresenceUser {
                    name: nickname,
                    color,
                },
            },
        );
        (id, sorted_users(room))
    }
    pub async fn update_room_color(
        &self,
        key: &str,
        id: u64,
        color: Option<String>,
    ) -> Vec<PresenceUser> {
        let mut presence = self.presence.write().await;
        if let Some(room) = presence.get_mut(key) {
            if let Some(identity) = room.get(&id).map(|connection| connection.identity.clone()) {
                for connection in room.values_mut() {
                    if connection.identity == identity {
                        connection.user.color = color.clone();
                    }
                }
            }
            sorted_users(room)
        } else {
            Vec::new()
        }
    }
    pub async fn leave_room(&self, key: &str, id: u64) -> Vec<PresenceUser> {
        let mut presence = self.presence.write().await;
        if let Some(room) = presence.get_mut(key) {
            room.remove(&id);
            let users = sorted_users(room);
            let empty = room.is_empty();
            if empty {
                presence.remove(key);
            }
            users
        } else {
            Vec::new()
        }
    }
}

fn sorted_users(room: &HashMap<u64, PresenceConnection>) -> Vec<PresenceUser> {
    let mut by_identity: HashMap<&str, PresenceUser> = HashMap::new();
    for connection in room.values() {
        by_identity
            .entry(&connection.identity)
            .or_insert_with(|| connection.user.clone());
    }
    let mut users: Vec<PresenceUser> = by_identity.into_values().collect();
    users.sort_by_key(|value| value.name.to_lowercase());
    users
}

pub type SharedState = Arc<AppState>;