Fix parsing flowsets (#23)

* Fix parsing flowsets

* Fix fmt

Author: mikhailzagurskiy

src/lib.rs

@@ -312,6 +312,16 @@ mod tests {
         assert_yaml_snapshot!(NetflowParser::default().parse_bytes(&packet));
     }
 
+    #[test]
+    fn it_parses_v9_many_flows() {
+        let packet = [
+            0, 9, 0, 3, 0, 0, 9, 9, 0, 1, 2, 3, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 16, 1, 2, 0,
+            2, 0, 1, 0, 4, 0, 8, 0, 4, 1, 2, 0, 12, 9, 2, 3, 4, 9, 9, 9, 8, 9, 2, 3, 4, 9, 9,
+            9, 8,
+        ];
+        assert_yaml_snapshot!(NetflowParser::default().parse_bytes(&packet));
+    }
+
     #[test]
     fn it_parses_v9_options_template() {
         let packet = [

src/snapshots/netflow_parser__tests__it_parses_v9_many_flows.snap

@@ -0,0 +1,38 @@
+---
+source: src/lib.rs
+expression: "NetflowParser::default().parse_bytes(&packet)"
+---
+- V9:
+    header:
+      version: 9
+      count: 3
+      sys_up_time: 2313
+      unix_secs: 66051
+      sequence_number: 1
+      source_id: 1
+    flowsets:
+      - flow_set_id: 0
+        template:
+          length: 16
+          template_id: 258
+          field_count: 2
+          fields:
+            - field_type_number: 1
+              field_type: InBytes
+              field_length: 4
+            - field_type_number: 8
+              field_type: Ipv4SrcAddr
+              field_length: 4
+      - flow_set_id: 258
+        data:
+          length: 12
+          data_fields:
+            - InBytes:
+                DataNumber: 151126788
+              Ipv4SrcAddr:
+                Ip4Addr: 9.9.9.8
+            - InBytes:
+                DataNumber: 151126788
+              Ipv4SrcAddr:
+                Ip4Addr: 9.9.9.8
+

src/variable_versions/v9.rs

@@ -41,10 +41,37 @@ pub struct V9 {
     /// V9 Header
     pub header: Header,
     /// Flowsets
-    #[nom(Count = "header.count", Parse = "{ |i| FlowSet::parse(i, parser) }")]
+    #[nom(Parse = "{ |i| parse_flowsets(i, parser, header.count as usize) }")]
     pub flowsets: Vec<FlowSet>,
 }
 
+fn parse_flowsets<'a, 'b>(
+    i: &'a [u8],
+    parser: &'b mut V9Parser,
+    mut count: usize,
+) -> IResult<&'a [u8], Vec<FlowSet>> {
+    let mut flowsets = vec![];
+    let mut remaining = i;
+
+    // Header.count represents total number of records in data + records in templates
+    while count > 0 {
+        let (i, flowset) = FlowSet::parse(remaining, parser)?;
+        remaining = i;
+
+        if flowset.template.is_some() || flowset.options_template.is_some() {
+            count = count.saturating_sub(1);
+        } else if let Some(data) = flowset.data.as_ref() {
+            count = count.saturating_sub(data.data_fields.len());
+        } else if let Some(_) = flowset.options_data.as_ref() {
+            count = count.saturating_sub(1);
+        }
+
+        flowsets.push(flowset)
+    }
+
+    Ok((remaining, flowsets))
+}
+
 #[derive(Debug, PartialEq, Clone, Copy, Serialize, Nom)]
 pub struct Header {
     /// The version of NetFlow records exported in this packet; for Version 9, this value is 9