Skip to content

Commit 4374d6a

Browse files
mikemiles-devMichael Mileusnich
and
Michael Mileusnich
authored
Lib cleanup (#40)
Co-authored-by: Michael Mileusnich <[email protected]>
1 parent 50176f3 commit 4374d6a

File tree

2 files changed

+59
-27
lines changed

2 files changed

+59
-27
lines changed

src/lib.rs

Lines changed: 21 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -70,17 +70,18 @@
7070
//!
7171
//! ```cargo run --example netflow_udp_listener_single_threaded```
7272
73+
pub mod netflow_header;
7374
pub mod protocol;
7475
pub mod static_versions;
7576
pub mod variable_versions;
7677

7778
use serde::Serialize;
79+
80+
use netflow_header::{NetflowHeader, NetflowVersion};
7881
use static_versions::{v5::V5, v7::V7};
7982
use variable_versions::ipfix::{IPFix, IPFixParser};
8083
use variable_versions::v9::{V9Parser, V9};
8184

82-
use nom_derive::{Nom, Parse};
83-
8485
#[derive(Debug, Clone, Serialize)]
8586
pub struct NetflowPacketError {
8687
pub error_message: String,
@@ -127,13 +128,6 @@ struct ParsedNetflow {
127128
netflow_packet: NetflowPacketResult,
128129
}
129130

130-
/// Struct is used simply to match how to handle the result of the packet
131-
#[derive(Nom)]
132-
struct NetflowHeader {
133-
/// Netflow Version
134-
version: u16,
135-
}
136-
137131
/// Trait provided for all static parser versions
138132
trait NetflowByteParserStatic {
139133
fn parse_bytes(packet: &[u8]) -> Result<ParsedNetflow, Box<dyn std::error::Error>>;
@@ -159,13 +153,17 @@ impl NetflowParser {
159153
&'a mut self,
160154
packet: &'a [u8],
161155
) -> Result<ParsedNetflow, Box<dyn std::error::Error>> {
162-
match NetflowHeader::parse_be(packet) {
163-
Ok((i, netflow_header)) if netflow_header.version == 5 => V5::parse_bytes(i),
164-
Ok((i, netflow_header)) if netflow_header.version == 7 => V7::parse_bytes(i),
165-
Ok((i, netflow_header)) if netflow_header.version == 9 => {
156+
match NetflowHeader::parse_header(packet) {
157+
Ok((i, netflow_header)) if netflow_header.version == NetflowVersion::V5 => {
158+
V5::parse_bytes(i)
159+
}
160+
Ok((i, netflow_header)) if netflow_header.version == NetflowVersion::V7 => {
161+
V7::parse_bytes(i)
162+
}
163+
Ok((i, netflow_header)) if netflow_header.version == NetflowVersion::V9 => {
166164
self.v9_parser.parse_bytes(i)
167165
}
168-
Ok((i, netflow_header)) if netflow_header.version == 10 => {
166+
Ok((i, netflow_header)) if netflow_header.version == NetflowVersion::IPFix => {
169167
self.ipfix_parser.parse_bytes(i)
170168
}
171169
_ => Err("Not Supported".to_string().into()),
@@ -196,22 +194,18 @@ impl NetflowParser {
196194
if packet.is_empty() {
197195
return vec![];
198196
}
199-
match self.parse_by_version(packet) {
200-
Ok(parsed_netflow) => {
197+
self.parse_by_version(packet)
198+
.map(|parsed_netflow| {
201199
let mut parsed = vec![parsed_netflow.netflow_packet];
202-
if !parsed_netflow.remaining.is_empty() {
203-
parsed.append(&mut self.parse_bytes(parsed_netflow.remaining.as_slice()));
204-
}
200+
parsed.append(&mut self.parse_bytes(parsed_netflow.remaining.as_slice()));
205201
parsed
206-
}
207-
Err(parsed_error) => {
208-
let netflow_packet_error = NetflowPacketError {
209-
error_message: parsed_error.to_string(),
202+
})
203+
.unwrap_or_else(|e| {
204+
vec![NetflowPacketResult::Error(NetflowPacketError {
205+
error_message: e.to_string(),
210206
bytes: packet.to_vec(),
211-
};
212-
vec![NetflowPacketResult::Error(netflow_packet_error)]
213-
}
214-
}
207+
})]
208+
})
215209
}
216210
}
217211

src/netflow_header.rs

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
use nom::number::complete::be_u16;
2+
use nom::IResult;
3+
use nom_derive::{Nom, Parse};
4+
5+
/// Struct is used simply to match how to handle the result of the packet
6+
#[derive(Nom)]
7+
pub struct NetflowHeader {
8+
/// Netflow Version
9+
#[nom(Map = "NetflowVersion::from", Parse = "be_u16")]
10+
pub version: NetflowVersion,
11+
}
12+
13+
impl NetflowHeader {
14+
pub fn parse_header(packet: &[u8]) -> IResult<&[u8], NetflowHeader> {
15+
NetflowHeader::parse_be(packet)
16+
}
17+
}
18+
19+
#[derive(PartialEq)]
20+
pub enum NetflowVersion {
21+
V5,
22+
V7,
23+
V9,
24+
IPFix,
25+
Unsupported,
26+
}
27+
28+
impl From<u16> for NetflowVersion {
29+
fn from(version: u16) -> Self {
30+
match version {
31+
5 => NetflowVersion::V5,
32+
7 => NetflowVersion::V7,
33+
9 => NetflowVersion::V9,
34+
10 => NetflowVersion::IPFix,
35+
_ => NetflowVersion::Unsupported,
36+
}
37+
}
38+
}

0 commit comments

Comments
 (0)