add methods to parse specific pppoe headers

2 files changed, 57 insertions, 1 deletions

diff --git a/src/error.rs b/src/error.rs
index 8560022..e6d1e12 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -13,6 +13,8 @@ pub enum ParseError {
     InvalidPppoeType(u8),
     InvalidPppoeCode(u8),
 
+    UnexpectedCode(u8),
+
     PayloadLengthOutOfBound {
         actual_packet_length: u16,
         payload_length: u16,
diff --git a/src/header.rs b/src/header.rs
index 8f48888..9a676b2 100644
--- a/src/header.rs
+++ b/src/header.rs
@@ -40,6 +40,13 @@ pub struct Header<'a>(&'a mut [u8]);
 
 impl<'a> Header<'a> {
     pub fn from_buffer(buffer: &mut [u8]) -> Result<Header, ParseError> {
+        Self::from_buffer_with_code(buffer, None)
+    }
+
+    pub fn from_buffer_with_code(
+        buffer: &mut [u8],
+        expected_code: Option<Code>,
+    ) -> Result<Header, ParseError> {
         Self::ensure_minimal_buffer_length(buffer)?;
         if buffer[0] != 0x11 {
             let version = buffer[0] >> 4;
@@ -51,7 +58,12 @@ impl<'a> Header<'a> {
             };
         }
 
-        Code::try_from(buffer[1])?;
+        let code = Code::try_from(buffer[1])?;
+        if let Some(expected_code) = expected_code {
+            if code != expected_code {
+                return Err(ParseError::UnexpectedCode(code as u8));
+            }
+        }
 
         let length = usize::from(NE::read_u16(&buffer[4..]));
         if length + 6 > buffer.len() {
@@ -68,6 +80,26 @@ impl<'a> Header<'a> {
         Ok(Header(buffer))
     }
 
+    pub fn padi_from_buffer(buffer: &mut [u8]) -> Result<Header, ParseError> {
+        Self::from_buffer_with_code(buffer, Some(Code::Padi))
+    }
+
+    pub fn pado_from_buffer(buffer: &mut [u8]) -> Result<Header, ParseError> {
+        Self::from_buffer_with_code(buffer, Some(Code::Pado))
+    }
+
+    pub fn padr_from_buffer(buffer: &mut [u8]) -> Result<Header, ParseError> {
+        Self::from_buffer_with_code(buffer, Some(Code::Padr))
+    }
+
+    pub fn pads_from_buffer(buffer: &mut [u8]) -> Result<Header, ParseError> {
+        Self::from_buffer_with_code(buffer, Some(Code::Pads))
+    }
+
+    pub fn padt_from_buffer(buffer: &mut [u8]) -> Result<Header, ParseError> {
+        Self::from_buffer_with_code(buffer, Some(Code::Padt))
+    }
+
     fn ensure_minimal_buffer_length(buffer: &mut [u8]) -> Result<(), ParseError> {
         if buffer.len() < 6 {
             return Err(ParseError::BufferTooSmall(buffer.len()));
@@ -486,4 +518,26 @@ mod tests {
             _ => false,
         });
     }
+
+    #[test]
+    fn pppoe_code_conditional_parsing() {
+        let buffer = &mut [0u8; 200];
+        let codes = [Code::Padi, Code::Pado, Code::Padr, Code::Pads, Code::Padt];
+        let session_ids = [0, 0, 0, 1, 1];
+
+        for (i, code) in codes.iter().cloned().enumerate() {
+            Header::create_packet(buffer, code, session_ids[i])
+                .unwrap()
+                .add_tag(Tag::ServiceName(b"abc"))
+                .unwrap();
+            for (j, code) in codes.iter().cloned().enumerate() {
+                let header = Header::from_buffer_with_code(buffer, Some(code));
+                if j == i {
+                    assert!(header.is_ok())
+                } else {
+                    assert!(header.is_err())
+                }
+            }
+        }
+    }
 }