瀏覽代碼

Make PkgLength return something useful

Because we constantly shave bits off the stream as we successfully parse
stuff, we can't use the raw PkgLength for anything useful. Instead, we
work out when in the stream we should stop parsing, and provide that through
utility methods on `PkgLength`. We still need to be able to parse raw
lengths for the tests though, so we don't have to constuct streams that are
thousands of bytes long for the longer lengths.
Isaac Woods 5 年之前
父節點
當前提交
aaa815f84c
共有 1 個文件被更改,包括 65 次插入6 次删除
  1. 65 6
      aml_parser/src/pkg_length.rs

+ 65 - 6
aml_parser/src/pkg_length.rs

@@ -5,9 +5,41 @@ use crate::{
 use bit_field::BitField;
 use log::trace;
 
-pub type PkgLength = u32;
+#[derive(Clone, Copy, PartialEq, Eq, Debug)]
+pub struct PkgLength {
+    pub raw_length: u32,
+    /// The distance from the end of the structure this `PkgLength` refers to, and the end of the
+    /// stream.
+    pub end_offset: u32,
+}
+
+impl PkgLength {
+    pub fn from_raw_length(stream: &[u8], raw_length: u32) -> PkgLength {
+        PkgLength { raw_length, end_offset: stream.len() as u32 - raw_length }
+    }
+
+    /// Returns `true` if the given stream is still within the structure this `PkgLength` refers
+    /// to.
+    pub fn still_parsing(&self, stream: &[u8]) -> bool {
+        stream.len() as u32 > self.end_offset
+    }
+}
 
 pub fn pkg_length<'a>() -> impl Parser<'a, PkgLength> {
+    move |input: &'a [u8]| {
+        let (new_input, raw_length) = raw_pkg_length().parse(input)?;
+
+        /*
+         * NOTE: we use the original input here, because `raw_length` includes the length of the
+         * `PkgLength`.
+         */
+        Ok((new_input, PkgLength::from_raw_length(input, raw_length)))
+    }
+}
+
+/// Parses a `PkgLength` and returns the *raw length*. If you want an instance of `PkgLength`, use
+/// `pkg_length` instead.
+pub fn raw_pkg_length<'a>() -> impl Parser<'a, u32> {
     /*
      * PkgLength := PkgLeadByte |
      * <PkgLeadByte ByteData> |
@@ -17,6 +49,8 @@ pub fn pkg_length<'a>() -> impl Parser<'a, PkgLength> {
      * The length encoded by the PkgLength includes the number of bytes used to encode it.
      */
     move |input: &'a [u8]| {
+        let starting_len = input.len() as u32;
+
         let (new_input, lead_byte) = take().parse(input)?;
         let byte_count = lead_byte.get_bits(6..8);
 
@@ -52,18 +86,43 @@ mod tests {
     use super::*;
     use crate::{test_utils::*, AmlError};
 
+    fn test_correct_pkglength(stream: &[u8], expected_raw_length: u32, expected_leftover: &[u8]) {
+        check_ok!(
+            pkg_length().parse(stream),
+            PkgLength::from_raw_length(stream, expected_raw_length),
+            &expected_leftover
+        );
+    }
+
+    #[test]
+    fn test_raw_pkg_length() {
+        check_ok!(raw_pkg_length().parse(&[0b01000101, 0x14]), 325, &[]);
+        check_ok!(raw_pkg_length().parse(&[0b01000111, 0x14, 0x46]), 327, &[0x46]);
+        check_ok!(raw_pkg_length().parse(&[0b10000111, 0x14, 0x46]), 287047, &[]);
+    }
+
     #[test]
     fn test_pkg_length() {
         check_err!(pkg_length().parse(&[]), AmlError::UnexpectedEndOfStream, &[]);
-        check_ok!(pkg_length().parse(&[0x00]), 0, &[]);
-        check_ok!(pkg_length().parse(&[0x05, 0xf5, 0x7f]), 5, &[0xf5, 0x7f]);
-        check_ok!(pkg_length().parse(&[0b01000101, 0x14]), 325, &[]);
-        check_ok!(pkg_length().parse(&[0b01000111, 0x14, 0x46]), 327, &[0x46]);
-        check_ok!(pkg_length().parse(&[0b10000111, 0x14, 0x46]), 287047, &[]);
+        test_correct_pkglength(&[0x00], 0, &[]);
+        test_correct_pkglength(
+            &[0x05, 0xf5, 0x7f, 0x3e, 0x54, 0x03],
+            5,
+            &[0xf5, 0x7f, 0x3e, 0x54, 0x03],
+        );
         check_err!(
             pkg_length().parse(&[0b11000000, 0xff, 0x4f]),
             AmlError::UnexpectedEndOfStream,
             &[0b11000000, 0xff, 0x4f]
         );
     }
+
+    #[test]
+    #[should_panic]
+    fn not_enough_stream() {
+        /*
+         * TODO: Ideally, this shouldn't panic the parser, but return a `UnexpectedEndOfStream`.
+         */
+        test_correct_pkglength(&[0x05, 0xf5], 5, &[0xf5]);
+    }
 }