1 use { 2 super::{ 3 define, 4 define::SchemaVersion, 5 errors::{DigestError, DigestErrorValue}, 6 }, 7 bytes::BytesMut, 8 bytesio::bytes_reader::BytesReader, 9 hmac::{Hmac, Mac, NewMac}, 10 sha2::Sha256, 11 }; 12 13 pub struct DigestProcessor { 14 reader: BytesReader, 15 key: BytesMut, 16 } 17 18 impl DigestProcessor { new(data: BytesMut, key: BytesMut) -> Self19 pub fn new(data: BytesMut, key: BytesMut) -> Self { 20 Self { 21 reader: BytesReader::new(data), 22 key, 23 } 24 } 25 26 /* return validate digest and schema version*/ read_digest(&mut self) -> Result<(BytesMut, SchemaVersion), DigestError>27 pub fn read_digest(&mut self) -> Result<(BytesMut, SchemaVersion), DigestError> { 28 if let Ok(digest) = self.generate_and_validate(SchemaVersion::Schema0) { 29 return Ok((digest, SchemaVersion::Schema0)); 30 } 31 32 let digest = self.generate_and_validate(SchemaVersion::Schema1)?; 33 Ok((digest, SchemaVersion::Schema1)) 34 } 35 generate_and_fill_digest(&mut self) -> Result<Vec<u8>, DigestError>36 pub fn generate_and_fill_digest(&mut self) -> Result<Vec<u8>, DigestError> { 37 let (left_part, _, right_part) = self.cook_raw_message(SchemaVersion::Schema0)?; 38 let raw_message = [left_part.clone(), right_part.clone()].concat(); 39 let computed_digest = self.make_digest(raw_message)?; 40 41 let result = [left_part, computed_digest, right_part].concat(); 42 43 Ok(result) 44 } 45 generate_digest(&mut self) -> Result<BytesMut, DigestError>46 pub fn generate_digest(&mut self) -> Result<BytesMut, DigestError> { 47 let (left_part, _, right_part) = self.cook_raw_message(SchemaVersion::Schema0)?; 48 let raw_message = [left_part, right_part].concat(); 49 let digest = self.make_digest(raw_message)?; 50 51 Ok(digest) 52 } 53 find_digest_offset(&mut self, version: SchemaVersion) -> Result<usize, DigestError>54 fn find_digest_offset(&mut self, version: SchemaVersion) -> Result<usize, DigestError> { 55 let mut digest_offset: usize = 0; 56 57 match version { 58 SchemaVersion::Schema0 => { 59 digest_offset += self.reader.get(772)? as usize; 60 digest_offset += self.reader.get(773)? as usize; 61 digest_offset += self.reader.get(774)? as usize; 62 digest_offset += self.reader.get(775)? as usize; 63 64 digest_offset %= 728; 65 digest_offset += 776; 66 } 67 SchemaVersion::Schema1 => { 68 digest_offset += self.reader.get(8)? as usize; 69 digest_offset += self.reader.get(9)? as usize; 70 digest_offset += self.reader.get(10)? as usize; 71 digest_offset += self.reader.get(11)? as usize; 72 73 digest_offset %= 728; 74 digest_offset += 12; 75 } 76 SchemaVersion::Unknown => { 77 return Err(DigestError { 78 value: DigestErrorValue::UnknowSchema, 79 }); 80 } 81 } 82 83 Ok(digest_offset) 84 } 85 /* 86 +-----------------------------------------------------------+ 87 | 764 bytes | 88 * +--------------+-----------------------+--------------------+ 89 * | left part | digest data (32 bytes)| right part | 90 * +--------------+-----------------------+--------------------+ 91 * | 92 / 93 digest offset 94 pice together the left part and right part to get the raw message. 95 */ cook_raw_message( &mut self, version: SchemaVersion, ) -> Result<(BytesMut, BytesMut, BytesMut), DigestError>96 fn cook_raw_message( 97 &mut self, 98 version: SchemaVersion, 99 ) -> Result<(BytesMut, BytesMut, BytesMut), DigestError> { 100 let digest_offset: usize = self.find_digest_offset(version)?; 101 102 let mut new_reader = BytesReader::new(self.reader.get_remaining_bytes()); 103 104 let left_part = new_reader.read_bytes(digest_offset)?; 105 let digest_data = new_reader.read_bytes(define::RTMP_DIGEST_LENGTH)?; 106 let right_part = new_reader.extract_remaining_bytes(); 107 108 Ok((left_part, digest_data, right_part)) 109 } make_digest(&mut self, raw_message: Vec<u8>) -> Result<BytesMut, DigestError>110 pub fn make_digest(&mut self, raw_message: Vec<u8>) -> Result<BytesMut, DigestError> { 111 let mut mac = Hmac::<Sha256>::new_from_slice(&self.key[..]).unwrap(); 112 mac.update(&raw_message); 113 let result = mac.finalize().into_bytes(); 114 115 if result.len() != define::RTMP_DIGEST_LENGTH { 116 return Err(DigestError { 117 value: DigestErrorValue::DigestLengthNotCorrect, 118 }); 119 } 120 121 let mut rv = BytesMut::new(); 122 rv.extend_from_slice(result.as_slice()); 123 124 Ok(rv) 125 } 126 generate_and_validate(&mut self, version: SchemaVersion) -> Result<BytesMut, DigestError>127 fn generate_and_validate(&mut self, version: SchemaVersion) -> Result<BytesMut, DigestError> { 128 let (left_part, digest_data, right_part) = self.cook_raw_message(version)?; 129 let raw_message = [left_part, right_part].concat(); 130 131 let computed_digest = self.make_digest(raw_message)?; 132 133 if digest_data == computed_digest { 134 return Ok(digest_data); 135 } 136 137 Err(DigestError { 138 value: DigestErrorValue::CannotGenerate, 139 }) 140 } 141 } 142