diff --git a/README.md b/README.md index cff9d2f..4a83d4b 100644 --- a/README.md +++ b/README.md @@ -258,6 +258,35 @@ The parser also automatically validates: - Each occurrence is decoded as a distinct value in the record - Descriptor order is significant and preserved +### Cumulative Decoded Output Limits + +A small NetFlow v9 or IPFIX message can expand into many decoded field values, +especially when a template contains zero-width or very small fields. The parser +therefore applies two cumulative limits to each message, including pending data +replayed when a template arrives: + +- 65,536 decoded field values +- 4 MiB of decoded field content + +The content-byte limit counts field contents, not IPFIX variable-length prefixes. +Both limits must be greater than zero. A message that exceeds either limit is +rejected with `NetflowError::DecodedOutputLimitExceeded`; no packet from that +message is returned. + +```rust +use netflow_parser::NetflowParser; + +let parser = NetflowParser::builder() + .with_max_decoded_field_values_per_message(32_768) + .with_max_decoded_field_payload_bytes_per_message(2 * 1024 * 1024) + .build() + .expect("valid limits"); +``` + +Use the `with_v9_*` and `with_ipfix_*` variants to configure the protocols +independently. These limits complement the per-FlowSet record limit; they bound +the combined decoded output of all Sets or FlowSets in one message. + ### Template TTL (Time-to-Live) > **Note:** Only time-based TTL is supported. See [RELEASES.md](RELEASES.md) for details. diff --git a/SECURITY.md b/SECURITY.md index 0e7d394..3330c5f 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -124,6 +124,10 @@ let parser = NetflowParser::builder() // Limit fields per template (DoS protection) .with_max_field_count(5000) + // Bound cumulative decoded output from each message + .with_max_decoded_field_values_per_message(65_536) + .with_max_decoded_field_payload_bytes_per_message(4 * 1024 * 1024) + // Limit error sample size (prevents memory exhaustion) .with_max_error_sample_size(256) @@ -154,6 +158,7 @@ The parser includes several DoS mitigations: - **Template Field Count Limit:** Default 10,000 fields per template - **Template Total Size Validation:** Maximum 65,535 bytes per template +- **Cumulative Decoded Output:** Defaults to 65,536 field values and 4 MiB of field content per message - **Error Sample Size Limit:** Default 256 bytes to prevent memory exhaustion - **LRU Template Cache:** Prevents unbounded cache growth diff --git a/benches/hot_path_bench.rs b/benches/hot_path_bench.rs index f5a1750..db5b155 100644 --- a/benches/hot_path_bench.rs +++ b/benches/hot_path_bench.rs @@ -1,6 +1,8 @@ use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; -use netflow_parser::NetflowParser; +use netflow_parser::scoped_parser::AutoScopedParser; +use netflow_parser::{NetflowPacket, NetflowParser}; use std::hint::black_box; +use std::net::SocketAddr; fn v9_template_packet() -> Vec { vec![ @@ -167,55 +169,179 @@ fn ipfix_data_packet(flow_count: u16) -> Vec { packet } -fn bench_warm_v9_data_hot_path(c: &mut Criterion) { - let template = v9_template_packet(); - let mut group = c.benchmark_group("Hot Path V9 Data"); - - for flow_count in [100u16, 500, 1000] { - let data = v9_data_packet(flow_count); - group.throughput(Throughput::Bytes(data.len() as u64)); - group.bench_with_input(BenchmarkId::from_parameter(flow_count), &data, |b, pkt| { - let mut parser = NetflowParser::default(); - let template_result = parser.parse_bytes(&template); - assert!(template_result.error.is_none()); - assert_eq!(template_result.packets.len(), 1); - - b.iter(|| { - let result = parser.parse_bytes(black_box(pkt)); - black_box(result.packets.len()); - }); - }); +#[derive(Clone, Copy)] +enum Protocol { + V9, + Ipfix, +} + +impl Protocol { + fn name(self) -> &'static str { + match self { + Self::V9 => "v9", + Self::Ipfix => "ipfix", + } + } + + fn template_packet(self) -> Vec { + match self { + Self::V9 => v9_template_packet(), + Self::Ipfix => ipfix_template_packet(), + } } - group.finish(); + fn data_packet(self, flow_count: u16) -> Vec { + match self { + Self::V9 => v9_data_packet(flow_count), + Self::Ipfix => ipfix_data_packet(flow_count), + } + } + + fn assert_decoded_records(self, packets: &[NetflowPacket], expected: usize) { + assert_eq!(packets.len(), 1, "fixture must decode one outer packet"); + let actual: usize = match (self, &packets[0]) { + (Self::V9, NetflowPacket::V9(packet)) => packet + .flowsets + .iter() + .map(|flowset| match &flowset.body { + netflow_parser::variable_versions::v9::FlowSetBody::Data(data) => { + data.fields.len() + } + _ => 0, + }) + .sum(), + (Self::Ipfix, NetflowPacket::IPFix(packet)) => packet + .flowsets + .iter() + .map(|flowset| match &flowset.body { + netflow_parser::variable_versions::ipfix::FlowSetBody::Data(data) => { + data.fields.len() + } + _ => 0, + }) + .sum(), + _ => panic!("fixture decoded as the wrong protocol"), + }; + assert_eq!(actual, expected, "fixture decoded-record count mismatch"); + } } -fn bench_warm_ipfix_data_hot_path(c: &mut Criterion) { - let template = ipfix_template_packet(); - let mut group = c.benchmark_group("Hot Path IPFIX Data"); - - for flow_count in [100u16, 500, 1000] { - let data = ipfix_data_packet(flow_count); - group.throughput(Throughput::Bytes(data.len() as u64)); - group.bench_with_input(BenchmarkId::from_parameter(flow_count), &data, |b, pkt| { - let mut parser = NetflowParser::default(); - let template_result = parser.parse_bytes(&template); - assert!(template_result.error.is_none()); - assert_eq!(template_result.packets.len(), 1); - - b.iter(|| { - let result = parser.parse_bytes(black_box(pkt)); - black_box(result.packets.len()); - }); - }); +#[derive(Clone, Copy)] +enum Scenario { + DirectParse, + DirectIterator, + AutoParse, + AutoIterator, +} + +impl Scenario { + fn name(self) -> &'static str { + match self { + Self::DirectParse => "direct/parse", + Self::DirectIterator => "direct/iterator", + Self::AutoParse => "auto/parse", + Self::AutoIterator => "auto/iterator", + } } +} - group.finish(); +fn bench_warmed_hot_paths(c: &mut Criterion) { + let source = SocketAddr::from(([192, 0, 2, 1], 2055)); + + for protocol in [Protocol::V9, Protocol::Ipfix] { + let template = protocol.template_packet(); + for scenario in [ + Scenario::DirectParse, + Scenario::DirectIterator, + Scenario::AutoParse, + Scenario::AutoIterator, + ] { + let mut group = + c.benchmark_group(format!("Hot Path/{}/{}", protocol.name(), scenario.name())); + + for flow_count in [1u16, 1000] { + let data = protocol.data_packet(flow_count); + group.throughput(Throughput::Elements(u64::from(flow_count))); + group.bench_with_input( + BenchmarkId::from_parameter(flow_count), + &data, + |b, packet| match scenario { + Scenario::DirectParse => { + let mut parser = NetflowParser::default(); + assert!(parser.parse_bytes(&template).is_ok()); + let result = parser.parse_bytes(packet); + assert!(result.is_ok()); + protocol.assert_decoded_records( + &result.packets, + usize::from(flow_count), + ); + b.iter(|| { + drop(black_box( + parser.parse_bytes(black_box(packet.as_slice())), + )); + }); + } + Scenario::DirectIterator => { + let mut parser = NetflowParser::default(); + assert!(parser.parse_bytes(&template).is_ok()); + let packets = parser + .iter_packets(packet) + .map(Result::unwrap) + .collect::>(); + protocol.assert_decoded_records(&packets, usize::from(flow_count)); + b.iter(|| { + for result in parser.iter_packets(black_box(packet.as_slice())) + { + black_box(result.unwrap()); + } + }); + } + Scenario::AutoParse => { + let mut parser = AutoScopedParser::new(); + assert!(parser.parse_from_source(source, &template).is_ok()); + let result = parser.parse_from_source(source, packet); + assert!(result.is_ok()); + protocol.assert_decoded_records( + &result.packets, + usize::from(flow_count), + ); + b.iter(|| { + drop(black_box( + parser.parse_from_source( + source, + black_box(packet.as_slice()), + ), + )); + }); + } + Scenario::AutoIterator => { + let mut parser = AutoScopedParser::new(); + assert!(parser.parse_from_source(source, &template).is_ok()); + let packets = parser + .iter_packets_from_source(source, packet) + .unwrap() + .map(Result::unwrap) + .collect::>(); + protocol.assert_decoded_records(&packets, usize::from(flow_count)); + b.iter(|| { + let iterator = parser + .iter_packets_from_source( + source, + black_box(packet.as_slice()), + ) + .unwrap(); + for result in iterator { + black_box(result.unwrap()); + } + }); + } + }, + ); + } + group.finish(); + } + } } -criterion_group!( - benches, - bench_warm_v9_data_hot_path, - bench_warm_ipfix_data_hot_path -); +criterion_group!(benches, bench_warmed_hot_paths); criterion_main!(benches); diff --git a/src/lib.rs b/src/lib.rs index 0cde9a6..5a27f51 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -58,7 +58,9 @@ pub use variable_versions::enterprise_registry::{EnterpriseFieldDef, EnterpriseF pub use variable_versions::metrics::{CacheInfo, CacheMetrics, ParserCacheInfo}; pub use variable_versions::ttl::TtlConfig; pub use variable_versions::{ - Config, ConfigError, DEFAULT_MAX_RECORDS_PER_FLOWSET, NoTemplateInfo, PendingFlowsConfig, + Config, ConfigError, DEFAULT_MAX_DECODED_FIELD_PAYLOAD_BYTES_PER_MESSAGE, + DEFAULT_MAX_DECODED_FIELD_VALUES_PER_MESSAGE, DEFAULT_MAX_RECORDS_PER_FLOWSET, + DecodedOutputLimit, DecodedOutputLimits, NoTemplateInfo, PendingFlowsConfig, }; // Rust-idiomatic naming aliases @@ -575,6 +577,69 @@ impl NetflowParserBuilder { self } + /// Sets the cumulative decoded field-value limit for both v9 and IPFIX messages. + /// + /// The default is 65,536. Zero is rejected by [`NetflowParserBuilder::build`]. + #[must_use = "builder methods consume self and return a new builder; the return value must be used"] + pub fn with_max_decoded_field_values_per_message(mut self, count: usize) -> Self { + self.v9_config.max_decoded_field_values_per_message = count; + self.ipfix_config.max_decoded_field_values_per_message = count; + self + } + + /// Sets the cumulative decoded field-value limit for v9 messages. + /// + /// The default is 65,536. Zero is rejected by [`NetflowParserBuilder::build`]. + #[must_use = "builder methods consume self and return a new builder; the return value must be used"] + pub fn with_v9_max_decoded_field_values_per_message(mut self, count: usize) -> Self { + self.v9_config.max_decoded_field_values_per_message = count; + self + } + + /// Sets the cumulative decoded field-value limit for IPFIX messages. + /// + /// The default is 65,536. Zero is rejected by [`NetflowParserBuilder::build`]. + #[must_use = "builder methods consume self and return a new builder; the return value must be used"] + pub fn with_ipfix_max_decoded_field_values_per_message(mut self, count: usize) -> Self { + self.ipfix_config.max_decoded_field_values_per_message = count; + self + } + + /// Sets the cumulative decoded field-content-byte limit for both protocols. + /// + /// The default is 4 MiB. IPFIX variable-length prefixes are excluded. + /// Zero is rejected by [`NetflowParserBuilder::build`]. + #[must_use = "builder methods consume self and return a new builder; the return value must be used"] + pub fn with_max_decoded_field_payload_bytes_per_message(mut self, bytes: usize) -> Self { + self.v9_config.max_decoded_field_payload_bytes_per_message = bytes; + self.ipfix_config + .max_decoded_field_payload_bytes_per_message = bytes; + self + } + + /// Sets the cumulative decoded field-content-byte limit for v9 messages. + /// + /// The default is 4 MiB. Zero is rejected by [`NetflowParserBuilder::build`]. + #[must_use = "builder methods consume self and return a new builder; the return value must be used"] + pub fn with_v9_max_decoded_field_payload_bytes_per_message(mut self, bytes: usize) -> Self { + self.v9_config.max_decoded_field_payload_bytes_per_message = bytes; + self + } + + /// Sets the cumulative decoded field-content-byte limit for IPFIX messages. + /// + /// The default is 4 MiB. Variable-length prefixes are excluded. + /// Zero is rejected by [`NetflowParserBuilder::build`]. + #[must_use = "builder methods consume self and return a new builder; the return value must be used"] + pub fn with_ipfix_max_decoded_field_payload_bytes_per_message( + mut self, + bytes: usize, + ) -> Self { + self.ipfix_config + .max_decoded_field_payload_bytes_per_message = bytes; + self + } + /// Registers a custom enterprise field definition for both V9 and IPFIX parsers. /// /// This allows library users to define their own enterprise-specific fields without @@ -880,6 +945,7 @@ impl NetflowParserBuilder { /// /// Returns an error if: /// - Template cache size is 0 + /// - A decoded-output limit is 0 /// - Parser initialization fails /// /// # Examples @@ -988,6 +1054,18 @@ pub enum NetflowError { /// Description of the partial parse result message: String, }, + + /// A complete v9/IPFIX message would exceed a configured decoded-output limit. + DecodedOutputLimitExceeded { + /// Protocol of the rejected message. + protocol: TemplateProtocol, + /// Which cumulative limit was exceeded. + limit: DecodedOutputLimit, + /// Configured finite limit. + configured: usize, + /// Cumulative amount required at the rejection boundary. + attempted: usize, + }, } impl std::fmt::Display for NetflowError { @@ -1031,6 +1109,15 @@ impl std::fmt::Display for NetflowError { NetflowError::Partial { message } => { write!(f, "Partial parse error: {}", message) } + NetflowError::DecodedOutputLimitExceeded { + protocol, + limit, + configured, + attempted, + } => write!( + f, + "Decoded output limit exceeded for {protocol:?}: {limit:?} attempted {attempted}, configured {configured}" + ), } } } diff --git a/src/tests.rs b/src/tests.rs index 4cbf9c3..4789120 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -504,6 +504,10 @@ mod base_tests { max_error_sample_size: 256, max_records_per_flowset: crate::variable_versions::config::DEFAULT_MAX_RECORDS_PER_FLOWSET, + max_decoded_field_values_per_message: + crate::DEFAULT_MAX_DECODED_FIELD_VALUES_PER_MESSAGE, + max_decoded_field_payload_bytes_per_message: + crate::DEFAULT_MAX_DECODED_FIELD_PAYLOAD_BYTES_PER_MESSAGE, ttl_config: None, enterprise_registry: Arc::new(EnterpriseFieldRegistry::new()), pending_flows_config: None, @@ -529,6 +533,10 @@ mod base_tests { max_error_sample_size: 256, max_records_per_flowset: crate::variable_versions::config::DEFAULT_MAX_RECORDS_PER_FLOWSET, + max_decoded_field_values_per_message: + crate::DEFAULT_MAX_DECODED_FIELD_VALUES_PER_MESSAGE, + max_decoded_field_payload_bytes_per_message: + crate::DEFAULT_MAX_DECODED_FIELD_PAYLOAD_BYTES_PER_MESSAGE, ttl_config: None, enterprise_registry: Arc::new(EnterpriseFieldRegistry::new()), pending_flows_config: None, @@ -577,6 +585,10 @@ mod base_tests { max_error_sample_size: 256, max_records_per_flowset: crate::variable_versions::config::DEFAULT_MAX_RECORDS_PER_FLOWSET, + max_decoded_field_values_per_message: + crate::DEFAULT_MAX_DECODED_FIELD_VALUES_PER_MESSAGE, + max_decoded_field_payload_bytes_per_message: + crate::DEFAULT_MAX_DECODED_FIELD_PAYLOAD_BYTES_PER_MESSAGE, ttl_config: None, enterprise_registry: Arc::new(EnterpriseFieldRegistry::new()), pending_flows_config: None, @@ -625,6 +637,10 @@ mod base_tests { max_error_sample_size: 256, max_records_per_flowset: crate::variable_versions::config::DEFAULT_MAX_RECORDS_PER_FLOWSET, + max_decoded_field_values_per_message: + crate::DEFAULT_MAX_DECODED_FIELD_VALUES_PER_MESSAGE, + max_decoded_field_payload_bytes_per_message: + crate::DEFAULT_MAX_DECODED_FIELD_PAYLOAD_BYTES_PER_MESSAGE, ttl_config: None, enterprise_registry: Arc::new(EnterpriseFieldRegistry::new()), pending_flows_config: None, @@ -672,6 +688,10 @@ mod base_tests { max_error_sample_size: 256, max_records_per_flowset: crate::variable_versions::config::DEFAULT_MAX_RECORDS_PER_FLOWSET, + max_decoded_field_values_per_message: + crate::DEFAULT_MAX_DECODED_FIELD_VALUES_PER_MESSAGE, + max_decoded_field_payload_bytes_per_message: + crate::DEFAULT_MAX_DECODED_FIELD_PAYLOAD_BYTES_PER_MESSAGE, ttl_config: None, enterprise_registry: Arc::new(EnterpriseFieldRegistry::new()), pending_flows_config: None, @@ -719,6 +739,10 @@ mod base_tests { max_error_sample_size: 256, max_records_per_flowset: crate::variable_versions::config::DEFAULT_MAX_RECORDS_PER_FLOWSET, + max_decoded_field_values_per_message: + crate::DEFAULT_MAX_DECODED_FIELD_VALUES_PER_MESSAGE, + max_decoded_field_payload_bytes_per_message: + crate::DEFAULT_MAX_DECODED_FIELD_PAYLOAD_BYTES_PER_MESSAGE, ttl_config: None, enterprise_registry: Arc::new(EnterpriseFieldRegistry::new()), pending_flows_config: None, @@ -769,6 +793,10 @@ mod base_tests { max_error_sample_size: 256, max_records_per_flowset: crate::variable_versions::config::DEFAULT_MAX_RECORDS_PER_FLOWSET, + max_decoded_field_values_per_message: + crate::DEFAULT_MAX_DECODED_FIELD_VALUES_PER_MESSAGE, + max_decoded_field_payload_bytes_per_message: + crate::DEFAULT_MAX_DECODED_FIELD_PAYLOAD_BYTES_PER_MESSAGE, ttl_config: None, enterprise_registry: Arc::new(EnterpriseFieldRegistry::new()), pending_flows_config: None, @@ -809,6 +837,10 @@ mod base_tests { max_error_sample_size: 256, max_records_per_flowset: crate::variable_versions::config::DEFAULT_MAX_RECORDS_PER_FLOWSET, + max_decoded_field_values_per_message: + crate::DEFAULT_MAX_DECODED_FIELD_VALUES_PER_MESSAGE, + max_decoded_field_payload_bytes_per_message: + crate::DEFAULT_MAX_DECODED_FIELD_PAYLOAD_BYTES_PER_MESSAGE, ttl_config: None, enterprise_registry: Arc::new(EnterpriseFieldRegistry::new()), pending_flows_config: None, @@ -844,6 +876,10 @@ mod base_tests { max_error_sample_size: 256, max_records_per_flowset: crate::variable_versions::config::DEFAULT_MAX_RECORDS_PER_FLOWSET, + max_decoded_field_values_per_message: + crate::DEFAULT_MAX_DECODED_FIELD_VALUES_PER_MESSAGE, + max_decoded_field_payload_bytes_per_message: + crate::DEFAULT_MAX_DECODED_FIELD_PAYLOAD_BYTES_PER_MESSAGE, ttl_config: None, enterprise_registry: Arc::new(EnterpriseFieldRegistry::new()), pending_flows_config: None, @@ -883,6 +919,10 @@ mod base_tests { max_error_sample_size: 256, max_records_per_flowset: crate::variable_versions::config::DEFAULT_MAX_RECORDS_PER_FLOWSET, + max_decoded_field_values_per_message: + crate::DEFAULT_MAX_DECODED_FIELD_VALUES_PER_MESSAGE, + max_decoded_field_payload_bytes_per_message: + crate::DEFAULT_MAX_DECODED_FIELD_PAYLOAD_BYTES_PER_MESSAGE, ttl_config: None, enterprise_registry: Arc::new(EnterpriseFieldRegistry::new()), pending_flows_config: None, diff --git a/src/variable_versions/config.rs b/src/variable_versions/config.rs index 634e394..cb224de 100644 --- a/src/variable_versions/config.rs +++ b/src/variable_versions/config.rs @@ -1,6 +1,10 @@ //! Parser configuration, traits, and constants for V9 and IPFIX parsers. use super::metrics::CacheMetricsInner; +use super::output_budget::{ + DEFAULT_MAX_DECODED_FIELD_PAYLOAD_BYTES_PER_MESSAGE, + DEFAULT_MAX_DECODED_FIELD_VALUES_PER_MESSAGE, +}; use super::pending_flows::{PendingFlowCache, PendingFlowsConfig}; use crate::template_store::TemplateStore; use crate::variable_versions::enterprise_registry::EnterpriseFieldRegistry; @@ -45,6 +49,10 @@ pub struct Config { /// Maximum number of data records to parse per flowset. Default: 1,024. /// This prevents CPU-bound DoS from maliciously large flowsets. pub max_records_per_flowset: usize, + /// Maximum decoded field values returned by one message. Default: 65,536. + pub max_decoded_field_values_per_message: usize, + /// Maximum decoded field content bytes returned by one message. Default: 4 MiB. + pub max_decoded_field_payload_bytes_per_message: usize, /// Optional TTL configuration for template expiration. pub ttl_config: Option, /// Registry of custom enterprise-specific field definitions for IPFIX. @@ -92,6 +100,10 @@ pub enum ConfigError { EmptyAllowedVersions, /// Max records per flowset must be greater than 0 InvalidRecordsPerFlowset(usize), + /// Decoded field-value message limit must be greater than 0. + InvalidDecodedFieldValueLimit(usize), + /// Decoded field-payload-byte message limit must be greater than 0. + InvalidDecodedFieldPayloadByteLimit(usize), /// Pending flow max_total_bytes must be >= max_entry_size_bytes InvalidPendingTotalBytes { max_total_bytes: usize, @@ -167,6 +179,20 @@ impl std::fmt::Display for ConfigError { count ) } + ConfigError::InvalidDecodedFieldValueLimit(count) => { + write!( + f, + "Invalid max decoded field values per message: {}. Must be greater than 0.", + count + ) + } + ConfigError::InvalidDecodedFieldPayloadByteLimit(bytes) => { + write!( + f, + "Invalid max decoded field payload bytes per message: {}. Must be greater than 0.", + bytes + ) + } ConfigError::EmptyAllowedVersions => { write!( f, @@ -205,6 +231,9 @@ impl Default for Config { max_template_total_size: usize::from(u16::MAX), max_error_sample_size: 256, max_records_per_flowset: DEFAULT_MAX_RECORDS_PER_FLOWSET, + max_decoded_field_values_per_message: DEFAULT_MAX_DECODED_FIELD_VALUES_PER_MESSAGE, + max_decoded_field_payload_bytes_per_message: + DEFAULT_MAX_DECODED_FIELD_PAYLOAD_BYTES_PER_MESSAGE, ttl_config: None, enterprise_registry: Arc::new(EnterpriseFieldRegistry::new()), pending_flows_config: None, @@ -254,6 +283,12 @@ impl Config { if self.max_records_per_flowset == 0 { return Err(ConfigError::InvalidRecordsPerFlowset(0)); } + if self.max_decoded_field_values_per_message == 0 { + return Err(ConfigError::InvalidDecodedFieldValueLimit(0)); + } + if self.max_decoded_field_payload_bytes_per_message == 0 { + return Err(ConfigError::InvalidDecodedFieldPayloadByteLimit(0)); + } if self.max_error_sample_size == 0 { return Err(ConfigError::InvalidErrorSampleSize(0)); } @@ -279,6 +314,7 @@ pub(crate) trait ParserFields { fn set_max_template_total_size_field(&mut self, size: usize); fn set_max_error_sample_size_field(&mut self, size: usize); fn set_max_records_per_flowset_field(&mut self, count: usize); + fn set_decoded_output_limits_fields(&mut self, values: usize, payload_bytes: usize); fn set_ttl_config_field(&mut self, config: Option); /// Apply an enterprise registry update. Default is a no-op (V9 has no registry). fn set_enterprise_registry(&mut self, _registry: Arc) {} @@ -308,6 +344,12 @@ pub trait ParserConfig: ParserFields { if config.max_records_per_flowset == 0 { return Err(ConfigError::InvalidRecordsPerFlowset(0)); } + if config.max_decoded_field_values_per_message == 0 { + return Err(ConfigError::InvalidDecodedFieldValueLimit(0)); + } + if config.max_decoded_field_payload_bytes_per_message == 0 { + return Err(ConfigError::InvalidDecodedFieldPayloadByteLimit(0)); + } if config.max_error_sample_size == 0 { return Err(ConfigError::InvalidErrorSampleSize(0)); } @@ -326,6 +368,10 @@ pub trait ParserConfig: ParserFields { self.set_max_template_total_size_field(config.max_template_total_size); self.set_max_error_sample_size_field(config.max_error_sample_size); self.set_max_records_per_flowset_field(config.max_records_per_flowset); + self.set_decoded_output_limits_fields( + config.max_decoded_field_values_per_message, + config.max_decoded_field_payload_bytes_per_message, + ); self.set_ttl_config_field(config.ttl_config); self.set_enterprise_registry(config.enterprise_registry); // Safety: validate_config above already verified pending_flows_config, diff --git a/src/variable_versions/ipfix/parser.rs b/src/variable_versions/ipfix/parser.rs index f71f9cc..b5b9a73 100644 --- a/src/variable_versions/ipfix/parser.rs +++ b/src/variable_versions/ipfix/parser.rs @@ -21,15 +21,17 @@ use crate::variable_versions::enterprise_registry::EnterpriseFieldRegistry; use crate::variable_versions::field_value::FieldValue; use crate::variable_versions::lazy_lru::LazyLruCache; use crate::variable_versions::metrics::CacheMetricsInner; +use crate::variable_versions::output_budget::{PendingOutputError, PendingOutputPreflight}; use crate::variable_versions::template_events::TemplateProtocol; use crate::variable_versions::ttl::{TemplateWithTtl, TtlConfig}; use crate::variable_versions::v9::{ DATA_TEMPLATE_V9_ID, Data as V9Data, OPTIONS_TEMPLATE_V9_ID, OptionsData as V9OptionsData, OptionsTemplate as V9OptionsTemplate, Template as V9Template, }; +use crate::variable_versions::wire::RecordBodyKind; use crate::variable_versions::{ - Config, ConfigError, ParserConfig, ParserFields, PendingFlowCache, PendingFlowEntry, - PendingFlowsConfig, + Config, ConfigError, DecodedOutputBudget, ParserConfig, ParserFields, PendingFlowCache, + PendingFlowEntry, PendingFlowsConfig, }; use crate::{NetflowError, NetflowPacket, ParsedNetflow}; @@ -41,6 +43,24 @@ use nom_derive::Parse; use std::num::NonZeroUsize; use std::sync::Arc; +const MIN_STORE_REPLAY_TRIGGER_MESSAGE_LENGTH: u16 = 20; +const MIN_TEMPLATE_REPLAY_TRIGGER_MESSAGE_LENGTH: u16 = 28; + +enum IpfixPendingReplayOutcome { + Replayed { new_header_length: u16 }, + TemporarilyDoesNotFit, + Failed, +} + +impl From for IpfixPendingReplayOutcome { + fn from(error: PendingOutputError) -> Self { + match error { + PendingOutputError::TemporarilyDoesNotFit => Self::TemporarilyDoesNotFit, + PendingOutputError::NeverFits | PendingOutputError::Invalid => Self::Failed, + } + } +} + impl Default for IPFixParser { fn default() -> Self { // Safe to unwrap because DEFAULT_MAX_TEMPLATE_CACHE_SIZE is non-zero @@ -50,6 +70,10 @@ impl Default for IPFixParser { max_template_total_size: usize::from(u16::MAX), max_error_sample_size: 256, max_records_per_flowset: DEFAULT_MAX_RECORDS_PER_FLOWSET, + max_decoded_field_values_per_message: + crate::DEFAULT_MAX_DECODED_FIELD_VALUES_PER_MESSAGE, + max_decoded_field_payload_bytes_per_message: + crate::DEFAULT_MAX_DECODED_FIELD_PAYLOAD_BYTES_PER_MESSAGE, ttl_config: None, enterprise_registry: Arc::new(EnterpriseFieldRegistry::new()), pending_flows_config: None, @@ -65,6 +89,14 @@ impl Default for IPFixParser { } impl IPFixParser { + pub(crate) fn start_decoded_output_message(&mut self) { + self.decoded_output_budget.reset(); + } + + pub(crate) fn decoded_output_limit_was_exceeded(&self) -> bool { + self.decoded_output_budget.is_exceeded() + } + /// Validates a configuration without allocating parser internals. pub fn validate_config(config: &Config) -> Result<(), ConfigError> { config.validate() @@ -76,11 +108,18 @@ impl IPFixParser { /// * `config` - Configuration struct containing max_template_cache_size and optional ttl_config /// /// # Errors - /// Returns `ConfigError` if `max_template_cache_size` is 0 + /// Returns `ConfigError` if the template cache size or either decoded-output + /// limit is zero. pub fn try_new(config: Config) -> Result { let cache_size = NonZeroUsize::new(config.max_template_cache_size).ok_or( ConfigError::InvalidCacheSize(config.max_template_cache_size), )?; + if config.max_decoded_field_values_per_message == 0 { + return Err(ConfigError::InvalidDecodedFieldValueLimit(0)); + } + if config.max_decoded_field_payload_bytes_per_message == 0 { + return Err(ConfigError::InvalidDecodedFieldPayloadByteLimit(0)); + } let pending_flows = config .pending_flows_config @@ -98,6 +137,10 @@ impl IPFixParser { max_template_total_size: config.max_template_total_size, max_error_sample_size: config.max_error_sample_size, max_records_per_flowset: config.max_records_per_flowset, + decoded_output_budget: DecodedOutputBudget::new( + config.max_decoded_field_values_per_message, + config.max_decoded_field_payload_bytes_per_message, + ), enterprise_registry: config.enterprise_registry, metrics: CacheMetricsInner::new(), pending_flows, @@ -424,6 +467,9 @@ impl ParserFields for IPFixParser { fn set_max_records_per_flowset_field(&mut self, count: usize) { self.max_records_per_flowset = count; } + fn set_decoded_output_limits_fields(&mut self, values: usize, payload_bytes: usize) { + self.decoded_output_budget.set_limits(values, payload_bytes); + } fn set_ttl_config_field(&mut self, config: Option) { self.ttl_config = config; } @@ -490,11 +536,21 @@ impl IPFixParser { remaining, } } - Err(e) => ParsedNetflow::Error { - error: NetflowError::Partial { - message: format!("IPFIX parse error: {}", e), - }, - }, + Err(e) => { + let error = if let Some(exceeded) = self.decoded_output_budget.take_exceeded() { + NetflowError::DecodedOutputLimitExceeded { + protocol: TemplateProtocol::Ipfix, + limit: exceeded.limit, + configured: exceeded.configured, + attempted: exceeded.attempted, + } + } else { + NetflowError::Partial { + message: format!("IPFIX parse error: {}", e), + } + }; + ParsedNetflow::Error { error } + } } } @@ -628,28 +684,36 @@ impl IPFixParser { ) { for &template_id in learned { let entries = cache.drain(template_id, &mut self.metrics); - let total_entries = entries.len(); - for (processed, entry) in entries.iter().enumerate() { + let mut entries = entries.into_iter(); + while let Some(entry) = entries.next() { // Bound flowset count, consistent with V9 replay. if ipfix.flowsets.len() >= u16::MAX as usize { - let remaining = (total_entries - processed) as u64; - self.metrics.record_pending_replay_failed_n(remaining); + let mut retained = Vec::with_capacity(entries.len().saturating_add(1)); + retained.push(entry); + retained.extend(entries); + cache.restore_replay_suffix(template_id, retained); break; } - let flowset_length = - u16::try_from(entry.raw_data.len().saturating_add(4)).unwrap_or(u16::MAX); - let Some(new_header_length) = ipfix.header.length.checked_add(flowset_length) - else { - // Count this entry plus all remaining as failed. - let remaining = (total_entries - processed) as u64; - self.metrics.record_pending_replay_failed_n(remaining); - break; - }; - if self.try_replay_ipfix_flow(&mut ipfix.flowsets, template_id, entry) { - self.metrics.record_pending_replayed(); - ipfix.header.length = new_header_length; - } else { - self.metrics.record_pending_replay_failed(); + match self.try_replay_ipfix_flow( + &mut ipfix.flowsets, + template_id, + &entry, + ipfix.header.length, + ) { + IpfixPendingReplayOutcome::Replayed { new_header_length } => { + self.metrics.record_pending_replayed(); + ipfix.header.length = new_header_length; + } + IpfixPendingReplayOutcome::Failed => { + self.metrics.record_pending_replay_failed(); + } + IpfixPendingReplayOutcome::TemporarilyDoesNotFit => { + let mut retained = Vec::with_capacity(entries.len().saturating_add(1)); + retained.push(entry); + retained.extend(entries); + cache.restore_replay_suffix(template_id, retained); + break; + } } } } @@ -661,105 +725,300 @@ impl IPFixParser { flowsets: &mut Vec, template_id: u16, entry: &PendingFlowEntry, - ) -> bool { + current_header_length: u16, + ) -> IpfixPendingReplayOutcome { // Use peek_valid_template to avoid false LRU promotion on failed parse. - // Promote only after successful replay. + // Preserve the established lookup priority and promote only after a + // complete entry has been materialized within the message budget. + self.try_replay_ipfix_data(flowsets, template_id, entry, current_header_length) + .or_else(|| { + self.try_replay_ipfix_options_data( + flowsets, + template_id, + entry, + current_header_length, + ) + }) + .or_else(|| { + self.try_replay_embedded_v9_data( + flowsets, + template_id, + entry, + current_header_length, + ) + }) + .or_else(|| { + self.try_replay_embedded_v9_options_data( + flowsets, + template_id, + entry, + current_header_length, + ) + }) + .unwrap_or(IpfixPendingReplayOutcome::Failed) + } - // Try IPFIX templates - if let Some(template) = crate::variable_versions::peek_valid_template( + fn try_replay_ipfix_data( + &mut self, + flowsets: &mut Vec, + template_id: u16, + entry: &PendingFlowEntry, + current_header_length: u16, + ) -> Option { + let template = crate::variable_versions::peek_valid_template( &mut self.templates, &template_id, &self.ttl_config, &mut self.metrics, - ) && let Ok((_, data)) = Data::parse_with_registry( + )?; + let preflight = crate::variable_versions::output_budget::measure_variable_output( &entry.raw_data, - &template, - &self.enterprise_registry, + template.get_fields(), self.max_records_per_flowset, - ) { - // Don't record_hit() — the original flowset already recorded - // a miss. Replay success is tracked via record_pending_replayed(). - self.templates.promote(&template_id); - flowsets.push(FlowSet { - header: FlowSetHeader { - header_id: template_id, - length: u16::try_from(entry.raw_data.len().saturating_add(4)) - .unwrap_or(u16::MAX), - }, - body: FlowSetBody::Data(data), - }); - return true; - } + |field| field.field_length, + RecordBodyKind::Ipfix, + ); + let (preflight, flowset_length, new_header_length) = match self + .validate_ipfix_pending_replay( + preflight, + entry.raw_data.len(), + current_header_length, + ) { + Ok(result) => result, + Err(PendingOutputError::Invalid) => return None, + Err(error) => return Some(error.into()), + }; + let (mut data, padding_len) = + match self + .decoded_output_budget + .materialize_pending(Some(preflight), |budget| { + Data::parse_with_registry_and_budget( + &entry.raw_data, + &template, + &self.enterprise_registry, + self.max_records_per_flowset, + budget, + ) + }) { + Ok(result) => result, + Err(PendingOutputError::Invalid) => return None, + Err(error) => return Some(error.into()), + }; + data.padding = entry.raw_data[entry.raw_data.len() - padding_len..].to_vec(); + self.templates.promote(&template_id); + flowsets.push(Self::replayed_flowset( + template_id, + flowset_length, + FlowSetBody::Data(data), + )); + Some(IpfixPendingReplayOutcome::Replayed { new_header_length }) + } - // Try IPFIX options templates - if let Some(template) = crate::variable_versions::peek_valid_template( + fn try_replay_ipfix_options_data( + &mut self, + flowsets: &mut Vec, + template_id: u16, + entry: &PendingFlowEntry, + current_header_length: u16, + ) -> Option { + let template = crate::variable_versions::peek_valid_template( &mut self.ipfix_options_templates, &template_id, &self.ttl_config, &mut self.metrics, - ) && let Ok((_, data)) = OptionsData::parse_with_registry( + )?; + let preflight = crate::variable_versions::output_budget::measure_variable_output( &entry.raw_data, - &template, - &self.enterprise_registry, + template.get_fields(), self.max_records_per_flowset, - ) { - self.ipfix_options_templates.promote(&template_id); - flowsets.push(FlowSet { - header: FlowSetHeader { - header_id: template_id, - length: u16::try_from(entry.raw_data.len().saturating_add(4)) - .unwrap_or(u16::MAX), - }, - body: FlowSetBody::OptionsData(data), - }); - return true; - } + |field| field.field_length, + RecordBodyKind::Ipfix, + ); + let (preflight, flowset_length, new_header_length) = match self + .validate_ipfix_pending_replay( + preflight, + entry.raw_data.len(), + current_header_length, + ) { + Ok(result) => result, + Err(PendingOutputError::Invalid) => return None, + Err(error) => return Some(error.into()), + }; + let (mut data, padding_len) = + match self + .decoded_output_budget + .materialize_pending(Some(preflight), |budget| { + OptionsData::parse_with_registry_and_budget( + &entry.raw_data, + &template, + &self.enterprise_registry, + self.max_records_per_flowset, + budget, + ) + }) { + Ok(result) => result, + Err(PendingOutputError::Invalid) => return None, + Err(error) => return Some(error.into()), + }; + data.padding = entry.raw_data[entry.raw_data.len() - padding_len..].to_vec(); + self.ipfix_options_templates.promote(&template_id); + flowsets.push(Self::replayed_flowset( + template_id, + flowset_length, + FlowSetBody::OptionsData(data), + )); + Some(IpfixPendingReplayOutcome::Replayed { new_header_length }) + } - // Try V9 templates - if let Some(template) = crate::variable_versions::peek_valid_template( + fn try_replay_embedded_v9_data( + &mut self, + flowsets: &mut Vec, + template_id: u16, + entry: &PendingFlowEntry, + current_header_length: u16, + ) -> Option { + let template = crate::variable_versions::peek_valid_template( &mut self.v9_templates, &template_id, &self.ttl_config, &mut self.metrics, - ) && let Ok((_, data)) = - V9Data::parse_with_limit(&entry.raw_data, &template, self.max_records_per_flowset) - { - self.v9_templates.promote(&template_id); - flowsets.push(FlowSet { - header: FlowSetHeader { - header_id: template_id, - length: u16::try_from(entry.raw_data.len().saturating_add(4)) - .unwrap_or(u16::MAX), - }, - body: FlowSetBody::V9Data(data), - }); - return true; - } + )?; + let preflight = V9Data::decoded_output_preflight( + &entry.raw_data, + &template, + self.max_records_per_flowset, + RecordBodyKind::Ipfix, + ); + let (preflight, flowset_length, new_header_length) = match self + .validate_ipfix_pending_replay( + preflight, + entry.raw_data.len(), + current_header_length, + ) { + Ok(result) => result, + Err(PendingOutputError::Invalid) => return None, + Err(error) => return Some(error.into()), + }; + let (mut data, padding_len) = + match self + .decoded_output_budget + .materialize_pending(Some(preflight), |budget| { + V9Data::parse_with_budget( + &entry.raw_data, + &template, + self.max_records_per_flowset, + budget, + ) + }) { + Ok(result) => result, + Err(PendingOutputError::Invalid) => return None, + Err(error) => return Some(error.into()), + }; + data.padding = entry.raw_data[entry.raw_data.len() - padding_len..].to_vec(); + self.v9_templates.promote(&template_id); + flowsets.push(Self::replayed_flowset( + template_id, + flowset_length, + FlowSetBody::V9Data(data), + )); + Some(IpfixPendingReplayOutcome::Replayed { new_header_length }) + } - // Try V9 options templates - if let Some(template) = crate::variable_versions::peek_valid_template( + fn try_replay_embedded_v9_options_data( + &mut self, + flowsets: &mut Vec, + template_id: u16, + entry: &PendingFlowEntry, + current_header_length: u16, + ) -> Option { + let template = crate::variable_versions::peek_valid_template( &mut self.v9_options_templates, &template_id, &self.ttl_config, &mut self.metrics, - ) && let Ok((_, data)) = V9OptionsData::parse_with_limit( + )?; + let preflight = V9OptionsData::decoded_output_preflight( &entry.raw_data, &template, self.max_records_per_flowset, - ) { - self.v9_options_templates.promote(&template_id); - flowsets.push(FlowSet { - header: FlowSetHeader { - header_id: template_id, - length: u16::try_from(entry.raw_data.len().saturating_add(4)) - .unwrap_or(u16::MAX), - }, - body: FlowSetBody::V9OptionsData(data), - }); - return true; + RecordBodyKind::Ipfix, + ); + let (preflight, flowset_length, new_header_length) = match self + .validate_ipfix_pending_replay( + preflight, + entry.raw_data.len(), + current_header_length, + ) { + Ok(result) => result, + Err(PendingOutputError::Invalid) => return None, + Err(error) => return Some(error.into()), + }; + let data = + match self + .decoded_output_budget + .materialize_pending(Some(preflight), |budget| { + V9OptionsData::parse_with_budget( + &entry.raw_data, + &template, + self.max_records_per_flowset, + budget, + ) + }) { + Ok((data, _)) => data, + Err(PendingOutputError::Invalid) => return None, + Err(error) => return Some(error.into()), + }; + self.v9_options_templates.promote(&template_id); + flowsets.push(Self::replayed_flowset( + template_id, + flowset_length, + FlowSetBody::V9OptionsData(data), + )); + Some(IpfixPendingReplayOutcome::Replayed { new_header_length }) + } + + fn validate_ipfix_pending_replay( + &self, + preflight: Option, + body_length: usize, + current_header_length: u16, + ) -> Result<(PendingOutputPreflight, u16, u16), PendingOutputError> { + let preflight = self + .decoded_output_budget + .validate_pending_full_budget(preflight)?; + let flowset_length = body_length + .checked_add(4) + .and_then(|length| u16::try_from(length).ok()) + .ok_or(PendingOutputError::NeverFits)?; + + // A store-backed parser can restore a template from an empty Data Set + // (16-byte message header plus 4-byte Set header). Without a store, + // replay requires a valid one-field Template Set, which needs 28 bytes. + let minimum_trigger_length = if self.template_store.is_some() { + MIN_STORE_REPLAY_TRIGGER_MESSAGE_LENGTH + } else { + MIN_TEMPLATE_REPLAY_TRIGGER_MESSAGE_LENGTH + }; + if minimum_trigger_length.checked_add(flowset_length).is_none() { + return Err(PendingOutputError::NeverFits); } - false + self.decoded_output_budget + .validate_pending_remaining(preflight)?; + let new_header_length = current_header_length + .checked_add(flowset_length) + .ok_or(PendingOutputError::TemporarilyDoesNotFit)?; + Ok((preflight, flowset_length, new_header_length)) + } + + fn replayed_flowset(template_id: u16, flowset_length: u16, body: FlowSetBody) -> FlowSet { + FlowSet { + header: FlowSetHeader { + header_id: template_id, + length: flowset_length, + }, + body, + } } /// Returns a sorted, deduplicated list of all available template IDs. @@ -1202,11 +1461,12 @@ impl FlowSetBody { if template.get_fields().is_empty() { return Ok((i, FlowSetBody::Empty)); } - let (i, data) = Data::parse_with_registry( + let (i, data) = Data::parse_with_registry_and_budget( i, &template, &parser.enterprise_registry, parser.max_records_per_flowset, + &mut parser.decoded_output_budget, )?; return Ok((i, FlowSetBody::Data(data))); } @@ -1222,11 +1482,12 @@ impl FlowSetBody { if template.get_fields().is_empty() { return Ok((i, FlowSetBody::Empty)); } - let (i, data) = OptionsData::parse_with_registry( + let (i, data) = OptionsData::parse_with_registry_and_budget( i, &template, &parser.enterprise_registry, parser.max_records_per_flowset, + &mut parser.decoded_output_budget, )?; return Ok((i, FlowSetBody::OptionsData(data))); } @@ -1239,8 +1500,12 @@ impl FlowSetBody { &mut parser.metrics, ) { parser.metrics.record_hit(); - let (i, data) = - V9Data::parse_with_limit(i, &template, parser.max_records_per_flowset)?; + let (i, data) = V9Data::parse_with_budget( + i, + &template, + parser.max_records_per_flowset, + &mut parser.decoded_output_budget, + )?; return Ok((i, FlowSetBody::V9Data(data))); } @@ -1252,10 +1517,11 @@ impl FlowSetBody { &mut parser.metrics, ) { parser.metrics.record_hit(); - let (i, data) = V9OptionsData::parse_with_limit( + let (i, data) = V9OptionsData::parse_with_budget( i, &template, parser.max_records_per_flowset, + &mut parser.decoded_output_budget, )?; return Ok((i, FlowSetBody::V9OptionsData(data))); } @@ -1271,11 +1537,12 @@ impl FlowSetBody { if template.get_fields().is_empty() { return Ok((i, FlowSetBody::Empty)); } - let (i, data) = Data::parse_with_registry( + let (i, data) = Data::parse_with_registry_and_budget( i, &template, &parser.enterprise_registry, parser.max_records_per_flowset, + &mut parser.decoded_output_budget, )?; return Ok((i, FlowSetBody::Data(data))); } @@ -1284,26 +1551,32 @@ impl FlowSetBody { if template.get_fields().is_empty() { return Ok((i, FlowSetBody::Empty)); } - let (i, data) = OptionsData::parse_with_registry( + let (i, data) = OptionsData::parse_with_registry_and_budget( i, &template, &parser.enterprise_registry, parser.max_records_per_flowset, + &mut parser.decoded_output_budget, )?; return Ok((i, FlowSetBody::OptionsData(data))); } if let Some(template) = parser.fetch_v9_template_from_store(id) { parser.metrics.record_hit(); - let (i, data) = - V9Data::parse_with_limit(i, &template, parser.max_records_per_flowset)?; + let (i, data) = V9Data::parse_with_budget( + i, + &template, + parser.max_records_per_flowset, + &mut parser.decoded_output_budget, + )?; return Ok((i, FlowSetBody::V9Data(data))); } if let Some(template) = parser.fetch_v9_options_template_from_store(id) { parser.metrics.record_hit(); - let (i, data) = V9OptionsData::parse_with_limit( + let (i, data) = V9OptionsData::parse_with_budget( i, &template, parser.max_records_per_flowset, + &mut parser.decoded_output_budget, )?; return Ok((i, FlowSetBody::V9OptionsData(data))); } @@ -1366,15 +1639,41 @@ fn collect_varlen_field_lengths(fields: &[TemplateField]) -> Vec { } impl Data { - /// Parse Data using the enterprise registry to resolve custom enterprise fields - pub(super) fn parse_with_registry<'a>( + pub(super) fn parse_with_registry_and_budget<'a>( i: &'a [u8], template: &Template, registry: &EnterpriseFieldRegistry, max_records: usize, + budget: &mut DecodedOutputBudget, + ) -> IResult<&'a [u8], Self> { + let template_field_lengths = collect_varlen_field_lengths(template.get_fields()); + let (i, fields) = FieldParser::parse_with_registry_and_budget( + i, + template, + registry, + max_records, + budget, + )?; + Ok(( + i, + Self { + fields, + padding: vec![], + template_field_lengths, + }, + )) + } + + /// Parse one data body with explicit finite output limits. + pub fn parse_with_limits<'a>( + i: &'a [u8], + template: &Template, + limits: crate::DecodedOutputLimits, ) -> IResult<&'a [u8], Self> { let template_field_lengths = collect_varlen_field_lengths(template.get_fields()); - let (i, fields) = FieldParser::parse_with_registry(i, template, registry, max_records)?; + let mut budget = limits.budget(); + let (i, fields) = + FieldParser::parse_with_budget(i, template, limits.max_records(), &mut budget)?; Ok(( i, Self { @@ -1387,15 +1686,41 @@ impl Data { } impl OptionsData { - /// Parse OptionsData using the enterprise registry to resolve custom enterprise fields - pub(super) fn parse_with_registry<'a>( + pub(super) fn parse_with_registry_and_budget<'a>( i: &'a [u8], template: &OptionsTemplate, registry: &EnterpriseFieldRegistry, max_records: usize, + budget: &mut DecodedOutputBudget, + ) -> IResult<&'a [u8], Self> { + let template_field_lengths = collect_varlen_field_lengths(template.get_fields()); + let (i, fields) = FieldParser::parse_with_registry_and_budget( + i, + template, + registry, + max_records, + budget, + )?; + Ok(( + i, + Self { + fields, + padding: vec![], + template_field_lengths, + }, + )) + } + + /// Parse one options-data body with explicit finite output limits. + pub fn parse_with_limits<'a>( + i: &'a [u8], + template: &OptionsTemplate, + limits: crate::DecodedOutputLimits, ) -> IResult<&'a [u8], Self> { let template_field_lengths = collect_varlen_field_lengths(template.get_fields()); - let (i, fields) = FieldParser::parse_with_registry(i, template, registry, max_records)?; + let mut budget = limits.budget(); + let (i, fields) = + FieldParser::parse_with_budget(i, template, limits.max_records(), &mut budget)?; Ok(( i, Self { @@ -1416,6 +1741,7 @@ impl<'a> FieldParser { mut i: &'a [u8], template: &T, max_records: usize, + budget: &mut DecodedOutputBudget, parse_field: F, ) -> IResult<&'a [u8], Vec>> where @@ -1439,15 +1765,56 @@ impl<'a> FieldParser { } }) .sum(); - // template_fields is non-empty (checked above) and each contributes >= 1 byte, - // so template_size is always > 0 here. - let estimated_records = (i.len() / template_size).min(max_records); + if template_size == 0 { + return Err(nom::Err::Error(nom::error::Error::new( + i, + nom::error::ErrorKind::Verify, + ))); + } + let field_count = template_fields.len(); + let all_fixed = template_fields + .iter() + .all(|field| field.field_length != u16::MAX); + let fixed_payload_per_record = all_fixed.then_some(template_size); + let (remaining_values, remaining_payload) = budget.remaining(); + let mut estimated_records = (i.len() / template_size) + .min(max_records) + .min(remaining_values / field_count); + if let Some(payload_per_record) = fixed_payload_per_record { + estimated_records = estimated_records.min(remaining_payload / payload_per_record); + } let mut res = Vec::with_capacity(estimated_records); // Try to parse as much as we can, but if it fails, just return what we have so far. while !i.is_empty() && res.len() < max_records { + if let Some(payload_per_record) = fixed_payload_per_record + && i.len() < payload_per_record + { + break; + } let before = i; - let mut vec = Vec::with_capacity(template_fields.len()); + let checkpoint = budget.checkpoint(); + let record_payload = if let Some(payload_per_record) = fixed_payload_per_record { + payload_per_record + } else { + let Some((_, payload)) = + crate::variable_versions::output_budget::scan_variable_record( + i, + template_fields, + |field| field.field_length, + ) + else { + break; + }; + payload + }; + if budget.reserve(field_count, record_payload).is_err() { + return Err(nom::Err::Error(nom::error::Error::new( + i, + nom::error::ErrorKind::TooLarge, + ))); + } + let mut vec = Vec::with_capacity(field_count); for field in template_fields.iter() { match parse_field(field, i) { Ok((remaining, field_value)) => { @@ -1455,6 +1822,7 @@ impl<'a> FieldParser { i = remaining; } Err(_) => { + budget.rollback(checkpoint); i = before; return Ok((i, res)); } @@ -1463,6 +1831,7 @@ impl<'a> FieldParser { // Guard against infinite loops: if no bytes were consumed after // parsing a full record, stop to prevent CPU-bound DoS. if std::ptr::eq(i, before) { + budget.rollback(checkpoint); break; } res.push(vec); @@ -1478,19 +1847,34 @@ impl<'a> FieldParser { template: &T, max_records: usize, ) -> IResult<&'a [u8], Vec>> { - Self::parse_inner(i, template, max_records, |field, input| { + let mut budget = DecodedOutputBudget::new( + crate::DEFAULT_MAX_DECODED_FIELD_VALUES_PER_MESSAGE, + crate::DEFAULT_MAX_DECODED_FIELD_PAYLOAD_BYTES_PER_MESSAGE, + ); + Self::parse_inner(i, template, max_records, &mut budget, |field, input| { + field.parse_as_field_value(input) + }) + } + + fn parse_with_budget( + i: &'a [u8], + template: &T, + max_records: usize, + budget: &mut DecodedOutputBudget, + ) -> IResult<&'a [u8], Vec>> { + Self::parse_inner(i, template, max_records, budget, |field, input| { field.parse_as_field_value(input) }) } - /// Same as parse but uses the enterprise registry to resolve custom enterprise fields - fn parse_with_registry( + fn parse_with_registry_and_budget( i: &'a [u8], template: &T, registry: &EnterpriseFieldRegistry, max_records: usize, + budget: &mut DecodedOutputBudget, ) -> IResult<&'a [u8], Vec>> { - Self::parse_inner(i, template, max_records, |field, input| { + Self::parse_inner(i, template, max_records, budget, |field, input| { field.parse_as_field_value_with_registry(input, registry) }) } diff --git a/src/variable_versions/ipfix/types.rs b/src/variable_versions/ipfix/types.rs index 3ad1445..7a5fbcb 100644 --- a/src/variable_versions/ipfix/types.rs +++ b/src/variable_versions/ipfix/types.rs @@ -6,6 +6,7 @@ use super::lookup::IPFixField; use crate::template_store::TemplateStore; +use crate::variable_versions::DecodedOutputBudget; use crate::variable_versions::PendingFlowCache; use crate::variable_versions::enterprise_registry::EnterpriseFieldRegistry; use crate::variable_versions::field_value::FieldValue; @@ -49,6 +50,7 @@ pub struct IPFixParser { pub(crate) max_template_total_size: usize, pub(crate) max_error_sample_size: usize, pub(crate) max_records_per_flowset: usize, + pub(crate) decoded_output_budget: DecodedOutputBudget, pub(crate) enterprise_registry: Arc, pub(crate) metrics: CacheMetricsInner, pub(crate) pending_flows: Option, @@ -66,7 +68,11 @@ pub struct IPFixParser { /// A parsed IPFIX message containing a header and a list of flowsets. #[derive(Nom, Debug, PartialEq, Clone, Serialize)] -#[nom(ExtraArgs(parser: &mut IPFixParser))] +#[nom( + ExtraArgs(parser: &mut IPFixParser), + PreExec = "parser.start_decoded_output_message();", + PostExec = "if parser.decoded_output_limit_was_exceeded() { return Err(nom::Err::Error(nom::error::Error::new(i, nom::error::ErrorKind::TooLarge))); }" +)] pub struct IPFix { /// IPFix Header pub header: Header, diff --git a/src/variable_versions/metrics.rs b/src/variable_versions/metrics.rs index 0b7e837..bca266b 100644 --- a/src/variable_versions/metrics.rs +++ b/src/variable_versions/metrics.rs @@ -124,12 +124,6 @@ impl CacheMetricsInner { self.pending_replay_failed = self.pending_replay_failed.saturating_add(1); } - /// Record multiple pending flows that failed to replay at once - #[inline] - pub(crate) fn record_pending_replay_failed_n(&mut self, n: u64) { - self.pending_replay_failed = self.pending_replay_failed.saturating_add(n); - } - /// Record a successful read-through against the secondary template store. #[inline] pub(crate) fn record_template_store_restored(&mut self) { diff --git a/src/variable_versions/mod.rs b/src/variable_versions/mod.rs index 6c65cda..95b9fdc 100644 --- a/src/variable_versions/mod.rs +++ b/src/variable_versions/mod.rs @@ -75,10 +75,12 @@ pub mod field_value; pub mod ipfix; pub(crate) mod lazy_lru; pub mod metrics; +pub(crate) mod output_budget; pub(crate) mod pending_flows; pub mod template_events; pub mod ttl; pub mod v9; +pub(crate) mod wire; // Re-export public types to preserve existing import paths pub use config::ParserConfig; @@ -87,10 +89,15 @@ pub use config::{ Config, ConfigError, DEFAULT_MAX_RECORDS_PER_FLOWSET, DEFAULT_MAX_TEMPLATE_CACHE_SIZE, MAX_FIELD_COUNT, }; +pub use output_budget::{ + DEFAULT_MAX_DECODED_FIELD_PAYLOAD_BYTES_PER_MESSAGE, + DEFAULT_MAX_DECODED_FIELD_VALUES_PER_MESSAGE, DecodedOutputLimit, DecodedOutputLimits, +}; pub use pending_flows::PendingFlowsConfig; // Re-export crate-internal types for use by sibling modules pub(crate) use config::TemplateId; +pub(crate) use output_budget::{DecodedOutputBudget, PendingReplayOutcome}; pub(crate) use pending_flows::{PendingFlowCache, PendingFlowEntry}; use crate::variable_versions::metrics::CacheMetricsInner; diff --git a/src/variable_versions/output_budget.rs b/src/variable_versions/output_budget.rs new file mode 100644 index 0000000..6306c2e --- /dev/null +++ b/src/variable_versions/output_budget.rs @@ -0,0 +1,566 @@ +//! Cumulative decoded-output accounting shared by NetFlow v9 and IPFIX. + +use crate::variable_versions::wire::{ + RecordBodyKind, minimum_record_size, record_body_is_complete, +}; + +/// Default maximum number of decoded field values returned by one message. +pub const DEFAULT_MAX_DECODED_FIELD_VALUES_PER_MESSAGE: usize = 65_536; + +/// Default maximum number of decoded field payload bytes returned by one message. +pub const DEFAULT_MAX_DECODED_FIELD_PAYLOAD_BYTES_PER_MESSAGE: usize = 4 * 1024 * 1024; + +/// The cumulative decoded-output limit that rejected a message. +#[non_exhaustive] +#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)] +pub enum DecodedOutputLimit { + /// Number of materialized field values. + FieldValues, + /// Sum of materialized field content bytes. + FieldPayloadBytes, +} + +/// Finite limits for advanced one-body `Data` and `OptionsData` parsing. +/// +/// Use [`DecodedOutputLimits::new`] to override the same finite defaults used +/// by the stateful parser. All three values must be greater than zero. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DecodedOutputLimits { + max_records: usize, + max_field_values: usize, + max_field_payload_bytes: usize, +} + +impl DecodedOutputLimits { + /// Validate and create a bounded one-body parsing policy. + pub fn new( + max_records: usize, + max_field_values: usize, + max_field_payload_bytes: usize, + ) -> Result { + if max_records == 0 { + return Err(crate::ConfigError::InvalidRecordsPerFlowset(0)); + } + if max_field_values == 0 { + return Err(crate::ConfigError::InvalidDecodedFieldValueLimit(0)); + } + if max_field_payload_bytes == 0 { + return Err(crate::ConfigError::InvalidDecodedFieldPayloadByteLimit(0)); + } + Ok(Self { + max_records, + max_field_values, + max_field_payload_bytes, + }) + } + + pub(crate) fn max_records(self) -> usize { + self.max_records + } + + pub(crate) fn budget(self) -> DecodedOutputBudget { + DecodedOutputBudget::new(self.max_field_values, self.max_field_payload_bytes) + } +} + +impl Default for DecodedOutputLimits { + fn default() -> Self { + Self { + max_records: crate::DEFAULT_MAX_RECORDS_PER_FLOWSET, + max_field_values: DEFAULT_MAX_DECODED_FIELD_VALUES_PER_MESSAGE, + max_field_payload_bytes: DEFAULT_MAX_DECODED_FIELD_PAYLOAD_BYTES_PER_MESSAGE, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct OutputBudgetExceeded { + pub(crate) limit: DecodedOutputLimit, + pub(crate) configured: usize, + pub(crate) attempted: usize, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum PendingOutputError { + TemporarilyDoesNotFit, + NeverFits, + Invalid, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum PendingReplayOutcome { + Replayed, + TemporarilyDoesNotFit, + Failed, +} + +/// Exact allocation-free framing and decoded-cost result for one queued body. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct PendingOutputPreflight { + field_values: usize, + field_payload_bytes: usize, + remainder_len: usize, +} + +impl PendingOutputPreflight { + #[inline] + fn cost(self) -> (usize, usize) { + (self.field_values, self.field_payload_bytes) + } +} + +impl From for PendingReplayOutcome { + fn from(error: PendingOutputError) -> Self { + match error { + PendingOutputError::TemporarilyDoesNotFit => Self::TemporarilyDoesNotFit, + PendingOutputError::NeverFits | PendingOutputError::Invalid => Self::Failed, + } + } +} + +/// Per-message accounting state. It is reset before every independently parsed +/// v9/IPFIX message and deliberately performs no allocation. +#[derive(Debug, Clone)] +pub(crate) struct DecodedOutputBudget { + max_values: usize, + max_payload_bytes: usize, + used_values: usize, + used_payload_bytes: usize, + exceeded: Option, +} + +impl DecodedOutputBudget { + pub(crate) fn new(max_values: usize, max_payload_bytes: usize) -> Self { + Self { + max_values, + max_payload_bytes, + used_values: 0, + used_payload_bytes: 0, + exceeded: None, + } + } + + #[inline] + pub(crate) fn reset(&mut self) { + self.used_values = 0; + self.used_payload_bytes = 0; + self.exceeded = None; + } + + pub(crate) fn set_limits(&mut self, max_values: usize, max_payload_bytes: usize) { + self.max_values = max_values; + self.max_payload_bytes = max_payload_bytes; + self.reset(); + } + + #[inline] + pub(crate) fn checkpoint(&self) -> (usize, usize) { + (self.used_values, self.used_payload_bytes) + } + + #[inline] + pub(crate) fn rollback(&mut self, checkpoint: (usize, usize)) { + self.used_values = checkpoint.0; + self.used_payload_bytes = checkpoint.1; + self.exceeded = None; + } + + #[inline] + pub(crate) fn reserve( + &mut self, + values: usize, + payload_bytes: usize, + ) -> Result<(), OutputBudgetExceeded> { + let Some(attempted_values) = self.used_values.checked_add(values) else { + return Err(self.record_exceeded(DecodedOutputLimit::FieldValues, usize::MAX)); + }; + if attempted_values > self.max_values { + return Err(self.record_exceeded(DecodedOutputLimit::FieldValues, attempted_values)); + } + + let Some(attempted_payload) = self.used_payload_bytes.checked_add(payload_bytes) else { + return Err(self.record_exceeded(DecodedOutputLimit::FieldPayloadBytes, usize::MAX)); + }; + if attempted_payload > self.max_payload_bytes { + return Err( + self.record_exceeded(DecodedOutputLimit::FieldPayloadBytes, attempted_payload) + ); + } + + self.used_values = attempted_values; + self.used_payload_bytes = attempted_payload; + Ok(()) + } + + #[inline] + fn record_exceeded( + &mut self, + limit: DecodedOutputLimit, + attempted: usize, + ) -> OutputBudgetExceeded { + let configured = match limit { + DecodedOutputLimit::FieldValues => self.max_values, + DecodedOutputLimit::FieldPayloadBytes => self.max_payload_bytes, + }; + let exceeded = OutputBudgetExceeded { + limit, + configured, + attempted, + }; + self.exceeded.get_or_insert(exceeded); + exceeded + } + + #[inline] + pub(crate) fn take_exceeded(&mut self) -> Option { + self.exceeded.take() + } + + #[inline] + pub(crate) fn is_exceeded(&self) -> bool { + self.exceeded.is_some() + } + + #[inline] + pub(crate) fn used(&self) -> (usize, usize) { + (self.used_values, self.used_payload_bytes) + } + + #[inline] + pub(crate) fn remaining(&self) -> (usize, usize) { + ( + self.max_values.saturating_sub(self.used_values), + self.max_payload_bytes + .saturating_sub(self.used_payload_bytes), + ) + } + + /// Preflight one queued body, materialize it only when it fits the current + /// message, then commit its actual decoded cost. + pub(crate) fn materialize_pending<'a, T, E>( + &mut self, + preflight: Option, + parse: impl FnOnce(&mut DecodedOutputBudget) -> Result<(&'a [u8], T), E>, + ) -> Result<(T, usize), PendingOutputError> { + let preflight = self.validate_pending_full_budget(preflight)?; + self.validate_pending_remaining(preflight)?; + let measured = preflight.cost(); + + let mut scratch = Self::new(self.max_values, self.max_payload_bytes); + let (remaining, value) = + parse(&mut scratch).map_err(|_| PendingOutputError::Invalid)?; + let actual = scratch.used(); + if actual != measured || remaining.len() != preflight.remainder_len { + return Err(PendingOutputError::Invalid); + } + if self.reserve(actual.0, actual.1).is_err() { + return Err(PendingOutputError::Invalid); + } + Ok((value, preflight.remainder_len)) + } + + pub(crate) fn validate_pending_full_budget( + &self, + preflight: Option, + ) -> Result { + let Some(preflight) = preflight else { + return Err(PendingOutputError::Invalid); + }; + let measured = preflight.cost(); + if measured.0 > self.max_values || measured.1 > self.max_payload_bytes { + return Err(PendingOutputError::NeverFits); + } + Ok(preflight) + } + + pub(crate) fn validate_pending_remaining( + &self, + preflight: PendingOutputPreflight, + ) -> Result<(), PendingOutputError> { + let measured = preflight.cost(); + let remaining = self.remaining(); + if measured.0 > remaining.0 || measured.1 > remaining.1 { + return Err(PendingOutputError::TemporarilyDoesNotFit); + } + Ok(()) + } +} + +pub(crate) fn measure_fixed_output( + input: &[u8], + fields: &[T], + max_records: usize, + width: impl Fn(&T) -> u16, + body_kind: RecordBodyKind, +) -> Option { + if fields.is_empty() { + return None; + } + measure_fixed_widths( + input, + fields.len(), + fields.iter().map(width), + max_records, + body_kind, + ) +} + +pub(crate) fn measure_fixed_widths( + input: &[u8], + field_count: usize, + widths: impl IntoIterator, + max_records: usize, + body_kind: RecordBodyKind, +) -> Option { + if field_count == 0 { + return None; + } + let payload_per_record = minimum_record_size(widths, None)?; + if payload_per_record == 0 { + return None; + } + let records = (input.len() / payload_per_record).min(max_records); + let consumed = records.checked_mul(payload_per_record)?; + let remainder = input.get(consumed..)?; + if !record_body_is_complete(records, remainder, payload_per_record, body_kind) { + return None; + } + Some(PendingOutputPreflight { + field_values: records.checked_mul(field_count)?, + field_payload_bytes: consumed, + remainder_len: remainder.len(), + }) +} + +pub(crate) fn measure_variable_output( + mut input: &[u8], + fields: &[T], + max_records: usize, + width: impl Fn(&T) -> u16, + body_kind: RecordBodyKind, +) -> Option { + if fields.is_empty() { + return None; + } + + let minimum_record_size = minimum_record_size(fields.iter().map(&width), Some(u16::MAX))?; + if minimum_record_size == 0 { + return None; + } + + let mut records = 0usize; + let mut payload_bytes = 0usize; + while !input.is_empty() && records < max_records { + let Some((remaining, record_payload)) = scan_variable_record(input, fields, &width) + else { + break; + }; + input = remaining; + records = records.checked_add(1)?; + payload_bytes = payload_bytes.checked_add(record_payload)?; + } + if !record_body_is_complete(records, input, minimum_record_size, body_kind) { + return None; + } + Some(PendingOutputPreflight { + field_values: records.checked_mul(fields.len())?, + field_payload_bytes: payload_bytes, + remainder_len: input.len(), + }) +} + +/// Scan one complete variable-width record without allocating. The returned +/// payload excludes RFC 7011 length prefixes. `None` means the bytes form only +/// an incomplete trailing record/padding or the template cannot make progress. +pub(crate) fn scan_variable_record<'a, T>( + input: &'a [u8], + fields: &[T], + width: impl Fn(&T) -> u16, +) -> Option<(&'a [u8], usize)> { + let mut remaining = input; + let mut payload_bytes = 0usize; + for field in fields { + let content_length = match width(field) { + u16::MAX => { + let (&first, rest) = remaining.split_first()?; + remaining = rest; + if first == 255 { + let bytes = remaining.get(..2)?; + remaining = &remaining[2..]; + usize::from(u16::from_be_bytes([bytes[0], bytes[1]])) + } else { + usize::from(first) + } + } + fixed => usize::from(fixed), + }; + remaining = remaining.get(content_length..)?; + payload_bytes = payload_bytes.checked_add(content_length)?; + } + (remaining.len() < input.len()).then_some((remaining, payload_bytes)) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::cell::Cell; + + fn preflight(field_values: usize, field_payload_bytes: usize) -> PendingOutputPreflight { + PendingOutputPreflight { + field_values, + field_payload_bytes, + remainder_len: 0, + } + } + + #[test] + fn pending_preflight_does_not_materialize_when_current_message_is_full() { + let mut budget = DecodedOutputBudget::new(2, 2); + budget.reserve(1, 1).unwrap(); + let called = Cell::new(false); + + let result = + budget.materialize_pending(Some(preflight(2, 2)), |_: &mut DecodedOutputBudget| { + called.set(true); + Ok::<_, ()>((&[][..], ())) + }); + + assert_eq!(result, Err(PendingOutputError::TemporarilyDoesNotFit)); + assert!(!called.get()); + } + + #[test] + fn pending_preflight_does_not_materialize_an_entry_that_can_never_fit() { + for cost in [(3, 1), (1, 3)] { + let mut budget = DecodedOutputBudget::new(2, 2); + let called = Cell::new(false); + + let result = budget.materialize_pending( + Some(preflight(cost.0, cost.1)), + |_: &mut DecodedOutputBudget| { + called.set(true); + Ok::<_, ()>((&[][..], ())) + }, + ); + + assert_eq!(result, Err(PendingOutputError::NeverFits)); + assert!(!called.get()); + } + } + + #[test] + fn pending_materialization_rejects_a_preflight_cost_mismatch() { + let mut budget = DecodedOutputBudget::new(3, 3); + + let result = budget.materialize_pending(Some(preflight(1, 1)), |scratch| { + scratch.reserve(2, 1).unwrap(); + Ok::<_, ()>((&[][..], ())) + }); + + assert_eq!(result, Err(PendingOutputError::Invalid)); + assert_eq!(budget.used(), (0, 0)); + } + + #[test] + fn pending_materialization_rejects_a_preflight_boundary_mismatch() { + let mut budget = DecodedOutputBudget::new(3, 3); + let expected = PendingOutputPreflight { + field_values: 1, + field_payload_bytes: 1, + remainder_len: 1, + }; + + let result = budget.materialize_pending(Some(expected), |scratch| { + scratch.reserve(1, 1).unwrap(); + Ok::<_, ()>((&[][..], ())) + }); + + assert_eq!(result, Err(PendingOutputError::Invalid)); + assert_eq!(budget.used(), (0, 0)); + } + + #[test] + fn wire_preflight_counts_both_zero_length_prefixes_as_progress() { + let widths = [u16::MAX, 1]; + for input in [&[0, 7][..], &[255, 0, 0, 7][..]] { + assert_eq!( + scan_variable_record(input, &widths, |width| *width), + Some((&[][..], 1)) + ); + assert_eq!( + measure_variable_output( + input, + &widths, + 1, + |width| *width, + RecordBodyKind::Ipfix, + ) + .map(PendingOutputPreflight::cost), + Some((2, 1)), + ); + } + } + + #[test] + fn exact_limits_succeed_and_one_over_is_reported() { + let mut budget = DecodedOutputBudget::new(2, 4); + assert!(budget.reserve(2, 4).is_ok()); + assert_eq!(budget.used(), (2, 4)); + + let error = budget.reserve(1, 0).unwrap_err(); + assert_eq!(error.limit, DecodedOutputLimit::FieldValues); + assert_eq!(error.configured, 2); + assert_eq!(error.attempted, 3); + assert_eq!(budget.used(), (2, 4)); + } + + #[test] + fn checked_overflow_is_a_limit_failure() { + let mut budget = DecodedOutputBudget::new(usize::MAX - 1, usize::MAX - 1); + assert!(budget.reserve(usize::MAX - 2, 0).is_ok()); + let error = budget.reserve(4, 0).unwrap_err(); + assert_eq!(error.attempted, usize::MAX); + } + + #[test] + fn value_overflow_is_rejected_at_the_maximum_limit() { + let mut budget = DecodedOutputBudget::new(usize::MAX, usize::MAX); + assert!(budget.reserve(usize::MAX - 1, 0).is_ok()); + + let error = budget.reserve(2, 0).unwrap_err(); + assert_eq!(error.limit, DecodedOutputLimit::FieldValues); + assert_eq!(error.configured, usize::MAX); + assert_eq!(error.attempted, usize::MAX); + assert_eq!(budget.used(), (usize::MAX - 1, 0)); + } + + #[test] + fn payload_overflow_is_rejected_at_the_maximum_limit() { + let mut budget = DecodedOutputBudget::new(usize::MAX, usize::MAX); + assert!(budget.reserve(0, usize::MAX - 1).is_ok()); + + let error = budget.reserve(0, 2).unwrap_err(); + assert_eq!(error.limit, DecodedOutputLimit::FieldPayloadBytes); + assert_eq!(error.configured, usize::MAX); + assert_eq!(error.attempted, usize::MAX); + assert_eq!(budget.used(), (0, usize::MAX - 1)); + } + + #[test] + fn rollback_restores_both_counters_and_failure_state() { + let mut budget = DecodedOutputBudget::new(2, 2); + let checkpoint = budget.checkpoint(); + assert!(budget.reserve(2, 2).is_ok()); + assert!(budget.reserve(1, 0).is_err()); + budget.rollback(checkpoint); + assert_eq!(budget.used(), (0, 0)); + assert!(budget.take_exceeded().is_none()); + } + + #[test] + fn public_limits_reject_zero_values() { + assert!(DecodedOutputLimits::new(0, 1, 1).is_err()); + assert!(DecodedOutputLimits::new(1, 0, 1).is_err()); + assert!(DecodedOutputLimits::new(1, 1, 0).is_err()); + } +} diff --git a/src/variable_versions/pending_flows.rs b/src/variable_versions/pending_flows.rs index c50138f..474c851 100644 --- a/src/variable_versions/pending_flows.rs +++ b/src/variable_versions/pending_flows.rs @@ -337,6 +337,32 @@ impl PendingFlowCache { result } + /// Restore an unprocessed FIFO suffix after replay reaches a cumulative + /// message budget. The entries came from `drain` immediately before this + /// call, so reinserting them cannot exceed the cache's prior byte or key + /// limits and must not change pending-flow metrics. + pub(crate) fn restore_replay_suffix( + &mut self, + template_id: u16, + entries: Vec, + ) { + if entries.is_empty() { + return; + } + let restored_bytes = entries + .iter() + .map(|entry| entry.raw_data.len()) + .fold(0usize, usize::saturating_add); + let displaced = self.cache.push(template_id, entries); + debug_assert!( + displaced.is_none(), + "drained pending key was unexpectedly replaced" + ); + self.total_bytes = self.total_bytes.saturating_add(restored_bytes); + #[cfg(debug_assertions)] + self.debug_verify_total_bytes(); + } + /// Remove expired entries from a single template's vector. /// If all entries expire, the key is removed from the cache. fn prune_expired_for_template( diff --git a/src/variable_versions/v9/parser.rs b/src/variable_versions/v9/parser.rs index 5a1042d..404b772 100644 --- a/src/variable_versions/v9/parser.rs +++ b/src/variable_versions/v9/parser.rs @@ -21,11 +21,13 @@ use crate::variable_versions::enterprise_registry::EnterpriseFieldRegistry; use crate::variable_versions::field_value::FieldValue; use crate::variable_versions::lazy_lru::LazyLruCache; use crate::variable_versions::metrics::CacheMetricsInner; +use crate::variable_versions::output_budget::PendingOutputError; use crate::variable_versions::template_events::TemplateProtocol; use crate::variable_versions::ttl::{TemplateWithTtl, TtlConfig}; +use crate::variable_versions::wire::RecordBodyKind; use crate::variable_versions::{ - Config, ConfigError, ParserConfig, ParserFields, PendingFlowCache, PendingFlowEntry, - PendingFlowsConfig, + Config, ConfigError, DecodedOutputBudget, ParserConfig, ParserFields, PendingFlowCache, + PendingFlowEntry, PendingFlowsConfig, PendingReplayOutcome, }; use crate::{NetflowError, NetflowPacket, ParsedNetflow}; @@ -48,6 +50,7 @@ pub struct V9Parser { pub(crate) max_template_total_size: usize, pub(crate) max_error_sample_size: usize, pub(crate) max_records_per_flowset: usize, + pub(crate) decoded_output_budget: DecodedOutputBudget, pub(crate) metrics: CacheMetricsInner, pub(crate) pending_flows: Option, /// Optional secondary-tier template store. See [`crate::template_store`]. @@ -71,6 +74,10 @@ impl Default for V9Parser { max_template_total_size: usize::from(u16::MAX), max_error_sample_size: 256, max_records_per_flowset: DEFAULT_MAX_RECORDS_PER_FLOWSET, + max_decoded_field_values_per_message: + crate::DEFAULT_MAX_DECODED_FIELD_VALUES_PER_MESSAGE, + max_decoded_field_payload_bytes_per_message: + crate::DEFAULT_MAX_DECODED_FIELD_PAYLOAD_BYTES_PER_MESSAGE, ttl_config: None, enterprise_registry: Arc::new(EnterpriseFieldRegistry::new()), pending_flows_config: None, @@ -86,6 +93,10 @@ impl Default for V9Parser { } impl V9Parser { + pub(crate) fn start_decoded_output_message(&mut self) { + self.decoded_output_budget.reset(); + } + /// Validates a configuration without allocating parser internals. pub fn validate_config(config: &Config) -> Result<(), ConfigError> { config.validate() @@ -97,11 +108,18 @@ impl V9Parser { /// * `config` - Configuration struct containing max_template_cache_size and optional ttl_config /// /// # Errors - /// Returns `ConfigError` if `max_template_cache_size` is 0 + /// Returns `ConfigError` if the template cache size or either decoded-output + /// limit is zero. pub fn try_new(config: Config) -> Result { let cache_size = NonZeroUsize::new(config.max_template_cache_size).ok_or( ConfigError::InvalidCacheSize(config.max_template_cache_size), )?; + if config.max_decoded_field_values_per_message == 0 { + return Err(ConfigError::InvalidDecodedFieldValueLimit(0)); + } + if config.max_decoded_field_payload_bytes_per_message == 0 { + return Err(ConfigError::InvalidDecodedFieldPayloadByteLimit(0)); + } let pending_flows = config .pending_flows_config @@ -117,6 +135,10 @@ impl V9Parser { max_template_total_size: config.max_template_total_size, max_error_sample_size: config.max_error_sample_size, max_records_per_flowset: config.max_records_per_flowset, + decoded_output_budget: DecodedOutputBudget::new( + config.max_decoded_field_values_per_message, + config.max_decoded_field_payload_bytes_per_message, + ), metrics: CacheMetricsInner::new(), pending_flows, template_store: config.template_store, @@ -353,6 +375,9 @@ impl ParserFields for V9Parser { fn set_max_records_per_flowset_field(&mut self, count: usize) { self.max_records_per_flowset = count; } + fn set_decoded_output_limits_fields(&mut self, values: usize, payload_bytes: usize) { + self.decoded_output_budget.set_limits(values, payload_bytes); + } fn set_ttl_config_field(&mut self, config: Option) { self.ttl_config = config; } @@ -414,11 +439,21 @@ impl V9Parser { remaining, } } - Err(e) => ParsedNetflow::Error { - error: NetflowError::Partial { - message: format!("V9 parse error: {}", e), - }, - }, + Err(e) => { + let error = if let Some(exceeded) = self.decoded_output_budget.take_exceeded() { + NetflowError::DecodedOutputLimitExceeded { + protocol: TemplateProtocol::V9, + limit: exceeded.limit, + configured: exceeded.configured, + attempted: exceeded.attempted, + } + } else { + NetflowError::Partial { + message: format!("V9 parse error: {}", e), + } + }; + ParsedNetflow::Error { error } + } } } @@ -513,18 +548,27 @@ impl V9Parser { ) { for &template_id in learned { let entries = cache.drain(template_id, &mut self.metrics); - let total_entries = entries.len(); - for (processed, entry) in entries.iter().enumerate() { + let mut entries = entries.into_iter(); + while let Some(entry) = entries.next() { if v9.flowsets.len() >= u16::MAX as usize { - // Count this entry plus all remaining as failed, then break. - let remaining = (total_entries - processed) as u64; - self.metrics.record_pending_replay_failed_n(remaining); + let mut retained = Vec::with_capacity(entries.len().saturating_add(1)); + retained.push(entry); + retained.extend(entries); + cache.restore_replay_suffix(template_id, retained); break; } - if self.try_replay_v9_flow(&mut v9.flowsets, template_id, entry) { - self.metrics.record_pending_replayed(); - } else { - self.metrics.record_pending_replay_failed(); + match self.try_replay_v9_flow(&mut v9.flowsets, template_id, &entry) { + PendingReplayOutcome::Replayed => self.metrics.record_pending_replayed(), + PendingReplayOutcome::Failed => { + self.metrics.record_pending_replay_failed(); + } + PendingReplayOutcome::TemporarilyDoesNotFit => { + let mut retained = Vec::with_capacity(entries.len().saturating_add(1)); + retained.push(entry); + retained.extend(entries); + cache.restore_replay_suffix(template_id, retained); + break; + } } } } @@ -537,29 +581,50 @@ impl V9Parser { flowsets: &mut Vec, template_id: u16, entry: &PendingFlowEntry, - ) -> bool { + ) -> PendingReplayOutcome { // Try regular template (peek to avoid false LRU promotion on failed parse) if let Some(template) = crate::variable_versions::peek_valid_template( &mut self.templates, &template_id, &self.ttl_config, &mut self.metrics, - ) && let Ok((_, data)) = - Data::parse_with_limit(&entry.raw_data, &template, self.max_records_per_flowset) - { - // Don't record_hit() here — the original flowset already - // recorded a miss. Replay success is tracked separately - // via record_pending_replayed() in the caller. - self.templates.promote(&template_id); - flowsets.push(FlowSet { - header: FlowSetHeader { - flowset_id: template_id, - length: u16::try_from(entry.raw_data.len().saturating_add(4)) - .unwrap_or(u16::MAX), - }, - body: FlowSetBody::Data(data), - }); - return true; + ) { + let preflight = Data::decoded_output_preflight( + &entry.raw_data, + &template, + self.max_records_per_flowset, + RecordBodyKind::NetFlowV9, + ); + match self + .decoded_output_budget + .materialize_pending(preflight, |budget| { + Data::parse_with_budget( + &entry.raw_data, + &template, + self.max_records_per_flowset, + budget, + ) + }) { + Ok((mut data, padding_len)) => { + data.padding = + entry.raw_data[entry.raw_data.len() - padding_len..].to_vec(); + // Don't record_hit() here — the original flowset already + // recorded a miss. Replay success is tracked separately + // via record_pending_replayed() in the caller. + self.templates.promote(&template_id); + flowsets.push(FlowSet { + header: FlowSetHeader { + flowset_id: template_id, + length: u16::try_from(entry.raw_data.len().saturating_add(4)) + .unwrap_or(u16::MAX), + }, + body: FlowSetBody::Data(data), + }); + return PendingReplayOutcome::Replayed; + } + Err(PendingOutputError::Invalid) => {} + Err(error) => return error.into(), + } } // Try options template (peek to avoid false LRU promotion on failed parse) if let Some(template) = crate::variable_versions::peek_valid_template( @@ -567,11 +632,27 @@ impl V9Parser { &template_id, &self.ttl_config, &mut self.metrics, - ) && let Ok((_, options_data)) = OptionsData::parse_with_limit( - &entry.raw_data, - &template, - self.max_records_per_flowset, ) { + let preflight = OptionsData::decoded_output_preflight( + &entry.raw_data, + &template, + self.max_records_per_flowset, + RecordBodyKind::NetFlowV9, + ); + let options_data = + match self + .decoded_output_budget + .materialize_pending(preflight, |budget| { + OptionsData::parse_with_budget( + &entry.raw_data, + &template, + self.max_records_per_flowset, + budget, + ) + }) { + Ok((data, _)) => data, + Err(error) => return error.into(), + }; self.options_templates.promote(&template_id); flowsets.push(FlowSet { header: FlowSetHeader { @@ -581,9 +662,9 @@ impl V9Parser { }, body: FlowSetBody::OptionsData(options_data), }); - return true; + return PendingReplayOutcome::Replayed; } - false + PendingReplayOutcome::Failed } /// Returns a sorted, deduplicated list of all available template IDs. @@ -711,8 +792,12 @@ impl FlowSetBody { &mut parser.metrics, ) { parser.metrics.record_hit(); - let (i, data) = - Data::parse_with_limit(i, &template, parser.max_records_per_flowset)?; + let (i, data) = Data::parse_with_budget( + i, + &template, + parser.max_records_per_flowset, + &mut parser.decoded_output_budget, + )?; return Ok((i, FlowSetBody::Data(data))); } @@ -724,10 +809,11 @@ impl FlowSetBody { &mut parser.metrics, ) { parser.metrics.record_hit(); - let (i, options_data) = OptionsData::parse_with_limit( + let (i, options_data) = OptionsData::parse_with_budget( i, &template, parser.max_records_per_flowset, + &mut parser.decoded_output_budget, )?; return Ok((i, FlowSetBody::OptionsData(options_data))); } @@ -738,16 +824,21 @@ impl FlowSetBody { // served from the hot path. if let Some(template) = parser.fetch_template_from_store(id) { parser.metrics.record_hit(); - let (i, data) = - Data::parse_with_limit(i, &template, parser.max_records_per_flowset)?; + let (i, data) = Data::parse_with_budget( + i, + &template, + parser.max_records_per_flowset, + &mut parser.decoded_output_budget, + )?; return Ok((i, FlowSetBody::Data(data))); } if let Some(template) = parser.fetch_options_template_from_store(id) { parser.metrics.record_hit(); - let (i, options_data) = OptionsData::parse_with_limit( + let (i, options_data) = OptionsData::parse_with_budget( i, &template, parser.max_records_per_flowset, + &mut parser.decoded_output_budget, )?; return Ok((i, FlowSetBody::OptionsData(options_data))); } @@ -1023,9 +1114,23 @@ impl FlowSetParser { impl<'a> FieldParser { #[inline] pub(super) fn parse( + input: &'a [u8], + template: &Template, + max_records: usize, + ) -> IResult<&'a [u8], Vec>> { + let mut budget = DecodedOutputBudget::new( + crate::DEFAULT_MAX_DECODED_FIELD_VALUES_PER_MESSAGE, + crate::DEFAULT_MAX_DECODED_FIELD_PAYLOAD_BYTES_PER_MESSAGE, + ); + Self::parse_with_budget(input, template, max_records, &mut budget) + } + + #[inline] + pub(super) fn parse_with_budget( mut input: &'a [u8], template: &Template, max_records: usize, + budget: &mut DecodedOutputBudget, ) -> IResult<&'a [u8], Vec>> { let template_fields = &template.fields; // Estimate per-record size for capacity pre-allocation. @@ -1047,11 +1152,19 @@ impl<'a> FieldParser { // Calculate how many complete records we can parse based on input length let record_count = (input.len() / template_total_size).min(max_records); - let mut res = Vec::with_capacity(record_count); let field_count = template_fields.len(); + let (remaining_values, remaining_payload) = budget.remaining(); + let bounded_capacity = record_count + .min(remaining_values / field_count) + .min(remaining_payload / template_total_size); + let mut res = Vec::with_capacity(bounded_capacity); for _ in 0..record_count { let before = input; + let checkpoint = budget.checkpoint(); + if budget.reserve(field_count, template_total_size).is_err() { + return Err(nom::Err::Error(NomError::new(input, ErrorKind::TooLarge))); + } let mut record = Vec::with_capacity(field_count); for template_field in template_fields { @@ -1061,6 +1174,7 @@ impl<'a> FieldParser { record.push((template_field.field_type, field_value)); } Err(_) => { + budget.rollback(checkpoint); input = before; return Ok((input, res)); } @@ -1070,6 +1184,7 @@ impl<'a> FieldParser { // Guard against infinite loops: if no bytes were consumed after // parsing a full record, stop to prevent CPU-bound DoS. if std::ptr::eq(input, before) { + budget.rollback(checkpoint); break; } res.push(record); diff --git a/src/variable_versions/v9/types.rs b/src/variable_versions/v9/types.rs index 4382402..93f5120 100644 --- a/src/variable_versions/v9/types.rs +++ b/src/variable_versions/v9/types.rs @@ -7,6 +7,8 @@ use super::lookup::{ScopeFieldType, V9Field}; use crate::variable_versions::field_value::FieldValue; +use crate::DecodedOutputLimits; +use crate::variable_versions::DecodedOutputBudget; use nom::bytes::complete::take; use nom::combinator::{complete, map_res}; use nom_derive::{Nom, Parse}; @@ -25,7 +27,10 @@ use super::parser::V9Parser; /// A parsed NetFlow V9 packet containing a header and a list of flowsets. #[derive(Debug, PartialEq, Clone, Serialize, Nom)] -#[nom(ExtraArgs(parser: &mut V9Parser))] +#[nom( + ExtraArgs(parser: &mut V9Parser), + PreExec = "parser.start_decoded_output_message();" +)] pub struct V9 { /// V9 Header pub header: Header, @@ -180,22 +185,117 @@ impl OptionsData { template: &OptionsTemplate, max_records: usize, ) -> nom::IResult<&'a [u8], Self> { - let mut fields = Vec::new(); + let mut budget = DecodedOutputBudget::new( + crate::DEFAULT_MAX_DECODED_FIELD_VALUES_PER_MESSAGE, + crate::DEFAULT_MAX_DECODED_FIELD_PAYLOAD_BYTES_PER_MESSAGE, + ); + Self::parse_with_budget(i, template, max_records, &mut budget) + } + + pub(crate) fn parse_with_budget<'a>( + i: &'a [u8], + template: &OptionsTemplate, + max_records: usize, + budget: &mut DecodedOutputBudget, + ) -> nom::IResult<&'a [u8], Self> { + let values_per_record = template + .scope_fields + .len() + .saturating_add(template.option_fields.len()); + let payload_per_record = template + .scope_fields + .iter() + .map(|field| usize::from(field.field_length)) + .chain( + template + .option_fields + .iter() + .map(|field| usize::from(field.field_length)), + ) + .fold(0usize, usize::saturating_add); + if values_per_record == 0 || payload_per_record == 0 { + return Err(nom::Err::Error(nom::error::Error::new( + i, + nom::error::ErrorKind::Verify, + ))); + } + + let (remaining_values, remaining_payload) = budget.remaining(); + let bounded_capacity = max_records + .min(i.len() / payload_per_record) + .min(remaining_values / values_per_record) + .min(remaining_payload / payload_per_record); + let mut fields = Vec::with_capacity(bounded_capacity); let mut remaining = i; while !remaining.is_empty() && fields.len() < max_records { + if remaining.len() < payload_per_record { + break; + } + let checkpoint = budget.checkpoint(); + if budget + .reserve(values_per_record, payload_per_record) + .is_err() + { + return Err(nom::Err::Error(nom::error::Error::new( + remaining, + nom::error::ErrorKind::TooLarge, + ))); + } match complete(|i| OptionsDataFields::parse(i, template))(remaining) { Ok((i, record)) => { if std::ptr::eq(i, remaining) { + budget.rollback(checkpoint); break; } remaining = i; fields.push(record); } - Err(_) => break, + Err(_) => { + budget.rollback(checkpoint); + break; + } } } Ok((remaining, Self { fields })) } + + pub(crate) fn decoded_output_preflight( + input: &[u8], + template: &OptionsTemplate, + max_records: usize, + body_kind: crate::variable_versions::wire::RecordBodyKind, + ) -> Option { + let field_count = template + .scope_fields + .len() + .checked_add(template.option_fields.len())?; + crate::variable_versions::output_budget::measure_fixed_widths( + input, + field_count, + template + .scope_fields + .iter() + .map(|field| field.field_length) + .chain( + template + .option_fields + .iter() + .map(|field| field.field_length), + ), + max_records, + body_kind, + ) + } + + /// Parse one options-data body with explicit finite output limits. + pub fn parse_with_limits<'a>( + i: &'a [u8], + template: &OptionsTemplate, + limits: DecodedOutputLimits, + ) -> nom::IResult<&'a [u8], Self> { + let mut budget = limits.budget(); + Self::parse_with_budget(i, template, limits.max_records(), &mut budget) + } } pub(crate) struct ScopeParser; @@ -244,13 +344,13 @@ impl Data { } } - /// Parse data records with a configurable maximum record limit. - pub(crate) fn parse_with_limit<'a>( + pub(crate) fn parse_with_budget<'a>( i: &'a [u8], template: &Template, max_records: usize, + budget: &mut DecodedOutputBudget, ) -> nom::IResult<&'a [u8], Self> { - let (i, fields) = FieldParser::parse(i, template, max_records)?; + let (i, fields) = FieldParser::parse_with_budget(i, template, max_records, budget)?; Ok(( i, Self { @@ -259,6 +359,31 @@ impl Data { }, )) } + + pub(crate) fn decoded_output_preflight( + input: &[u8], + template: &Template, + max_records: usize, + body_kind: crate::variable_versions::wire::RecordBodyKind, + ) -> Option { + crate::variable_versions::output_budget::measure_fixed_output( + input, + &template.fields, + max_records, + |field| field.field_length, + body_kind, + ) + } + + /// Parse one data body with explicit finite output limits. + pub fn parse_with_limits<'a>( + i: &'a [u8], + template: &Template, + limits: DecodedOutputLimits, + ) -> nom::IResult<&'a [u8], Self> { + let mut budget = limits.budget(); + Self::parse_with_budget(i, template, limits.max_records(), &mut budget) + } } pub(crate) struct FlowSetParser; diff --git a/src/variable_versions/wire.rs b/src/variable_versions/wire.rs new file mode 100644 index 0000000..f337110 --- /dev/null +++ b/src/variable_versions/wire.rs @@ -0,0 +1,51 @@ +//! Shared checked arithmetic for v9/IPFIX wire records. + +/// Returns the minimum number of bytes one record must consume. +/// +/// IPFIX variable-length fields consume at least their one-byte length prefix. +/// Pass `None` for fixed-width-only formats such as NetFlow v9. +pub(crate) fn minimum_record_size( + widths: impl IntoIterator, + variable_length_marker: Option, +) -> Option { + widths.into_iter().try_fold(0usize, |total, width| { + let width = if variable_length_marker == Some(width) { + 1 + } else { + usize::from(width) + }; + total.checked_add(width) + }) +} + +/// Padding is unambiguous only when it is shorter than a complete record. +pub(crate) fn is_short_padding(padding_len: usize, minimum_record_size: usize) -> bool { + padding_len < minimum_record_size +} + +/// NetFlow v9 additionally limits alignment padding to three bytes. +pub(crate) fn is_v9_padding(padding_len: usize, record_size: usize) -> bool { + padding_len <= 3 && is_short_padding(padding_len, record_size) +} + +#[derive(Clone, Copy)] +pub(crate) enum RecordBodyKind { + NetFlowV9, + Ipfix, +} + +/// Require at least one decoded record and only protocol-valid padding. +pub(crate) fn record_body_is_complete( + decoded_records: usize, + remainder: &[u8], + minimum_record_size: usize, + kind: RecordBodyKind, +) -> bool { + if decoded_records == 0 { + return false; + } + match kind { + RecordBodyKind::NetFlowV9 => is_v9_padding(remainder.len(), minimum_record_size), + RecordBodyKind::Ipfix => is_short_padding(remainder.len(), minimum_record_size), + } +} diff --git a/tests/decoded_output_allocation.rs b/tests/decoded_output_allocation.rs new file mode 100644 index 0000000..3f3f7d8 --- /dev/null +++ b/tests/decoded_output_allocation.rs @@ -0,0 +1,157 @@ +//! Fresh-process allocation gates for adversarial decoded output. + +#![cfg(feature = "parse_unknown_fields")] + +use netflow_parser::{NetflowPacket, NetflowParser, PendingFlowsConfig}; +use std::alloc::{GlobalAlloc, Layout, System}; +use std::sync::atomic::{AtomicUsize, Ordering}; + +struct CountingAllocator; + +static LIVE_BYTES: AtomicUsize = AtomicUsize::new(0); +static PEAK_BYTES: AtomicUsize = AtomicUsize::new(0); + +fn record_live(live: usize) { + let mut peak = PEAK_BYTES.load(Ordering::Relaxed); + while live > peak { + match PEAK_BYTES.compare_exchange_weak(peak, live, Ordering::Relaxed, Ordering::Relaxed) + { + Ok(_) => break, + Err(actual) => peak = actual, + } + } +} + +unsafe impl GlobalAlloc for CountingAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let ptr = unsafe { System.alloc(layout) }; + if !ptr.is_null() { + let live = LIVE_BYTES.fetch_add(layout.size(), Ordering::Relaxed) + layout.size(); + record_live(live); + } + ptr + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + LIVE_BYTES.fetch_sub(layout.size(), Ordering::Relaxed); + unsafe { System.dealloc(ptr, layout) }; + } + + unsafe fn realloc(&self, ptr: *mut u8, old: Layout, new_size: usize) -> *mut u8 { + let new_ptr = unsafe { System.realloc(ptr, old, new_size) }; + if !new_ptr.is_null() { + let live = if new_size >= old.size() { + LIVE_BYTES.fetch_add(new_size - old.size(), Ordering::Relaxed) + new_size + - old.size() + } else { + LIVE_BYTES.fetch_sub(old.size() - new_size, Ordering::Relaxed) + - (old.size() - new_size) + }; + record_live(live); + } + new_ptr + } +} + +#[global_allocator] +static ALLOCATOR: CountingAllocator = CountingAllocator; + +fn reset_peak() -> usize { + let live = LIVE_BYTES.load(Ordering::Relaxed); + PEAK_BYTES.store(live, Ordering::Relaxed); + live +} + +fn peak_delta(baseline: usize) -> usize { + PEAK_BYTES.load(Ordering::Relaxed).saturating_sub(baseline) +} + +fn v9_message(flowsets: &[Vec]) -> Vec { + let mut packet = vec![0, 9, 0, flowsets.len() as u8]; + packet.extend_from_slice(&[0; 12]); + packet.extend_from_slice(&1u32.to_be_bytes()); + for flowset in flowsets { + packet.extend_from_slice(flowset); + } + packet +} + +fn wide_template() -> Vec { + let mut fields: Vec<(u16, u16)> = (1000..1064).map(|field| (field, 0)).collect(); + fields.push((1, 1)); + let length = 8 + fields.len() * 4; + let mut set = Vec::with_capacity(length); + set.extend_from_slice(&0u16.to_be_bytes()); + set.extend_from_slice(&(length as u16).to_be_bytes()); + set.extend_from_slice(&256u16.to_be_bytes()); + set.extend_from_slice(&(fields.len() as u16).to_be_bytes()); + for (field_type, field_length) in fields { + set.extend_from_slice(&field_type.to_be_bytes()); + set.extend_from_slice(&field_length.to_be_bytes()); + } + set +} + +fn data(body_len: usize) -> Vec { + let mut set = Vec::with_capacity(body_len + 4); + set.extend_from_slice(&256u16.to_be_bytes()); + set.extend_from_slice(&((body_len + 4) as u16).to_be_bytes()); + set.resize(body_len + 4, 1); + set +} + +#[test] +fn default_limit_stays_bounded_and_pending_shortage_is_preflight_only() { + let template = wide_template(); + let mut parser = NetflowParser::default(); + assert!( + parser + .parse_bytes(&v9_message(std::slice::from_ref(&template))) + .is_ok() + ); + + let baseline = reset_peak(); + let result = parser.parse_bytes(&v9_message(&[data(1000)])); + assert!(result.is_ok(), "{:?}", result.error); + let NetflowPacket::V9(packet) = &result.packets[0] else { + panic!("expected v9 packet") + }; + let records = packet + .flowsets + .iter() + .map(|flowset| match &flowset.body { + netflow_parser::variable_versions::v9::FlowSetBody::Data(data) => data.fields.len(), + _ => 0, + }) + .sum::(); + assert_eq!(records, 1000); + let bounded_peak = peak_delta(baseline); + eprintln!("default_bound_peak_requested_bytes={bounded_peak}"); + assert!( + bounded_peak <= 16 * 1024 * 1024, + "default-bound parse peaked {} requested bytes above baseline", + bounded_peak + ); + drop(result); + + let mut pending = NetflowParser::builder() + .with_v9_pending_flows(PendingFlowsConfig::default()) + .build() + .unwrap(); + assert!(pending.parse_bytes(&v9_message(&[data(1000)])).is_ok()); + assert_eq!(pending.v9_cache_info().pending_flow_count, 1); + + // Nine current records consume 585 values, leaving too little room for + // the queued 65,000-value entry. Replay must stop after wire-only preflight. + let baseline = reset_peak(); + let refresh = pending.parse_bytes(&v9_message(&[template, data(9)])); + assert!(refresh.is_ok(), "{:?}", refresh.error); + assert_eq!(pending.v9_cache_info().pending_flow_count, 1); + let preflight_peak = peak_delta(baseline); + eprintln!("pending_preflight_peak_requested_bytes={preflight_peak}"); + assert!( + preflight_peak < 1024 * 1024, + "temporarily blocked replay allocated {} requested bytes", + preflight_peak + ); +} diff --git a/tests/decoded_output_default_regression.rs b/tests/decoded_output_default_regression.rs new file mode 100644 index 0000000..9df63f0 --- /dev/null +++ b/tests/decoded_output_default_regression.rs @@ -0,0 +1,55 @@ +#![cfg(feature = "parse_unknown_fields")] + +use netflow_parser::NetflowParser; + +fn v9_message(flowset: &[u8]) -> Vec { + let mut packet = vec![ + 0, 9, // version + 0, 1, // one FlowSet + 0, 0, 0, 0, // system uptime + 0, 0, 0, 0, // export time + 0, 0, 0, 1, // sequence + 0, 0, 0, 1, // source ID + ]; + packet.extend_from_slice(flowset); + packet +} + +fn v9_template(template_id: u16, fields: &[(u16, u16)]) -> Vec { + let length = 8 + fields.len() * 4; + let mut flowset = Vec::with_capacity(length); + flowset.extend_from_slice(&0u16.to_be_bytes()); + flowset.extend_from_slice(&(length as u16).to_be_bytes()); + flowset.extend_from_slice(&template_id.to_be_bytes()); + flowset.extend_from_slice(&(fields.len() as u16).to_be_bytes()); + for (field_type, field_length) in fields { + flowset.extend_from_slice(&field_type.to_be_bytes()); + flowset.extend_from_slice(&field_length.to_be_bytes()); + } + flowset +} + +fn v9_data(template_id: u16, body: &[u8]) -> Vec { + let mut flowset = Vec::with_capacity(body.len() + 4); + flowset.extend_from_slice(&template_id.to_be_bytes()); + flowset.extend_from_slice(&((body.len() + 4) as u16).to_be_bytes()); + flowset.extend_from_slice(body); + flowset +} + +#[test] +fn default_limits_reject_cumulative_decoded_output_amplification() { + // Each one-byte wire record materializes 65 field values: 64 zero-width + // unknown fields and one one-byte field. Across 1,009 records this creates + // 65,585 values, just beyond the default cumulative limit of 65,536. + let mut fields: Vec<(u16, u16)> = (1000..1064).map(|field| (field, 0)).collect(); + fields.push((1, 1)); + + let mut parser = NetflowParser::default(); + let template = parser.parse_bytes(&v9_message(&v9_template(256, &fields))); + assert!(template.is_ok(), "template failed: {:?}", template.error); + + let result = parser.parse_bytes(&v9_message(&v9_data(256, &[1; 1009]))); + assert!(result.packets.is_empty()); + assert!(result.error.is_some()); +} diff --git a/tests/decoded_output_limits.rs b/tests/decoded_output_limits.rs new file mode 100644 index 0000000..97576e4 --- /dev/null +++ b/tests/decoded_output_limits.rs @@ -0,0 +1,1109 @@ +use netflow_parser::variable_versions::{ipfix, v9}; +use netflow_parser::{ + Config, ConfigError, DecodedOutputLimit, DecodedOutputLimits, InMemoryTemplateStore, + IpfixField, NetflowError, NetflowPacket, NetflowParser, PendingFlowsConfig, + TemplateProtocol, V9Field, +}; +use std::sync::Arc; + +fn v9_message(flowsets: &[Vec]) -> Vec { + let mut packet = vec![ + 0, + 9, // version + 0, + flowsets.len() as u8, // current parser's declared count + 0, + 0, + 0, + 0, // sys_uptime + 0, + 0, + 0, + 0, // unix_secs + 0, + 0, + 0, + 1, // sequence + 0, + 0, + 0, + 1, // source_id + ]; + for flowset in flowsets { + packet.extend_from_slice(flowset); + } + packet +} + +fn v9_template(template_id: u16, fields: &[(u16, u16)]) -> Vec { + let length = 8 + fields.len() * 4; + let mut set = Vec::with_capacity(length); + set.extend_from_slice(&0u16.to_be_bytes()); + set.extend_from_slice(&(length as u16).to_be_bytes()); + set.extend_from_slice(&template_id.to_be_bytes()); + set.extend_from_slice(&(fields.len() as u16).to_be_bytes()); + for (field_type, field_length) in fields { + set.extend_from_slice(&field_type.to_be_bytes()); + set.extend_from_slice(&field_length.to_be_bytes()); + } + set +} + +fn v9_options_template(template_id: u16) -> Vec { + let mut set = Vec::new(); + set.extend_from_slice(&1u16.to_be_bytes()); + set.extend_from_slice(&18u16.to_be_bytes()); + set.extend_from_slice(&template_id.to_be_bytes()); + set.extend_from_slice(&4u16.to_be_bytes()); + set.extend_from_slice(&4u16.to_be_bytes()); + set.extend_from_slice(&1u16.to_be_bytes()); + set.extend_from_slice(&1u16.to_be_bytes()); + set.extend_from_slice(&2u16.to_be_bytes()); + set.extend_from_slice(&1u16.to_be_bytes()); + set +} + +fn v9_data(template_id: u16, body: &[u8]) -> Vec { + let mut set = Vec::with_capacity(body.len() + 4); + set.extend_from_slice(&template_id.to_be_bytes()); + set.extend_from_slice(&((body.len() + 4) as u16).to_be_bytes()); + set.extend_from_slice(body); + set +} + +fn ipfix_message(sets: &[Vec]) -> Vec { + let length = 16 + sets.iter().map(Vec::len).sum::(); + let mut packet = Vec::with_capacity(length); + packet.extend_from_slice(&10u16.to_be_bytes()); + packet.extend_from_slice(&(length as u16).to_be_bytes()); + packet.extend_from_slice(&0u32.to_be_bytes()); + packet.extend_from_slice(&1u32.to_be_bytes()); + packet.extend_from_slice(&1u32.to_be_bytes()); + for set in sets { + packet.extend_from_slice(set); + } + packet +} + +fn ipfix_template(template_id: u16, field_type: u16, field_length: u16) -> Vec { + ipfix_template_fields(template_id, &[(field_type, field_length)]) +} + +fn ipfix_template_fields(template_id: u16, fields: &[(u16, u16)]) -> Vec { + let mut set = Vec::new(); + set.extend_from_slice(&2u16.to_be_bytes()); + set.extend_from_slice(&((8 + fields.len() * 4) as u16).to_be_bytes()); + set.extend_from_slice(&template_id.to_be_bytes()); + set.extend_from_slice(&(fields.len() as u16).to_be_bytes()); + for (field_type, field_length) in fields { + set.extend_from_slice(&field_type.to_be_bytes()); + set.extend_from_slice(&field_length.to_be_bytes()); + } + set +} + +fn ipfix_options_template(template_id: u16) -> Vec { + let mut set = Vec::new(); + set.extend_from_slice(&3u16.to_be_bytes()); + set.extend_from_slice(&18u16.to_be_bytes()); + set.extend_from_slice(&template_id.to_be_bytes()); + set.extend_from_slice(&2u16.to_be_bytes()); + set.extend_from_slice(&1u16.to_be_bytes()); + set.extend_from_slice(&1u16.to_be_bytes()); + set.extend_from_slice(&1u16.to_be_bytes()); + set.extend_from_slice(&2u16.to_be_bytes()); + set.extend_from_slice(&1u16.to_be_bytes()); + set +} + +fn ipfix_data(template_id: u16, body: &[u8]) -> Vec { + let mut set = Vec::with_capacity(body.len() + 4); + set.extend_from_slice(&template_id.to_be_bytes()); + set.extend_from_slice(&((body.len() + 4) as u16).to_be_bytes()); + set.extend_from_slice(body); + set +} + +fn ipfix_replay_boundary_body() -> Vec { + // The replayed Set is 65,508 bytes: it fits after a 20-byte store-backed + // trigger, but not after the smallest 28-byte on-wire Template message. + let mut body = Vec::with_capacity(65_504); + body.push(255); + body.extend_from_slice(&65_501u16.to_be_bytes()); + body.resize(65_504, b'x'); + body +} + +fn assert_limit( + error: Option, + protocol: TemplateProtocol, + limit: DecodedOutputLimit, + configured: usize, + attempted: usize, +) { + assert!(matches!( + error, + Some(NetflowError::DecodedOutputLimitExceeded { + protocol: actual_protocol, + limit: actual_limit, + configured: actual_configured, + attempted: actual_attempted, + }) if actual_protocol == protocol + && actual_limit == limit + && actual_configured == configured + && actual_attempted == attempted + )); +} + +#[test] +fn v9_value_budget_is_cumulative_across_flowsets_and_exact() { + let template = v9_template(256, &[(1, 1)]); + + let mut exact = NetflowParser::builder() + .with_v9_max_decoded_field_values_per_message(4) + .build() + .unwrap(); + assert!( + exact + .parse_bytes(&v9_message(std::slice::from_ref(&template))) + .is_ok() + ); + let result = + exact.parse_bytes(&v9_message(&[v9_data(256, &[1, 2]), v9_data(256, &[3, 4])])); + assert!(result.is_ok(), "{:?}", result.error); + + let mut one_over = NetflowParser::builder() + .with_v9_max_decoded_field_values_per_message(3) + .build() + .unwrap(); + assert!(one_over.parse_bytes(&v9_message(&[template])).is_ok()); + let result = + one_over.parse_bytes(&v9_message(&[v9_data(256, &[1, 2]), v9_data(256, &[3, 4])])); + assert!(result.packets.is_empty()); + assert_limit( + result.error, + TemplateProtocol::V9, + DecodedOutputLimit::FieldValues, + 3, + 4, + ); +} + +#[test] +fn v9_options_budget_counts_scope_and_option_values() { + let mut parser = NetflowParser::builder() + .with_v9_max_decoded_field_values_per_message(1) + .build() + .unwrap(); + assert!( + parser + .parse_bytes(&v9_message(&[v9_options_template(300)])) + .is_ok() + ); + let result = parser.parse_bytes(&v9_message(&[v9_data(300, &[1, 2])])); + assert_limit( + result.error, + TemplateProtocol::V9, + DecodedOutputLimit::FieldValues, + 1, + 2, + ); +} + +#[test] +fn ipfix_value_budget_is_cumulative_across_sets_and_exact() { + let template = ipfix_template(256, 1, 1); + let mut parser = NetflowParser::builder() + .with_ipfix_max_decoded_field_values_per_message(3) + .build() + .unwrap(); + assert!(parser.parse_bytes(&ipfix_message(&[template])).is_ok()); + + let result = parser.parse_bytes(&ipfix_message(&[ + ipfix_data(256, &[1, 2]), + ipfix_data(256, &[3, 4]), + ])); + assert!(result.packets.is_empty()); + assert_limit( + result.error, + TemplateProtocol::Ipfix, + DecodedOutputLimit::FieldValues, + 3, + 4, + ); +} + +#[test] +fn output_budgets_reset_between_messages() { + let mut v9_parser = NetflowParser::builder() + .with_v9_max_decoded_field_values_per_message(1) + .build() + .unwrap(); + assert!( + v9_parser + .parse_bytes(&v9_message(&[v9_template(256, &[(1, 1)])])) + .is_ok() + ); + for value in [1, 2] { + let result = v9_parser.parse_bytes(&v9_message(&[v9_data(256, &[value])])); + assert!(result.is_ok(), "{:?}", result.error); + } + + let mut ipfix_parser = NetflowParser::builder() + .with_ipfix_max_decoded_field_values_per_message(1) + .build() + .unwrap(); + assert!( + ipfix_parser + .parse_bytes(&ipfix_message(&[ipfix_template(256, 1, 1)])) + .is_ok() + ); + for value in [1, 2] { + let result = ipfix_parser.parse_bytes(&ipfix_message(&[ipfix_data(256, &[value])])); + assert!(result.is_ok(), "{:?}", result.error); + } +} + +#[test] +fn direct_v9_message_parser_resets_the_output_budget() { + let mut config = Config::default(); + config.max_decoded_field_values_per_message = 1; + let mut parser = v9::V9Parser::try_new(config).unwrap(); + + let template = v9_message(&[v9_template(256, &[(1, 1)])]); + assert!(v9::V9::parse(&template[2..], &mut parser).is_ok()); + + for value in [1, 2] { + let message = v9_message(&[v9_data(256, &[value])]); + let (_, packet) = v9::V9::parse(&message[2..], &mut parser).unwrap(); + let records = packet + .flowsets + .iter() + .filter_map(|flowset| match &flowset.body { + v9::FlowSetBody::Data(data) => Some(data.fields.len()), + _ => None, + }) + .sum::(); + assert_eq!(records, 1); + } + + let over_limit = v9_message(&[v9_data(256, &[3, 4])]); + assert!(v9::V9::parse(&over_limit[2..], &mut parser).is_err()); +} + +#[test] +fn direct_ipfix_message_parser_resets_and_enforces_the_output_budget() { + let mut config = Config::default(); + config.max_decoded_field_values_per_message = 1; + let mut parser = ipfix::IPFixParser::try_new(config).unwrap(); + + let template = ipfix_message(&[ipfix_template(256, 1, 1)]); + assert!(ipfix::IPFix::parse(&template[2..], &mut parser).is_ok()); + + for value in [1, 2] { + let message = ipfix_message(&[ipfix_data(256, &[value])]); + let (_, packet) = ipfix::IPFix::parse(&message[2..], &mut parser).unwrap(); + let records = packet + .flowsets + .iter() + .filter_map(|flowset| match &flowset.body { + ipfix::FlowSetBody::Data(data) => Some(data.fields.len()), + _ => None, + }) + .sum::(); + assert_eq!(records, 1); + } + + let over_limit = ipfix_message(&[ipfix_data(256, &[3, 4])]); + assert!(ipfix::IPFix::parse(&over_limit[2..], &mut parser).is_err()); +} + +#[test] +fn ipfix_variable_field_payload_budget_excludes_length_prefix() { + // interfaceName (82) is a String; 65535 selects RFC 7011 variable length. + let template = ipfix_template(256, 82, u16::MAX); + let mut exact = NetflowParser::builder() + .with_ipfix_max_decoded_field_payload_bytes_per_message(4) + .build() + .unwrap(); + assert!( + exact + .parse_bytes(&ipfix_message(std::slice::from_ref(&template))) + .is_ok() + ); + assert!( + exact + .parse_bytes(&ipfix_message(&[ipfix_data( + 256, + &[4, b't', b'e', b's', b't'] + )])) + .is_ok() + ); + + let mut one_over = NetflowParser::builder() + .with_ipfix_max_decoded_field_payload_bytes_per_message(4) + .build() + .unwrap(); + assert!(one_over.parse_bytes(&ipfix_message(&[template])).is_ok()); + let result = one_over.parse_bytes(&ipfix_message(&[ipfix_data( + 256, + &[5, b't', b'e', b's', b't', b'!'], + )])); + assert_limit( + result.error, + TemplateProtocol::Ipfix, + DecodedOutputLimit::FieldPayloadBytes, + 4, + 5, + ); +} + +#[test] +fn v9_payload_budget_is_cumulative_across_flowsets() { + let template = v9_template(256, &[(1, 2)]); + let mut parser = NetflowParser::builder() + .with_v9_max_decoded_field_payload_bytes_per_message(3) + .build() + .unwrap(); + assert!(parser.parse_bytes(&v9_message(&[template])).is_ok()); + let result = + parser.parse_bytes(&v9_message(&[v9_data(256, &[0, 1]), v9_data(256, &[0, 2])])); + assert_limit( + result.error, + TemplateProtocol::V9, + DecodedOutputLimit::FieldPayloadBytes, + 3, + 4, + ); +} + +#[test] +#[cfg(feature = "parse_unknown_fields")] +fn default_value_budget_rejects_many_tiny_materialized_fields() { + let mut fields: Vec<(u16, u16)> = (1000..1064).map(|field| (field, 0)).collect(); + fields.push((1, 1)); + let mut parser = NetflowParser::default(); + assert!( + parser + .parse_bytes(&v9_message(&[v9_template(256, &fields)])) + .is_ok() + ); + let result = parser.parse_bytes(&v9_message(&[v9_data(256, &[1; 1009])])); + assert_limit( + result.error, + TemplateProtocol::V9, + DecodedOutputLimit::FieldValues, + netflow_parser::DEFAULT_MAX_DECODED_FIELD_VALUES_PER_MESSAGE, + 65_585, + ); +} + +#[test] +fn large_variable_value_is_rejected_before_payload_materialization() { + let template = ipfix_template(256, 82, u16::MAX); + let mut parser = NetflowParser::builder() + .with_ipfix_max_decoded_field_payload_bytes_per_message(4095) + .build() + .unwrap(); + assert!(parser.parse_bytes(&ipfix_message(&[template])).is_ok()); + + let mut body = Vec::with_capacity(4099); + body.push(255); + body.extend_from_slice(&4096u16.to_be_bytes()); + body.resize(4099, b'x'); + let result = parser.parse_bytes(&ipfix_message(&[ipfix_data(256, &body)])); + assert_limit( + result.error, + TemplateProtocol::Ipfix, + DecodedOutputLimit::FieldPayloadBytes, + 4095, + 4096, + ); +} + +#[test] +fn ipfix_native_and_embedded_options_count_every_value() { + for template in [ipfix_options_template(300), { + let mut embedded = v9_options_template(300); + embedded[0..2].copy_from_slice(&1u16.to_be_bytes()); + embedded + }] { + let mut parser = NetflowParser::builder() + .with_ipfix_max_decoded_field_values_per_message(1) + .build() + .unwrap(); + assert!(parser.parse_bytes(&ipfix_message(&[template])).is_ok()); + let result = parser.parse_bytes(&ipfix_message(&[ipfix_data(300, &[1, 2])])); + assert_limit( + result.error, + TemplateProtocol::Ipfix, + DecodedOutputLimit::FieldValues, + 1, + 2, + ); + } +} + +#[test] +fn ipfix_embedded_v9_data_uses_the_message_budget() { + let mut parser = NetflowParser::builder() + .with_ipfix_max_decoded_field_values_per_message(1) + .build() + .unwrap(); + let template = v9_template(256, &[(1, 1)]); + assert!(parser.parse_bytes(&ipfix_message(&[template])).is_ok()); + let result = parser.parse_bytes(&ipfix_message(&[ipfix_data(256, &[1, 2])])); + assert_limit( + result.error, + TemplateProtocol::Ipfix, + DecodedOutputLimit::FieldValues, + 1, + 2, + ); +} + +#[test] +fn pending_variable_preflight_uses_the_parser_padding_boundary() { + let mut parser = NetflowParser::builder() + .with_ipfix_pending_flows(PendingFlowsConfig::default()) + .with_ipfix_max_decoded_field_values_per_message(2) + .build() + .unwrap(); + + // One complete one-byte varlen + one-byte fixed record, then one byte of + // legal Set padding that is not a complete second record. + assert!( + parser + .parse_bytes(&ipfix_message(&[ipfix_data(256, &[1, b'x', 7, 0])])) + .is_ok() + ); + assert_eq!(parser.ipfix_cache_info().pending_flow_count, 1); + + let template = ipfix_template_fields(256, &[(82, u16::MAX), (1, 1)]); + let replay = parser.parse_bytes(&ipfix_message(&[template])); + assert!(replay.is_ok(), "{:?}", replay.error); + assert_eq!(parser.ipfix_cache_info().pending_flow_count, 0); + let NetflowPacket::IPFix(packet) = &replay.packets[0] else { + panic!("expected IPFIX packet") + }; + assert!(packet.flowsets.len() >= 2, "pending flow was not appended"); +} + +#[test] +fn pending_replay_retains_a_temporarily_non_fitting_suffix() { + let mut parser = NetflowParser::builder() + .with_v9_pending_flows(PendingFlowsConfig::default()) + .with_v9_max_decoded_field_values_per_message(2) + .build() + .unwrap(); + + assert!( + parser + .parse_bytes(&v9_message(&[v9_data(256, &[1, 2])])) + .is_ok() + ); + assert_eq!(parser.v9_cache_info().pending_flow_count, 1); + + let template = v9_template(256, &[(1, 1)]); + let current = parser.parse_bytes(&v9_message(&[template.clone(), v9_data(256, &[3])])); + assert!(current.is_ok(), "{:?}", current.error); + assert_eq!(parser.v9_cache_info().pending_flow_count, 1); + + let replay = parser.parse_bytes(&v9_message(&[template])); + assert!(replay.is_ok(), "{:?}", replay.error); + assert_eq!(parser.v9_cache_info().pending_flow_count, 0); + let NetflowPacket::V9(packet) = &replay.packets[0] else { + panic!("expected v9 packet") + }; + assert!(packet.flowsets.len() >= 2, "pending flow was not appended"); +} + +#[test] +fn pending_replay_appends_the_largest_fitting_fifo_prefix() { + let mut parser = NetflowParser::builder() + .with_v9_pending_flows(PendingFlowsConfig::default()) + .with_v9_max_decoded_field_values_per_message(2) + .build() + .unwrap(); + + for value in [1, 2] { + assert!( + parser + .parse_bytes(&v9_message(&[v9_data(256, &[value])])) + .is_ok() + ); + } + assert_eq!(parser.v9_cache_info().pending_flow_count, 2); + + let template = v9_template(256, &[(1, 1)]); + let first = parser.parse_bytes(&v9_message(&[template.clone(), v9_data(256, &[3])])); + assert!(first.is_ok(), "{:?}", first.error); + let NetflowPacket::V9(packet) = &first.packets[0] else { + panic!("expected v9 packet") + }; + assert_eq!( + packet + .flowsets + .iter() + .filter(|flowset| matches!(flowset.body, v9::FlowSetBody::Data(_))) + .count(), + 2 + ); + let info = parser.v9_cache_info(); + assert_eq!(info.pending_flow_count, 1); + assert_eq!(info.metrics.pending_replayed, 1); + + let second = parser.parse_bytes(&v9_message(&[template])); + assert!(second.is_ok(), "{:?}", second.error); + let info = parser.v9_cache_info(); + assert_eq!(info.pending_flow_count, 0); + assert_eq!(info.metrics.pending_replayed, 2); +} + +#[test] +fn pending_replay_drops_incomplete_fixed_entry_and_continues_fifo() { + let mut parser = NetflowParser::builder() + .with_v9_pending_flows(PendingFlowsConfig::default()) + .build() + .unwrap(); + + assert!( + parser + .parse_bytes(&v9_message(&[v9_data(256, &[1])])) + .is_ok() + ); + assert!( + parser + .parse_bytes(&v9_message(&[v9_data(256, &[2, 3])])) + .is_ok() + ); + assert_eq!(parser.v9_cache_info().pending_flow_count, 2); + + let replay = parser.parse_bytes(&v9_message(&[v9_template(256, &[(1, 2)])])); + assert!(replay.is_ok(), "{:?}", replay.error); + let info = parser.v9_cache_info(); + assert_eq!(info.pending_flow_count, 0); + assert_eq!(info.metrics.pending_replay_failed, 1); + assert_eq!(info.metrics.pending_replayed, 1); + + let NetflowPacket::V9(packet) = &replay.packets[0] else { + panic!("expected v9 packet") + }; + let data = packet + .flowsets + .iter() + .filter_map(|flowset| match &flowset.body { + v9::FlowSetBody::Data(data) => Some(data), + _ => None, + }) + .collect::>(); + assert_eq!(data.len(), 1); + assert_eq!(data[0].fields.len(), 1); +} + +#[test] +fn pending_replay_drops_an_entry_that_can_never_fit_and_continues_fifo() { + let mut v9_parser = NetflowParser::builder() + .with_v9_pending_flows(PendingFlowsConfig::default()) + .with_v9_max_decoded_field_values_per_message(1) + .build() + .unwrap(); + + for body in [&[1, 2][..], &[3][..]] { + assert!( + v9_parser + .parse_bytes(&v9_message(&[v9_data(256, body)])) + .is_ok() + ); + } + let replay = v9_parser.parse_bytes(&v9_message(&[v9_template(256, &[(1, 1)])])); + assert!(replay.is_ok(), "{:?}", replay.error); + let info = v9_parser.v9_cache_info(); + assert_eq!(info.pending_flow_count, 0); + assert_eq!(info.metrics.pending_replay_failed, 1); + assert_eq!(info.metrics.pending_replayed, 1); + + let NetflowPacket::V9(packet) = &replay.packets[0] else { + panic!("expected v9 packet") + }; + assert_eq!( + packet + .flowsets + .iter() + .filter(|flowset| matches!(flowset.body, v9::FlowSetBody::Data(_))) + .count(), + 1 + ); + + let mut ipfix_parser = NetflowParser::builder() + .with_ipfix_pending_flows(PendingFlowsConfig::default()) + .with_ipfix_max_decoded_field_values_per_message(1) + .build() + .unwrap(); + + for body in [&[1, 2][..], &[3][..]] { + assert!( + ipfix_parser + .parse_bytes(&ipfix_message(&[ipfix_data(256, body)])) + .is_ok() + ); + } + let replay = ipfix_parser.parse_bytes(&ipfix_message(&[ipfix_template(256, 1, 1)])); + assert!(replay.is_ok(), "{:?}", replay.error); + let info = ipfix_parser.ipfix_cache_info(); + assert_eq!(info.pending_flow_count, 0); + assert_eq!(info.metrics.pending_replay_failed, 1); + assert_eq!(info.metrics.pending_replayed, 1); + + let NetflowPacket::IPFix(packet) = &replay.packets[0] else { + panic!("expected IPFIX packet") + }; + assert_eq!( + packet + .flowsets + .iter() + .filter(|flowset| matches!(flowset.body, ipfix::FlowSetBody::Data(_))) + .count(), + 1 + ); +} + +#[test] +fn ipfix_pending_replay_classifies_entries_before_framing_pressure() { + let template = ipfix_template(256, 82, u16::MAX); + + let mut parser = NetflowParser::builder() + .with_ipfix_pending_flows(PendingFlowsConfig::default()) + .build() + .unwrap(); + + // This valid variable-length record can fill a standalone IPFIX message, + // but no message containing the Set that triggers replay has room for it. + let mut permanently_too_large = Vec::with_capacity(65_515); + permanently_too_large.push(255); + permanently_too_large.extend_from_slice(&65_512u16.to_be_bytes()); + permanently_too_large.resize(65_515, b'x'); + assert!( + parser + .parse_bytes(&ipfix_message(&[ipfix_data(256, &permanently_too_large)])) + .is_ok() + ); + assert!( + parser + .parse_bytes(&ipfix_message(&[ipfix_data(256, &[1, b'y'])])) + .is_ok() + ); + + let replay = parser.parse_bytes(&ipfix_message(std::slice::from_ref(&template))); + assert!(replay.is_ok(), "{:?}", replay.error); + let info = parser.ipfix_cache_info(); + assert_eq!(info.pending_flow_count, 0); + assert_eq!(info.metrics.pending_replay_failed, 1); + assert_eq!(info.metrics.pending_replayed, 1); + + let mut parser = NetflowParser::builder() + .with_ipfix_pending_flows(PendingFlowsConfig::default()) + .build() + .unwrap(); + + // The first entry is malformed. The second is valid but cannot fit while + // the current message is full, so only the valid suffix must be retained. + for body in [&[5, b'x'][..], &[1, b'y'][..]] { + assert!( + parser + .parse_bytes(&ipfix_message(&[ipfix_data(256, body)])) + .is_ok() + ); + } + let full_message = ipfix_message(&[template.clone(), ipfix_data(4, &vec![0; 65_503])]); + assert_eq!(full_message.len(), usize::from(u16::MAX)); + let blocked = parser.parse_bytes(&full_message); + assert!(blocked.is_ok(), "{:?}", blocked.error); + let info = parser.ipfix_cache_info(); + assert_eq!(info.pending_flow_count, 1); + assert_eq!(info.metrics.pending_replay_failed, 1); + assert_eq!(info.metrics.pending_replayed, 0); + + let replay = parser.parse_bytes(&ipfix_message(&[template])); + assert!(replay.is_ok(), "{:?}", replay.error); + let info = parser.ipfix_cache_info(); + assert_eq!(info.pending_flow_count, 0); + assert_eq!(info.metrics.pending_replay_failed, 1); + assert_eq!(info.metrics.pending_replayed, 1); +} + +#[test] +fn ipfix_no_store_drops_entry_that_cannot_fit_with_template_trigger() { + let mut parser = NetflowParser::builder() + .with_ipfix_pending_flows(PendingFlowsConfig::default()) + .build() + .unwrap(); + + let body = ipfix_replay_boundary_body(); + assert!( + parser + .parse_bytes(&ipfix_message(&[ipfix_data(256, &body)])) + .is_ok() + ); + + let replay = parser.parse_bytes(&ipfix_message(&[ipfix_template(256, 82, u16::MAX)])); + assert!(replay.is_ok(), "{:?}", replay.error); + let info = parser.ipfix_cache_info(); + assert_eq!(info.pending_flow_count, 0); + assert_eq!(info.metrics.pending_replay_failed, 1); + assert_eq!(info.metrics.pending_replayed, 0); +} + +#[test] +fn ipfix_store_restoration_replays_at_twenty_byte_boundary() { + let store = Arc::new(InMemoryTemplateStore::new()); + let mut parser = NetflowParser::builder() + .with_template_store(store.clone()) + .with_ipfix_pending_flows(PendingFlowsConfig::default()) + .build() + .unwrap(); + + let body = ipfix_replay_boundary_body(); + assert!( + parser + .parse_bytes(&ipfix_message(&[ipfix_data(256, &body)])) + .is_ok() + ); + + let mut writer = NetflowParser::builder() + .with_template_store(store) + .build() + .unwrap(); + assert!( + writer + .parse_bytes(&ipfix_message(&[ipfix_template(256, 82, u16::MAX)])) + .is_ok() + ); + + let replay = parser.parse_bytes(&ipfix_message(&[ipfix_data(256, &[])])); + assert!(replay.is_ok(), "{:?}", replay.error); + let info = parser.ipfix_cache_info(); + assert_eq!(info.pending_flow_count, 0); + assert_eq!(info.metrics.pending_replay_failed, 0); + assert_eq!(info.metrics.pending_replayed, 1); +} + +#[test] +fn pending_replay_rejects_max_record_truncated_prefixes() { + let mut v9_parser = NetflowParser::builder() + .with_v9_pending_flows(PendingFlowsConfig::default()) + .with_v9_max_records_per_flowset(1) + .build() + .unwrap(); + assert!( + v9_parser + .parse_bytes(&v9_message(&[v9_data(256, &[1, 2])])) + .is_ok() + ); + let replay = v9_parser.parse_bytes(&v9_message(&[v9_template(256, &[(1, 1)])])); + assert!(replay.is_ok(), "{:?}", replay.error); + let info = v9_parser.v9_cache_info(); + assert_eq!(info.pending_flow_count, 0); + assert_eq!(info.metrics.pending_replay_failed, 1); + assert_eq!(info.metrics.pending_replayed, 0); + + let mut ipfix_parser = NetflowParser::builder() + .with_ipfix_pending_flows(PendingFlowsConfig::default()) + .with_ipfix_max_records_per_flowset(1) + .build() + .unwrap(); + assert!( + ipfix_parser + .parse_bytes(&ipfix_message(&[ipfix_data(256, &[1, b'x', 1, b'y'])])) + .is_ok() + ); + let replay = ipfix_parser.parse_bytes(&ipfix_message(&[ipfix_template(256, 82, u16::MAX)])); + assert!(replay.is_ok(), "{:?}", replay.error); + let info = ipfix_parser.ipfix_cache_info(); + assert_eq!(info.pending_flow_count, 0); + assert_eq!(info.metrics.pending_replay_failed, 1); + assert_eq!(info.metrics.pending_replayed, 0); +} + +#[test] +fn pending_replay_preserves_valid_fixed_and_variable_padding() { + let mut v9_parser = NetflowParser::builder() + .with_v9_pending_flows(PendingFlowsConfig::default()) + .build() + .unwrap(); + assert!( + v9_parser + .parse_bytes(&v9_message(&[v9_data(256, &[1, 2, 3, 4, 0, 0])])) + .is_ok() + ); + let replay = v9_parser.parse_bytes(&v9_message(&[v9_template(256, &[(1, 4)])])); + assert!(replay.is_ok(), "{:?}", replay.error); + let NetflowPacket::V9(packet) = &replay.packets[0] else { + panic!("expected v9 packet") + }; + let data = packet + .flowsets + .iter() + .find_map(|flowset| match &flowset.body { + v9::FlowSetBody::Data(data) => Some(data), + _ => None, + }) + .expect("expected replayed v9 data"); + assert_eq!(data.fields.len(), 1); + assert_eq!(data.padding, [0, 0]); + + let mut ipfix_parser = NetflowParser::builder() + .with_ipfix_pending_flows(PendingFlowsConfig::default()) + .build() + .unwrap(); + // One variable-length byte, one fixed byte, then legal short padding. + assert!( + ipfix_parser + .parse_bytes(&ipfix_message(&[ipfix_data(256, &[1, b'x', 7, 0])])) + .is_ok() + ); + let replay = ipfix_parser.parse_bytes(&ipfix_message(&[ipfix_template_fields( + 256, + &[(82, u16::MAX), (1, 1)], + )])); + assert!(replay.is_ok(), "{:?}", replay.error); + let info = ipfix_parser.ipfix_cache_info(); + assert_eq!(info.pending_flow_count, 0); + assert_eq!(info.metrics.pending_replay_failed, 0); + assert_eq!(info.metrics.pending_replayed, 1); + let NetflowPacket::IPFix(packet) = &replay.packets[0] else { + panic!("expected IPFIX packet") + }; + let data = packet + .flowsets + .iter() + .filter_map(|flowset| match &flowset.body { + ipfix::FlowSetBody::Data(data) => Some(data), + _ => None, + }) + .collect::>(); + assert_eq!(data.len(), 1); + for replayed in data { + assert_eq!(replayed.fields.len(), 1); + assert_eq!(replayed.fields[0].len(), 2); + assert_eq!(replayed.padding, [0]); + } +} + +#[test] +fn pending_replay_rejects_every_incomplete_body_kind() { + let mut v9_parser = NetflowParser::builder() + .with_v9_pending_flows(PendingFlowsConfig::default()) + .build() + .unwrap(); + assert!( + v9_parser + .parse_bytes(&v9_message(&[v9_data(300, &[1])])) + .is_ok() + ); + let replay = v9_parser.parse_bytes(&v9_message(&[v9_options_template(300)])); + assert!(replay.is_ok(), "{:?}", replay.error); + assert_eq!(v9_parser.v9_cache_info().metrics.pending_replay_failed, 1); + + for template in [ + ipfix_template(300, 1, 2), + ipfix_options_template(300), + v9_template(300, &[(1, 2)]), + v9_options_template(300), + ] { + let mut parser = NetflowParser::builder() + .with_ipfix_pending_flows(PendingFlowsConfig::default()) + .build() + .unwrap(); + assert!( + parser + .parse_bytes(&ipfix_message(&[ipfix_data(300, &[1])])) + .is_ok() + ); + let replay = parser.parse_bytes(&ipfix_message(&[template])); + assert!(replay.is_ok(), "{:?}", replay.error); + let info = parser.ipfix_cache_info(); + assert_eq!(info.pending_flow_count, 0); + assert_eq!(info.metrics.pending_replay_failed, 1); + assert_eq!(info.metrics.pending_replayed, 0); + } +} + +#[test] +fn ipfix_pending_replay_continues_after_a_same_id_template_does_not_match() { + let mut parser = NetflowParser::builder() + .with_ipfix_pending_flows(PendingFlowsConfig::default()) + .build() + .unwrap(); + + assert!( + parser + .parse_bytes(&ipfix_message(&[ipfix_data(300, &[1, 2])])) + .is_ok() + ); + assert_eq!(parser.ipfix_cache_info().pending_flow_count, 1); + + // The first lookup candidate needs three bytes and cannot decode the + // queued body. The later same-ID options template is the matching owner. + let replay = parser.parse_bytes(&ipfix_message(&[ + ipfix_template(300, 1, 3), + ipfix_options_template(300), + ])); + assert!(replay.is_ok(), "{:?}", replay.error); + assert_eq!(parser.ipfix_cache_info().pending_flow_count, 0); + assert_eq!(parser.ipfix_cache_info().metrics.pending_replayed, 1); + + let NetflowPacket::IPFix(packet) = &replay.packets[0] else { + panic!("expected IPFIX packet") + }; + assert!( + packet + .flowsets + .iter() + .any(|flowset| { matches!(flowset.body, ipfix::FlowSetBody::OptionsData(_)) }) + ); +} + +#[test] +fn v9_pending_replay_continues_after_a_same_id_template_does_not_match() { + let mut parser = NetflowParser::builder() + .with_v9_pending_flows(PendingFlowsConfig::default()) + .build() + .unwrap(); + + assert!( + parser + .parse_bytes(&v9_message(&[v9_data(300, &[1, 2])])) + .is_ok() + ); + assert_eq!(parser.v9_cache_info().pending_flow_count, 1); + + let replay = parser.parse_bytes(&v9_message(&[ + v9_template(300, &[(1, 3)]), + v9_options_template(300), + ])); + assert!(replay.is_ok(), "{:?}", replay.error); + assert_eq!(parser.v9_cache_info().pending_flow_count, 0); + assert_eq!(parser.v9_cache_info().metrics.pending_replayed, 1); + + let NetflowPacket::V9(packet) = &replay.packets[0] else { + panic!("expected v9 packet") + }; + assert!( + packet + .flowsets + .iter() + .any(|flowset| matches!(flowset.body, v9::FlowSetBody::OptionsData(_))) + ); +} + +#[test] +fn zero_output_limits_are_rejected() { + assert!( + NetflowParser::builder() + .with_max_decoded_field_values_per_message(0) + .build() + .is_err() + ); + assert!( + NetflowParser::builder() + .with_max_decoded_field_payload_bytes_per_message(0) + .build() + .is_err() + ); + + let mut config = Config::default(); + config.max_decoded_field_values_per_message = 0; + assert_eq!( + v9::V9Parser::try_new(config).unwrap_err(), + ConfigError::InvalidDecodedFieldValueLimit(0) + ); + + let mut config = Config::default(); + config.max_decoded_field_payload_bytes_per_message = 0; + assert_eq!( + v9::V9Parser::try_new(config).unwrap_err(), + ConfigError::InvalidDecodedFieldPayloadByteLimit(0) + ); + + let mut config = Config::default(); + config.max_decoded_field_values_per_message = 0; + assert_eq!( + ipfix::IPFixParser::try_new(config).unwrap_err(), + ConfigError::InvalidDecodedFieldValueLimit(0) + ); + + let mut config = Config::default(); + config.max_decoded_field_payload_bytes_per_message = 0; + assert_eq!( + ipfix::IPFixParser::try_new(config).unwrap_err(), + ConfigError::InvalidDecodedFieldPayloadByteLimit(0) + ); +} + +#[test] +fn low_level_data_and_options_companions_apply_explicit_bounds() { + let limits = DecodedOutputLimits::new(16, 1, 16).unwrap(); + + let v9_template = v9::Template { + template_id: 256, + field_count: 1, + fields: vec![v9::TemplateField { + field_type_number: 1, + field_type: V9Field::from(1), + field_length: 1, + }], + }; + assert!(v9::Data::parse_with_limits(&[1], &v9_template, limits).is_ok()); + assert!(v9::Data::parse_with_limits(&[1, 2], &v9_template, limits).is_err()); + + let v9_options = v9::OptionsTemplate { + template_id: 300, + options_scope_length: 4, + options_length: 4, + scope_fields: vec![v9::OptionsTemplateScopeField { + field_type_number: 1, + field_type: v9::lookup::ScopeFieldType::from(1), + field_length: 1, + }], + option_fields: vec![v9::TemplateField { + field_type_number: 1, + field_type: V9Field::from(1), + field_length: 1, + }], + }; + assert!(v9::OptionsData::parse_with_limits(&[1, 2], &v9_options, limits).is_err()); + + let ipfix_data = ipfix::Template { + template_id: 256, + field_count: 1, + fields: vec![ipfix::TemplateField { + field_type_number: 1, + field_length: 1, + enterprise_number: None, + field_type: IpfixField::new(1, None), + }], + }; + assert!(ipfix::Data::parse_with_limits(&[1, 2], &ipfix_data, limits).is_err()); + + let ipfix_options = ipfix::OptionsTemplate { + template_id: 300, + field_count: 2, + scope_field_count: 1, + fields: vec![ + ipfix::TemplateField { + field_type_number: 1, + field_length: 1, + enterprise_number: None, + field_type: IpfixField::new(1, None), + }, + ipfix::TemplateField { + field_type_number: 2, + field_length: 1, + enterprise_number: None, + field_type: IpfixField::new(2, None), + }, + ], + }; + assert!(ipfix::OptionsData::parse_with_limits(&[1, 2], &ipfix_options, limits).is_err()); +} diff --git a/tests/hot_path_allocations.rs b/tests/hot_path_allocations.rs new file mode 100644 index 0000000..beeabc7 --- /dev/null +++ b/tests/hot_path_allocations.rs @@ -0,0 +1,324 @@ +// Deterministic warmed allocation measurements for the common parser APIs. + +use netflow_parser::scoped_parser::AutoScopedParser; +use netflow_parser::{NetflowPacket, NetflowParser}; +use std::alloc::{GlobalAlloc, Layout, System}; +use std::hint::black_box; +use std::net::SocketAddr; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + +struct CountingAllocator; + +static TRACKING: AtomicBool = AtomicBool::new(false); +static ALLOCATION_CALLS: AtomicUsize = AtomicUsize::new(0); +static REQUESTED_BYTES: AtomicUsize = AtomicUsize::new(0); +static LIVE_BYTES: AtomicUsize = AtomicUsize::new(0); + +fn record_allocation(size: usize) { + if TRACKING.load(Ordering::Relaxed) { + ALLOCATION_CALLS.fetch_add(1, Ordering::Relaxed); + REQUESTED_BYTES.fetch_add(size, Ordering::Relaxed); + } +} + +unsafe impl GlobalAlloc for CountingAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let ptr = unsafe { System.alloc(layout) }; + if !ptr.is_null() { + LIVE_BYTES.fetch_add(layout.size(), Ordering::Relaxed); + record_allocation(layout.size()); + } + ptr + } + + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + let ptr = unsafe { System.alloc_zeroed(layout) }; + if !ptr.is_null() { + LIVE_BYTES.fetch_add(layout.size(), Ordering::Relaxed); + record_allocation(layout.size()); + } + ptr + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + LIVE_BYTES.fetch_sub(layout.size(), Ordering::Relaxed); + unsafe { System.dealloc(ptr, layout) }; + } + + unsafe fn realloc(&self, ptr: *mut u8, old: Layout, new_size: usize) -> *mut u8 { + let new_ptr = unsafe { System.realloc(ptr, old, new_size) }; + if !new_ptr.is_null() { + if new_size >= old.size() { + LIVE_BYTES.fetch_add(new_size - old.size(), Ordering::Relaxed); + } else { + LIVE_BYTES.fetch_sub(old.size() - new_size, Ordering::Relaxed); + } + record_allocation(new_size); + } + new_ptr + } +} + +#[global_allocator] +static ALLOCATOR: CountingAllocator = CountingAllocator; + +const ITERATIONS: usize = 100; +const FIELDS: [(u16, u16); 6] = [(8, 4), (12, 4), (7, 2), (11, 2), (1, 4), (2, 4)]; + +#[derive(Clone, Copy)] +enum Protocol { + V9, + Ipfix, +} + +impl Protocol { + fn name(self) -> &'static str { + match self { + Self::V9 => "v9", + Self::Ipfix => "ipfix", + } + } + + fn template_packet(self) -> Vec { + match self { + Self::V9 => v9_template_packet(), + Self::Ipfix => ipfix_template_packet(), + } + } + + fn data_packet(self, flow_count: u16) -> Vec { + match self { + Self::V9 => v9_data_packet(flow_count), + Self::Ipfix => ipfix_data_packet(flow_count), + } + } + + fn assert_records(self, packets: &[NetflowPacket], expected: usize) { + assert_eq!(packets.len(), 1); + let records: usize = match (self, &packets[0]) { + (Self::V9, NetflowPacket::V9(packet)) => packet + .flowsets + .iter() + .map(|flowset| match &flowset.body { + netflow_parser::variable_versions::v9::FlowSetBody::Data(data) => { + data.fields.len() + } + _ => 0, + }) + .sum(), + (Self::Ipfix, NetflowPacket::IPFix(packet)) => packet + .flowsets + .iter() + .map(|flowset| match &flowset.body { + netflow_parser::variable_versions::ipfix::FlowSetBody::Data(data) => { + data.fields.len() + } + _ => 0, + }) + .sum(), + _ => panic!("wrong protocol"), + }; + assert_eq!(records, expected); + } +} + +#[derive(Clone, Copy)] +enum Scenario { + DirectParse, + DirectIterator, + AutoParse, + AutoIterator, +} + +impl Scenario { + fn name(self) -> &'static str { + match self { + Self::DirectParse => "direct/parse", + Self::DirectIterator => "direct/iterator", + Self::AutoParse => "auto/parse", + Self::AutoIterator => "auto/iterator", + } + } +} + +fn v9_template_packet() -> Vec { + let set_length = 8 + FIELDS.len() * 4; + let mut packet = vec![0, 9, 0, 1]; + packet.extend_from_slice(&[0; 12]); + packet.extend_from_slice(&1u32.to_be_bytes()); + packet.extend_from_slice(&0u16.to_be_bytes()); + packet.extend_from_slice(&(set_length as u16).to_be_bytes()); + packet.extend_from_slice(&256u16.to_be_bytes()); + packet.extend_from_slice(&(FIELDS.len() as u16).to_be_bytes()); + for (field, length) in FIELDS { + packet.extend_from_slice(&field.to_be_bytes()); + packet.extend_from_slice(&length.to_be_bytes()); + } + packet +} + +fn v9_data_packet(flow_count: u16) -> Vec { + let body = records(flow_count); + let mut packet = vec![0, 9, 0, 1]; + packet.extend_from_slice(&[0; 12]); + packet.extend_from_slice(&1u32.to_be_bytes()); + packet.extend_from_slice(&256u16.to_be_bytes()); + packet.extend_from_slice(&((body.len() + 4) as u16).to_be_bytes()); + packet.extend_from_slice(&body); + packet +} + +fn ipfix_template_packet() -> Vec { + let set_length = 8 + FIELDS.len() * 4; + let mut packet = Vec::with_capacity(16 + set_length); + packet.extend_from_slice(&10u16.to_be_bytes()); + packet.extend_from_slice(&((16 + set_length) as u16).to_be_bytes()); + packet.extend_from_slice(&0u32.to_be_bytes()); + packet.extend_from_slice(&1u32.to_be_bytes()); + packet.extend_from_slice(&1u32.to_be_bytes()); + packet.extend_from_slice(&2u16.to_be_bytes()); + packet.extend_from_slice(&(set_length as u16).to_be_bytes()); + packet.extend_from_slice(&256u16.to_be_bytes()); + packet.extend_from_slice(&(FIELDS.len() as u16).to_be_bytes()); + for (field, length) in FIELDS { + packet.extend_from_slice(&field.to_be_bytes()); + packet.extend_from_slice(&length.to_be_bytes()); + } + packet +} + +fn ipfix_data_packet(flow_count: u16) -> Vec { + let body = records(flow_count); + let mut packet = Vec::with_capacity(20 + body.len()); + packet.extend_from_slice(&10u16.to_be_bytes()); + packet.extend_from_slice(&((20 + body.len()) as u16).to_be_bytes()); + packet.extend_from_slice(&0u32.to_be_bytes()); + packet.extend_from_slice(&1u32.to_be_bytes()); + packet.extend_from_slice(&1u32.to_be_bytes()); + packet.extend_from_slice(&256u16.to_be_bytes()); + packet.extend_from_slice(&((body.len() + 4) as u16).to_be_bytes()); + packet.extend_from_slice(&body); + packet +} + +fn records(flow_count: u16) -> Vec { + let mut body = Vec::with_capacity(usize::from(flow_count) * 20); + for value in 0..flow_count { + body.extend_from_slice(&[10, 0, (value >> 8) as u8, value as u8]); + body.extend_from_slice(&[10, 0, 1, value as u8]); + body.extend_from_slice(&80u16.to_be_bytes()); + body.extend_from_slice(&443u16.to_be_bytes()); + body.extend_from_slice(&1280u32.to_be_bytes()); + body.extend_from_slice(&10u32.to_be_bytes()); + } + body +} + +fn measure(mut operation: impl FnMut()) -> (usize, usize, isize) { + operation(); + let baseline = LIVE_BYTES.load(Ordering::SeqCst); + ALLOCATION_CALLS.store(0, Ordering::SeqCst); + REQUESTED_BYTES.store(0, Ordering::SeqCst); + TRACKING.store(true, Ordering::SeqCst); + for _ in 0..ITERATIONS { + operation(); + } + TRACKING.store(false, Ordering::SeqCst); + let after = LIVE_BYTES.load(Ordering::SeqCst); + let calls = ALLOCATION_CALLS.load(Ordering::SeqCst); + let bytes = REQUESTED_BYTES.load(Ordering::SeqCst); + assert_eq!(calls % ITERATIONS, 0); + assert_eq!(bytes % ITERATIONS, 0); + ( + calls / ITERATIONS, + bytes / ITERATIONS, + after as isize - baseline as isize, + ) +} + +fn run(protocol: Protocol, scenario: Scenario, flow_count: u16) -> (usize, usize, isize) { + let template = protocol.template_packet(); + let packet = protocol.data_packet(flow_count); + let source = SocketAddr::from(([192, 0, 2, 1], 2055)); + + match scenario { + Scenario::DirectParse => { + let mut parser = NetflowParser::default(); + assert!(parser.parse_bytes(&template).is_ok()); + let check = parser.parse_bytes(&packet); + assert!(check.is_ok()); + protocol.assert_records(&check.packets, usize::from(flow_count)); + measure(|| drop(black_box(parser.parse_bytes(black_box(&packet))))) + } + Scenario::DirectIterator => { + let mut parser = NetflowParser::default(); + assert!(parser.parse_bytes(&template).is_ok()); + let check = parser + .iter_packets(&packet) + .map(Result::unwrap) + .collect::>(); + protocol.assert_records(&check, usize::from(flow_count)); + measure(|| { + for result in parser.iter_packets(black_box(&packet)) { + black_box(result.unwrap()); + } + }) + } + Scenario::AutoParse => { + let mut parser = AutoScopedParser::new(); + assert!(parser.parse_from_source(source, &template).is_ok()); + let check = parser.parse_from_source(source, &packet); + assert!(check.is_ok()); + protocol.assert_records(&check.packets, usize::from(flow_count)); + measure(|| { + drop(black_box( + parser.parse_from_source(source, black_box(&packet)), + )); + }) + } + Scenario::AutoIterator => { + let mut parser = AutoScopedParser::new(); + assert!(parser.parse_from_source(source, &template).is_ok()); + let check = parser + .iter_packets_from_source(source, &packet) + .unwrap() + .map(Result::unwrap) + .collect::>(); + protocol.assert_records(&check, usize::from(flow_count)); + measure(|| { + let iterator = parser + .iter_packets_from_source(source, black_box(&packet)) + .unwrap(); + for result in iterator { + black_box(result.unwrap()); + } + }) + } + } +} + +#[test] +fn warmed_common_hot_path_allocations() { + for protocol in [Protocol::V9, Protocol::Ipfix] { + for scenario in [ + Scenario::DirectParse, + Scenario::DirectIterator, + Scenario::AutoParse, + Scenario::AutoIterator, + ] { + for flow_count in [1, 1000] { + let (calls, bytes, live_delta) = run(protocol, scenario, flow_count); + println!( + "allocation_profile\t{}/{}/{}\t{}\t{}\t{}", + protocol.name(), + scenario.name(), + flow_count, + calls, + bytes, + live_delta + ); + assert_eq!(live_delta, 0); + } + } + } +}