1 use {
2     super::{
3         define,
4         errors::{SessionError, SessionErrorValue},
5     },
6     crate::{
7         amf0::Amf0ValueType,
8         chunk::{
9             define::{chunk_type, csid_type, CHUNK_SIZE},
10             packetizer::ChunkPacketizer,
11             unpacketizer::{ChunkUnpacketizer, UnpackResult},
12             ChunkInfo,
13         },
14         handshake::handshake::{ClientHandshakeState, SimpleHandshakeClient},
15         messages::{
16             define::{msg_type_id, RtmpMessageData},
17             parser::MessageParser,
18         },
19         netconnection::commands::{ConnectProperties, NetConnection},
20         netstream::writer::NetStreamWriter,
21         protocol_control_messages::writer::ProtocolControlMessagesWriter,
22         user_control_messages::writer::EventMessagesWriter,
23     },
24     bytes::BytesMut,
25     netio::{
26         bytes_writer::{AsyncBytesWriter, BytesWriter},
27         netio::NetworkIO,
28     },
29     std::{collections::HashMap, sync::Arc, time::Duration},
30     tokio::{net::TcpStream, sync::Mutex},
31 };
32 
33 enum ClientSessionState {
34     Handshake,
35     Connect,
36     CreateStream,
37     Play,
38     PublishingContent,
39 }
40 
41 enum ClientSessionPlayState {
42     Handshake,
43     Connect,
44     CreateStream,
45     Play,
46 }
47 
48 enum ClientSessionPublishState {
49     Handshake,
50     Connect,
51     CreateStream,
52     PublishingContent,
53 }
54 
55 enum ClientType {
56     Play,
57     Publish,
58 }
59 pub struct ClientSession {
60     packetizer: ChunkPacketizer,
61     unpacketizer: ChunkUnpacketizer,
62     handshaker: SimpleHandshakeClient,
63     io: Arc<Mutex<NetworkIO>>,
64 
65     play_state: ClientSessionPlayState,
66     publish_state: ClientSessionPublishState,
67     state: ClientSessionState,
68     client_type: ClientType,
69     stream_name: String,
70 }
71 
72 impl ClientSession {
73     fn new(
74         stream: TcpStream,
75         timeout: Duration,
76         client_type: ClientType,
77         stream_name: String,
78     ) -> Self {
79         let net_io = Arc::new(Mutex::new(NetworkIO::new(stream, timeout)));
80 
81         // let reader = BytesReader::new(BytesMut::new());
82 
83         Self {
84             io: Arc::clone(&net_io),
85 
86             packetizer: ChunkPacketizer::new(Arc::clone(&net_io)),
87             unpacketizer: ChunkUnpacketizer::new(),
88             handshaker: SimpleHandshakeClient::new(Arc::clone(&net_io)),
89 
90             play_state: ClientSessionPlayState::Handshake,
91             publish_state: ClientSessionPublishState::Handshake,
92             state: ClientSessionState::Handshake,
93             client_type: client_type,
94             stream_name: stream_name,
95         }
96     }
97 
98     pub async fn run(&mut self) -> Result<(), SessionError> {
99         loop {
100             match self.state {
101                 ClientSessionState::Handshake => {
102                     self.handshake().await?;
103                 }
104                 ClientSessionState::Connect => {
105                     self.send_connect(&(define::TRANSACTION_ID_CONNECT as f64))
106                         .await?;
107                 }
108                 ClientSessionState::CreateStream => {
109                     self.send_create_stream(&(define::TRANSACTION_ID_CREATE_STREAM as f64))
110                         .await?;
111                 }
112                 ClientSessionState::Play => {
113                     self.send_play(&0.0, &self.stream_name.clone(), &0.0, &0.0, &false)
114                         .await?;
115                 }
116                 ClientSessionState::PublishingContent => {
117                     self.send_publish(&0.0, &self.stream_name.clone(), &"live".to_string())
118                         .await?;
119                 }
120             }
121 
122             let data = self.io.lock().await.read().await?;
123             self.unpacketizer.extend_data(&data[..]);
124             let result = self.unpacketizer.read_chunk()?;
125 
126             match result {
127                 UnpackResult::ChunkInfo(chunk_info) => {
128                     let mut message_parser = MessageParser::new(chunk_info);
129                     let mut msg = message_parser.parse()?;
130 
131                     self.process_messages(&mut msg).await?;
132                 }
133                 _ => {}
134             }
135         }
136 
137         // Ok(())
138     }
139 
140     async fn handshake(&mut self) -> Result<(), SessionError> {
141         loop {
142             self.handshaker.handshake().await?;
143             if self.handshaker.state == ClientHandshakeState::Finish {
144                 break;
145             }
146 
147             let data = self.io.lock().await.read().await?;
148             self.handshaker.extend_data(&data[..]);
149         }
150         self.state = ClientSessionState::Connect;
151 
152         Ok(())
153     }
154 
155     pub async fn process_messages(
156         &mut self,
157         msg: &mut RtmpMessageData,
158     ) -> Result<(), SessionError> {
159         match msg {
160             RtmpMessageData::Amf0Command {
161                 command_name,
162                 transaction_id,
163                 command_object,
164                 others,
165             } => self.process_amf0_command_message(
166                 command_name,
167                 transaction_id,
168                 command_object,
169                 others,
170             )?,
171             RtmpMessageData::SetPeerBandwidth { properties } => {
172                 self.on_set_peer_bandwidth().await?
173             }
174             RtmpMessageData::SetChunkSize { chunk_size } => self.on_set_chunk_size(chunk_size)?,
175             RtmpMessageData::AudioData { data } => {}
176             RtmpMessageData::VideoData { data } => {}
177 
178             _ => {}
179         }
180         Ok(())
181     }
182 
183     pub fn process_amf0_command_message(
184         &mut self,
185         command_name: &Amf0ValueType,
186         transaction_id: &Amf0ValueType,
187         command_object: &Amf0ValueType,
188         others: &mut Vec<Amf0ValueType>,
189     ) -> Result<(), SessionError> {
190         let empty_cmd_name = &String::new();
191         let cmd_name = match command_name {
192             Amf0ValueType::UTF8String(str) => str,
193             _ => empty_cmd_name,
194         };
195 
196         let transaction_id = match transaction_id {
197             Amf0ValueType::Number(number) => number.clone() as u8,
198             _ => 0,
199         };
200 
201         let empty_cmd_obj: HashMap<String, Amf0ValueType> = HashMap::new();
202         let obj = match command_object {
203             Amf0ValueType::Object(obj) => obj,
204             // Amf0ValueType::Null =>
205             _ => &empty_cmd_obj,
206         };
207 
208         match cmd_name.as_str() {
209             "_reslut" => match transaction_id {
210                 define::TRANSACTION_ID_CONNECT => {
211                     self.on_result_connect()?;
212                 }
213                 define::TRANSACTION_ID_CREATE_STREAM => {
214                     self.on_result_create_stream()?;
215                 }
216                 _ => {}
217             },
218             "_error" => {
219                 self.on_error()?;
220             }
221             "onStatus" => {
222                 match others.remove(0) {
223                     Amf0ValueType::Object(obj) => self.on_status(&obj),
224                     _ => Err(SessionError {
225                         value: SessionErrorValue::Amf0ValueCountNotCorrect,
226                     }),
227                 };
228             }
229 
230             _ => {}
231         }
232 
233         Ok(())
234     }
235 
236     pub async fn send_connect(&mut self, transaction_id: &f64) -> Result<(), SessionError> {
237         let app_name = String::from("app");
238         let properties = ConnectProperties::new(app_name);
239 
240         let mut netconnection = NetConnection::new(BytesWriter::new());
241         let data = netconnection.connect(transaction_id, &properties)?;
242 
243         let mut chunk_info = ChunkInfo::new(
244             csid_type::COMMAND_AMF0_AMF3,
245             chunk_type::TYPE_0,
246             0,
247             data.len() as u32,
248             msg_type_id::COMMAND_AMF0,
249             0,
250             data,
251         );
252 
253         self.packetizer.write_chunk(&mut chunk_info).await?;
254         Ok(())
255     }
256 
257     pub async fn send_create_stream(&mut self, transaction_id: &f64) -> Result<(), SessionError> {
258         let mut netconnection = NetConnection::new(BytesWriter::new());
259         let data = netconnection.create_stream(transaction_id)?;
260 
261         let mut chunk_info = ChunkInfo::new(
262             csid_type::COMMAND_AMF0_AMF3,
263             chunk_type::TYPE_0,
264             0,
265             data.len() as u32,
266             msg_type_id::COMMAND_AMF0,
267             0,
268             data,
269         );
270 
271         self.packetizer.write_chunk(&mut chunk_info).await?;
272 
273         Ok(())
274     }
275 
276     pub async fn send_delete_stream(
277         &mut self,
278         transaction_id: &f64,
279         stream_id: &f64,
280     ) -> Result<(), SessionError> {
281         let mut netstream = NetStreamWriter::new(BytesWriter::new(), Arc::clone(&self.io));
282         netstream.delete_stream(transaction_id, stream_id).await?;
283 
284         Ok(())
285     }
286 
287     pub async fn send_publish(
288         &mut self,
289         transaction_id: &f64,
290         stream_name: &String,
291         stream_type: &String,
292     ) -> Result<(), SessionError> {
293         let mut netstream = NetStreamWriter::new(BytesWriter::new(), Arc::clone(&self.io));
294         netstream
295             .publish(transaction_id, stream_name, stream_type)
296             .await?;
297 
298         Ok(())
299     }
300 
301     pub async fn send_play(
302         &mut self,
303         transaction_id: &f64,
304         stream_name: &String,
305         start: &f64,
306         duration: &f64,
307         reset: &bool,
308     ) -> Result<(), SessionError> {
309         let mut netstream = NetStreamWriter::new(BytesWriter::new(), Arc::clone(&self.io));
310         netstream
311             .play(transaction_id, stream_name, start, duration, reset)
312             .await?;
313 
314         Ok(())
315     }
316 
317     pub async fn send_set_chunk_size(&mut self) -> Result<(), SessionError> {
318         let mut controlmessage =
319             ProtocolControlMessagesWriter::new(AsyncBytesWriter::new(self.io.clone()));
320         controlmessage.write_set_chunk_size(CHUNK_SIZE).await?;
321         Ok(())
322     }
323 
324     pub async fn send_window_acknowledgement_size(
325         &mut self,
326         window_size: u32,
327     ) -> Result<(), SessionError> {
328         let mut controlmessage =
329             ProtocolControlMessagesWriter::new(AsyncBytesWriter::new(self.io.clone()));
330         controlmessage
331             .write_window_acknowledgement_size(window_size)
332             .await?;
333         Ok(())
334     }
335 
336     pub async fn send_set_buffer_length(
337         &mut self,
338         stream_id: u32,
339         ms: u32,
340     ) -> Result<(), SessionError> {
341         let mut eventmessages = EventMessagesWriter::new(AsyncBytesWriter::new(self.io.clone()));
342         eventmessages.write_set_buffer_length(stream_id, ms).await?;
343 
344         Ok(())
345     }
346 
347     pub async fn send_audio(&mut self, data: BytesMut) -> Result<(), SessionError> {
348         let mut chunk_info = ChunkInfo::new(
349             csid_type::AUDIO,
350             chunk_type::TYPE_0,
351             0,
352             data.len() as u32,
353             msg_type_id::AUDIO,
354             0,
355             data,
356         );
357 
358         self.packetizer.write_chunk(&mut chunk_info).await?;
359 
360         Ok(())
361     }
362 
363     pub async fn send_video(&mut self, data: BytesMut) -> Result<(), SessionError> {
364         let mut chunk_info = ChunkInfo::new(
365             csid_type::VIDEO,
366             chunk_type::TYPE_0,
367             0,
368             data.len() as u32,
369             msg_type_id::VIDEO,
370             0,
371             data,
372         );
373 
374         self.packetizer.write_chunk(&mut chunk_info).await?;
375 
376         Ok(())
377     }
378 
379     pub fn on_result_connect(&mut self) -> Result<(), SessionError> {
380         self.state = ClientSessionState::CreateStream;
381         Ok(())
382     }
383 
384     pub fn on_result_create_stream(&mut self) -> Result<(), SessionError> {
385         match self.client_type {
386             ClientType::Play => {
387                 self.state = ClientSessionState::Play;
388             }
389             ClientType::Publish => {
390                 self.state = ClientSessionState::PublishingContent;
391             }
392         }
393         Ok(())
394     }
395 
396     pub fn on_set_chunk_size(&mut self, chunk_size: &mut u32) -> Result<(), SessionError> {
397         self.unpacketizer
398             .update_max_chunk_size(chunk_size.clone() as usize);
399         Ok(())
400     }
401 
402     pub async fn on_set_peer_bandwidth(&mut self) -> Result<(), SessionError> {
403         self.send_window_acknowledgement_size(250000).await?;
404         Ok(())
405     }
406     pub fn on_error(&mut self) -> Result<(), SessionError> {
407         Ok(())
408     }
409 
410     pub fn on_status(&mut self, obj: &HashMap<String, Amf0ValueType>) -> Result<(), SessionError> {
411         Ok(())
412     }
413 }
414