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