1 use alloc::string::String;
2 use core::convert::TryInto;
3 use core::fmt;
4 use core::marker::PhantomData;
5 
6 use crate::pod::{from_bytes, slice_from_bytes, Pod};
7 use crate::ReadRef;
8 
9 /// A newtype for byte slices.
10 ///
11 /// It has these important features:
12 /// - no methods that can panic, such as `Index`
13 /// - convenience methods for `Pod` types
14 /// - a useful `Debug` implementation
15 #[derive(Default, Clone, Copy, PartialEq, Eq)]
16 pub struct Bytes<'data>(pub &'data [u8]);
17 
18 impl<'data> fmt::Debug for Bytes<'data> {
fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result19     fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
20         debug_list_bytes(self.0, fmt)
21     }
22 }
23 
24 impl<'data> Bytes<'data> {
25     /// Return the length of the byte slice.
26     #[inline]
len(&self) -> usize27     pub fn len(&self) -> usize {
28         self.0.len()
29     }
30 
31     /// Return true if the byte slice is empty.
32     #[inline]
is_empty(&self) -> bool33     pub fn is_empty(&self) -> bool {
34         self.0.is_empty()
35     }
36 
37     /// Skip over the given number of bytes at the start of the byte slice.
38     ///
39     /// Modifies the byte slice to start after the bytes.
40     ///
41     /// Returns an error if there are too few bytes.
42     #[inline]
skip(&mut self, offset: usize) -> Result<(), ()>43     pub fn skip(&mut self, offset: usize) -> Result<(), ()> {
44         match self.0.get(offset..) {
45             Some(tail) => {
46                 self.0 = tail;
47                 Ok(())
48             }
49             None => {
50                 self.0 = &[];
51                 Err(())
52             }
53         }
54     }
55 
56     /// Return a reference to the given number of bytes at the start of the byte slice.
57     ///
58     /// Modifies the byte slice to start after the bytes.
59     ///
60     /// Returns an error if there are too few bytes.
61     #[inline]
read_bytes(&mut self, count: usize) -> Result<Bytes<'data>, ()>62     pub fn read_bytes(&mut self, count: usize) -> Result<Bytes<'data>, ()> {
63         match (self.0.get(..count), self.0.get(count..)) {
64             (Some(head), Some(tail)) => {
65                 self.0 = tail;
66                 Ok(Bytes(head))
67             }
68             _ => {
69                 self.0 = &[];
70                 Err(())
71             }
72         }
73     }
74 
75     /// Return a reference to the given number of bytes at the given offset of the byte slice.
76     ///
77     /// Returns an error if the offset is invalid or there are too few bytes.
78     #[inline]
read_bytes_at(mut self, offset: usize, count: usize) -> Result<Bytes<'data>, ()>79     pub fn read_bytes_at(mut self, offset: usize, count: usize) -> Result<Bytes<'data>, ()> {
80         self.skip(offset)?;
81         self.read_bytes(count)
82     }
83 
84     /// Return a reference to a `Pod` struct at the start of the byte slice.
85     ///
86     /// Modifies the byte slice to start after the bytes.
87     ///
88     /// Returns an error if there are too few bytes or the slice is incorrectly aligned.
89     #[inline]
read<T: Pod>(&mut self) -> Result<&'data T, ()>90     pub fn read<T: Pod>(&mut self) -> Result<&'data T, ()> {
91         match from_bytes(self.0) {
92             Ok((value, tail)) => {
93                 self.0 = tail;
94                 Ok(value)
95             }
96             Err(()) => {
97                 self.0 = &[];
98                 Err(())
99             }
100         }
101     }
102 
103     /// Return a reference to a `Pod` struct at the given offset of the byte slice.
104     ///
105     /// Returns an error if there are too few bytes or the offset is incorrectly aligned.
106     #[inline]
read_at<T: Pod>(mut self, offset: usize) -> Result<&'data T, ()>107     pub fn read_at<T: Pod>(mut self, offset: usize) -> Result<&'data T, ()> {
108         self.skip(offset)?;
109         self.read()
110     }
111 
112     /// Return a reference to a slice of `Pod` structs at the start of the byte slice.
113     ///
114     /// Modifies the byte slice to start after the bytes.
115     ///
116     /// Returns an error if there are too few bytes or the offset is incorrectly aligned.
117     #[inline]
read_slice<T: Pod>(&mut self, count: usize) -> Result<&'data [T], ()>118     pub fn read_slice<T: Pod>(&mut self, count: usize) -> Result<&'data [T], ()> {
119         match slice_from_bytes(self.0, count) {
120             Ok((value, tail)) => {
121                 self.0 = tail;
122                 Ok(value)
123             }
124             Err(()) => {
125                 self.0 = &[];
126                 Err(())
127             }
128         }
129     }
130 
131     /// Return a reference to a slice of `Pod` structs at the given offset of the byte slice.
132     ///
133     /// Returns an error if there are too few bytes or the offset is incorrectly aligned.
134     #[inline]
read_slice_at<T: Pod>(mut self, offset: usize, count: usize) -> Result<&'data [T], ()>135     pub fn read_slice_at<T: Pod>(mut self, offset: usize, count: usize) -> Result<&'data [T], ()> {
136         self.skip(offset)?;
137         self.read_slice(count)
138     }
139 
140     /// Read a null terminated string.
141     ///
142     /// Does not assume any encoding.
143     /// Reads past the null byte, but doesn't return it.
144     #[inline]
read_string(&mut self) -> Result<&'data [u8], ()>145     pub fn read_string(&mut self) -> Result<&'data [u8], ()> {
146         match memchr::memchr(b'\0', self.0) {
147             Some(null) => {
148                 // These will never fail.
149                 let bytes = self.read_bytes(null)?;
150                 self.skip(1)?;
151                 Ok(bytes.0)
152             }
153             None => {
154                 self.0 = &[];
155                 Err(())
156             }
157         }
158     }
159 
160     /// Read a null terminated string at an offset.
161     ///
162     /// Does not assume any encoding. Does not return the null byte.
163     #[inline]
read_string_at(mut self, offset: usize) -> Result<&'data [u8], ()>164     pub fn read_string_at(mut self, offset: usize) -> Result<&'data [u8], ()> {
165         self.skip(offset)?;
166         self.read_string()
167     }
168 }
169 
170 // Only for Debug impl of `Bytes`.
debug_list_bytes(bytes: &[u8], fmt: &mut fmt::Formatter<'_>) -> fmt::Result171 fn debug_list_bytes(bytes: &[u8], fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
172     let mut list = fmt.debug_list();
173     list.entries(bytes.iter().take(8).copied().map(DebugByte));
174     if bytes.len() > 8 {
175         list.entry(&DebugLen(bytes.len()));
176     }
177     list.finish()
178 }
179 
180 struct DebugByte(u8);
181 
182 impl fmt::Debug for DebugByte {
fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result183     fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
184         write!(fmt, "0x{:02x}", self.0)
185     }
186 }
187 
188 struct DebugLen(usize);
189 
190 impl fmt::Debug for DebugLen {
fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result191     fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
192         write!(fmt, "...; {}", self.0)
193     }
194 }
195 
196 /// A newtype for byte strings.
197 ///
198 /// For byte slices that are strings of an unknown encoding.
199 ///
200 /// Provides a `Debug` implementation that interprets the bytes as UTF-8.
201 #[derive(Default, Clone, Copy, PartialEq, Eq)]
202 pub(crate) struct ByteString<'data>(pub &'data [u8]);
203 
204 impl<'data> fmt::Debug for ByteString<'data> {
fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result205     fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
206         write!(fmt, "\"{}\"", String::from_utf8_lossy(self.0))
207     }
208 }
209 
210 #[allow(dead_code)]
211 #[inline]
align(offset: usize, size: usize) -> usize212 pub(crate) fn align(offset: usize, size: usize) -> usize {
213     (offset + (size - 1)) & !(size - 1)
214 }
215 
216 #[allow(dead_code)]
data_range( data: &[u8], data_address: u64, range_address: u64, size: u64, ) -> Option<&[u8]>217 pub(crate) fn data_range(
218     data: &[u8],
219     data_address: u64,
220     range_address: u64,
221     size: u64,
222 ) -> Option<&[u8]> {
223     let offset = range_address.checked_sub(data_address)?;
224     data.get(offset.try_into().ok()?..)?
225         .get(..size.try_into().ok()?)
226 }
227 
228 /// A table of zero-terminated strings.
229 ///
230 /// This is used for most file formats.
231 #[derive(Debug, Clone, Copy)]
232 pub struct StringTable<'data, R = &'data [u8]>
233 where
234     R: ReadRef<'data>,
235 {
236     data: Option<R>,
237     start: u64,
238     end: u64,
239     marker: PhantomData<&'data ()>,
240 }
241 
242 impl<'data, R: ReadRef<'data>> StringTable<'data, R> {
243     /// Interpret the given data as a string table.
new(data: R, start: u64, end: u64) -> Self244     pub fn new(data: R, start: u64, end: u64) -> Self {
245         StringTable {
246             data: Some(data),
247             start,
248             end,
249             marker: PhantomData,
250         }
251     }
252 
253     /// Return the string at the given offset.
get(&self, offset: u32) -> Result<&'data [u8], ()>254     pub fn get(&self, offset: u32) -> Result<&'data [u8], ()> {
255         match self.data {
256             Some(data) => {
257                 let r_start = self.start.checked_add(offset.into()).ok_or(())?;
258                 data.read_bytes_at_until(r_start..self.end, 0)
259             }
260             None => Err(()),
261         }
262     }
263 }
264 
265 impl<'data, R: ReadRef<'data>> Default for StringTable<'data, R> {
default() -> Self266     fn default() -> Self {
267         StringTable {
268             data: None,
269             start: 0,
270             end: 0,
271             marker: PhantomData,
272         }
273     }
274 }
275 
276 #[cfg(test)]
277 mod tests {
278     use super::*;
279     use crate::pod::bytes_of;
280 
281     #[test]
bytes()282     fn bytes() {
283         let x = u32::to_be(0x0123_4567);
284         let data = Bytes(bytes_of(&x));
285 
286         let mut bytes = data;
287         assert_eq!(bytes.skip(0), Ok(()));
288         assert_eq!(bytes, data);
289 
290         let mut bytes = data;
291         assert_eq!(bytes.skip(4), Ok(()));
292         assert_eq!(bytes, Bytes(&[]));
293 
294         let mut bytes = data;
295         assert_eq!(bytes.skip(5), Err(()));
296         assert_eq!(bytes, Bytes(&[]));
297 
298         let mut bytes = data;
299         assert_eq!(bytes.read_bytes(0), Ok(Bytes(&[])));
300         assert_eq!(bytes, data);
301 
302         let mut bytes = data;
303         assert_eq!(bytes.read_bytes(4), Ok(data));
304         assert_eq!(bytes, Bytes(&[]));
305 
306         let mut bytes = data;
307         assert_eq!(bytes.read_bytes(5), Err(()));
308         assert_eq!(bytes, Bytes(&[]));
309 
310         assert_eq!(data.read_bytes_at(0, 0), Ok(Bytes(&[])));
311         assert_eq!(data.read_bytes_at(4, 0), Ok(Bytes(&[])));
312         assert_eq!(data.read_bytes_at(0, 4), Ok(data));
313         assert_eq!(data.read_bytes_at(1, 4), Err(()));
314 
315         let mut bytes = data;
316         assert_eq!(bytes.read::<u16>(), Ok(&u16::to_be(0x0123)));
317         assert_eq!(bytes, Bytes(&[0x45, 0x67]));
318         assert_eq!(data.read_at::<u16>(2), Ok(&u16::to_be(0x4567)));
319         assert_eq!(data.read_at::<u16>(3), Err(()));
320         assert_eq!(data.read_at::<u16>(4), Err(()));
321 
322         let mut bytes = data;
323         assert_eq!(bytes.read::<u32>(), Ok(&x));
324         assert_eq!(bytes, Bytes(&[]));
325 
326         let mut bytes = data;
327         assert_eq!(bytes.read::<u64>(), Err(()));
328         assert_eq!(bytes, Bytes(&[]));
329 
330         let mut bytes = data;
331         assert_eq!(bytes.read_slice::<u8>(0), Ok(&[][..]));
332         assert_eq!(bytes, data);
333 
334         let mut bytes = data;
335         assert_eq!(bytes.read_slice::<u8>(4), Ok(data.0));
336         assert_eq!(bytes, Bytes(&[]));
337 
338         let mut bytes = data;
339         assert_eq!(bytes.read_slice::<u8>(5), Err(()));
340         assert_eq!(bytes, Bytes(&[]));
341 
342         assert_eq!(data.read_slice_at::<u8>(0, 0), Ok(&[][..]));
343         assert_eq!(data.read_slice_at::<u8>(4, 0), Ok(&[][..]));
344         assert_eq!(data.read_slice_at::<u8>(0, 4), Ok(data.0));
345         assert_eq!(data.read_slice_at::<u8>(1, 4), Err(()));
346 
347         let data = Bytes(&[0x01, 0x02, 0x00, 0x04]);
348 
349         let mut bytes = data;
350         assert_eq!(bytes.read_string(), Ok(&data.0[..2]));
351         assert_eq!(bytes.0, &data.0[3..]);
352 
353         let mut bytes = data;
354         bytes.skip(3).unwrap();
355         assert_eq!(bytes.read_string(), Err(()));
356         assert_eq!(bytes.0, &[]);
357 
358         assert_eq!(data.read_string_at(0), Ok(&data.0[..2]));
359         assert_eq!(data.read_string_at(1), Ok(&data.0[1..2]));
360         assert_eq!(data.read_string_at(2), Ok(&[][..]));
361         assert_eq!(data.read_string_at(3), Err(()));
362     }
363 
364     #[test]
bytes_debug()365     fn bytes_debug() {
366         assert_eq!(format!("{:?}", Bytes(&[])), "[]");
367         assert_eq!(format!("{:?}", Bytes(&[0x01])), "[0x01]");
368         assert_eq!(
369             format!(
370                 "{:?}",
371                 Bytes(&[0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08])
372             ),
373             "[0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08]"
374         );
375         assert_eq!(
376             format!(
377                 "{:?}",
378                 Bytes(&[0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09])
379             ),
380             "[0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, ...; 9]"
381         );
382     }
383 }
384