khadhroony-bobobot/kb_lib/src/rpc_ws.rs

// file: kb_lib/src/rpc_ws.rs

//! Generic JSON-RPC 2.0 WebSocket helpers.
//!
//! This module provides generic JSON-RPC request and incoming-message parsing
//! helpers for WebSocket-based Solana RPC communication.
//!
//! At this stage, the top-level envelopes are typed while the method-specific
//! payloads remain as `serde_json::Value`. Later versions can progressively
//! replace selected payloads with official Solana RPC client types.

/// Generic JSON-RPC 2.0 request sent over WebSocket.
#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct KbJsonRpcWsRequest {
    /// JSON-RPC version, expected to be `"2.0"`.
    pub jsonrpc: std::string::String,
    /// Client request identifier.
    pub id: serde_json::Value,
    /// RPC method name.
    pub method: std::string::String,
    /// Ordered method parameters.
    pub params: std::vec::Vec<serde_json::Value>,
}

impl KbJsonRpcWsRequest {
    /// Creates a new JSON-RPC request with a numeric identifier.
    pub fn new_with_u64_id(
        id: u64,
        method: std::string::String,
        params: std::vec::Vec<serde_json::Value>,
    ) -> Self {
        Self {
            jsonrpc: "2.0".to_string(),
            id: serde_json::Value::from(id),
            method,
            params,
        }
    }

    /// Converts the request into a JSON value.
    pub fn to_value(&self) -> Result<serde_json::Value, crate::KbError> {
        let value_result = serde_json::to_value(self);
        match value_result {
            Ok(value) => Ok(value),
            Err(error) => Err(crate::KbError::Json(format!(
                "cannot serialize websocket json-rpc request '{}': {error}",
                self.method
            ))),
        }
    }

    /// Serializes the request into a compact JSON string.
    pub fn to_json_string(&self) -> Result<std::string::String, crate::KbError> {
        let text_result = serde_json::to_string(self);
        match text_result {
            Ok(text) => Ok(text),
            Err(error) => Err(crate::KbError::Json(format!(
                "cannot serialize websocket json-rpc request '{}': {error}",
                self.method
            ))),
        }
    }
}

/// JSON-RPC 2.0 success response.
#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct KbJsonRpcWsSuccessResponse {
    /// JSON-RPC version, expected to be `"2.0"`.
    pub jsonrpc: std::string::String,
    /// Result payload.
    pub result: serde_json::Value,
    /// Request identifier echoed by the server.
    pub id: serde_json::Value,
}

/// JSON-RPC 2.0 error object.
#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct KbJsonRpcWsErrorObject {
    /// Numeric JSON-RPC error code.
    pub code: i64,
    /// Human-readable error message.
    pub message: std::string::String,
    /// Optional server-provided structured payload.
    pub data: std::option::Option<serde_json::Value>,
}

/// JSON-RPC 2.0 error response.
#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct KbJsonRpcWsErrorResponse {
    /// JSON-RPC version, expected to be `"2.0"`.
    pub jsonrpc: std::string::String,
    /// Error payload.
    pub error: KbJsonRpcWsErrorObject,
    /// Request identifier echoed by the server.
    pub id: serde_json::Value,
}

/// JSON-RPC 2.0 notification parameter object.
#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct KbJsonRpcWsNotificationParams {
    /// Method-specific result payload.
    pub result: serde_json::Value,
    /// Active subscription identifier.
    pub subscription: u64,
}

/// JSON-RPC 2.0 notification message.
#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct KbJsonRpcWsNotification {
    /// JSON-RPC version, expected to be `"2.0"`.
    pub jsonrpc: std::string::String,
    /// Notification method name such as `slotNotification`.
    pub method: std::string::String,
    /// Notification payload.
    pub params: KbJsonRpcWsNotificationParams,
}

/// Parsed incoming JSON-RPC WebSocket message.
#[derive(Clone, Debug, PartialEq)]
pub enum KbJsonRpcWsIncomingMessage {
    /// JSON-RPC success response.
    SuccessResponse(KbJsonRpcWsSuccessResponse),
    /// JSON-RPC error response.
    ErrorResponse(KbJsonRpcWsErrorResponse),
    /// JSON-RPC notification.
    Notification(KbJsonRpcWsNotification),
}

impl KbJsonRpcWsIncomingMessage {
    /// Returns a short human-readable kind label.
    pub fn kind_name(&self) -> &'static str {
        match self {
            Self::SuccessResponse(_) => "success_response",
            Self::ErrorResponse(_) => "error_response",
            Self::Notification(_) => "notification",
        }
    }
}

/// Returns `true` when the text looks like a JSON object payload.
///
/// This is intentionally conservative and only checks for a leading `{` after
/// trimming left-side whitespace.
pub fn kb_is_probable_json_rpc_object_text(text: &str) -> bool {
    let trimmed = text.trim_start();
    trimmed.starts_with('{')
}

/// Parses a raw text message into a JSON-RPC incoming message.
///
/// This parser accepts only server-originating incoming message shapes:
/// success responses, error responses, and notifications.
pub fn parse_kb_json_rpc_ws_incoming_text(
    text: &str,
) -> Result<KbJsonRpcWsIncomingMessage, crate::KbError> {
    let value_result = serde_json::from_str::<serde_json::Value>(text);
    let value = match value_result {
        Ok(value) => value,
        Err(error) => {
            return Err(crate::KbError::Json(format!(
                "cannot parse websocket json-rpc text: {error}"
            )));
        }
    };

    parse_kb_json_rpc_ws_incoming_value(&value)
}

/// Parses a JSON value into a JSON-RPC incoming message.
///
/// This parser accepts only server-originating incoming message shapes:
/// success responses, error responses, and notifications.
pub fn parse_kb_json_rpc_ws_incoming_value(
    value: &serde_json::Value,
) -> Result<KbJsonRpcWsIncomingMessage, crate::KbError> {
    let object = match value.as_object() {
        Some(object) => object,
        None => {
            return Err(crate::KbError::Json(
                "json-rpc websocket payload must be a JSON object".to_string(),
            ));
        }
    };
    let jsonrpc_value_option = object.get("jsonrpc");
    let jsonrpc_value = match jsonrpc_value_option {
        Some(jsonrpc_value) => jsonrpc_value,
        None => {
            return Err(crate::KbError::Json(
                "json-rpc websocket payload is missing 'jsonrpc'".to_string(),
            ));
        }
    };
    let jsonrpc_string_option = jsonrpc_value.as_str();
    let jsonrpc_string = match jsonrpc_string_option {
        Some(jsonrpc_string) => jsonrpc_string,
        None => {
            return Err(crate::KbError::Json(
                "json-rpc websocket field 'jsonrpc' must be a string".to_string(),
            ));
        }
    };
    if jsonrpc_string != "2.0" {
        return Err(crate::KbError::Json(format!(
            "unsupported json-rpc version '{}'",
            jsonrpc_string
        )));
    }
    let has_method = object.contains_key("method");
    let has_params = object.contains_key("params");
    let has_result = object.contains_key("result");
    let has_error = object.contains_key("error");
    let has_id = object.contains_key("id");
    if has_method && has_params && !has_id {
        let notification_result = serde_json::from_value::<KbJsonRpcWsNotification>(value.clone());
        let notification = match notification_result {
            Ok(notification) => notification,
            Err(error) => {
                return Err(crate::KbError::Json(format!(
                    "cannot parse websocket json-rpc notification: {error}"
                )));
            }
        };
        return Ok(KbJsonRpcWsIncomingMessage::Notification(notification));
    }
    if has_id && has_result && !has_error {
        let response_result = serde_json::from_value::<KbJsonRpcWsSuccessResponse>(value.clone());
        let response = match response_result {
            Ok(response) => response,
            Err(error) => {
                return Err(crate::KbError::Json(format!(
                    "cannot parse websocket json-rpc success response: {error}"
                )));
            }
        };
        return Ok(KbJsonRpcWsIncomingMessage::SuccessResponse(response));
    }
    if has_id && has_error && !has_result {
        let response_result = serde_json::from_value::<KbJsonRpcWsErrorResponse>(value.clone());
        let response = match response_result {
            Ok(response) => response,
            Err(error) => {
                return Err(crate::KbError::Json(format!(
                    "cannot parse websocket json-rpc error response: {error}"
                )));
            }
        };
        return Ok(KbJsonRpcWsIncomingMessage::ErrorResponse(response));
    }
    Err(crate::KbError::Json(
        "unsupported websocket json-rpc message shape".to_string(),
    ))
}

#[cfg(test)]
mod tests {
    #[test]
    fn request_serialization_contains_expected_fields() {
        let request = crate::KbJsonRpcWsRequest::new_with_u64_id(
            7,
            "slotSubscribe".to_string(),
            std::vec::Vec::new(),
        );
        let value = request
            .to_value()
            .expect("request value serialization must succeed");
        assert_eq!(
            value["jsonrpc"],
            serde_json::Value::String("2.0".to_string())
        );
        assert_eq!(value["id"], serde_json::Value::from(7u64));
        assert_eq!(
            value["method"],
            serde_json::Value::String("slotSubscribe".to_string())
        );
        assert_eq!(value["params"], serde_json::Value::Array(vec![]));
    }

    #[test]
    fn parse_success_response_works() {
        let text = r#"{"jsonrpc":"2.0","result":42,"id":3}"#;
        let parsed = crate::parse_kb_json_rpc_ws_incoming_text(text).expect("parse must succeed");
        match parsed {
            crate::KbJsonRpcWsIncomingMessage::SuccessResponse(response) => {
                assert_eq!(response.jsonrpc, "2.0");
                assert_eq!(response.result, serde_json::Value::from(42u64));
                assert_eq!(response.id, serde_json::Value::from(3u64));
            }
            other => {
                panic!("unexpected parsed message: {other:?}");
            }
        }
    }

    #[test]
    fn parse_error_response_works() {
        let text =
            r#"{"jsonrpc":"2.0","error":{"code":-32601,"message":"Method not found"},"id":9}"#;
        let parsed = crate::parse_kb_json_rpc_ws_incoming_text(text).expect("parse must succeed");
        match parsed {
            crate::KbJsonRpcWsIncomingMessage::ErrorResponse(response) => {
                assert_eq!(response.jsonrpc, "2.0");
                assert_eq!(response.error.code, -32601);
                assert_eq!(response.error.message, "Method not found");
                assert_eq!(response.error.data, None);
                assert_eq!(response.id, serde_json::Value::from(9u64));
            }
            other => {
                panic!("unexpected parsed message: {other:?}");
            }
        }
    }

    #[test]
    fn parse_notification_works() {
        let text = r#"{"jsonrpc":"2.0","method":"slotNotification","params":{"result":{"parent":1,"root":2,"slot":3},"subscription":17}}"#;
        let parsed = crate::parse_kb_json_rpc_ws_incoming_text(text).expect("parse must succeed");
        match parsed {
            crate::KbJsonRpcWsIncomingMessage::Notification(notification) => {
                assert_eq!(notification.jsonrpc, "2.0");
                assert_eq!(notification.method, "slotNotification");
                assert_eq!(notification.params.subscription, 17);
                assert_eq!(
                    notification.params.result["slot"],
                    serde_json::Value::from(3u64)
                );
            }
            other => {
                panic!("unexpected parsed message: {other:?}");
            }
        }
    }

    #[test]
    fn probable_json_rpc_object_text_detects_object() {
        assert!(crate::kb_is_probable_json_rpc_object_text(
            " {\"jsonrpc\":\"2.0\"}"
        ));
        assert!(!crate::kb_is_probable_json_rpc_object_text("hello"));
    }
}