rialo-types 0.12.2

// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

//! WebSocket-related types for the Rialo network.
//!
//! This module contains types that are used for WebSocket REX connections
//! and need to be deserialized in the SVM.

use serde::{Deserialize, Serialize};

use crate::rex_info::RexId;

/// Size of BLS12-381 public key in bytes.
/// This is duplicated here to avoid dependency on the feature-gated `admin` module.
pub const BLS12381_PUBLIC_KEY_SIZE: usize = 96;

/// The state of a WebSocket connection account.
///
/// This struct tracks an active WebSocket connection that was established by a TEE.
/// The connection is identified by the `rex_id` of the Connect operation that created it.
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct WebSocketConnection {
    /// The RexId that created this connection (from the Connect operation).
    pub rex_id: RexId,

    /// The authority key (96 bytes BLS12-381) of the validator currently running this WebSocket connection.
    /// All Read operations for this connection will be routed to this validator.
    /// This matches the `authority_key` field in `ValidatorInfo`.
    #[serde(with = "serde_big_array::BigArray")]
    pub assigned_validator: [u8; BLS12381_PUBLIC_KEY_SIZE],

    /// The round when the connection was established.
    pub established_at_timestamp: Option<u64>,

    /// The WebSocket URL that the connection is connected to.
    pub url: String,

    /// The current status of the connection.
    pub status: ConnectionStatus,
}

impl Default for WebSocketConnection {
    fn default() -> Self {
        Self {
            rex_id: RexId::default(),
            assigned_validator: [0u8; BLS12381_PUBLIC_KEY_SIZE],
            established_at_timestamp: Some(0),
            url: String::new(),
            status: ConnectionStatus::default(),
        }
    }
}

/// The status of a WebSocket connection.
#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
pub enum ConnectionStatus {
    /// The connection is active and can receive data.
    #[default]
    Active,

    /// The connection failed to establish.
    Failed,

    /// The connection has been closed and is no longer usable.
    Closed,
}

impl ConnectionStatus {
    /// Returns `true` if the connection is active.
    pub fn is_active(&self) -> bool {
        matches!(self, ConnectionStatus::Active)
    }

    /// Returns `true` if the connection is closed.
    pub fn is_closed(&self) -> bool {
        matches!(self, ConnectionStatus::Closed)
    }

    /// Returns `true` if the connection has failed.
    pub fn is_failed(&self) -> bool {
        matches!(self, ConnectionStatus::Failed)
    }
}

#[cfg(test)]
mod tests {
    use rialo_s_pubkey::Pubkey;

    use super::*;

    #[test]
    fn test_connection_status_default() {
        let status = ConnectionStatus::default();
        assert!(status.is_active());
        assert!(!status.is_closed());
    }

    #[test]
    fn test_connection_status_active() {
        let status = ConnectionStatus::Active;
        assert!(status.is_active());
        assert!(!status.is_closed());
    }

    #[test]
    fn test_connection_status_closed() {
        let status = ConnectionStatus::Closed;
        assert!(!status.is_active());
        assert!(status.is_closed());
    }

    #[test]
    fn test_websocket_connection_serde_roundtrip() {
        let connection = WebSocketConnection {
            rex_id: RexId::new(Pubkey::default(), 42u64),
            assigned_validator: [0u8; BLS12381_PUBLIC_KEY_SIZE], // 96 bytes authority key
            established_at_timestamp: Some(100),
            url: "wss://example.com/stream".to_string(),
            status: ConnectionStatus::Active,
        };

        // Serialize to JSON
        let json = serde_json::to_string(&connection).expect("Failed to serialize");

        // Deserialize back
        let deserialized: WebSocketConnection =
            serde_json::from_str(&json).expect("Failed to deserialize");

        assert_eq!(connection, deserialized);
    }

    #[test]
    fn test_websocket_connection_serde_roundtrip_closed() {
        // Create a sample authority key with some non-zero bytes for variety
        let mut authority_key = [0u8; BLS12381_PUBLIC_KEY_SIZE];
        authority_key[0] = 0xAB;
        authority_key[95] = 0xCD;

        let connection = WebSocketConnection {
            rex_id: RexId::new(Pubkey::default(), 123u64),
            assigned_validator: authority_key,
            established_at_timestamp: Some(200),
            url: "wss://api.example.com/ws".to_string(),
            status: ConnectionStatus::Closed,
        };

        // Serialize to JSON
        let json = serde_json::to_string(&connection).expect("Failed to serialize");

        // Deserialize back
        let deserialized: WebSocketConnection =
            serde_json::from_str(&json).expect("Failed to deserialize");

        assert_eq!(connection, deserialized);
    }

    #[test]
    fn test_connection_status_serde_roundtrip() {
        for status in [ConnectionStatus::Active, ConnectionStatus::Closed] {
            let json = serde_json::to_string(&status).expect("Failed to serialize");
            let deserialized: ConnectionStatus =
                serde_json::from_str(&json).expect("Failed to deserialize");
            assert_eq!(status, deserialized);
        }
    }
}