1use core::fmt;
2
3use futures_util::stream::SplitStream;
4use serde::{Deserialize, Serialize};
5use serde_json::Value;
6use tokio::{net::TcpStream, sync::MutexGuard};
7use tokio_tungstenite::{
8 MaybeTlsStream, WebSocketStream,
9 tungstenite::{
10 http::{HeaderMap, HeaderValue},
11 protocol::Message,
12 },
13};
14use ustr::Ustr;
15
16use crate::{
17 Result, UA,
18 error::{Error, TradingViewError},
19 utils::format_packet,
20};
21use std::sync::LazyLock;
22
23pub static WEBSOCKET_HEADERS: LazyLock<HeaderMap<HeaderValue>> = LazyLock::new(|| {
24 let mut headers = HeaderMap::new();
25 headers.insert("Origin", "https://www.tradingview.com/".parse().unwrap());
26 headers.insert("User-Agent", UA.parse().unwrap());
27 headers
28});
29
30#[derive(Debug, Clone, PartialEq, Eq, Hash, Copy)]
35pub enum TradingViewDataEvent {
36 OnChartData,
37 OnChartDataUpdate,
38 OnQuoteData,
39 OnQuoteCompleted,
40 OnSeriesLoading,
41 OnSeriesCompleted,
42 OnSymbolResolved,
43 OnReplayOk,
44 OnReplayPoint,
45 OnReplayInstanceId,
46 OnReplayResolutions,
47 OnReplayDataEnd,
48 OnStudyLoading,
49 OnStudyCompleted,
50 OnError(TradingViewError),
51 UnknownEvent(Ustr),
52}
53
54impl From<String> for TradingViewDataEvent {
55 fn from(s: String) -> Self {
56 match s.as_str() {
57 "timescale_update" => TradingViewDataEvent::OnChartData,
58 "du" => TradingViewDataEvent::OnChartDataUpdate,
59
60 "qsd" => TradingViewDataEvent::OnQuoteData,
61 "quote_completed" => TradingViewDataEvent::OnQuoteCompleted,
62
63 "series_loading" => TradingViewDataEvent::OnSeriesLoading,
64 "series_completed" => TradingViewDataEvent::OnSeriesCompleted,
65
66 "symbol_resolved" => TradingViewDataEvent::OnSymbolResolved,
67
68 "replay_ok" => TradingViewDataEvent::OnReplayOk,
69 "replay_point" => TradingViewDataEvent::OnReplayPoint,
70 "replay_instance_id" => TradingViewDataEvent::OnReplayInstanceId,
71 "replay_resolutions" => TradingViewDataEvent::OnReplayResolutions,
72 "replay_data_end" => TradingViewDataEvent::OnReplayDataEnd,
73
74 "study_loading" => TradingViewDataEvent::OnStudyLoading,
75 "study_completed" => TradingViewDataEvent::OnStudyCompleted,
76
77 "symbol_error" => TradingViewDataEvent::OnError(TradingViewError::SymbolError),
78 "series_error" => TradingViewDataEvent::OnError(TradingViewError::SeriesError),
79 "critical_error" => TradingViewDataEvent::OnError(TradingViewError::CriticalError),
80 "study_error" => TradingViewDataEvent::OnError(TradingViewError::StudyError),
81 "protocol_error" => TradingViewDataEvent::OnError(TradingViewError::ProtocolError),
82 "replay_error" => TradingViewDataEvent::OnError(TradingViewError::ReplayError),
83
84 s => TradingViewDataEvent::UnknownEvent(s.into()),
85 }
86 }
87}
88
89impl From<Ustr> for TradingViewDataEvent {
90 fn from(s: Ustr) -> Self {
91 TradingViewDataEvent::from(s.to_string())
92 }
93}
94
95#[derive(Debug, Clone, PartialEq, Serialize)]
97pub struct SocketMessageSer {
98 pub m: Value,
99 pub p: Value,
100}
101
102#[derive(Debug, Clone, PartialEq, Deserialize)]
104pub struct SocketMessageDe {
105 pub m: Ustr,
106 pub p: Vec<Value>,
107 #[serde(default)]
108 pub t: u64, #[serde(default)]
110 pub t_ms: u64, }
112
113impl SocketMessageSer {
114 pub fn new<M, P>(m: M, p: P) -> Self
115 where
116 M: Serialize,
117 P: Serialize,
118 {
119 let m = serde_json::to_value(m).expect("Failed to serialize Socket Message");
120 let p = serde_json::to_value(p).expect("Failed to serialize Socket Message");
121 SocketMessageSer { m, p }
122 }
123
124 pub fn to_message(&self) -> Result<Message> {
125 let msg = format_packet(self)?;
126 Ok(msg)
127 }
128}
129
130#[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)]
134#[serde(rename_all = "camelCase")]
135pub struct SocketServerInfo {
136 #[serde(rename = "session_id")]
137 pub session_id: Ustr,
138 pub timestamp: i64,
139 pub timestamp_ms: i64,
140 pub release: Ustr,
141 #[serde(rename = "studies_metadata_hash")]
142 pub studies_metadata_hash: Ustr,
143 #[serde(rename = "auth_scheme_vsn")]
144 pub auth_scheme_vsn: i64,
145 pub protocol: Ustr,
146 pub via: Ustr,
147 #[serde(rename = "javastudies")]
148 pub sjavastudies: Vec<Ustr>,
149}
150
151impl fmt::Display for SocketServerInfo {
152 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
153 write!(
154 f,
155 "SocketServerInfo {{ session_id: {}, timestamp: {}, timestamp_ms: {}, release: {}, studies_metadata_hash: {}, auth_scheme_vsn: {}, protocol: {}, via: {}, sjavastudies: {:?} }}",
156 self.session_id,
157 self.timestamp,
158 self.timestamp_ms,
159 self.release,
160 self.studies_metadata_hash,
161 self.auth_scheme_vsn,
162 self.protocol,
163 self.via,
164 self.sjavastudies
165 )
166 }
167}
168
169#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
170#[serde(untagged)]
171pub enum SocketMessage<T> {
172 SocketServerInfo(SocketServerInfo),
173 SocketMessage(T),
174 Heartbeat(u64),
175 Other(Value),
176 Unknown(String),
177}
178
179impl<T> SocketMessage<T> {
180 pub fn heartbeat_echo(&self) -> Option<String> {
181 match self {
182 SocketMessage::Heartbeat(counter) => {
183 let payload = format!("~h~{counter}");
184 Some(format!("~m~{}~m~{payload}", payload.len()))
185 }
186 _ => None,
187 }
188 }
189}
190
191#[derive(Default, Clone, Debug, PartialEq, Serialize, Deserialize, Copy, Eq)]
196pub enum DataServer {
197 #[default]
198 Data,
199 ProData,
200 WidgetData,
201 MobileData,
202}
203
204impl std::fmt::Display for DataServer {
205 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
206 match *self {
207 DataServer::Data => write!(f, "data"),
208 DataServer::ProData => write!(f, "prodata"),
209 DataServer::WidgetData => write!(f, "widgetdata"),
210 DataServer::MobileData => write!(f, "mobile-data"),
211 }
212 }
213}
214
215pub trait Socket {
222 fn event_loop(
223 &self,
224 read: MutexGuard<SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>>,
225 ) -> impl Future<Output = Result<()>> + Send;
226
227 fn handle_raw_messages(&self, raw: Message) -> impl Future<Output = Result<()>> + Send;
228
229 fn handle_parsed_messages(
230 &self,
231 messages: Vec<SocketMessage<SocketMessageDe>>,
232 raw: &Message,
233 ) -> impl Future<Output = Result<()>> + Send;
234
235 fn handle_message_data(
236 &self,
237 message: SocketMessageDe,
238 ) -> impl Future<Output = Result<()>> + Send;
239
240 fn handle_error(&self, error: Error, context: Ustr) -> impl Future<Output = Result<()>> + Send;
241}
242
243#[cfg(test)]
248mod tests {
249 use super::*;
250
251 #[test]
252 fn test_study_loading_maps_to_on_study_loading() {
253 let event = TradingViewDataEvent::from("study_loading".to_string());
254 assert_eq!(event, TradingViewDataEvent::OnStudyLoading);
255 }
256
257 #[test]
258 fn test_series_loading_is_not_study_loading() {
259 let event = TradingViewDataEvent::from("series_loading".to_string());
260 assert_eq!(event, TradingViewDataEvent::OnSeriesLoading);
261 }
262
263 #[test]
264 fn test_study_completed_maps_correctly() {
265 let event = TradingViewDataEvent::from("study_completed".to_string());
266 assert_eq!(event, TradingViewDataEvent::OnStudyCompleted);
267 }
268
269 #[test]
270 fn test_series_completed_maps_correctly() {
271 let event = TradingViewDataEvent::from("series_completed".to_string());
272 assert_eq!(event, TradingViewDataEvent::OnSeriesCompleted);
273 }
274
275 #[test]
276 fn test_all_study_and_series_events_are_distinct() {
277 let loading = TradingViewDataEvent::from("study_loading".to_string());
278 let completed = TradingViewDataEvent::from("study_completed".to_string());
279 let series_loading = TradingViewDataEvent::from("series_loading".to_string());
280 let series_completed = TradingViewDataEvent::from("series_completed".to_string());
281
282 assert_ne!(loading, series_loading);
284 assert_ne!(completed, series_completed);
285 assert_ne!(loading, completed);
286 assert_ne!(series_loading, series_completed);
287
288 assert_eq!(loading, TradingViewDataEvent::OnStudyLoading);
290 assert_eq!(completed, TradingViewDataEvent::OnStudyCompleted);
291 assert_eq!(series_loading, TradingViewDataEvent::OnSeriesLoading);
292 assert_eq!(series_completed, TradingViewDataEvent::OnSeriesCompleted);
293 }
294
295 #[test]
296 fn test_ustr_from_maps_correctly() {
297 let event: TradingViewDataEvent = ustr::ustr("study_loading").into();
298 assert_eq!(event, TradingViewDataEvent::OnStudyLoading);
299 }
300}