1 use uuid::Uuid;
2 
3 use {
4     super::{
5         common::Common,
6         define,
7         define::SessionType,
8         errors::{SessionError, SessionErrorValue},
9     },
10     //crate::utils::print::print,
11     crate::{
12         amf0::Amf0ValueType,
13         channels::define::ChannelEventProducer,
14         chunk::{
15             define::CHUNK_SIZE,
16             unpacketizer::{ChunkUnpacketizer, UnpackResult},
17         },
18         handshake::{define::ClientHandshakeState, handshake_client::SimpleHandshakeClient},
19         messages::{define::RtmpMessageData, parser::MessageParser},
20         netconnection::writer::{ConnectProperties, NetConnection},
21         netstream::writer::NetStreamWriter,
22         protocol_control_messages::writer::ProtocolControlMessagesWriter,
23         user_control_messages::writer::EventMessagesWriter,
24     },
25     bytesio::{bytes_writer::AsyncBytesWriter, bytesio::BytesIO},
26     std::{collections::HashMap, sync::Arc},
27     tokio::{net::TcpStream, sync::Mutex},
28 };
29 
30 #[allow(dead_code)]
31 enum ClientSessionState {
32     Handshake,
33     Connect,
34     CreateStream,
35     Play,
36     PublishingContent,
37     StartPublish,
38     WaitStateChange,
39 }
40 
41 #[allow(dead_code)]
42 enum ClientSessionPlayState {
43     Handshake,
44     Connect,
45     CreateStream,
46     Play,
47 }
48 
49 #[allow(dead_code)]
50 enum ClientSessionPublishState {
51     Handshake,
52     Connect,
53     CreateStream,
54     PublishingContent,
55 }
56 #[allow(dead_code)]
57 pub enum ClientType {
58     Play,
59     Publish,
60 }
61 pub struct ClientSession {
62     io: Arc<Mutex<BytesIO>>,
63     common: Common,
64     handshaker: SimpleHandshakeClient,
65     unpacketizer: ChunkUnpacketizer,
66     app_name: String,
67     stream_name: String,
68     /* Used to mark the subscriber's the data producer
69     in channels and delete it from map when unsubscribe
70     is called. */
71     session_id: Uuid,
72     state: ClientSessionState,
73     client_type: ClientType,
74     sub_app_name: Option<String>,
75     sub_stream_name: Option<String>,
76 }
77 
78 impl ClientSession {
79     pub fn new(
80         stream: TcpStream,
81         client_type: ClientType,
82         app_name: String,
83         stream_name: String,
84         event_producer: ChannelEventProducer,
85     ) -> Self {
86         let remote_addr = if let Ok(addr) = stream.peer_addr() {
87             log::info!("server session: {}", addr.to_string());
88             Some(addr)
89         } else {
90             None
91         };
92 
93         let net_io = Arc::new(Mutex::new(BytesIO::new(stream)));
94         let subscriber_id = Uuid::new_v4();
95 
96         Self {
97             io: Arc::clone(&net_io),
98             common: Common::new(
99                 Arc::clone(&net_io),
100                 event_producer,
101                 SessionType::Client,
102                 remote_addr,
103             ),
104             handshaker: SimpleHandshakeClient::new(Arc::clone(&net_io)),
105             unpacketizer: ChunkUnpacketizer::new(),
106             app_name,
107             stream_name,
108             client_type,
109             state: ClientSessionState::Handshake,
110             session_id: subscriber_id,
111             sub_app_name: None,
112             sub_stream_name: None,
113         }
114     }
115 
116     pub async fn run(&mut self) -> Result<(), SessionError> {
117         loop {
118             match self.state {
119                 ClientSessionState::Handshake => {
120                     log::info!("[C -> S] handshake...");
121                     self.handshake().await?;
122                     continue;
123                 }
124                 ClientSessionState::Connect => {
125                     log::info!("[C -> S] connect...");
126                     self.send_connect(&(define::TRANSACTION_ID_CONNECT as f64))
127                         .await?;
128                     self.state = ClientSessionState::WaitStateChange;
129                 }
130                 ClientSessionState::CreateStream => {
131                     log::info!("[C -> S] CreateStream...");
132                     self.send_create_stream(&(define::TRANSACTION_ID_CREATE_STREAM as f64))
133                         .await?;
134                     self.state = ClientSessionState::WaitStateChange;
135                 }
136                 ClientSessionState::Play => {
137                     log::info!("[C -> S] Play...");
138                     self.send_play(&0.0, &self.stream_name.clone(), &0.0, &0.0, &false)
139                         .await?;
140                     self.state = ClientSessionState::WaitStateChange;
141                 }
142                 ClientSessionState::PublishingContent => {
143                     log::info!("[C -> S] PublishingContent...");
144                     self.send_publish(&0.0, &self.stream_name.clone(), &"live".to_string())
145                         .await?;
146                     self.state = ClientSessionState::WaitStateChange;
147                 }
148                 ClientSessionState::StartPublish => {
149                     log::info!("[C -> S] StartPublish...");
150                     self.common.send_channel_data().await?;
151                 }
152                 ClientSessionState::WaitStateChange => {}
153             }
154 
155             let data = self.io.lock().await.read().await?;
156             self.unpacketizer.extend_data(&data[..]);
157 
158             loop {
159                 let result = self.unpacketizer.read_chunks();
160 
161                 if let Ok(rv) = result {
162                     if let UnpackResult::Chunks(chunks) = rv {
163                         for chunk_info in chunks.iter() {
164                             let mut msg = MessageParser::new(chunk_info.clone()).parse()?;
165 
166                             let timestamp = chunk_info.message_header.timestamp;
167                             self.process_messages(&mut msg, &timestamp).await?;
168                         }
169                     }
170                 } else {
171                     break;
172                 }
173             }
174         }
175     }
176 
177     async fn handshake(&mut self) -> Result<(), SessionError> {
178         loop {
179             self.handshaker.handshake().await?;
180             if self.handshaker.state == ClientHandshakeState::Finish {
181                 log::info!("handshake finish");
182                 break;
183             }
184 
185             let data = self.io.lock().await.read().await?;
186             self.handshaker.extend_data(&data[..]);
187         }
188 
189         self.state = ClientSessionState::Connect;
190 
191         Ok(())
192     }
193 
194     pub async fn process_messages(
195         &mut self,
196         msg: &mut RtmpMessageData,
197         timestamp: &u32,
198     ) -> Result<(), SessionError> {
199         match msg {
200             RtmpMessageData::Amf0Command {
201                 command_name,
202                 transaction_id,
203                 command_object,
204                 others,
205             } => {
206                 log::info!("[C <- S] on_amf0_command_message...");
207                 self.on_amf0_command_message(command_name, transaction_id, command_object, others)
208                     .await?
209             }
210             RtmpMessageData::SetPeerBandwidth { .. } => {
211                 log::info!("[C <- S] on_set_peer_bandwidth...");
212                 self.on_set_peer_bandwidth().await?
213             }
214             RtmpMessageData::WindowAcknowledgementSize { .. } => {
215                 log::info!("[C <- S] on_windows_acknowledgement_size...");
216             }
217             RtmpMessageData::SetChunkSize { chunk_size } => {
218                 log::info!("[C <- S] on_set_chunk_size...");
219                 self.on_set_chunk_size(chunk_size)?;
220             }
221             RtmpMessageData::StreamBegin { stream_id } => {
222                 log::info!("[C <- S] on_stream_begin...");
223                 self.on_stream_begin(stream_id)?;
224             }
225             RtmpMessageData::StreamIsRecorded { stream_id } => {
226                 log::info!("[C <- S] on_stream_is_recorded...");
227                 self.on_stream_is_recorded(stream_id)?;
228             }
229             RtmpMessageData::AudioData { data } => self.common.on_audio_data(data, timestamp)?,
230             RtmpMessageData::VideoData { data } => self.common.on_video_data(data, timestamp)?,
231             RtmpMessageData::AmfData { raw_data } => {
232                 self.common.on_meta_data(raw_data, timestamp)?;
233             }
234 
235             _ => {}
236         }
237         Ok(())
238     }
239 
240     pub async fn on_amf0_command_message(
241         &mut self,
242         command_name: &Amf0ValueType,
243         transaction_id: &Amf0ValueType,
244         command_object: &Amf0ValueType,
245         others: &mut Vec<Amf0ValueType>,
246     ) -> Result<(), SessionError> {
247         let empty_cmd_name = &String::new();
248         let cmd_name = match command_name {
249             Amf0ValueType::UTF8String(str) => str,
250             _ => empty_cmd_name,
251         };
252 
253         let transaction_id = match transaction_id {
254             Amf0ValueType::Number(number) => *number as u8,
255             _ => 0,
256         };
257 
258         let empty_cmd_obj: HashMap<String, Amf0ValueType> = HashMap::new();
259         let _ = match command_object {
260             Amf0ValueType::Object(obj) => obj,
261             // Amf0ValueType::Null =>
262             _ => &empty_cmd_obj,
263         };
264 
265         match cmd_name.as_str() {
266             "_result" => match transaction_id {
267                 define::TRANSACTION_ID_CONNECT => {
268                     log::info!("[C <- S] on_result_connect...");
269                     self.on_result_connect().await?;
270                 }
271                 define::TRANSACTION_ID_CREATE_STREAM => {
272                     log::info!("[C <- S] on_result_create_stream...");
273                     self.on_result_create_stream()?;
274                 }
275                 _ => {}
276             },
277             "_error" => {
278                 self.on_error()?;
279             }
280             "onStatus" => {
281                 match others.remove(0) {
282                     Amf0ValueType::Object(obj) => self.on_status(&obj).await?,
283                     _ => {
284                         return Err(SessionError {
285                             value: SessionErrorValue::Amf0ValueCountNotCorrect,
286                         })
287                     }
288                 };
289             }
290 
291             _ => {}
292         }
293 
294         Ok(())
295     }
296 
297     pub async fn send_connect(&mut self, transaction_id: &f64) -> Result<(), SessionError> {
298         self.send_set_chunk_size().await?;
299 
300         let mut netconnection = NetConnection::new(Arc::clone(&self.io));
301         let mut properties = ConnectProperties::new_none();
302 
303         let url = format!("rtmp://localhost:1935/{app_name}", app_name = self.app_name);
304         properties.app = Some(self.app_name.clone());
305         properties.tc_url = Some(url.clone());
306 
307         match self.client_type {
308             ClientType::Play => {
309                 properties.flash_ver = Some("flashVerFMLE/3.0 (compatible; FMSc/1.0)".to_string());
310                 properties.swf_url = Some(url.clone());
311             }
312             ClientType::Publish => {
313                 properties.fpad = Some(false);
314                 properties.capabilities = Some(15_f64);
315                 properties.audio_codecs = Some(3191_f64);
316                 properties.video_codecs = Some(252_f64);
317                 properties.video_function = Some(1_f64);
318             }
319         }
320 
321         netconnection
322             .write_connect(transaction_id, &properties)
323             .await?;
324 
325         Ok(())
326     }
327 
328     pub async fn send_create_stream(&mut self, transaction_id: &f64) -> Result<(), SessionError> {
329         let mut netconnection = NetConnection::new(Arc::clone(&self.io));
330         netconnection.write_create_stream(transaction_id).await?;
331 
332         Ok(())
333     }
334 
335     pub async fn send_delete_stream(
336         &mut self,
337         transaction_id: &f64,
338         stream_id: &f64,
339     ) -> Result<(), SessionError> {
340         let mut netstream = NetStreamWriter::new(Arc::clone(&self.io));
341         netstream
342             .write_delete_stream(transaction_id, stream_id)
343             .await?;
344 
345         Ok(())
346     }
347 
348     pub async fn send_publish(
349         &mut self,
350         transaction_id: &f64,
351         stream_name: &String,
352         stream_type: &String,
353     ) -> Result<(), SessionError> {
354         let mut netstream = NetStreamWriter::new(Arc::clone(&self.io));
355         netstream
356             .write_publish(transaction_id, stream_name, stream_type)
357             .await?;
358 
359         Ok(())
360     }
361 
362     pub async fn send_play(
363         &mut self,
364         transaction_id: &f64,
365         stream_name: &String,
366         start: &f64,
367         duration: &f64,
368         reset: &bool,
369     ) -> Result<(), SessionError> {
370         let mut netstream = NetStreamWriter::new(Arc::clone(&self.io));
371         netstream
372             .write_play(transaction_id, stream_name, start, duration, reset)
373             .await?;
374 
375         Ok(())
376     }
377 
378     pub async fn send_set_chunk_size(&mut self) -> Result<(), SessionError> {
379         let mut controlmessage =
380             ProtocolControlMessagesWriter::new(AsyncBytesWriter::new(self.io.clone()));
381         controlmessage.write_set_chunk_size(CHUNK_SIZE).await?;
382         Ok(())
383     }
384 
385     pub async fn send_window_acknowledgement_size(
386         &mut self,
387         window_size: u32,
388     ) -> Result<(), SessionError> {
389         let mut controlmessage =
390             ProtocolControlMessagesWriter::new(AsyncBytesWriter::new(self.io.clone()));
391         controlmessage
392             .write_window_acknowledgement_size(window_size)
393             .await?;
394         Ok(())
395     }
396 
397     pub async fn send_set_buffer_length(
398         &mut self,
399         stream_id: u32,
400         ms: u32,
401     ) -> Result<(), SessionError> {
402         let mut eventmessages = EventMessagesWriter::new(AsyncBytesWriter::new(self.io.clone()));
403         eventmessages.write_set_buffer_length(stream_id, ms).await?;
404 
405         Ok(())
406     }
407 
408     pub async fn on_result_connect(&mut self) -> Result<(), SessionError> {
409         let mut controlmessage =
410             ProtocolControlMessagesWriter::new(AsyncBytesWriter::new(self.io.clone()));
411         controlmessage.write_acknowledgement(3107).await?;
412 
413         let mut netstream = NetStreamWriter::new(Arc::clone(&self.io));
414         netstream
415             .write_release_stream(&(define::TRANSACTION_ID_CONNECT as f64), &self.stream_name)
416             .await?;
417         netstream
418             .write_fcpublish(&(define::TRANSACTION_ID_CONNECT as f64), &self.stream_name)
419             .await?;
420 
421         self.state = ClientSessionState::CreateStream;
422 
423         Ok(())
424     }
425 
426     pub fn on_result_create_stream(&mut self) -> Result<(), SessionError> {
427         match self.client_type {
428             ClientType::Play => {
429                 self.state = ClientSessionState::Play;
430             }
431             ClientType::Publish => {
432                 self.state = ClientSessionState::PublishingContent;
433             }
434         }
435         Ok(())
436     }
437 
438     pub fn on_set_chunk_size(&mut self, chunk_size: &mut u32) -> Result<(), SessionError> {
439         self.unpacketizer
440             .update_max_chunk_size(*chunk_size as usize);
441         Ok(())
442     }
443 
444     pub fn on_stream_is_recorded(&mut self, stream_id: &mut u32) -> Result<(), SessionError> {
445         log::trace!("stream is recorded stream_id is {}", stream_id);
446         Ok(())
447     }
448 
449     pub fn on_stream_begin(&mut self, stream_id: &mut u32) -> Result<(), SessionError> {
450         log::trace!("stream is begin stream_id is {}", stream_id);
451         Ok(())
452     }
453 
454     pub async fn on_set_peer_bandwidth(&mut self) -> Result<(), SessionError> {
455         self.send_window_acknowledgement_size(250000).await?;
456         Ok(())
457     }
458 
459     pub fn on_error(&mut self) -> Result<(), SessionError> {
460         Ok(())
461     }
462 
463     pub async fn on_status(
464         &mut self,
465         obj: &HashMap<String, Amf0ValueType>,
466     ) -> Result<(), SessionError> {
467         if let Some(Amf0ValueType::UTF8String(code_info)) = obj.get("code") {
468             match &code_info[..] {
469                 "NetStream.Publish.Start" => {
470                     self.state = ClientSessionState::StartPublish;
471                     //subscribe from local session and publish to remote rtmp server
472                     if let (Some(app_name), Some(stream_name)) =
473                         (&self.sub_app_name, &self.sub_stream_name)
474                     {
475                         self.common
476                             .subscribe_from_channels(
477                                 app_name.clone(),
478                                 stream_name.clone(),
479                                 self.session_id,
480                             )
481                             .await?;
482                     } else {
483                         self.common
484                             .subscribe_from_channels(
485                                 self.app_name.clone(),
486                                 self.stream_name.clone(),
487                                 self.session_id,
488                             )
489                             .await?;
490                     }
491                 }
492                 "NetStream.Publish.Reset" => {}
493 
494                 "NetStream.Play.Start" => {
495                     //pull from remote rtmp server and publish to local session
496                     self.common
497                         .publish_to_channels(
498                             self.app_name.clone(),
499                             self.stream_name.clone(),
500                             self.session_id,
501                         )
502                         .await?
503                 }
504                 _ => {}
505             }
506         }
507         log::trace!("{}", obj.len());
508         Ok(())
509     }
510 
511     pub fn subscribe(&mut self, session: &ClientSession) {
512         self.sub_app_name = Some(session.app_name.clone());
513         self.sub_stream_name = Some(session.stream_name.clone());
514     }
515 }
516