vaultwarden/src/api/notifications.rs

use std::{
    net::IpAddr,
    sync::{Arc, LazyLock},
    time::{Duration, Instant},
};

use chrono::{NaiveDateTime, Utc};
use rmpv::Value;
use rocket::{Route, futures::StreamExt};
use rocket_ws::{Message, WebSocket};
use tokio::sync::mpsc::Sender;

use crate::{
    CONFIG, Error,
    auth::{ClientIp, WsAccessTokenHeader},
    db::{
        DbConn,
        models::{AuthRequestId, Cipher, CollectionId, Device, DeviceId, Folder, PushId, Send as DbSend, User, UserId},
    },
};

use super::{
    push::push_auth_request, push::push_auth_response, push_cipher_update, push_folder_update, push_logout,
    push_send_update, push_user_update,
};

pub static WS_USERS: LazyLock<Arc<WebSocketUsers>> = LazyLock::new(|| {
    Arc::new(WebSocketUsers {
        map: Arc::new(dashmap::DashMap::new()),
    })
});

pub static WS_ANONYMOUS_SUBSCRIPTIONS: LazyLock<Arc<AnonymousWebSocketSubscriptions>> = LazyLock::new(|| {
    Arc::new(AnonymousWebSocketSubscriptions {
        map: Arc::new(dashmap::DashMap::new()),
        connections: Arc::new(dashmap::DashMap::new()),
    })
});

/// The anonymous hub needs no authentication, so bound how much a single client can hold open.
/// One connection is needed per pending login request, several at once are only expected behind NAT.
const MAX_ANONYMOUS_CONNECTIONS_PER_IP: u32 = 25;

/// How often a Ping is sent to the client.
const WS_PING_INTERVAL: Duration = Duration::from_secs(15);

/// Close the connection if nothing, not even a Pong, was received from the client for this long.
/// Otherwise half-open connections (client gone without a FIN, e.g. behind a proxy or NAT) are kept forever.
/// Same as the default `ClientTimeoutInterval` of ASP.NET Core SignalR, which the official server uses.
const WS_CLIENT_TIMEOUT: Duration = Duration::from_secs(30);

static NOTIFICATIONS_DISABLED: LazyLock<bool> = LazyLock::new(|| !CONFIG.enable_websocket() && !CONFIG.push_enabled());

pub fn routes() -> Vec<Route> {
    if CONFIG.enable_websocket() {
        routes![websockets_hub, anonymous_websockets_hub]
    } else {
        info!("WebSocket are disabled, realtime sync functionality will not work!");
        routes![]
    }
}

#[derive(FromForm, Debug)]
struct WsAccessToken {
    access_token: Option<String>,
}

struct WSEntryMapGuard {
    users: Arc<WebSocketUsers>,
    user_uuid: UserId,
    entry_uuid: uuid::Uuid,
    addr: IpAddr,
}

impl WSEntryMapGuard {
    fn new(users: Arc<WebSocketUsers>, user_uuid: UserId, entry_uuid: uuid::Uuid, addr: IpAddr) -> Self {
        Self {
            users,
            user_uuid,
            entry_uuid,
            addr,
        }
    }
}

impl Drop for WSEntryMapGuard {
    fn drop(&mut self) {
        info!("Closing WS connection from {}", self.addr);
        if let Some(mut entry) = self.users.map.get_mut(self.user_uuid.as_ref()) {
            entry.retain(|(uuid, _)| uuid != &self.entry_uuid);
        }
    }
}

struct WSAnonymousEntryMapGuard {
    subscriptions: Arc<AnonymousWebSocketSubscriptions>,
    token: String,
    entry_uuid: uuid::Uuid,
    addr: IpAddr,
}

impl WSAnonymousEntryMapGuard {
    fn new(
        subscriptions: Arc<AnonymousWebSocketSubscriptions>,
        token: String,
        entry_uuid: uuid::Uuid,
        addr: IpAddr,
    ) -> Self {
        Self {
            subscriptions,
            token,
            entry_uuid,
            addr,
        }
    }
}

impl Drop for WSAnonymousEntryMapGuard {
    fn drop(&mut self) {
        info!("Closing WS connection from {}", self.addr);
        if let Some(mut entry) = self.subscriptions.map.get_mut(&self.token) {
            entry.retain(|(uuid, _)| uuid != &self.entry_uuid);
        }
        self.subscriptions.map.remove_if(&self.token, |_, senders| senders.is_empty());
        self.subscriptions.release(self.addr);
    }
}

#[expect(tail_expr_drop_order)]
#[get("/hub?<data..>")]
fn websockets_hub<'r>(
    ws: WebSocket,
    data: WsAccessToken,
    ip: ClientIp,
    header_token: WsAccessTokenHeader,
) -> Result<rocket_ws::Stream!['r], Error> {
    info!("Accepting Rocket WS connection from {}", ip.ip);

    let token = if let Some(token) = data.access_token {
        token
    } else if let Some(token) = header_token.access_token {
        token
    } else {
        err_code!("Invalid claim", 401)
    };

    let Ok(claims) = crate::auth::decode_login(&token) else {
        err_code!("Invalid token", 401)
    };

    let (mut rx, guard) = {
        let users = Arc::clone(&WS_USERS);

        // Add a channel to send messages to this client to the map
        let entry_uuid = uuid::Uuid::new_v4();
        let (tx, rx) = tokio::sync::mpsc::channel::<Message>(100);
        users.map.entry(claims.sub.to_string()).or_default().push((entry_uuid, tx));

        // Once the guard goes out of scope, the connection will have been closed and the entry will be deleted from the map
        (rx, WSEntryMapGuard::new(users, claims.sub, entry_uuid, ip.ip))
    };

    Ok({
        rocket_ws::Stream! { ws => {
            let mut ws = ws;
            let _guard = guard;
            let mut interval = tokio::time::interval(WS_PING_INTERVAL);
            let mut last_received = Instant::now();
            loop {
                tokio::select! {
                    res = ws.next() =>  {
                        match res {
                            Some(Ok(message)) => {
                                // Any message, including a Pong, means the client is still there
                                last_received = Instant::now();

                                match message {
                                    // Respond to any pings
                                    Message::Ping(ping) => yield Message::Pong(ping),
                                    Message::Pong(_) => {/* Ignored */},

                                    // We should receive an initial message with the protocol and version, and we will reply to it
                                    Message::Text(ref message) => {
                                        let msg = message.strip_suffix(RECORD_SEPARATOR as char).unwrap_or(message);

                                        if serde_json::from_str(msg).ok() == Some(INITIAL_MESSAGE) {
                                            yield Message::binary(INITIAL_RESPONSE);
                                        }
                                    }

                                    // Prevent sending anything back when a `Close` Message is received.
                                    // Just break the loop
                                    Message::Close(_) => break,

                                    // Just echo anything else the client sends
                                    _ => yield message,
                                }
                            }
                            _ => break,
                        }
                    }

                    res = rx.recv() => {
                        match res {
                            Some(res) => yield res,
                            None => break,
                        }
                    }

                    _ = interval.tick() => {
                        // The client stopped responding without closing the connection, drop it
                        if last_received.elapsed() > WS_CLIENT_TIMEOUT {
                            break;
                        }
                        yield Message::Ping(create_ping());
                    }
                }
            }
        }}
    })
}

#[expect(tail_expr_drop_order)]
#[get("/anonymous-hub?<token..>")]
fn anonymous_websockets_hub<'r>(ws: WebSocket, token: String, ip: ClientIp) -> Result<rocket_ws::Stream!['r], Error> {
    info!("Accepting Anonymous Rocket WS connection from {}", ip.ip);

    let (mut rx, guard) = {
        let subscriptions = Arc::clone(&WS_ANONYMOUS_SUBSCRIPTIONS);

        if !subscriptions.try_reserve(ip.ip) {
            err_code!("Too many connections", 429)
        }

        // Add a channel to send messages to this client to the map.
        // Clients reconnect with the same token while a login request is still pending, so keep
        // every subscriber instead of replacing, otherwise the older one takes the newer one down.
        let (tx, rx) = tokio::sync::mpsc::channel::<Message>(100);
        let entry_uuid = uuid::Uuid::new_v4();
        subscriptions.map.entry(token.clone()).or_default().push((entry_uuid, tx));

        // Once the guard goes out of scope, the connection will have been closed and the entry will be deleted from the map
        (rx, WSAnonymousEntryMapGuard::new(subscriptions, token, entry_uuid, ip.ip))
    };

    Ok({
        rocket_ws::Stream! { ws => {
            let mut ws = ws;
            let _guard = guard;
            let mut interval = tokio::time::interval(WS_PING_INTERVAL);
            let mut last_received = Instant::now();
            loop {
                tokio::select! {
                    res = ws.next() =>  {
                        match res {
                            Some(Ok(message)) => {
                                // Any message, including a Pong, means the client is still there
                                last_received = Instant::now();

                                match message {
                                    // Respond to any pings
                                    Message::Ping(ping) => yield Message::Pong(ping),
                                    Message::Pong(_) => {/* Ignored */},

                                    // We should receive an initial message with the protocol and version, and we will reply to it
                                    Message::Text(ref message) => {
                                        let msg = message.strip_suffix(RECORD_SEPARATOR as char).unwrap_or(message);

                                        if serde_json::from_str(msg).ok() == Some(INITIAL_MESSAGE) {
                                            yield Message::binary(INITIAL_RESPONSE);
                                        }
                                    }

                                    // Prevent sending anything back when a `Close` Message is received.
                                    // Just break the loop
                                    Message::Close(_) => break,

                                    // Just echo anything else the client sends
                                    _ => yield message,
                                }
                            }
                            _ => break,
                        }
                    }

                    res = rx.recv() => {
                        match res {
                            Some(res) => yield res,
                            None => break,
                        }
                    }

                    _ = interval.tick() => {
                        // The client stopped responding without closing the connection, drop it
                        if last_received.elapsed() > WS_CLIENT_TIMEOUT {
                            break;
                        }
                        yield Message::Ping(create_ping());
                    }
                }
            }
        }}
    })
}

//
// Websockets server
//

fn serialize(val: &Value) -> Vec<u8> {
    use rmpv::encode::write_value;

    let mut buf = Vec::new();
    write_value(&mut buf, val).expect("Error encoding MsgPack");

    // Add size bytes at the start
    // Extracted from BinaryMessageFormat.js
    let mut size: usize = buf.len();
    let mut len_buf: Vec<u8> = Vec::new();

    loop {
        #[expect(clippy::cast_possible_truncation, reason = "masked to 7 bits, fits u8")]
        let mut size_part = (size & 0x7f) as u8;
        size >>= 7;

        if size > 0 {
            size_part |= 0x80;
        }

        len_buf.push(size_part);

        if size == 0 {
            break;
        }
    }

    len_buf.append(&mut buf);
    len_buf
}

fn serialize_date(date: NaiveDateTime) -> Value {
    let seconds: i64 = date.and_utc().timestamp();
    let nanos: i64 = date.and_utc().timestamp_subsec_nanos().into();
    let timestamp = (nanos << 34) | seconds;

    let bs = timestamp.to_be_bytes();

    // -1 is Timestamp
    // https://github.com/msgpack/msgpack/blob/master/spec.md#timestamp-extension-type
    Value::Ext(-1, bs.to_vec())
}

fn convert_option<T: Into<Value>>(option: Option<T>) -> Value {
    match option {
        Some(a) => a.into(),
        None => Value::Nil,
    }
}

const RECORD_SEPARATOR: u8 = 0x1e;
const INITIAL_RESPONSE: [u8; 3] = [0x7b, 0x7d, RECORD_SEPARATOR]; // {, }, <RS>

#[derive(Deserialize, Copy, Clone, Eq, PartialEq)]
struct InitialMessage<'a> {
    protocol: &'a str,
    version: i32,
}

static INITIAL_MESSAGE: InitialMessage<'static> = InitialMessage {
    protocol: "messagepack",
    version: 1,
};

// We attach the UUID to the sender so we can differentiate them when we need to remove them from the Vec
type UserSenders = (uuid::Uuid, Sender<Message>);
#[derive(Clone)]
pub struct WebSocketUsers {
    map: Arc<dashmap::DashMap<String, Vec<UserSenders>>>,
}

impl WebSocketUsers {
    async fn send_update(&self, user_id: &UserId, data: &[u8]) {
        if let Some(user) = self.map.get(user_id.as_ref()).map(|v| v.clone()) {
            for (_, sender) in &user {
                if let Err(e) = sender.send(Message::binary(data)).await {
                    error!("Error sending WS update {e}");
                }
            }
        }
    }

    // NOTE: The last modified date needs to be updated before calling these methods
    pub async fn send_user_update(&self, ut: UpdateType, user: &User, push_uuid: Option<&PushId>, conn: &DbConn) {
        // Skip any processing if both WebSockets and Push are not active
        if *NOTIFICATIONS_DISABLED {
            return;
        }
        let data = create_update(
            vec![("UserId".into(), user.uuid.to_string().into()), ("Date".into(), serialize_date(user.updated_at))],
            ut,
            None,
        );

        if CONFIG.enable_websocket() {
            self.send_update(&user.uuid, &data).await;
        }

        if CONFIG.push_enabled() {
            push_user_update(ut, user, push_uuid, conn).await;
        }
    }

    pub async fn send_logout(&self, user: &User, acting_device: Option<&Device>, conn: &DbConn) {
        // Skip any processing if both WebSockets and Push are not active
        if *NOTIFICATIONS_DISABLED {
            return;
        }
        let acting_device_id = acting_device.map(|d| d.uuid.clone());
        let data = create_update(
            vec![("UserId".into(), user.uuid.to_string().into()), ("Date".into(), serialize_date(user.updated_at))],
            UpdateType::LogOut,
            acting_device_id,
        );

        if CONFIG.enable_websocket() {
            self.send_update(&user.uuid, &data).await;
        }

        if CONFIG.push_enabled() {
            push_logout(user, acting_device, conn).await;
        }
    }

    pub async fn send_folder_update(&self, ut: UpdateType, folder: &Folder, device: &Device, conn: &DbConn) {
        // Skip any processing if both WebSockets and Push are not active
        if *NOTIFICATIONS_DISABLED {
            return;
        }
        let data = create_update(
            vec![
                ("Id".into(), folder.uuid.to_string().into()),
                ("UserId".into(), folder.user_uuid.to_string().into()),
                ("RevisionDate".into(), serialize_date(folder.updated_at)),
            ],
            ut,
            Some(device.uuid.clone()),
        );

        if CONFIG.enable_websocket() {
            self.send_update(&folder.user_uuid, &data).await;
        }

        if CONFIG.push_enabled() {
            push_folder_update(ut, folder, device, conn).await;
        }
    }

    pub async fn send_cipher_update(
        &self,
        ut: UpdateType,
        cipher: &Cipher,
        user_ids: &[UserId],
        device: &Device,
        collection_uuids: Option<Vec<CollectionId>>,
        conn: &DbConn,
    ) {
        // Skip any processing if both WebSockets and Push are not active
        if *NOTIFICATIONS_DISABLED {
            return;
        }
        let org_id = convert_option(cipher.organization_uuid.as_deref());
        // Depending if there are collections provided or not, we need to have different values for the following variables.
        // The user_uuid should be `null`, and the revision date should be set to now, else the clients won't sync the collection change.
        let (user_id, collection_uuids, revision_date) = if let Some(collection_uuids) = collection_uuids {
            (
                Value::Nil,
                Value::Array(collection_uuids.into_iter().map(|v| v.to_string().into()).collect::<Vec<Value>>()),
                serialize_date(Utc::now().naive_utc()),
            )
        } else {
            (convert_option(cipher.user_uuid.as_deref()), Value::Nil, serialize_date(cipher.updated_at))
        };

        let data = create_update(
            vec![
                ("Id".into(), cipher.uuid.to_string().into()),
                ("UserId".into(), user_id),
                ("OrganizationId".into(), org_id),
                ("CollectionIds".into(), collection_uuids),
                ("RevisionDate".into(), revision_date),
            ],
            ut,
            Some(device.uuid.clone()), // Acting device id (unique device/app uuid)
        );

        if CONFIG.enable_websocket() {
            for uuid in user_ids {
                self.send_update(uuid, &data).await;
            }
        }

        if CONFIG.push_enabled() && user_ids.len() == 1 {
            push_cipher_update(ut, cipher, device, conn).await;
        }
    }

    pub async fn send_send_update(
        &self,
        ut: UpdateType,
        send: &DbSend,
        user_ids: &[UserId],
        device: &Device,
        conn: &DbConn,
    ) {
        // Skip any processing if both WebSockets and Push are not active
        if *NOTIFICATIONS_DISABLED {
            return;
        }
        let user_id = convert_option(send.user_uuid.as_deref());

        let data = create_update(
            vec![
                ("Id".into(), send.uuid.to_string().into()),
                ("UserId".into(), user_id),
                ("RevisionDate".into(), serialize_date(send.revision_date)),
            ],
            ut,
            None,
        );

        if CONFIG.enable_websocket() {
            for uuid in user_ids {
                self.send_update(uuid, &data).await;
            }
        }
        if CONFIG.push_enabled() && user_ids.len() == 1 {
            push_send_update(ut, send, device, conn).await;
        }
    }

    pub async fn send_auth_request(&self, user_id: &UserId, auth_request_uuid: &str, device: &Device, conn: &DbConn) {
        // Skip any processing if both WebSockets and Push are not active
        if *NOTIFICATIONS_DISABLED {
            return;
        }
        let data = create_update(
            vec![("Id".into(), auth_request_uuid.to_owned().into()), ("UserId".into(), user_id.to_string().into())],
            UpdateType::AuthRequest,
            Some(device.uuid.clone()),
        );
        if CONFIG.enable_websocket() {
            self.send_update(user_id, &data).await;
        }

        if CONFIG.push_enabled() {
            push_auth_request(user_id, auth_request_uuid, device, conn).await;
        }
    }

    pub async fn send_auth_response(
        &self,
        user_id: &UserId,
        auth_request_id: &AuthRequestId,
        device: &Device,
        conn: &DbConn,
    ) {
        // Skip any processing if both WebSockets and Push are not active
        if *NOTIFICATIONS_DISABLED {
            return;
        }
        let data = create_update(
            vec![("Id".into(), auth_request_id.to_string().into()), ("UserId".into(), user_id.to_string().into())],
            UpdateType::AuthRequestResponse,
            Some(device.uuid.clone()),
        );
        if CONFIG.enable_websocket() {
            self.send_update(user_id, &data).await;
        }

        if CONFIG.push_enabled() {
            push_auth_response(user_id, auth_request_id, device, conn).await;
        }
    }
}

#[derive(Clone)]
pub struct AnonymousWebSocketSubscriptions {
    map: Arc<dashmap::DashMap<String, Vec<UserSenders>>>,
    connections: Arc<dashmap::DashMap<IpAddr, u32>>,
}

impl AnonymousWebSocketSubscriptions {
    /// Takes a connection slot for this address, returns false when it already reached the limit.
    fn try_reserve(&self, addr: IpAddr) -> bool {
        let mut count = self.connections.entry(addr).or_insert(0);
        if *count >= MAX_ANONYMOUS_CONNECTIONS_PER_IP {
            return false;
        }
        *count += 1;
        true
    }

    /// Releases a slot taken by `try_reserve`.
    fn release(&self, addr: IpAddr) {
        let empty = if let Some(mut count) = self.connections.get_mut(&addr) {
            *count = count.saturating_sub(1);
            *count == 0
        } else {
            false
        };
        // Only remove once the guard above is dropped, otherwise this deadlocks.
        if empty {
            self.connections.remove_if(&addr, |_, count| *count == 0);
        }
    }

    async fn send_update(&self, token: &str, data: &[u8]) {
        // Clone the senders so the map isn't kept locked while sending.
        let senders = self.map.get(token).map(|v| v.clone()).unwrap_or_default();
        for (_, sender) in senders {
            if let Err(e) = sender.send(Message::binary(data)).await {
                error!("Error sending WS update {e}");
            }
        }
    }

    pub async fn send_auth_response(&self, user_id: &UserId, auth_request_id: &AuthRequestId) {
        if !CONFIG.enable_websocket() {
            return;
        }
        let data = create_anonymous_update(
            vec![("Id".into(), auth_request_id.to_string().into()), ("UserId".into(), user_id.to_string().into())],
            UpdateType::AuthRequestResponse,
            user_id,
        );
        self.send_update(auth_request_id, &data).await;
    }
}

/* Message Structure
[
    1, // MessageType.Invocation
    {}, // Headers (map)
    null, // InvocationId
    "ReceiveMessage", // Target
    [ // Arguments
        {
            "ContextId": acting_device_id || Nil,
            "Type": ut as i32,
            "Payload": {}
        }
    ]
]
*/
fn create_update(payload: Vec<(Value, Value)>, ut: UpdateType, acting_device_id: Option<DeviceId>) -> Vec<u8> {
    use rmpv::Value as V;

    let value = V::Array(vec![
        1.into(),
        V::Map(vec![]),
        V::Nil,
        "ReceiveMessage".into(),
        V::Array(vec![V::Map(vec![
            ("ContextId".into(), acting_device_id.map_or(V::Nil, |v| v.to_string().into())),
            ("Type".into(), (ut as i32).into()),
            ("Payload".into(), payload.into()),
        ])]),
    ]);

    serialize(&value)
}

fn create_anonymous_update(payload: Vec<(Value, Value)>, ut: UpdateType, user_id: &UserId) -> Vec<u8> {
    use rmpv::Value as V;

    let value = V::Array(vec![
        1.into(),
        V::Map(vec![]),
        V::Nil,
        // This word is misspelled, but upstream has this too
        // https://github.com/bitwarden/server/blob/dff9f1cf538198819911cf2c20f8cda3307701c5/src/Notifications/HubHelpers.cs#L86
        // https://github.com/bitwarden/clients/blob/9612a4ac45063e372a6fbe87eb253c7cb3c588fb/libs/common/src/auth/services/anonymous-hub.service.ts#L45
        "AuthRequestResponseRecieved".into(),
        V::Array(vec![V::Map(vec![
            ("Type".into(), (ut as i32).into()),
            ("Payload".into(), payload.into()),
            ("UserId".into(), user_id.to_string().into()),
        ])]),
    ]);

    serialize(&value)
}

fn create_ping() -> Vec<u8> {
    serialize(&Value::Array(vec![6.into()]))
}

// https://github.com/bitwarden/server/blob/375af7c43b10d9da03525d41452f95de3f921541/src/Core/Enums/PushType.cs
#[derive(Copy, Clone, Eq, PartialEq)]
pub enum UpdateType {
    SyncCipherUpdate = 0,
    SyncCipherCreate = 1,
    SyncLoginDelete = 2,
    SyncFolderDelete = 3,
    SyncCiphers = 4,

    SyncVault = 5,
    SyncOrgKeys = 6,
    SyncFolderCreate = 7,
    SyncFolderUpdate = 8,
    // SyncCipherDelete = 9, // Redirects to `SyncLoginDelete` on upstream
    SyncSettings = 10,

    LogOut = 11,

    SyncSendCreate = 12,
    SyncSendUpdate = 13,
    SyncSendDelete = 14,

    AuthRequest = 15,
    AuthRequestResponse = 16,

    // SyncOrganizations = 17, // Not supported
    // SyncOrganizationStatusChanged = 18, // Not supported
    // SyncOrganizationCollectionSettingChanged = 19, // Not supported

    // Notification = 20, // Not supported
    // NotificationStatus = 21, // Not supported

    // RefreshSecurityTasks = 22, // Not supported
    None = 100,
}

pub type Notify<'a> = &'a rocket::State<Arc<WebSocketUsers>>;
pub type AnonymousNotify<'a> = &'a rocket::State<Arc<AnonymousWebSocketSubscriptions>>;