Add ipv4only/ipv6only logic to actions

This commit is contained in:
ppom 2025-07-31 12:00:00 +02:00
commit b927ba4fdf
No known key found for this signature in database
8 changed files with 201 additions and 27 deletions

View file

@ -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<BTreeSet<Arc<Pattern>>>,
#[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());

View file

@ -32,6 +32,8 @@ pub enum Duplicate {
pub struct Filter {
#[serde(skip)]
longuest_action_duration: TimeDelta,
#[serde(skip)]
has_ip: bool,
regex: Vec<String>,
#[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| {

View file

@ -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;

View file

@ -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 <ip> test",
@ -714,7 +719,6 @@ mod patternip_tests {
ip,
"line: {line}"
);
println!("line ok: {line}");
}));
}
}

View file

@ -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))

View file

@ -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();

View file

@ -248,6 +248,7 @@ mod tests {
"f1",
"a1",
&patterns,
0,
)],
vec!["test <az>"],
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 <az>"],
Some(3),

View file

@ -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 <az> >> {}", &bed.out_file)],
@ -149,6 +154,7 @@ async fn three_matches_then_action_then_delayed_action() {
"test",
"a2",
&bed.az_patterns,
0,
),
],
vec!["test <az>"],
@ -271,6 +277,7 @@ async fn one_match_one_action() {
"test",
"a1",
&bed.az_patterns,
0,
)],
vec!["test <az>"],
None,
@ -317,6 +324,7 @@ async fn one_match_one_delayed_action() {
"test",
"a1",
&bed.az_patterns,
0,
)],
vec!["test <az>"],
None,
@ -382,6 +390,7 @@ async fn one_db_match_one_runtime_match_one_action() {
"test",
"a1",
&bed.az_patterns,
0,
)],
vec!["test <az>"],
Some(2),
@ -448,6 +457,7 @@ async fn one_outdated_db_match() {
"test",
"a1",
&bed.az_patterns,
0,
)],
vec!["test <az>"],
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 <az> >> {}", &bed.out_file)],
@ -504,6 +515,7 @@ async fn trigger_unmatched_pattern() {
"test",
"a2",
&bed.az_patterns,
0,
),
],
vec!["test <az>"],
@ -566,6 +578,7 @@ async fn trigger_matched_pattern() {
"test",
"a1",
&bed.az_patterns,
0,
),
Action::new(
vec!["sh", "-c", &format!("echo a2 <az> >> {}", &bed.out_file)],
@ -575,6 +588,7 @@ async fn trigger_matched_pattern() {
"test",
"a2",
&bed.az_patterns,
0,
),
],
vec!["test <az>"],
@ -644,6 +658,7 @@ async fn multiple_triggers() {
"test",
"a1",
&bed.az_patterns,
0,
),
Action::new(
vec!["sh", "-c", &format!("echo a2 <az> >> {}", &bed.out_file)],
@ -653,6 +668,7 @@ async fn multiple_triggers() {
"test",
"a2",
&bed.az_patterns,
0,
),
],
vec!["test <az>"],
@ -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 <ip> >> {}", &bed.out_file)],
None,
false,
"test",
"test",
"a4",
&bed.ip_patterns,
4,
),
Action::new(
vec!["sh", "-c", &format!("echo ipv6 <ip> >> {}", &bed.out_file)],
None,
false,
"test",
"test",
"a6",
&bed.ip_patterns,
6,
),
Action::new(
vec!["sh", "-c", &format!("echo any <ip> >> {}", &bed.out_file)],
Some("20ms"),
false,
"test",
"test",
"ax",
&bed.ip_patterns,
0,
),
],
vec!["test <ip>"],
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(),
);
}