scuffle_amf0/
decode.rs

1use std::borrow::Cow;
2use std::io::{Cursor, Seek, SeekFrom};
3
4use byteorder::{BigEndian, ReadBytesExt};
5use num_traits::FromPrimitive;
6
7use super::{Amf0Marker, Amf0ReadError, Amf0Value};
8
9/// An AMF0 Decoder.
10///
11/// This decoder takes a reference to a byte slice and reads the AMF0 data from
12/// it. All returned objects are references to the original byte slice. Making
13/// it very cheap to use.
14pub struct Amf0Decoder<'a> {
15    cursor: Cursor<&'a [u8]>,
16}
17
18impl<'a> Amf0Decoder<'a> {
19    /// Create a new AMF0 decoder.
20    pub const fn new(buff: &'a [u8]) -> Self {
21        Self {
22            cursor: Cursor::new(buff),
23        }
24    }
25
26    /// Check if the decoder has reached the end of the AMF0 data.
27    pub const fn is_empty(&self) -> bool {
28        self.cursor.get_ref().len() == self.cursor.position() as usize
29    }
30
31    fn read_bytes(&mut self, len: usize) -> Result<&'a [u8], Amf0ReadError> {
32        let pos = self.cursor.position();
33        self.cursor.seek(SeekFrom::Current(len as i64))?;
34        Ok(&self.cursor.get_ref()[pos as usize..pos as usize + len])
35    }
36
37    /// Read all the encoded values from the decoder.
38    pub fn decode_all(&mut self) -> Result<Vec<Amf0Value<'a>>, Amf0ReadError> {
39        let mut results = vec![];
40
41        while !self.is_empty() {
42            results.push(self.decode()?);
43        }
44
45        Ok(results)
46    }
47
48    /// Read the next encoded value from the decoder.
49    pub fn decode(&mut self) -> Result<Amf0Value<'a>, Amf0ReadError> {
50        let marker = self.cursor.read_u8()?;
51        let marker = Amf0Marker::from_u8(marker).ok_or(Amf0ReadError::UnknownMarker(marker))?;
52
53        match marker {
54            Amf0Marker::Number => Ok(Amf0Value::Number(self.read_number()?)),
55            Amf0Marker::Boolean => Ok(Amf0Value::Boolean(self.read_bool()?)),
56            Amf0Marker::String => Ok(Amf0Value::String(self.read_string()?)),
57            Amf0Marker::Object => Ok(Amf0Value::Object(self.read_object()?.into())),
58            Amf0Marker::Null => Ok(Amf0Value::Null),
59            Amf0Marker::EcmaArray => Ok(Amf0Value::Object(self.read_ecma_array()?.into())),
60            Amf0Marker::LongString => Ok(Amf0Value::LongString(self.read_long_string()?)),
61            _ => Err(Amf0ReadError::UnsupportedType(marker)),
62        }
63    }
64
65    /// Read the next encoded value from the decoder and check if it matches the
66    /// specified marker.
67    pub fn decode_with_type(&mut self, specified_marker: Amf0Marker) -> Result<Amf0Value<'a>, Amf0ReadError> {
68        let marker = self.cursor.read_u8()?;
69        self.cursor.seek(SeekFrom::Current(-1))?; // seek back to the original position
70
71        let marker = Amf0Marker::from_u8(marker).ok_or(Amf0ReadError::UnknownMarker(marker))?;
72        if marker != specified_marker {
73            return Err(Amf0ReadError::WrongType(specified_marker, marker));
74        }
75
76        self.decode()
77    }
78
79    fn read_number(&mut self) -> Result<f64, Amf0ReadError> {
80        Ok(self.cursor.read_f64::<BigEndian>()?)
81    }
82
83    fn read_bool(&mut self) -> Result<bool, Amf0ReadError> {
84        Ok(self.cursor.read_u8()? > 0)
85    }
86
87    fn read_string(&mut self) -> Result<Cow<'a, str>, Amf0ReadError> {
88        let l = self.cursor.read_u16::<BigEndian>()?;
89        let bytes = self.read_bytes(l as usize)?;
90
91        Ok(Cow::Borrowed(std::str::from_utf8(bytes)?))
92    }
93
94    fn is_read_object_eof(&mut self) -> Result<bool, Amf0ReadError> {
95        let pos = self.cursor.position();
96        let marker = self.cursor.read_u24::<BigEndian>().map(Amf0Marker::from_u32);
97
98        match marker {
99            Ok(Some(Amf0Marker::ObjectEnd)) => Ok(true),
100            _ => {
101                self.cursor.seek(SeekFrom::Start(pos))?;
102                Ok(false)
103            }
104        }
105    }
106
107    fn read_object(&mut self) -> Result<Vec<(Cow<'a, str>, Amf0Value<'a>)>, Amf0ReadError> {
108        let mut properties = Vec::new();
109
110        loop {
111            let is_eof = self.is_read_object_eof()?;
112
113            if is_eof {
114                break;
115            }
116
117            let key = self.read_string()?;
118            let val = self.decode()?;
119
120            properties.push((key, val));
121        }
122
123        Ok(properties)
124    }
125
126    fn read_ecma_array(&mut self) -> Result<Vec<(Cow<'a, str>, Amf0Value<'a>)>, Amf0ReadError> {
127        let len = self.cursor.read_u32::<BigEndian>()?;
128
129        let mut properties = Vec::new();
130
131        for _ in 0..len {
132            let key = self.read_string()?;
133            let val = self.decode()?;
134            properties.push((key, val));
135        }
136
137        // Sometimes the object end marker is present and sometimes it is not.
138        // If it is there just read it, if not then we are done.
139        self.is_read_object_eof().ok(); // ignore the result
140
141        Ok(properties)
142    }
143
144    fn read_long_string(&mut self) -> Result<Cow<'a, str>, Amf0ReadError> {
145        let l = self.cursor.read_u32::<BigEndian>()?;
146
147        let buff = self.read_bytes(l as usize)?;
148        let val = std::str::from_utf8(buff)?;
149
150        Ok(Cow::Borrowed(val))
151    }
152}
153
154impl<'a> Iterator for Amf0Decoder<'a> {
155    type Item = Result<Amf0Value<'a>, Amf0ReadError>;
156
157    fn next(&mut self) -> Option<Self::Item> {
158        if self.is_empty() {
159            return None;
160        }
161
162        Some(self.decode())
163    }
164}
165
166#[cfg(test)]
167#[cfg_attr(all(test, coverage_nightly), coverage(off))]
168mod tests {
169    use super::*;
170
171    #[test]
172    fn test_reader_bool() {
173        let amf0_bool = vec![0x01, 0x01]; // true
174        let mut amf_reader = Amf0Decoder::new(&amf0_bool);
175        let value = amf_reader.decode_with_type(Amf0Marker::Boolean).unwrap();
176        assert_eq!(value, Amf0Value::Boolean(true));
177    }
178
179    #[test]
180    fn test_reader_number() {
181        let mut amf0_number = vec![0x00];
182        amf0_number.extend_from_slice(&772.161_f64.to_be_bytes());
183
184        let mut amf_reader = Amf0Decoder::new(&amf0_number);
185        let value = amf_reader.decode_with_type(Amf0Marker::Number).unwrap();
186        assert_eq!(value, Amf0Value::Number(772.161));
187    }
188
189    #[test]
190    fn test_reader_string() {
191        let mut amf0_string = vec![0x02, 0x00, 0x0b]; // 11 bytes
192        amf0_string.extend_from_slice(b"Hello World");
193
194        let mut amf_reader = Amf0Decoder::new(&amf0_string);
195        let value = amf_reader.decode_with_type(Amf0Marker::String).unwrap();
196        assert_eq!(value, Amf0Value::String(Cow::Borrowed("Hello World")));
197    }
198
199    #[test]
200    fn test_reader_long_string() {
201        let mut amf0_string = vec![0x0c, 0x00, 0x00, 0x00, 0x0b]; // 11 bytes
202        amf0_string.extend_from_slice(b"Hello World");
203
204        let mut amf_reader = Amf0Decoder::new(&amf0_string);
205        let value = amf_reader.decode_with_type(Amf0Marker::LongString).unwrap();
206        assert_eq!(value, Amf0Value::LongString(Cow::Borrowed("Hello World")));
207    }
208
209    #[test]
210    fn test_reader_object() {
211        let mut amf0_object = vec![0x03, 0x00, 0x04]; // 1 property with 4 bytes
212        amf0_object.extend_from_slice(b"test");
213        amf0_object.extend_from_slice(&[0x05]); // null
214        amf0_object.extend_from_slice(&[0x00, 0x00, 0x09]); // object end (0x00 0x00 0x09)
215
216        let mut amf_reader = Amf0Decoder::new(&amf0_object);
217        let value = amf_reader.decode_with_type(Amf0Marker::Object).unwrap();
218
219        assert_eq!(value, Amf0Value::Object(vec![("test".into(), Amf0Value::Null)].into()));
220    }
221
222    #[test]
223    fn test_reader_ecma_array() {
224        let mut amf0_object = vec![0x08, 0x00, 0x00, 0x00, 0x01]; // 1 property
225        amf0_object.extend_from_slice(&[0x00, 0x04]); // 4 bytes
226        amf0_object.extend_from_slice(b"test");
227        amf0_object.extend_from_slice(&[0x05]); // null
228
229        let mut amf_reader = Amf0Decoder::new(&amf0_object);
230        let value = amf_reader.decode_with_type(Amf0Marker::EcmaArray).unwrap();
231
232        assert_eq!(value, Amf0Value::Object(vec![("test".into(), Amf0Value::Null)].into()));
233    }
234
235    #[test]
236    fn test_reader_multi_value() {
237        let mut amf0_multi = vec![0x00];
238        amf0_multi.extend_from_slice(&772.161_f64.to_be_bytes());
239        amf0_multi.extend_from_slice(&[0x01, 0x01]); // true
240        amf0_multi.extend_from_slice(&[0x02, 0x00, 0x0b]); // 11 bytes
241        amf0_multi.extend_from_slice(b"Hello World");
242        amf0_multi.extend_from_slice(&[0x03, 0x00, 0x04]); // 1 property with 4 bytes
243        amf0_multi.extend_from_slice(b"test");
244        amf0_multi.extend_from_slice(&[0x05]); // null
245        amf0_multi.extend_from_slice(&[0x00, 0x00, 0x09]); // object end (0x00 0x00 0x09)
246
247        let mut amf_reader = Amf0Decoder::new(&amf0_multi);
248        let values = amf_reader.decode_all().unwrap();
249
250        assert_eq!(values.len(), 4);
251
252        assert_eq!(values[0], Amf0Value::Number(772.161));
253        assert_eq!(values[1], Amf0Value::Boolean(true));
254        assert_eq!(values[2], Amf0Value::String(Cow::Borrowed("Hello World")));
255        assert_eq!(values[3], Amf0Value::Object(vec![("test".into(), Amf0Value::Null)].into()));
256    }
257
258    #[test]
259    fn test_reader_iterator() {
260        let mut amf0_multi = vec![0x00];
261        amf0_multi.extend_from_slice(&772.161_f64.to_be_bytes());
262        amf0_multi.extend_from_slice(&[0x01, 0x01]); // true
263        amf0_multi.extend_from_slice(&[0x02, 0x00, 0x0b]); // 11 bytes
264        amf0_multi.extend_from_slice(b"Hello World");
265
266        let amf_reader = Amf0Decoder::new(&amf0_multi);
267        let values = amf_reader.collect::<Result<Vec<_>, _>>().unwrap();
268
269        assert_eq!(values.len(), 3);
270
271        assert_eq!(values[0], Amf0Value::Number(772.161));
272        assert_eq!(values[1], Amf0Value::Boolean(true));
273        assert_eq!(values[2], Amf0Value::String(Cow::Borrowed("Hello World")));
274    }
275
276    #[test]
277    fn test_reader_invalid_marker() {
278        let amf0_unsupported_marker = vec![Amf0Marker::Unsupported as u8];
279        let mut amf_reader = Amf0Decoder::new(&amf0_unsupported_marker);
280        let result = amf_reader.decode();
281
282        assert!(matches!(result, Err(Amf0ReadError::UnsupportedType(Amf0Marker::Unsupported))));
283    }
284}