rialo_types/
websocket.rs

1// Copyright (c) Subzero Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4//! WebSocket-related types for the Rialo network.
5//!
6//! This module contains types that are used for WebSocket REX connections
7//! and need to be deserialized in the SVM.
8
9use serde::{Deserialize, Serialize};
10
11use crate::rex_info::RexId;
12
13/// Size of BLS12-381 public key in bytes.
14/// This is duplicated here to avoid dependency on the feature-gated `admin` module.
15pub const BLS12381_PUBLIC_KEY_SIZE: usize = 96;
16
17/// The state of a WebSocket connection account.
18///
19/// This struct tracks an active WebSocket connection that was established by a TEE.
20/// The connection is identified by the `rex_id` of the Connect operation that created it.
21#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
22pub struct WebSocketConnection {
23    /// The RexId that created this connection (from the Connect operation).
24    pub rex_id: RexId,
25
26    /// The authority key (96 bytes BLS12-381) of the validator currently running this WebSocket connection.
27    /// All Read operations for this connection will be routed to this validator.
28    /// This matches the `authority_key` field in `ValidatorInfo`.
29    #[serde(with = "serde_big_array::BigArray")]
30    pub assigned_validator: [u8; BLS12381_PUBLIC_KEY_SIZE],
31
32    /// The round when the connection was established.
33    pub established_at_timestamp: Option<u64>,
34
35    /// The WebSocket URL that the connection is connected to.
36    pub url: String,
37
38    /// The current status of the connection.
39    pub status: ConnectionStatus,
40}
41
42impl Default for WebSocketConnection {
43    fn default() -> Self {
44        Self {
45            rex_id: RexId::default(),
46            assigned_validator: [0u8; BLS12381_PUBLIC_KEY_SIZE],
47            established_at_timestamp: Some(0),
48            url: String::new(),
49            status: ConnectionStatus::default(),
50        }
51    }
52}
53
54/// The status of a WebSocket connection.
55#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
56pub enum ConnectionStatus {
57    /// The connection is active and can receive data.
58    #[default]
59    Active,
60
61    /// The connection failed to establish.
62    Failed,
63
64    /// The connection has been closed and is no longer usable.
65    Closed,
66}
67
68impl ConnectionStatus {
69    /// Returns `true` if the connection is active.
70    pub fn is_active(&self) -> bool {
71        matches!(self, ConnectionStatus::Active)
72    }
73
74    /// Returns `true` if the connection is closed.
75    pub fn is_closed(&self) -> bool {
76        matches!(self, ConnectionStatus::Closed)
77    }
78
79    /// Returns `true` if the connection has failed.
80    pub fn is_failed(&self) -> bool {
81        matches!(self, ConnectionStatus::Failed)
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use rialo_s_pubkey::Pubkey;
88
89    use super::*;
90
91    #[test]
92    fn test_connection_status_default() {
93        let status = ConnectionStatus::default();
94        assert!(status.is_active());
95        assert!(!status.is_closed());
96    }
97
98    #[test]
99    fn test_connection_status_active() {
100        let status = ConnectionStatus::Active;
101        assert!(status.is_active());
102        assert!(!status.is_closed());
103    }
104
105    #[test]
106    fn test_connection_status_closed() {
107        let status = ConnectionStatus::Closed;
108        assert!(!status.is_active());
109        assert!(status.is_closed());
110    }
111
112    #[test]
113    fn test_websocket_connection_serde_roundtrip() {
114        let connection = WebSocketConnection {
115            rex_id: RexId::new(Pubkey::default(), 42u64),
116            assigned_validator: [0u8; BLS12381_PUBLIC_KEY_SIZE], // 96 bytes authority key
117            established_at_timestamp: Some(100),
118            url: "wss://example.com/stream".to_string(),
119            status: ConnectionStatus::Active,
120        };
121
122        // Serialize to JSON
123        let json = serde_json::to_string(&connection).expect("Failed to serialize");
124
125        // Deserialize back
126        let deserialized: WebSocketConnection =
127            serde_json::from_str(&json).expect("Failed to deserialize");
128
129        assert_eq!(connection, deserialized);
130    }
131
132    #[test]
133    fn test_websocket_connection_serde_roundtrip_closed() {
134        // Create a sample authority key with some non-zero bytes for variety
135        let mut authority_key = [0u8; BLS12381_PUBLIC_KEY_SIZE];
136        authority_key[0] = 0xAB;
137        authority_key[95] = 0xCD;
138
139        let connection = WebSocketConnection {
140            rex_id: RexId::new(Pubkey::default(), 123u64),
141            assigned_validator: authority_key,
142            established_at_timestamp: Some(200),
143            url: "wss://api.example.com/ws".to_string(),
144            status: ConnectionStatus::Closed,
145        };
146
147        // Serialize to JSON
148        let json = serde_json::to_string(&connection).expect("Failed to serialize");
149
150        // Deserialize back
151        let deserialized: WebSocketConnection =
152            serde_json::from_str(&json).expect("Failed to deserialize");
153
154        assert_eq!(connection, deserialized);
155    }
156
157    #[test]
158    fn test_connection_status_serde_roundtrip() {
159        for status in [ConnectionStatus::Active, ConnectionStatus::Closed] {
160            let json = serde_json::to_string(&status).expect("Failed to serialize");
161            let deserialized: ConnectionStatus =
162                serde_json::from_str(&json).expect("Failed to deserialize");
163            assert_eq!(status, deserialized);
164        }
165    }
166}
rialo_types/websocket.rs

rialo_types/
websocket.rs
