push-packet 0.1.0

Packet-inspection and routing library for Linux, built on eBPF XDP and AF_XDP with aya.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104

use crate::{
    cast,
    rules::{Rule, RuleId},
};

/// A simple interface for storing [`Rule`]s and getting the next available [`RuleId`] with id
/// reclamation.
#[derive(Default)]
pub(crate) struct Filter {
    rules: Vec<Option<Rule>>,
    free: Vec<u32>,
}

impl Filter {
    pub fn next_rule_id(&self) -> RuleId {
        match self.free.last() {
            Some(&index) => RuleId(index),
            None => RuleId(cast::usize_to_rule_index(self.rules.len())),
        }
    }

    pub fn add(&mut self, rule: Rule) {
        match self.free.pop() {
            None => self.rules.push(Some(rule)),
            Some(index) => self.rules[cast::rule_index_to_usize(index)] = Some(rule),
        }
    }

    pub fn get(&self, rule_id: RuleId) -> Option<&Rule> {
        match self.rules.get(cast::rule_index_to_usize(rule_id.0)) {
            Some(Some(rule)) => Some(rule),
            _ => None,
        }
    }

    pub fn remove(&mut self, rule_id: RuleId) -> Option<Rule> {
        if cast::rule_index_to_usize(rule_id.0) >= self.rules.len() {
            return None;
        }
        match self.rules[cast::rule_index_to_usize(rule_id.0)].take() {
            None => None,
            Some(rule) => {
                self.free.push(rule_id.0);
                Some(rule)
            }
        }
    }

    pub fn iter_rules(&self) -> impl Iterator<Item = (RuleId, &Rule)> {
        self.rules.iter().enumerate().filter_map(|(index, rule)| {
            rule.as_ref()
                .map(|rule| (RuleId(cast::usize_to_rule_index(index)), rule))
        })
    }
}

#[cfg(test)]
mod tests {

    use crate::{
        filter::Filter,
        rules::{Action, Rule},
    };

    fn test_rule() -> Rule {
        Rule::source_cidr("127.0.0.1")
            .action(Action::Pass)
            .build()
            .unwrap()
    }

    #[test]
    fn filter_ids_increment() {
        let mut filter = Filter::default();
        let id = filter.next_rule_id();
        filter.add(test_rule());
        let id_2 = filter.next_rule_id();
        assert_eq!(id.0, 0);
        assert_eq!(id_2.0, 1);
    }

    #[test]
    fn filter_reclaims_ids() {
        let mut filter = Filter::default();
        let id = filter.next_rule_id();
        filter.add(test_rule());
        filter.remove(id);
        let id_2 = filter.next_rule_id();
        filter.add(test_rule());
        let id_3 = filter.next_rule_id();
        assert_eq!(id.0, 0);
        assert_eq!(id_2.0, 0);
        assert_eq!(id_3.0, 1);
    }

    #[test]
    fn filter_double_remove_returns_none() {
        let mut filter = Filter::default();
        let id = filter.next_rule_id();
        filter.add(test_rule());
        assert!(filter.remove(id).is_some());
        assert!(filter.remove(id).is_none());
    }
}