1 // Copyright 2024 Google LLC 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 use crate::internal_utils::*; 16 use crate::parser::mp4box::BoxSize; 17 18 #[derive(Debug)] 19 pub struct IBitStream<'a> { 20 pub data: &'a [u8], 21 pub bit_offset: usize, 22 } 23 24 impl IBitStream<'_> { read_bit(&mut self) -> AvifResult<u8>25 fn read_bit(&mut self) -> AvifResult<u8> { 26 let byte_offset = self.bit_offset / 8; 27 if byte_offset >= self.data.len() { 28 return Err(AvifError::BmffParseFailed("Not enough bits".into())); 29 } 30 let byte = self.data[byte_offset]; 31 let shift = 7 - (self.bit_offset % 8); 32 self.bit_offset += 1; 33 Ok((byte >> shift) & 0x01) 34 } 35 read(&mut self, n: usize) -> AvifResult<u32>36 pub(crate) fn read(&mut self, n: usize) -> AvifResult<u32> { 37 assert!(n <= 32); 38 let mut value: u32 = 0; 39 for _i in 0..n { 40 value <<= 1; 41 value |= self.read_bit()? as u32; 42 } 43 Ok(value) 44 } 45 read_bool(&mut self) -> AvifResult<bool>46 pub(crate) fn read_bool(&mut self) -> AvifResult<bool> { 47 let bit = self.read_bit()?; 48 Ok(bit == 1) 49 } 50 skip(&mut self, n: usize) -> AvifResult<()>51 pub(crate) fn skip(&mut self, n: usize) -> AvifResult<()> { 52 if checked_add!(self.bit_offset, n)? > checked_mul!(self.data.len(), 8)? { 53 return Err(AvifError::BmffParseFailed("Not enough bytes".into())); 54 } 55 self.bit_offset += n; 56 Ok(()) 57 } 58 skip_uvlc(&mut self) -> AvifResult<()>59 pub(crate) fn skip_uvlc(&mut self) -> AvifResult<()> { 60 // See the section 4.10.3. uvlc() of the AV1 specification. 61 let mut leading_zeros = 0u128; // leadingZeros 62 while !self.read_bool()? { 63 leading_zeros += 1; 64 } 65 if leading_zeros < 32 { 66 self.skip(leading_zeros as usize)?; // f(leadingZeros) value; 67 } 68 Ok(()) 69 } 70 remaining_bits(&self) -> AvifResult<usize>71 pub(crate) fn remaining_bits(&self) -> AvifResult<usize> { 72 checked_sub!(checked_mul!(self.data.len(), 8)?, self.bit_offset) 73 } 74 } 75 76 #[derive(Debug)] 77 pub struct IStream<'a> { 78 pub data: &'a [u8], 79 pub offset: usize, 80 } 81 82 impl IStream<'_> { create(data: &[u8]) -> IStream83 pub(crate) fn create(data: &[u8]) -> IStream { 84 IStream { data, offset: 0 } 85 } 86 check(&self, size: usize) -> AvifResult<()>87 fn check(&self, size: usize) -> AvifResult<()> { 88 if self.bytes_left()? < size { 89 return Err(AvifError::BmffParseFailed("".into())); 90 } 91 Ok(()) 92 } 93 sub_stream(&mut self, size: &BoxSize) -> AvifResult<IStream>94 pub(crate) fn sub_stream(&mut self, size: &BoxSize) -> AvifResult<IStream> { 95 let offset = self.offset; 96 checked_incr!( 97 self.offset, 98 match size { 99 BoxSize::FixedSize(size) => { 100 self.check(*size)?; 101 *size 102 } 103 BoxSize::UntilEndOfStream => self.bytes_left()?, 104 } 105 ); 106 Ok(IStream { 107 data: &self.data[offset..self.offset], 108 offset: 0, 109 }) 110 } 111 sub_bit_stream(&mut self, size: usize) -> AvifResult<IBitStream>112 pub(crate) fn sub_bit_stream(&mut self, size: usize) -> AvifResult<IBitStream> { 113 self.check(size)?; 114 let offset = self.offset; 115 checked_incr!(self.offset, size); 116 Ok(IBitStream { 117 data: &self.data[offset..self.offset], 118 bit_offset: 0, 119 }) 120 } 121 bytes_left(&self) -> AvifResult<usize>122 pub(crate) fn bytes_left(&self) -> AvifResult<usize> { 123 if self.data.len() < self.offset { 124 return Err(AvifError::UnknownError("".into())); 125 } 126 Ok(self.data.len() - self.offset) 127 } 128 has_bytes_left(&self) -> AvifResult<bool>129 pub(crate) fn has_bytes_left(&self) -> AvifResult<bool> { 130 Ok(self.bytes_left()? > 0) 131 } 132 get_slice(&mut self, size: usize) -> AvifResult<&[u8]>133 pub(crate) fn get_slice(&mut self, size: usize) -> AvifResult<&[u8]> { 134 self.check(size)?; 135 let offset_start = self.offset; 136 checked_incr!(self.offset, size); 137 Ok(&self.data[offset_start..offset_start + size]) 138 } 139 get_immutable_vec(&self, size: usize) -> AvifResult<Vec<u8>>140 pub(crate) fn get_immutable_vec(&self, size: usize) -> AvifResult<Vec<u8>> { 141 self.check(size)?; 142 Ok(self.data[self.offset..self.offset + size].to_vec()) 143 } 144 get_vec(&mut self, size: usize) -> AvifResult<Vec<u8>>145 fn get_vec(&mut self, size: usize) -> AvifResult<Vec<u8>> { 146 Ok(self.get_slice(size)?.to_vec()) 147 } 148 read_u8(&mut self) -> AvifResult<u8>149 pub(crate) fn read_u8(&mut self) -> AvifResult<u8> { 150 self.check(1)?; 151 let value = self.data[self.offset]; 152 checked_incr!(self.offset, 1); 153 Ok(value) 154 } 155 read_u16(&mut self) -> AvifResult<u16>156 pub(crate) fn read_u16(&mut self) -> AvifResult<u16> { 157 Ok(u16::from_be_bytes(self.get_slice(2)?.try_into().unwrap())) 158 } 159 read_u24(&mut self) -> AvifResult<u32>160 pub(crate) fn read_u24(&mut self) -> AvifResult<u32> { 161 Ok(self.read_uxx(3)? as u32) 162 } 163 read_u32(&mut self) -> AvifResult<u32>164 pub(crate) fn read_u32(&mut self) -> AvifResult<u32> { 165 Ok(u32::from_be_bytes(self.get_slice(4)?.try_into().unwrap())) 166 } 167 read_u64(&mut self) -> AvifResult<u64>168 pub(crate) fn read_u64(&mut self) -> AvifResult<u64> { 169 Ok(u64::from_be_bytes(self.get_slice(8)?.try_into().unwrap())) 170 } 171 read_i32(&mut self) -> AvifResult<i32>172 pub(crate) fn read_i32(&mut self) -> AvifResult<i32> { 173 Ok(self.read_u32()? as i32) 174 } 175 read_i16(&mut self) -> AvifResult<i16>176 pub(crate) fn read_i16(&mut self) -> AvifResult<i16> { 177 Ok(self.read_u16()? as i16) 178 } 179 skip_u32(&mut self) -> AvifResult<()>180 pub(crate) fn skip_u32(&mut self) -> AvifResult<()> { 181 self.skip(4) 182 } 183 skip_u64(&mut self) -> AvifResult<()>184 pub(crate) fn skip_u64(&mut self) -> AvifResult<()> { 185 self.skip(8) 186 } 187 read_fraction(&mut self) -> AvifResult<Fraction>188 pub(crate) fn read_fraction(&mut self) -> AvifResult<Fraction> { 189 Ok(Fraction(self.read_i32()?, self.read_u32()?)) 190 } 191 read_ufraction(&mut self) -> AvifResult<UFraction>192 pub(crate) fn read_ufraction(&mut self) -> AvifResult<UFraction> { 193 Ok(UFraction(self.read_u32()?, self.read_u32()?)) 194 } 195 196 // Reads size characters of a non-null-terminated string. read_string(&mut self, size: usize) -> AvifResult<String>197 pub(crate) fn read_string(&mut self, size: usize) -> AvifResult<String> { 198 Ok(String::from_utf8(self.get_vec(size)?).unwrap_or("".into())) 199 } 200 201 // Reads an xx-byte unsigner integer. read_uxx(&mut self, xx: u8) -> AvifResult<u64>202 pub(crate) fn read_uxx(&mut self, xx: u8) -> AvifResult<u64> { 203 let n: usize = xx.into(); 204 if n == 0 { 205 return Ok(0); 206 } 207 if n > 8 { 208 return Err(AvifError::NotImplemented); 209 } 210 let mut out = [0; 8]; 211 let start = out.len() - n; 212 out[start..].copy_from_slice(self.get_slice(n)?); 213 Ok(u64::from_be_bytes(out)) 214 } 215 216 // Reads a null-terminated string. read_c_string(&mut self) -> AvifResult<String>217 pub(crate) fn read_c_string(&mut self) -> AvifResult<String> { 218 self.check(1)?; 219 let null_position = self.data[self.offset..] 220 .iter() 221 .position(|&x| x == b'\0') 222 .ok_or(AvifError::BmffParseFailed("".into()))?; 223 let range = self.offset..self.offset + null_position; 224 self.offset += null_position + 1; 225 Ok(String::from_utf8(self.data[range].to_vec()).unwrap_or("".into())) 226 } 227 read_version_and_flags(&mut self) -> AvifResult<(u8, u32)>228 pub(crate) fn read_version_and_flags(&mut self) -> AvifResult<(u8, u32)> { 229 let version = self.read_u8()?; 230 let flags = self.read_u24()?; 231 Ok((version, flags)) 232 } 233 read_and_enforce_version_and_flags( &mut self, enforced_version: u8, ) -> AvifResult<(u8, u32)>234 pub(crate) fn read_and_enforce_version_and_flags( 235 &mut self, 236 enforced_version: u8, 237 ) -> AvifResult<(u8, u32)> { 238 let (version, flags) = self.read_version_and_flags()?; 239 if version != enforced_version { 240 return Err(AvifError::BmffParseFailed("".into())); 241 } 242 Ok((version, flags)) 243 } 244 skip(&mut self, size: usize) -> AvifResult<()>245 pub(crate) fn skip(&mut self, size: usize) -> AvifResult<()> { 246 self.check(size)?; 247 checked_incr!(self.offset, size); 248 Ok(()) 249 } 250 rewind(&mut self, size: usize) -> AvifResult<()>251 pub(crate) fn rewind(&mut self, size: usize) -> AvifResult<()> { 252 checked_decr!(self.offset, size); 253 Ok(()) 254 } 255 read_uleb128(&mut self) -> AvifResult<u32>256 pub(crate) fn read_uleb128(&mut self) -> AvifResult<u32> { 257 // See the section 4.10.5. of the AV1 specification. 258 let mut value: u64 = 0; 259 for i in 0..8 { 260 // leb128_byte contains 8 bits read from the bitstream. 261 let leb128_byte = self.read_u8()?; 262 // The bottom 7 bits are used to compute the variable value. 263 value |= u64::from(leb128_byte & 0x7F) << (i * 7); 264 // The most significant bit is used to indicate that there are more 265 // bytes to be read. 266 if (leb128_byte & 0x80) == 0 { 267 // It is a requirement of bitstream conformance that the value 268 // returned from the leb128 parsing process is less than or 269 // equal to (1 << 32)-1. 270 return u32_from_u64(value); 271 } 272 } 273 // It is a requirement of bitstream conformance that the most 274 // significant bit of leb128_byte is equal to 0 if i is equal to 7. 275 Err(AvifError::BmffParseFailed( 276 "uleb value did not terminate after 8 bytes".into(), 277 )) 278 } 279 } 280 281 #[cfg(test)] 282 mod tests { 283 use super::*; 284 285 #[test] read_uxx()286 fn read_uxx() { 287 let mut stream = IStream::create(&[1, 2, 3, 4, 5, 6, 7, 8]); 288 assert_eq!(stream.read_uxx(0), Ok(0)); 289 assert_eq!(stream.offset, 0); 290 assert_eq!(stream.read_uxx(1), Ok(1)); 291 assert_eq!(stream.offset, 1); 292 stream.offset = 0; 293 assert_eq!(stream.read_uxx(2), Ok(258)); 294 stream.offset = 0; 295 assert_eq!(stream.read_u16(), Ok(258)); 296 stream.offset = 0; 297 assert_eq!(stream.read_uxx(3), Ok(66051)); 298 stream.offset = 0; 299 assert_eq!(stream.read_u24(), Ok(66051)); 300 stream.offset = 0; 301 assert_eq!(stream.read_uxx(4), Ok(16909060)); 302 stream.offset = 0; 303 assert_eq!(stream.read_u32(), Ok(16909060)); 304 stream.offset = 0; 305 assert_eq!(stream.read_uxx(5), Ok(4328719365)); 306 stream.offset = 0; 307 assert_eq!(stream.read_uxx(6), Ok(1108152157446)); 308 stream.offset = 0; 309 assert_eq!(stream.read_uxx(7), Ok(283686952306183)); 310 stream.offset = 0; 311 assert_eq!(stream.read_uxx(8), Ok(72623859790382856)); 312 stream.offset = 0; 313 assert_eq!(stream.read_u64(), Ok(72623859790382856)); 314 stream.offset = 0; 315 assert_eq!(stream.read_uxx(9), Err(AvifError::NotImplemented)); 316 } 317 318 #[test] read_string()319 fn read_string() { 320 let bytes = "abcd\0e".as_bytes(); 321 assert_eq!(IStream::create(bytes).read_string(4), Ok("abcd".into())); 322 assert_eq!(IStream::create(bytes).read_string(5), Ok("abcd\0".into())); 323 assert_eq!(IStream::create(bytes).read_string(6), Ok("abcd\0e".into())); 324 assert!(matches!( 325 IStream::create(bytes).read_string(8), 326 Err(AvifError::BmffParseFailed(_)) 327 )); 328 assert_eq!(IStream::create(bytes).read_c_string(), Ok("abcd".into())); 329 } 330 } 331