From 9f274bdc15fb2c1f7b5012e5eec4071d5911ad4e Mon Sep 17 00:00:00 2001
From: Michael Mileusnich <mmileusnich@gmail.com>
Date: Thu, 1 Aug 2024 11:59:50 -0500
Subject: [PATCH] More cleanup

---
 src/parser.rs                                 | 54 ++++++-------------
 ...__tests__base_tests__it_creates_error.snap |  2 -
 ...__base_tests__it_parses_v5_incomplete.snap | 17 ++++++
 src/tests.rs                                  |  6 +++
 4 files changed, 40 insertions(+), 39 deletions(-)
 create mode 100644 src/snapshots/netflow_parser__tests__base_tests__it_parses_v5_incomplete.snap

diff --git a/src/parser.rs b/src/parser.rs
index 0458762..b5574de 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -7,35 +7,12 @@ use crate::{NetflowPacket, NetflowParser};
 
 use serde::Serialize;
 
-/// Struct is used simply to match how to handle the result of the packet
 #[derive(Nom)]
-pub struct NetflowHeader {
-    /// Netflow Version
+/// Generic Netflow Header for shared versions
+pub struct GenericNetflowHeader {
     pub version: u16,
 }
 
-/// Enum of supported Netflow Versions
-pub enum NetflowVersion<'a> {
-    V5(&'a [u8]),
-    V7(&'a [u8]),
-    V9(&'a [u8]),
-    IPFix(&'a [u8]),
-    Unknown(Vec<u8>),
-}
-
-/// Implementing From for NetflowVersion
-impl<'a> From<&'a [u8]> for NetflowVersion<'a> {
-    fn from(packet: &[u8]) -> NetflowVersion {
-        match NetflowHeader::parse_be(packet) {
-            Ok((i, header)) if header.version == 5 => NetflowVersion::V5(i),
-            Ok((i, header)) if header.version == 7 => NetflowVersion::V7(i),
-            Ok((i, header)) if header.version == 9 => NetflowVersion::V9(i),
-            Ok((i, header)) if header.version == 10 => NetflowVersion::IPFix(i),
-            _ => NetflowVersion::Unknown(packet.to_vec()),
-        }
-    }
-}
-
 #[derive(Debug, Clone)]
 pub struct ParsedNetflow {
     pub remaining: Vec<u8>,
@@ -58,6 +35,7 @@ pub enum NetflowParseError {
     V7(String),
     V9(String),
     IPFix(String),
+    Incomplete(String),
     UnknownVersion(Vec<u8>),
 }
 
@@ -67,24 +45,26 @@ impl NetflowParser {
         &'a mut self,
         packet: &'a [u8],
     ) -> Result<ParsedNetflow, NetflowParseError> {
-        match NetflowVersion::from(packet) {
-            NetflowVersion::V5(v5_packet) => V5::parse(v5_packet)
+        let (packet, version) = GenericNetflowHeader::parse(packet)
+            .map(|(remaining, header)| (remaining, header.version))
+            .map_err(|e| NetflowParseError::Incomplete(e.to_string()))?;
+
+        match version {
+            5 => V5::parse(packet)
                 .map(|(remaining, v5)| ParsedNetflow::new(remaining, NetflowPacket::V5(v5)))
                 .map_err(|e| NetflowParseError::V5(e.to_string())),
-            NetflowVersion::V7(v7_packet) => V7::parse(v7_packet)
+            7 => V7::parse(packet)
                 .map(|(remaining, v7)| ParsedNetflow::new(remaining, NetflowPacket::V7(v7)))
                 .map_err(|e| NetflowParseError::V7(e.to_string())),
-            NetflowVersion::V9(v9_packet) => V9::parse(v9_packet, &mut self.v9_parser)
+            9 => V9::parse(packet, &mut self.v9_parser)
                 .map(|(remaining, v9)| ParsedNetflow::new(remaining, NetflowPacket::V9(v9)))
                 .map_err(|e| NetflowParseError::V9(e.to_string())),
-            NetflowVersion::IPFix(ipfix_packet) => {
-                IPFix::parse(ipfix_packet, &mut self.ipfix_parser)
-                    .map(|(remaining, ipfix)| {
-                        ParsedNetflow::new(remaining, NetflowPacket::IPFix(ipfix))
-                    })
-                    .map_err(|e| NetflowParseError::IPFix(e.to_string()))
-            }
-            NetflowVersion::Unknown(e) => Err(NetflowParseError::UnknownVersion(e)),
+            10 => IPFix::parse(packet, &mut self.ipfix_parser)
+                .map(|(remaining, ipfix)| {
+                    ParsedNetflow::new(remaining, NetflowPacket::IPFix(ipfix))
+                })
+                .map_err(|e| NetflowParseError::IPFix(e.to_string())),
+            _ => Err(NetflowParseError::UnknownVersion(packet.to_vec())),
         }
     }
 }
diff --git a/src/snapshots/netflow_parser__tests__base_tests__it_creates_error.snap b/src/snapshots/netflow_parser__tests__base_tests__it_creates_error.snap
index f10eff2..d655cf2 100644
--- a/src/snapshots/netflow_parser__tests__base_tests__it_creates_error.snap
+++ b/src/snapshots/netflow_parser__tests__base_tests__it_creates_error.snap
@@ -5,8 +5,6 @@ expression: "NetflowParser::default().parse_bytes(&packet)"
 - Error:
     error:
       UnknownVersion:
-        - 12
-        - 13
         - 14
     remaining:
       - 12
diff --git a/src/snapshots/netflow_parser__tests__base_tests__it_parses_v5_incomplete.snap b/src/snapshots/netflow_parser__tests__base_tests__it_parses_v5_incomplete.snap
new file mode 100644
index 0000000..4145771
--- /dev/null
+++ b/src/snapshots/netflow_parser__tests__base_tests__it_parses_v5_incomplete.snap
@@ -0,0 +1,17 @@
+---
+source: src/tests.rs
+expression: "NetflowParser::default().parse_bytes(&packet)"
+---
+- Error:
+    error:
+      V5: Parsing requires 4 bytes/chars
+    remaining:
+      - 0
+      - 5
+      - 0
+      - 0
+      - 1
+      - 1
+      - 1
+      - 1
+
diff --git a/src/tests.rs b/src/tests.rs
index c6ac45a..e575d49 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -48,6 +48,12 @@ mod base_tests {
         assert_yaml_snapshot!(NetflowParser::default().parse_bytes(&packet));
     }
 
+    #[test]
+    fn it_parses_v5_incomplete() {
+        let packet = [0, 5, 0, 0, 1, 1, 1, 1];
+        assert_yaml_snapshot!(NetflowParser::default().parse_bytes(&packet));
+    }
+
     #[test]
     fn it_parses_v5_and_re_exports() {
         let packet = [