rusty_common/websocket/
auth.rs

1//! WebSocket authentication module
2//!
3//! Provides authentication abstractions for WebSocket connections.
4
5use crate::collections::FxHashMap;
6use simd_json::OwnedValue;
7use simd_json::prelude::*;
8use smartstring::alias::String;
9
10use super::{WebSocketError, WebSocketResult};
11use crate::auth::hmac::generate_hmac_signature;
12use crate::types::Exchange;
13
14/// WebSocket authentication method
15#[derive(Debug, Clone)]
16pub enum AuthMethod {
17    /// No authentication
18    None,
19
20    /// API key authentication
21    ApiKey {
22        /// The API key.
23        key: String,
24        /// The API secret.
25        secret: String,
26    },
27
28    /// JWT token authentication
29    Jwt {
30        /// The JWT token.
31        token: String,
32    },
33
34    /// Custom authentication with headers
35    Custom {
36        /// The custom headers.
37        headers: FxHashMap<String, String>,
38    },
39}
40
41/// WebSocket authentication provider
42pub trait WebSocketAuth: Send + Sync {
43    /// Create authentication message
44    fn create_auth_message(&self) -> WebSocketResult<Option<OwnedValue>>;
45
46    /// Get authentication headers
47    fn get_auth_headers(&self) -> WebSocketResult<FxHashMap<String, String>>;
48
49    /// Handle authentication response
50    fn handle_auth_response(&mut self, response: &OwnedValue) -> WebSocketResult<bool>;
51}
52
53/// Default WebSocket authenticator
54pub struct DefaultWebSocketAuth {
55    exchange: Exchange,
56    method: AuthMethod,
57}
58
59impl DefaultWebSocketAuth {
60    /// Create a new authenticator
61    #[must_use]
62    pub const fn new(exchange: Exchange, method: AuthMethod) -> Self {
63        Self { exchange, method }
64    }
65}
66
67impl WebSocketAuth for DefaultWebSocketAuth {
68    fn create_auth_message(&self) -> WebSocketResult<Option<OwnedValue>> {
69        match &self.method {
70            AuthMethod::None => Ok(None),
71
72            AuthMethod::ApiKey { key, secret } => {
73                match self.exchange {
74                    Exchange::Binance => {
75                        // Binance doesn't use WebSocket auth messages
76                        // Authentication is done via listen key
77                        Ok(None)
78                    }
79
80                    Exchange::Bybit => {
81                        // Bybit uses auth message
82                        let expires = crate::time::get_timestamp_ms() as i64 + 10000;
83                        let message = format!("GET/realtime{expires}");
84                        let signature = generate_hmac_signature(secret, &message)
85                            .map_err(|e| WebSocketError::AuthenticationError(e.to_string()))?;
86
87                        Ok(Some(simd_json::json!({
88                            "op": "auth",
89                            "args": [key, expires, signature]
90                        })))
91                    }
92
93                    Exchange::Coinbase => {
94                        // Coinbase requires signature in subscription message
95                        Ok(None)
96                    }
97
98                    Exchange::Upbit => {
99                        // Upbit uses JWT
100                        Ok(None)
101                    }
102
103                    Exchange::Bithumb => {
104                        // Bithumb auth is in subscription message
105                        Ok(None)
106                    }
107                }
108            }
109
110            AuthMethod::Jwt { token: _ } => {
111                match self.exchange {
112                    Exchange::Upbit => {
113                        // Upbit sends JWT in Authorization header
114                        Ok(None)
115                    }
116                    _ => Err(WebSocketError::AuthenticationError(
117                        "JWT not supported for this exchange".into(),
118                    )),
119                }
120            }
121
122            AuthMethod::Custom { .. } => Ok(None),
123        }
124    }
125
126    fn get_auth_headers(&self) -> WebSocketResult<FxHashMap<String, String>> {
127        let mut headers = FxHashMap::default();
128
129        match &self.method {
130            AuthMethod::None => {}
131
132            AuthMethod::ApiKey { .. } => {
133                // Most exchanges don't use headers for API key auth
134            }
135
136            AuthMethod::Jwt { token } => {
137                if self.exchange == Exchange::Upbit {
138                    headers.insert("Authorization".into(), format!("Bearer {token}").into());
139                }
140            }
141
142            AuthMethod::Custom { headers: custom } => {
143                for (k, v) in custom {
144                    headers.insert(k.clone(), v.clone());
145                }
146            }
147        }
148
149        Ok(headers)
150    }
151
152    fn handle_auth_response(&mut self, response: &OwnedValue) -> WebSocketResult<bool> {
153        // Parse common response patterns
154        if let Some(success) = response.get("success").and_then(|v| v.as_bool()) {
155            return Ok(success);
156        }
157
158        if let Some(status) = response.get("status").and_then(|v| v.as_str()) {
159            return Ok(status == "success" || status == "ok");
160        }
161
162        if let Some(result) = response.get("result").and_then(|v| v.as_str()) {
163            return Ok(result == "true" || result == "success");
164        }
165
166        // Exchange specific patterns
167        if self.exchange == Exchange::Bybit
168            && let Some(ret_code) = response.get("ret_code").and_then(|v| v.as_i64())
169        {
170            return Ok(ret_code == 0);
171        }
172
173        // If no clear indication, assume success if no error field
174        Ok(!response.contains_key("error") && !response.contains_key("code"))
175    }
176}
177
178/// Create authenticator for exchange
179pub fn create_authenticator(
180    exchange: Exchange,
181    api_key: Option<String>,
182    api_secret: Option<String>,
183) -> Box<dyn WebSocketAuth> {
184    let method = match (api_key, api_secret) {
185        (Some(key), Some(secret)) => AuthMethod::ApiKey { key, secret },
186        _ => AuthMethod::None,
187    };
188
189    Box::new(DefaultWebSocketAuth::new(exchange, method))
190}
rusty_common/websocket/auth.rs

rusty_common/websocket/
auth.rs
