From b927ba4fdfa354a6955a1c00a604cd636bff2633 Mon Sep 17 00:00:00 2001 From: ppom Date: Thu, 31 Jul 2025 12:00:00 +0200 Subject: [PATCH] Add ipv4only/ipv6only logic to actions --- src/concepts/action.rs | 57 +++++++++++++-------- src/concepts/filter.rs | 10 ++++ src/concepts/mod.rs | 2 +- src/concepts/pattern/ip/mod.rs | 6 ++- src/concepts/pattern/mod.rs | 6 ++- src/daemon/filter/mod.rs | 29 ++++++++++- src/daemon/filter/state.rs | 25 +++++++-- src/daemon/filter/tests.rs | 93 ++++++++++++++++++++++++++++++++++ 8 files changed, 201 insertions(+), 27 deletions(-) diff --git a/src/concepts/action.rs b/src/concepts/action.rs index e4e0fc9..3e9972a 100644 --- a/src/concepts/action.rs +++ b/src/concepts/action.rs @@ -5,7 +5,7 @@ use chrono::TimeDelta; use serde::{Deserialize, Serialize}; use tokio::process::Command; -use super::parse_duration::*; +use super::{parse_duration::*, PatternType}; use super::{Match, Pattern}; #[derive(Clone, Debug, Default, Deserialize, Serialize)] @@ -28,6 +28,11 @@ pub struct Action { #[serde(default = "set_false", skip_serializing_if = "is_false")] oneshot: bool, + #[serde(default = "set_false", skip_serializing_if = "is_false")] + ipv4only: bool, + #[serde(default = "set_false", skip_serializing_if = "is_false")] + ipv6only: bool, + #[serde(skip)] patterns: Arc>>, #[serde(skip)] @@ -63,6 +68,13 @@ impl Action { self.oneshot } + pub fn ipv4only(&self) -> bool { + self.ipv4only + } + pub fn ipv6only(&self) -> bool { + self.ipv6only + } + pub fn setup( &mut self, stream_name: &str, @@ -110,6 +122,22 @@ impl Action { return Err("cannot have `onexit: true`, without an `after` directive".into()); } + if self.ipv4only && self.ipv6only { + return Err("cannot have `ipv4only: true` and `ipv6only: true` in one action".into()); + } + if self + .patterns + .iter() + .all(|pattern| pattern.pattern_type() != PatternType::Ip) + { + if self.ipv4only { + return Err("it makes no sense to have an action with `ipv4only: true` when no pattern of type ip is defined on the filter".into()); + } + if self.ipv6only { + return Err("it makes no sense to have an action with `ipv6only: true` when no pattern of type ip is defined on the filter".into()); + } + } + Ok(()) } @@ -175,11 +203,14 @@ impl Action { filter_name: &str, name: &str, config_patterns: &super::Patterns, + ip_only: u8, ) -> Self { let mut action = Self { cmd: cmd.into_iter().map(|s| s.into()).collect(), after: after.map(|s| s.into()), on_exit, + ipv4only: ip_only == 4, + ipv6only: ip_only == 6, ..Default::default() }; action @@ -203,28 +234,14 @@ pub mod tests { use super::*; - fn default_action() -> Action { - Action { - cmd: Vec::new(), - name: "".into(), - filter_name: "".into(), - stream_name: "".into(), - after: None, - after_duration: None, - on_exit: false, - oneshot: false, - patterns: Arc::new(BTreeSet::default()), - } - } - pub fn ok_action() -> Action { - let mut action = default_action(); + let mut action = Action::default(); action.cmd = vec!["command".into()]; action } pub fn ok_action_with_after(d: String, name: &str) -> Action { - let mut action = default_action(); + let mut action = Action::default(); action.cmd = vec!["command".into()]; action.after = Some(d); action @@ -240,16 +257,16 @@ pub mod tests { let patterns = Arc::new(BTreeSet::default()); // No command - action = default_action(); + action = Action::default(); assert!(action.setup(&name, &name, &name, patterns.clone()).is_err()); // No command - action = default_action(); + action = Action::default(); action.cmd = vec!["".into()]; assert!(action.setup(&name, &name, &name, patterns.clone()).is_err()); // No command - action = default_action(); + action = Action::default(); action.cmd = vec!["".into(), "arg1".into()]; assert!(action.setup(&name, &name, &name, patterns.clone()).is_err()); diff --git a/src/concepts/filter.rs b/src/concepts/filter.rs index 410a997..98be40b 100644 --- a/src/concepts/filter.rs +++ b/src/concepts/filter.rs @@ -32,6 +32,8 @@ pub enum Duplicate { pub struct Filter { #[serde(skip)] longuest_action_duration: TimeDelta, + #[serde(skip)] + has_ip: bool, regex: Vec, #[serde(skip)] @@ -119,6 +121,10 @@ impl Filter { &self.patterns } + pub fn check_ip(&self) -> bool { + self.has_ip + } + pub fn setup( &mut self, stream_name: &str, @@ -211,6 +217,10 @@ impl Filter { for (key, action) in &mut self.actions { action.setup(stream_name, name, key, self.patterns.clone())?; } + self.has_ip = self + .actions + .values() + .any(|action| action.ipv4only() || action.ipv6only()); self.longuest_action_duration = self.actions.values().fold(TimeDelta::seconds(0), |acc, v| { diff --git a/src/concepts/mod.rs b/src/concepts/mod.rs index b4a785e..40a624c 100644 --- a/src/concepts/mod.rs +++ b/src/concepts/mod.rs @@ -9,7 +9,7 @@ pub use action::Action; pub use config::{Config, Patterns}; pub use filter::{Duplicate, Filter}; use parse_duration::parse_duration; -pub use pattern::Pattern; +pub use pattern::{Pattern, PatternType}; use serde::{Deserialize, Serialize}; pub use stream::Stream; diff --git a/src/concepts/pattern/ip/mod.rs b/src/concepts/pattern/ip/mod.rs index 00cf1c1..4165852 100644 --- a/src/concepts/pattern/ip/mod.rs +++ b/src/concepts/pattern/ip/mod.rs @@ -152,6 +152,10 @@ pub struct PatternIp { } impl PatternIp { + pub fn pattern_type(&self) -> PatternType { + self.pattern_type + } + /// Setup the IP-specific part of a Pattern. /// Returns an optional regex string if of type IP, else None /// Returns an error if one of: @@ -688,6 +692,7 @@ mod patternip_tests { "test", "a1", &bed.ip_patterns, + 0, )], vec![ "^borned test", @@ -714,7 +719,6 @@ mod patternip_tests { ip, "line: {line}" ); - println!("line ok: {line}"); })); } } diff --git a/src/concepts/pattern/mod.rs b/src/concepts/pattern/mod.rs index 9fbb610..21c8d31 100644 --- a/src/concepts/pattern/mod.rs +++ b/src/concepts/pattern/mod.rs @@ -5,7 +5,7 @@ use serde::{Deserialize, Serialize}; mod ip; -use ip::PatternIp; +pub use ip::{PatternIp, PatternType}; #[derive(Clone, Debug, Deserialize, Serialize)] #[cfg_attr(test, derive(Default))] @@ -47,6 +47,10 @@ impl Pattern { &self.name_with_braces } + pub fn pattern_type(&self) -> PatternType { + self.ip.pattern_type() + } + pub fn setup(&mut self, name: &str) -> Result<(), String> { self._setup(name) .map_err(|msg| format!("pattern {}: {}", name, msg)) diff --git a/src/daemon/filter/mod.rs b/src/daemon/filter/mod.rs index 290f3a1..484b1b5 100644 --- a/src/daemon/filter/mod.rs +++ b/src/daemon/filter/mod.rs @@ -14,7 +14,7 @@ use tokio::sync::Semaphore; use tracing::{error, info}; use crate::{ - concepts::{Action, Duplicate, Filter, Match, Pattern, Time}, + concepts::{Action, Duplicate, Filter, Match, Pattern, PatternType, Time}, protocol::{Order, PatternStatus}, treedb::Database, }; @@ -240,6 +240,30 @@ impl FilterManager { startup: bool, only_after: bool, ) { + // Testing if we have an IPv4 or IPv6 + let ip_type = if self.filter.check_ip() { + self.filter + .patterns() + .iter() + .zip(&m) + .find(|(p, _)| p.pattern_type() == PatternType::Ip) + .map(|(_, m)| -> _ { + // Using this dumb heuristic is ok, + // because we know we have a valid IP address. + if m.contains(':') { + PatternType::Ipv6 + } else if m.contains('.') { + PatternType::Ipv4 + } else { + PatternType::Regex + } + }) + .unwrap_or(PatternType::Regex) + } else { + PatternType::Regex + }; + + // Scheduling each action for action in self .filter .actions() @@ -248,6 +272,9 @@ impl FilterManager { .filter(|action| !startup || !action.oneshot()) // If only_after, keep only after actions .filter(|action| !only_after || action.after_duration().is_some()) + // If specific ip version, check it + .filter(|action| !action.ipv4only() || ip_type == PatternType::Ipv4) + .filter(|action| !action.ipv6only() || ip_type == PatternType::Ipv6) { let exec_time = t + action.after_duration().unwrap_or_default(); let m = m.clone(); diff --git a/src/daemon/filter/state.rs b/src/daemon/filter/state.rs index 8e3d7d3..9c44d71 100644 --- a/src/daemon/filter/state.rs +++ b/src/daemon/filter/state.rs @@ -248,6 +248,7 @@ mod tests { "f1", "a1", &patterns, + 0, )], vec!["test "], Some(3), @@ -374,9 +375,27 @@ mod tests { let patterns = Pattern::new_map("az", "[a-z]+").unwrap(); let filter = Filter::new_static( vec![ - Action::new(vec!["true"], None, false, "s1", "f1", "a1", &patterns), - Action::new(vec!["true"], Some("1s"), false, "s1", "f1", "a2", &patterns), - Action::new(vec!["true"], Some("3s"), false, "s1", "f1", "a3", &patterns), + Action::new(vec!["true"], None, false, "s1", "f1", "a1", &patterns, 0), + Action::new( + vec!["true"], + Some("1s"), + false, + "s1", + "f1", + "a2", + &patterns, + 0, + ), + Action::new( + vec!["true"], + Some("3s"), + false, + "s1", + "f1", + "a3", + &patterns, + 0, + ), ], vec!["test "], Some(3), diff --git a/src/daemon/filter/tests.rs b/src/daemon/filter/tests.rs index a926b2a..0565db2 100644 --- a/src/daemon/filter/tests.rs +++ b/src/daemon/filter/tests.rs @@ -124,6 +124,10 @@ impl TestBed2 { ); assert!(state.triggers.is_empty(), "triggers must be empty"); } + + pub async fn reset_out_file(&self) { + tokio::fs::write(&self.out_file, "").await.unwrap(); + } } #[tokio::test] @@ -140,6 +144,7 @@ async fn three_matches_then_action_then_delayed_action() { "test", "a1", &bed.az_patterns, + 0, ), Action::new( vec!["sh", "-c", &format!("echo a2 >> {}", &bed.out_file)], @@ -149,6 +154,7 @@ async fn three_matches_then_action_then_delayed_action() { "test", "a2", &bed.az_patterns, + 0, ), ], vec!["test "], @@ -271,6 +277,7 @@ async fn one_match_one_action() { "test", "a1", &bed.az_patterns, + 0, )], vec!["test "], None, @@ -317,6 +324,7 @@ async fn one_match_one_delayed_action() { "test", "a1", &bed.az_patterns, + 0, )], vec!["test "], None, @@ -382,6 +390,7 @@ async fn one_db_match_one_runtime_match_one_action() { "test", "a1", &bed.az_patterns, + 0, )], vec!["test "], Some(2), @@ -448,6 +457,7 @@ async fn one_outdated_db_match() { "test", "a1", &bed.az_patterns, + 0, )], vec!["test "], Some(2), @@ -495,6 +505,7 @@ async fn trigger_unmatched_pattern() { "test", "a1", &bed.az_patterns, + 0, ), Action::new( vec!["sh", "-c", &format!("echo a2 >> {}", &bed.out_file)], @@ -504,6 +515,7 @@ async fn trigger_unmatched_pattern() { "test", "a2", &bed.az_patterns, + 0, ), ], vec!["test "], @@ -566,6 +578,7 @@ async fn trigger_matched_pattern() { "test", "a1", &bed.az_patterns, + 0, ), Action::new( vec!["sh", "-c", &format!("echo a2 >> {}", &bed.out_file)], @@ -575,6 +588,7 @@ async fn trigger_matched_pattern() { "test", "a2", &bed.az_patterns, + 0, ), ], vec!["test "], @@ -644,6 +658,7 @@ async fn multiple_triggers() { "test", "a1", &bed.az_patterns, + 0, ), Action::new( vec!["sh", "-c", &format!("echo a2 >> {}", &bed.out_file)], @@ -653,6 +668,7 @@ async fn multiple_triggers() { "test", "a2", &bed.az_patterns, + 0, ), ], vec!["test "], @@ -745,3 +761,80 @@ async fn multiple_triggers() { ); } } + +#[tokio::test] +async fn ip_specific() { + let bed = TestBed::default(); + let filter = Filter::new_static( + vec![ + Action::new( + vec!["sh", "-c", &format!("echo ipv4 >> {}", &bed.out_file)], + None, + false, + "test", + "test", + "a4", + &bed.ip_patterns, + 4, + ), + Action::new( + vec!["sh", "-c", &format!("echo ipv6 >> {}", &bed.out_file)], + None, + false, + "test", + "test", + "a6", + &bed.ip_patterns, + 6, + ), + Action::new( + vec!["sh", "-c", &format!("echo any >> {}", &bed.out_file)], + Some("20ms"), + false, + "test", + "test", + "ax", + &bed.ip_patterns, + 0, + ), + ], + vec!["test "], + None, + None, + "test", + "test", + Duplicate::Extend, + &bed.ip_patterns, + ); + + let bed = bed.part2(filter, Local::now(), None).await; + + assert_eq!( + bed.manager.handle_line("test 1.2.3.4", Local::now()), + React::Trigger, + ); + + // Wait for action to execute + tokio::time::sleep(Duration::from_millis(70)).await; + + assert_eq!( + "ipv4 1.2.3.4\nany 1.2.3.4\n", + &read_to_string(&bed.out_file).unwrap(), + ); + + bed.reset_out_file().await; + + assert_eq!( + bed.manager + .handle_line("test 1:2:3:4:5:6:7:8", Local::now()), + React::Trigger, + ); + + // Wait for action to execute + tokio::time::sleep(Duration::from_millis(70)).await; + + assert_eq!( + "ipv6 1:2:3:4:5:6:7:8\nany 1:2:3:4:5:6:7:8\n", + &read_to_string(&bed.out_file).unwrap(), + ); +}