aria2-rust-pro/crates/aria2-rust-pro-rpc/src/websocket/notification.rs

use std::collections::BTreeMap;

use crate::model::{RpcMeta, RpcValue};

/// Maximum queued outbound WebSocket frames retained per connected session.
pub(super) const MAX_PENDING_WEBSOCKET_FRAMES_PER_SESSION: usize = 256;
/// Maximum bridged runtime-event broadcast frames retained before session fan-out.
pub(super) const MAX_PENDING_WEBSOCKET_BROADCAST_FRAMES: usize = 1024;

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
/// Notification kinds that can be bridged onto WebSocket sessions.
pub enum RpcNotificationKind {
    /// A download transitioned into the active state.
    DownloadStarted,
    /// A download was paused.
    DownloadPaused,
    /// A download was stopped or removed.
    DownloadStopped,
    /// A download completed successfully.
    DownloadComplete,
    /// A download ended in error.
    DownloadError,
    /// A download was removed.
    DownloadRemoved,
    /// A download entered the waiting queue.
    DownloadWaiting,
    /// A waiting download became active again.
    DownloadActive,
    /// A `BitTorrent` download completed.
    DownloadBtDownloadComplete,
    /// Synthetic notification for version polling.
    SystemVersion,
    /// Synthetic notification for method-list polling.
    SystemListMethods,
    /// Synthetic notification for notification-list polling.
    SystemListNotifications,
}

#[derive(Debug, Clone, PartialEq)]
/// Notification payload bridged to WebSocket clients.
pub struct RpcNotificationEvent {
    /// Logical notification kind.
    pub kind: RpcNotificationKind,
    /// Explicit method name override, if present.
    pub method: String,
    /// Download gid associated with the notification, if any.
    pub gid: Option<String>,
    /// Additional notification payload.
    pub payload: Option<RpcValue>,
    /// Supplemental metadata.
    pub meta: RpcMeta,
}

#[derive(Debug, Clone, PartialEq, Eq)]
/// Minimal WebSocket frame shapes used by the RPC server.
pub enum RpcWebSocketFrame {
    /// UTF-8 text frame.
    Text(String),
    /// Binary frame carrying JSON-RPC bytes.
    Binary(Vec<u8>),
    /// Ping control frame.
    Ping(Vec<u8>),
    /// Pong control frame.
    Pong(Vec<u8>),
    /// Close control frame.
    Close,
}

impl RpcNotificationKind {
    /// Returns the canonical aria2-compatible method name for this notification.
    #[must_use]
    pub const fn method_name(self) -> &'static str {
        match self {
            Self::DownloadStarted | Self::DownloadActive => "aria2.onDownloadStart",
            Self::DownloadPaused | Self::DownloadWaiting => "aria2.onDownloadPause",
            Self::DownloadStopped | Self::DownloadRemoved => "aria2.onDownloadStop",
            Self::DownloadComplete => "aria2.onDownloadComplete",
            Self::DownloadError => "aria2.onDownloadError",
            Self::DownloadBtDownloadComplete => "aria2.onBtDownloadComplete",
            Self::SystemVersion => "aria2.getVersion",
            Self::SystemListMethods => "system.listMethods",
            Self::SystemListNotifications => "system.listNotifications",
        }
    }
}

impl RpcNotificationEvent {
    #[must_use]
    /// Returns the effective WebSocket method name for the event.
    pub fn websocket_method_name(&self) -> &str {
        if self.method.is_empty() {
            self.kind.method_name()
        } else {
            &self.method
        }
    }

    #[must_use]
    /// Converts the event into a text WebSocket frame.
    pub fn to_websocket_frame(&self) -> RpcWebSocketFrame {
        RpcWebSocketFrame::Text(self.to_websocket_json())
    }

    #[must_use]
    /// Renders the event into a JSON-RPC notification string.
    pub fn to_websocket_json(&self) -> String {
        let method = self.websocket_method_name();
        let params = websocket_notification_params(self.gid.as_deref(), self.payload.as_ref());
        let params = rpc_value_to_json(&RpcValue::Array(params));
        format!(
            "{{\"jsonrpc\":\"2.0\",\"method\":\"{}\",\"params\":{}}}",
            escape_json(method),
            params
        )
    }
}

/// Builds the JSON-RPC `params` array for a bridged WebSocket notification event.
pub(super) fn websocket_notification_params(
    gid: Option<&str>,
    payload: Option<&RpcValue>,
) -> Vec<RpcValue> {
    if let Some(RpcValue::Array(items)) = payload {
        return items.clone();
    }

    let mut event_spec = BTreeMap::new();
    if let Some(gid) = gid {
        event_spec.insert("gid".to_owned(), RpcValue::String(gid.to_owned()));
    }

    if let Some(RpcValue::Object(map)) = payload {
        for (key, value) in map {
            event_spec.insert(key.clone(), value.clone());
        }
    }

    vec![RpcValue::Object(event_spec)]
}

/// Renders a transport-neutral RPC value into compact JSON text for WebSocket frames.
pub(super) fn rpc_value_to_json(value: &RpcValue) -> String {
    match value {
        RpcValue::Null => "null".to_owned(),
        RpcValue::Bool(value) => value.to_string(),
        RpcValue::Number(value) => value.to_string(),
        RpcValue::String(value) => format!("\"{}\"", escape_json(value)),
        RpcValue::Array(values) => {
            let items = values
                .iter()
                .map(rpc_value_to_json)
                .collect::<Vec<_>>()
                .join(",");
            format!("[{items}]")
        }
        RpcValue::Object(map) => {
            let members = map
                .iter()
                .map(|(key, value)| {
                    format!("\"{}\":{}", escape_json(key), rpc_value_to_json(value))
                })
                .collect::<Vec<_>>()
                .join(",");
            format!("{{{members}}}")
        }
    }
}

/// Escapes a string for safe embedding in generated JSON text.
pub(super) fn escape_json(value: &str) -> String {
    let mut escaped = String::with_capacity(value.len());
    for ch in value.chars() {
        match ch {
            '"' => escaped.push_str("\\\""),
            '\\' => escaped.push_str("\\\\"),
            '\n' => escaped.push_str("\\n"),
            '\r' => escaped.push_str("\\r"),
            '\t' => escaped.push_str("\\t"),
            '\u{08}' => escaped.push_str("\\b"),
            '\u{0C}' => escaped.push_str("\\f"),
            ch if ch.is_control() => {
                use std::fmt::Write as _;
                let _ = write!(escaped, "\\u{:04x}", u32::from(ch));
            }
            ch => escaped.push(ch),
        }
    }
    escaped
}