1 use byteorder::{BigEndian, LittleEndian}; 2 3 use super::errors::PackError; 4 use super::{ 5 chunk::{ChunkBasicHeader, ChunkHeader, ChunkInfo, ChunkMessageHeader}, 6 define::CHUNK_SIZE, 7 }; 8 use netio::bytes_writer::AsyncBytesWriter; 9 use std::collections::HashMap; 10 11 use std::sync::Arc; 12 13 use netio::netio::NetworkIO; 14 use tokio::sync::Mutex; 15 16 #[derive(Eq, PartialEq, Debug)] 17 pub enum PackResult { 18 Success, 19 NotEnoughBytes, 20 } 21 22 pub struct ChunkPacketizer { 23 csid_2_chunk_header: HashMap<u32, ChunkHeader>, 24 //https://doc.rust-lang.org/stable/rust-by-example/scope/lifetime/fn.html 25 //https://zhuanlan.zhihu.com/p/165976086 26 //chunk_info: ChunkInfo, 27 max_chunk_size: usize, 28 //bytes: Cursor<Vec<u8>>, 29 writer: AsyncBytesWriter, 30 } 31 32 impl ChunkPacketizer { 33 pub fn new(io: Arc<Mutex<NetworkIO>>) -> Self { 34 Self { 35 csid_2_chunk_header: HashMap::new(), 36 //chunk_info: ChunkInfo::new(), 37 writer: AsyncBytesWriter::new(io), 38 max_chunk_size: CHUNK_SIZE as usize, 39 } 40 } 41 fn zip_chunk_header(&mut self, chunk_info: &mut ChunkInfo) -> Result<PackResult, PackError> { 42 chunk_info.basic_header.format = 0; 43 44 let pre_header = self 45 .csid_2_chunk_header 46 .get_mut(&chunk_info.basic_header.chunk_stream_id); 47 48 match pre_header { 49 Some(val) => { 50 let cur_msg_header = &mut chunk_info.message_header; 51 let pre_msg_header = &val.message_header; 52 53 if cur_msg_header.msg_streamd_id == pre_msg_header.msg_streamd_id { 54 chunk_info.basic_header.format = 1; 55 cur_msg_header.timestamp -= pre_msg_header.timestamp; 56 57 if cur_msg_header.msg_type_id == pre_msg_header.msg_type_id 58 && cur_msg_header.msg_length == pre_msg_header.msg_length 59 { 60 chunk_info.basic_header.format = 2; 61 if chunk_info.message_header.timestamp == pre_msg_header.timestamp { 62 chunk_info.basic_header.format = 3; 63 } 64 } 65 } 66 } 67 68 None => {} 69 } 70 71 Ok(PackResult::Success) 72 } 73 74 fn write_basic_header(&mut self, fmt: u8, csid: u32) -> Result<(), PackError> { 75 if csid >= 64 + 255 { 76 self.writer.write_u8(fmt << 6 | 1)?; 77 self.writer.write_u16::<BigEndian>((csid - 64) as u16)?; 78 } else if csid >= 64 { 79 self.writer.write_u8(fmt << 6 | 0)?; 80 self.writer.write_u8((csid - 64) as u8)?; 81 } else { 82 self.writer.write_u8(fmt << 6 | csid as u8)?; 83 } 84 85 Ok(()) 86 } 87 88 fn write_message_header( 89 &mut self, 90 basic_header: &ChunkBasicHeader, 91 message_header: &ChunkMessageHeader, 92 ) -> Result<(), PackError> { 93 let timestamp = if message_header.timestamp >= 0xFFFFFF { 94 0xFFFFFF 95 } else { 96 message_header.timestamp 97 }; 98 99 match basic_header.format { 100 0 => { 101 self.writer.write_u24::<BigEndian>(timestamp)?; 102 self.writer 103 .write_u24::<BigEndian>(message_header.msg_length)?; 104 self.writer.write_u8(message_header.msg_type_id)?; 105 self.writer 106 .write_u32::<LittleEndian>(message_header.msg_streamd_id)?; 107 } 108 1 => { 109 self.writer.write_u24::<BigEndian>(timestamp)?; 110 self.writer 111 .write_u24::<BigEndian>(message_header.msg_length)?; 112 } 113 2 => { 114 self.writer.write_u24::<BigEndian>(timestamp)?; 115 } 116 _ => {} 117 } 118 119 Ok(()) 120 } 121 122 fn write_extened_timestamp(&mut self, timestamp: u32) -> Result<(), PackError> { 123 self.writer.write_u32::<BigEndian>(timestamp)?; 124 125 Ok(()) 126 } 127 128 pub async fn write_chunk(&mut self, chunk_info: &mut ChunkInfo) -> Result<(), PackError> { 129 self.zip_chunk_header(chunk_info)?; 130 131 let mut whole_payload_size = chunk_info.payload.len(); 132 133 self.write_basic_header( 134 chunk_info.basic_header.format, 135 chunk_info.basic_header.chunk_stream_id, 136 )?; 137 self.write_message_header(&chunk_info.basic_header, &chunk_info.message_header)?; 138 139 if chunk_info.message_header.is_extended_timestamp { 140 self.write_extened_timestamp(chunk_info.message_header.timestamp)?; 141 } 142 143 let mut cur_payload_size: usize; 144 145 while whole_payload_size > 0 { 146 cur_payload_size = if whole_payload_size > self.max_chunk_size { 147 self.max_chunk_size 148 } else { 149 whole_payload_size 150 }; 151 152 let payload_bytes = chunk_info.payload.split_to(cur_payload_size); 153 self.writer.write(&payload_bytes[0..])?; 154 155 whole_payload_size -= cur_payload_size; 156 157 if whole_payload_size > 0 { 158 self.write_basic_header(3, chunk_info.basic_header.chunk_stream_id)?; 159 if chunk_info.message_header.is_extended_timestamp { 160 self.write_extened_timestamp(chunk_info.message_header.timestamp)?; 161 } 162 } 163 } 164 self.writer.flush().await?; 165 166 Ok(()) 167 } 168 } 169