Skip to main content

push_packet/engine/linear/
mod.rs

1//! Defines the [`LinearEngine`]
2mod rules;
3use aya::{
4    Ebpf, include_bytes_aligned,
5    maps::{Array, MapData},
6};
7use push_packet_common::engine::linear::{CAPACITY, Ipv4Rule, Ipv6Rule};
8
9use crate::{
10    array_ext::ArrayExt,
11    ebpf::{array_owned, load_xdp_program},
12    engine::Engine,
13    error::Error,
14    loader::Loader,
15    rules::{AddressFamily, Rule, RuleId},
16};
17
18/// Name of the ipv4 map
19const IP_V4_MAP: &str = "LINEAR_MAP_V4";
20/// Name of the ipv6 map
21const IP_V6_MAP: &str = "LINEAR_MAP_V6";
22/// Name of the rule count map
23const RULE_COUNT_MAP: &str = "LINEAR_RULE_COUNT";
24
25#[derive(Default)]
26/// Loader for a [`LinearEngine`].
27pub struct LinearEngineLoader;
28
29impl Loader for LinearEngineLoader {
30    type Component = LinearEngine;
31    fn load(self, ebpf: &mut Ebpf) -> Result<Self::Component, Error> {
32        let ipv4_rules = array_owned::<Ipv4Rule>(ebpf, IP_V4_MAP)?;
33        let ipv6_rules = array_owned::<Ipv6Rule>(ebpf, IP_V6_MAP)?;
34        let rule_count = array_owned::<u32>(ebpf, RULE_COUNT_MAP)?;
35
36        load_xdp_program(ebpf, LinearEngine::EBPF_PROGRAM_NAME)?;
37
38        Ok(LinearEngine {
39            v4_count: 0,
40            v6_count: 0,
41            ipv4_rules,
42            ipv6_rules,
43            rule_count,
44        })
45    }
46}
47
48/// The `LinearEngine` is a simple rule-matching engine that processes rules in order. It has a max
49/// capacity of [`CAPACITY`], but stops early based on the number of rules populated.
50pub struct LinearEngine {
51    v4_count: u32,
52    v6_count: u32,
53    ipv4_rules: Array<MapData, Ipv4Rule>,
54    ipv6_rules: Array<MapData, Ipv6Rule>,
55    rule_count: Array<MapData, u32>,
56}
57
58impl LinearEngine {
59    fn update_counts(&mut self) -> Result<(), Error> {
60        self.rule_count
61            .set(0, self.v4_count, 0)
62            .map_err(|e| Error::map(RULE_COUNT_MAP, e))?;
63        self.rule_count
64            .set(1, self.v6_count, 0)
65            .map_err(|e| Error::map(RULE_COUNT_MAP, e))
66    }
67
68    fn add_ipv4_rule(&mut self, rule_id: RuleId, rule: &Rule) -> Result<(), Error> {
69        let rule: Ipv4Rule = rule.into();
70        self.ipv4_rules
71            .set(rule_id.0, rule, 0)
72            .map_err(|e| Error::map(IP_V4_MAP, e))?;
73        self.v4_count = self.v4_count.max(rule_id.0 + 1);
74        self.update_counts()
75    }
76
77    fn add_ipv6_rule(&mut self, rule_id: RuleId, rule: &Rule) -> Result<(), Error> {
78        let rule: Ipv6Rule = rule.into();
79        self.ipv6_rules
80            .set(rule_id.0, rule, 0)
81            .map_err(|e| Error::map(IP_V6_MAP, e))?;
82        self.v6_count = self.v6_count.max(rule_id.0 + 1);
83        self.update_counts()
84    }
85}
86
87impl Engine for LinearEngine {
88    type Loader = LinearEngineLoader;
89    const EBPF_PROGRAM_NAME: &'static str = "linear";
90    #[cfg(feature = "build-ebpf")]
91    const EBPF_BYTES: &'static [u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/linear"));
92
93    #[cfg(not(feature = "build-ebpf"))]
94    const EBPF_BYTES: &'static [u8] = include_bytes_aligned!("../../../ebpf-bin/linear");
95
96    fn capacity(&self) -> Option<usize> {
97        Some(CAPACITY)
98    }
99
100    fn add_rule(&mut self, rule_id: RuleId, rule: &Rule) -> Result<(), Error> {
101        match rule.address_family() {
102            AddressFamily::Ipv4 => self.add_ipv4_rule(rule_id, rule),
103            AddressFamily::Ipv6 => self.add_ipv6_rule(rule_id, rule),
104            AddressFamily::Any => {
105                self.add_ipv4_rule(rule_id, rule)?;
106                self.add_ipv6_rule(rule_id, rule)?;
107                Ok(())
108            }
109        }
110    }
111
112    fn remove_rule(&mut self, rule_id: RuleId, rule: &Rule) -> Result<(), Error> {
113        let rule_id = rule_id.0;
114        match rule.address_family() {
115            AddressFamily::Ipv4 => self.ipv4_rules.clear(rule_id, IP_V4_MAP)?,
116            AddressFamily::Ipv6 => self.ipv6_rules.clear(rule_id, IP_V6_MAP)?,
117            AddressFamily::Any => {
118                self.ipv4_rules.clear(rule_id, IP_V4_MAP)?;
119                self.ipv6_rules.clear(rule_id, IP_V6_MAP)?;
120            }
121        }
122        Ok(())
123    }
124}