eggress_core/
detect.rs

1use crate::ProtocolId;
2
3/// Result of a protocol detection attempt.
4#[derive(Debug, Clone, PartialEq, Eq)]
5pub enum DetectResult {
6    /// The protocol was matched with the given confidence level (0–100).
7    Match { confidence: u8 },
8    /// More data is needed; `minimum` is the smallest number of total prefix
9    /// bytes required before a definitive result can be given.
10    NeedMore { minimum: usize },
11    /// This prefix does not match the protocol.
12    NoMatch,
13}
14
15/// A trait for protocol detectors that can identify a protocol from the
16/// initial bytes of a stream.
17pub trait ProtocolDetector: Send + Sync {
18    /// Returns the unique identifier for the protocol this detector handles.
19    fn id(&self) -> ProtocolId;
20
21    /// Attempts to detect the protocol from the given byte prefix.
22    ///
23    /// # Arguments
24    /// * `prefix` - The bytes read from the start of the stream so far.
25    fn detect(&self, prefix: &[u8]) -> DetectResult;
26}
27
28/// A simple prefix-based detector useful for testing and simple protocols.
29pub struct PrefixDetector {
30    id: ProtocolId,
31    prefix: Vec<u8>,
32    min_length: usize,
33}
34
35impl PrefixDetector {
36    pub fn new(id: ProtocolId, prefix: Vec<u8>) -> Self {
37        let min_length = prefix.len();
38        Self {
39            id,
40            prefix,
41            min_length,
42        }
43    }
44
45    pub fn with_min_length(id: ProtocolId, prefix: Vec<u8>, min_length: usize) -> Self {
46        Self {
47            id,
48            prefix,
49            min_length,
50        }
51    }
52}
53
54impl ProtocolDetector for PrefixDetector {
55    fn id(&self) -> ProtocolId {
56        self.id
57    }
58
59    fn detect(&self, data: &[u8]) -> DetectResult {
60        if data.len() < self.min_length {
61            DetectResult::NeedMore {
62                minimum: self.min_length,
63            }
64        } else if data.starts_with(&self.prefix) {
65            DetectResult::Match { confidence: 100 }
66        } else {
67            DetectResult::NoMatch
68        }
69    }
70}
71
72#[cfg(test)]
73mod tests {
74    use super::*;
75
76    #[test]
77    fn test_prefix_detector_match() {
78        let detector = PrefixDetector::new(ProtocolId::Http, b"GET ".to_vec());
79        assert_eq!(
80            detector.detect(b"GET / HTTP/1.1"),
81            DetectResult::Match { confidence: 100 }
82        );
83    }
84
85    #[test]
86    fn test_prefix_detector_no_match() {
87        let detector = PrefixDetector::new(ProtocolId::Http, b"GET ".to_vec());
88        assert_eq!(detector.detect(b"POST /"), DetectResult::NoMatch);
89    }
90
91    #[test]
92    fn test_prefix_detector_need_more() {
93        let detector = PrefixDetector::new(ProtocolId::Http, b"GET ".to_vec());
94        assert_eq!(
95            detector.detect(b"GE"),
96            DetectResult::NeedMore { minimum: 4 }
97        );
98    }
99
100    #[test]
101    fn test_prefix_detector_empty_input() {
102        let detector = PrefixDetector::new(ProtocolId::Http, b"GET ".to_vec());
103        assert_eq!(detector.detect(b""), DetectResult::NeedMore { minimum: 4 });
104    }
105
106    #[test]
107    fn test_prefix_detector_exact_match() {
108        let detector = PrefixDetector::new(ProtocolId::Socks5, b"\x05".to_vec());
109        assert_eq!(
110            detector.detect(b"\x05"),
111            DetectResult::Match { confidence: 100 }
112        );
113    }
114
115    #[test]
116    fn test_prefix_detector_with_min_length() {
117        let detector = PrefixDetector::with_min_length(ProtocolId::Http, b"GET ".to_vec(), 8);
118        // Have 4 bytes which is the prefix length but min_length is 8
119        assert_eq!(
120            detector.detect(b"GET /"),
121            DetectResult::NeedMore { minimum: 8 }
122        );
123    }
124}
eggress_core/detect.rs

eggress_core/
detect.rs
