1 // Copyright 2023 The ChromiumOS Authors 2 // Use of this source code is governed by a BSD-style license that can be 3 // found in the LICENSE file. 4 5 use std::borrow::Cow; 6 use std::fmt::Debug; 7 use std::io::Cursor; 8 use std::io::Seek; 9 use std::io::SeekFrom; 10 11 #[allow(clippy::len_without_is_empty)] 12 pub trait Header: Sized { 13 /// Parse the NALU header, returning it. parse<T: AsRef<[u8]>>(cursor: &mut Cursor<T>) -> Result<Self, String>14 fn parse<T: AsRef<[u8]>>(cursor: &mut Cursor<T>) -> Result<Self, String>; 15 /// Whether this header type indicates EOS. is_end(&self) -> bool16 fn is_end(&self) -> bool; 17 /// The length of the header. len(&self) -> usize18 fn len(&self) -> usize; 19 } 20 21 #[derive(Debug)] 22 pub struct Nalu<'a, U> { 23 pub header: U, 24 /// The mapping that backs this NALU. Possibly shared with the other NALUs 25 /// in the Access Unit. 26 pub data: Cow<'a, [u8]>, 27 28 pub size: usize, 29 pub offset: usize, 30 } 31 32 impl<'a, U> Nalu<'a, U> 33 where 34 U: Debug + Header, 35 { 36 /// Find the next Annex B encoded NAL unit. next(cursor: &mut Cursor<&'a [u8]>) -> Result<Nalu<'a, U>, String>37 pub fn next(cursor: &mut Cursor<&'a [u8]>) -> Result<Nalu<'a, U>, String> { 38 let bitstream = cursor.clone().into_inner(); 39 let pos = usize::try_from(cursor.position()).map_err(|err| err.to_string())?; 40 41 // Find the start code for this NALU 42 let current_nalu_offset = match Nalu::<'a, U>::find_start_code(cursor, pos) { 43 Some(offset) => offset, 44 None => return Err("No NAL found".into()), 45 }; 46 47 let mut start_code_offset = pos + current_nalu_offset; 48 49 // If the preceding byte is 00, then we actually have a four byte SC, 50 // i.e. 00 00 00 01 Where the first 00 is the "zero_byte()" 51 if start_code_offset > 0 && cursor.get_ref()[start_code_offset - 1] == 00 { 52 start_code_offset -= 1; 53 } 54 55 // The NALU offset is its offset + 3 bytes to skip the start code. 56 let nalu_offset = pos + current_nalu_offset + 3; 57 58 // Set the bitstream position to the start of the current NALU 59 cursor.set_position(u64::try_from(nalu_offset).map_err(|err| err.to_string())?); 60 61 let hdr = U::parse(cursor)?; 62 63 // Find the start of the subsequent NALU. 64 let mut next_nalu_offset = match Nalu::<'a, U>::find_start_code(cursor, nalu_offset) { 65 Some(offset) => offset, 66 None => { 67 let cur_pos = cursor.position(); 68 let end_pos = cursor.seek(SeekFrom::End(0)).map_err(|err| err.to_string())?; 69 let _ = cursor.seek(SeekFrom::Start(cur_pos)).map_err(|err| err.to_string())?; 70 (end_pos - cur_pos) as usize 71 } // Whatever data is left must be part of the current NALU 72 }; 73 74 while next_nalu_offset > 0 && cursor.get_ref()[nalu_offset + next_nalu_offset - 1] == 00 { 75 // Discard trailing_zero_8bits 76 next_nalu_offset -= 1; 77 } 78 79 let nal_size = if hdr.is_end() { 80 // the NALU is comprised of only the header 81 hdr.len() 82 } else { 83 next_nalu_offset 84 }; 85 86 Ok(Nalu { 87 header: hdr, 88 data: Cow::from(&bitstream[start_code_offset..nalu_offset + nal_size]), 89 size: nal_size, 90 offset: nalu_offset - start_code_offset, 91 }) 92 } 93 } 94 95 impl<'a, U> Nalu<'a, U> 96 where 97 U: Debug, 98 { find_start_code(data: &mut Cursor<&'a [u8]>, offset: usize) -> Option<usize>99 fn find_start_code(data: &mut Cursor<&'a [u8]>, offset: usize) -> Option<usize> { 100 // discard all zeroes until the start code pattern is found 101 data.get_ref()[offset..].windows(3).position(|window| window == [0x00, 0x00, 0x01]) 102 } 103 into_owned(self) -> Nalu<'static, U>104 pub fn into_owned(self) -> Nalu<'static, U> { 105 Nalu { 106 header: self.header, 107 size: self.size, 108 offset: self.offset, 109 data: Cow::Owned(self.data.into_owned()), 110 } 111 } 112 } 113 114 impl<'a, U> AsRef<[u8]> for Nalu<'a, U> { as_ref(&self) -> &[u8]115 fn as_ref(&self) -> &[u8] { 116 &self.data[self.offset..self.offset + self.size] 117 } 118 } 119