Skip to main content

stackforge_core/flow/
icmp_state.rs

1use std::time::Duration;
2
3use crate::Packet;
4
5use super::config::FlowConfig;
6use super::state::ConversationStatus;
7
8/// ICMP/ICMPv6 conversation state.
9///
10/// Tracks ICMP-specific metadata for echo request/reply pairs and other ICMP types.
11/// Echo requests and replies are correlated using the ICMP identifier field.
12#[derive(Debug, Clone)]
13pub struct IcmpFlowState {
14    /// ICMP type (e.g., 8 for Echo Request, 0 for Echo Reply).
15    pub icmp_type: u8,
16    /// ICMP code.
17    pub icmp_code: u8,
18    /// ICMP identifier (for echo, timestamp, and other types that use it).
19    pub identifier: Option<u16>,
20    /// Number of echo requests (type 8 for ICMP, 128 for ICMPv6).
21    pub request_count: u64,
22    /// Number of echo replies (type 0 for ICMP, 129 for ICMPv6).
23    pub reply_count: u64,
24    /// Last sequence number seen in an echo packet.
25    pub last_seq: Option<u16>,
26    /// Conversation status.
27    pub status: ConversationStatus,
28}
29
30impl IcmpFlowState {
31    #[must_use]
32    pub fn new(icmp_type: u8, icmp_code: u8) -> Self {
33        Self {
34            icmp_type,
35            icmp_code,
36            identifier: None,
37            request_count: 0,
38            reply_count: 0,
39            last_seq: None,
40            status: ConversationStatus::Active,
41        }
42    }
43
44    /// Update state when a new ICMP packet is received.
45    ///
46    /// Increments request or reply count based on ICMP type, and updates
47    /// the identifier and sequence number fields if present.
48    pub fn process_packet(&mut self, packet: &Packet, buf: &[u8], icmp_type: u8, icmp_code: u8) {
49        // Update type/code on every packet (they should be consistent)
50        self.icmp_type = icmp_type;
51        self.icmp_code = icmp_code;
52
53        // Get ICMP layer bounds to extract fields
54        if let Some(icmp_layer) = crate::layer::LayerKind::Icmp
55            .try_into()
56            .ok()
57            .and_then(|kind| packet.get_layer(kind))
58        {
59            let icmp_start = icmp_layer.start;
60
61            // Extract identifier (bytes 4-5) if present
62            if buf.len() >= icmp_start + 6 {
63                self.identifier = Some(u16::from_be_bytes([
64                    buf[icmp_start + 4],
65                    buf[icmp_start + 5],
66                ]));
67            }
68
69            // Extract sequence number (bytes 6-7) if present
70            if buf.len() >= icmp_start + 8 {
71                self.last_seq = Some(u16::from_be_bytes([
72                    buf[icmp_start + 6],
73                    buf[icmp_start + 7],
74                ]));
75            }
76
77            // Count requests and replies based on ICMP type
78            match icmp_type {
79                8 => {
80                    // ICMP Echo Request
81                    self.request_count += 1;
82                },
83                0 => {
84                    // ICMP Echo Reply
85                    self.reply_count += 1;
86                },
87                128 => {
88                    // ICMPv6 Echo Request
89                    self.request_count += 1;
90                },
91                129 => {
92                    // ICMPv6 Echo Reply
93                    self.reply_count += 1;
94                },
95                _ => {
96                    // Other ICMP types: no counting
97                },
98            }
99        }
100
101        self.status = ConversationStatus::Active;
102    }
103
104    /// Check whether this flow has timed out.
105    #[must_use]
106    pub fn check_timeout(&self, last_seen: Duration, now: Duration, config: &FlowConfig) -> bool {
107        // ICMP uses UDP timeout
108        now.saturating_sub(last_seen) > config.udp_timeout
109    }
110}
111
112impl Default for IcmpFlowState {
113    fn default() -> Self {
114        Self::new(0, 0)
115    }
116}
117
118#[cfg(test)]
119mod tests {
120    use super::*;
121
122    #[test]
123    fn test_icmp_state_new() {
124        let state = IcmpFlowState::new(8, 0);
125        assert_eq!(state.icmp_type, 8);
126        assert_eq!(state.icmp_code, 0);
127        assert_eq!(state.request_count, 0);
128        assert_eq!(state.reply_count, 0);
129        assert_eq!(state.identifier, None);
130        assert_eq!(state.last_seq, None);
131    }
132
133    #[test]
134    fn test_icmp_timeout() {
135        let config = FlowConfig::default(); // 120s UDP timeout
136        let state = IcmpFlowState::new(8, 0);
137
138        // Not timed out
139        assert!(!state.check_timeout(Duration::from_secs(100), Duration::from_secs(200), &config));
140
141        // Timed out
142        assert!(state.check_timeout(Duration::from_secs(100), Duration::from_secs(300), &config));
143    }
144}