Skip to main content

http_request/request/socket/websocket/
struct.rs

1use crate::*;
2
3/// Represents different types of WebSocket connections.
4///
5/// This enum encapsulates both direct and proxy-based WebSocket connections,
6/// providing a unified interface for WebSocket operations.
7#[derive(Debug)]
8pub enum WebSocketConnectionType {
9    Direct(WebSocketStream<MaybeTlsStream<AsyncTcpStream>>),
10    Proxy(WebSocketStream<WebSocketProxyTunnelStream>),
11}
12
13/// Represents a WebSocket client connection.
14///
15/// This struct manages the WebSocket lifecycle including:
16/// - Connection state
17/// - Message sending/receiving
18/// - Configuration
19#[derive(Debug)]
20pub struct WebSocket {
21    /// The WebSocket server URL.
22    pub(crate) url: Arc<String>,
23    /// HTTP headers for the WebSocket handshake.
24    pub(crate) header: Arc<RequestHeaders>,
25    /// Configuration settings for the WebSocket connection.
26    pub(crate) config: ArcRwLock<WebSocketConfig>,
27    /// Atomic flag indicating connection status.
28    pub(crate) connected: Arc<AtomicBool>,
29    /// The underlying WebSocket connection.
30    pub(crate) connection: WebSocketConnection,
31}
32
33/// Clone implementation for WebSocket.
34///
35/// Creates a new WebSocket instance with cloned configuration but resets:
36/// - Connection status to false
37/// - Connection to None
38impl Clone for WebSocket {
39    fn clone(&self) -> Self {
40        Self {
41            url: self.url.clone(),
42            header: self.header.clone(),
43            config: self.config.clone(),
44            connected: Arc::new(AtomicBool::new(false)),
45            connection: Arc::new(AsyncMutex::new(None)),
46        }
47    }
48}
49
50/// Default implementation for WebSocket.
51///
52/// Creates a WebSocket with:
53/// - Empty URL
54/// - Empty headers
55/// - Default configuration
56/// - Disconnected state
57/// - No active connection
58impl Default for WebSocket {
59    #[inline(always)]
60    fn default() -> Self {
61        Self {
62            url: Arc::new(String::new()),
63            header: Arc::new(hash_map_xx_hash3_64()),
64            config: Arc::new(RwLock::new(WebSocketConfig::default())),
65            connected: Arc::new(AtomicBool::new(false)),
66            connection: Arc::new(AsyncMutex::new(None)),
67        }
68    }
69}
70
71/// Stream implementation for WebSocketConnectionType.
72///
73/// Allows polling for incoming WebSocket messages.
74/// Handles both direct and proxy connections uniformly.
75impl Stream for WebSocketConnectionType {
76    type Item =
77        Result<tokio_tungstenite::tungstenite::Message, tokio_tungstenite::tungstenite::Error>;
78
79    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
80        match &mut *self {
81            WebSocketConnectionType::Direct(stream) => Pin::new(stream).poll_next(cx),
82            WebSocketConnectionType::Proxy(stream) => Pin::new(stream).poll_next(cx),
83        }
84    }
85}
86
87/// Sink implementation for WebSocketConnectionType.
88///
89/// Allows sending WebSocket messages.
90/// Handles both direct and proxy connections uniformly.
91impl Sink<tokio_tungstenite::tungstenite::Message> for WebSocketConnectionType {
92    type Error = tokio_tungstenite::tungstenite::Error;
93
94    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
95        match &mut *self {
96            WebSocketConnectionType::Direct(stream) => Pin::new(stream).poll_ready(cx),
97            WebSocketConnectionType::Proxy(stream) => Pin::new(stream).poll_ready(cx),
98        }
99    }
100
101    fn start_send(
102        mut self: Pin<&mut Self>,
103        item: tokio_tungstenite::tungstenite::Message,
104    ) -> Result<(), Self::Error> {
105        match &mut *self {
106            WebSocketConnectionType::Direct(stream) => Pin::new(stream).start_send(item),
107            WebSocketConnectionType::Proxy(stream) => Pin::new(stream).start_send(item),
108        }
109    }
110
111    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
112        match &mut *self {
113            WebSocketConnectionType::Direct(stream) => Pin::new(stream).poll_flush(cx),
114            WebSocketConnectionType::Proxy(stream) => Pin::new(stream).poll_flush(cx),
115        }
116    }
117
118    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
119        match &mut *self {
120            WebSocketConnectionType::Direct(stream) => Pin::new(stream).poll_close(cx),
121            WebSocketConnectionType::Proxy(stream) => Pin::new(stream).poll_close(cx),
122        }
123    }
124}