rtmp/handshake/
digest.rs

1use {
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
13pub struct DigestProcessor {
14    reader: BytesReader,
15    key: BytesMut,
16}
17
18impl DigestProcessor {
19    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*/
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
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
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
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     */
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    }
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
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}
rtmp/handshake/digest.rs

rtmp/handshake/
digest.rs
