Skip to main content

retina_core/conntrack/conn/tcp_conn/
reassembly.rs

1use crate::conntrack::conn::conn_info::{ConnInfo, ConnState};
2use crate::conntrack::pdu::L4Pdu;
3use crate::protocols::packet::tcp::{ACK, FIN, RST, SYN};
4use crate::protocols::stream::ParserRegistry;
5use crate::subscription::{Subscription, Trackable};
6
7use anyhow::{bail, Result};
8use std::collections::VecDeque;
9
10/// Represents a uni-directional TCP flow
11#[derive(Debug)]
12pub(crate) struct TcpFlow {
13    /// Expected sequence number of next segment
14    pub(super) next_seq: Option<u32>,
15    /// Flow status for consumed control packets.
16    /// Matches TCP flag bits.
17    pub(super) consumed_flags: u8,
18    /// Out-of-order buffer
19    pub(crate) ooo_buf: OutOfOrderBuffer,
20}
21
22impl TcpFlow {
23    /// Creates a default TCP flow
24    #[inline]
25    pub(super) fn default(capacity: usize) -> Self {
26        TcpFlow {
27            next_seq: None,
28            consumed_flags: 0,
29            ooo_buf: OutOfOrderBuffer::new(capacity),
30        }
31    }
32
33    /// Creates a new TCP flow with given next sequence number, flags,
34    /// and out-of-order buffer
35    #[inline]
36    pub(super) fn new(capacity: usize, next_seq: u32, flags: u8) -> Self {
37        TcpFlow {
38            next_seq: Some(next_seq),
39            consumed_flags: flags,
40            ooo_buf: OutOfOrderBuffer::new(capacity),
41        }
42    }
43
44    /// Attempt to insert incoming data segment into flow.
45    /// Buffer future segments and drop old segments.
46    /// Shunts TcpStream if the incoming segment causes out-of-order buffer overflow
47    #[inline]
48    pub(super) fn insert_segment<T: Trackable>(
49        &mut self,
50        mut segment: L4Pdu,
51        info: &mut ConnInfo<T>,
52        subscription: &Subscription<T::Subscribed>,
53        registry: &ParserRegistry,
54    ) {
55        let length = segment.length() as u32;
56        let cur_seq = segment.seq_no();
57
58        if let Some(next_seq) = self.next_seq {
59            if next_seq == cur_seq {
60                // Segment is the next expected segment in the sequence
61                self.consumed_flags |= segment.flags();
62                if segment.flags() & RST != 0 {
63                    info.consume_pdu(segment, subscription, registry);
64                    return;
65                }
66                let mut expected_seq = cur_seq.wrapping_add(length);
67                if segment.flags() & FIN != 0 {
68                    expected_seq = cur_seq.wrapping_add(1);
69                }
70                info.consume_pdu(segment, subscription, registry);
71                self.flush_ooo_buffer::<T>(expected_seq, info, subscription, registry);
72            } else if wrapping_lt(next_seq, cur_seq) {
73                // Segment comes after the next expected segment
74                self.buffer_ooo_seg(segment, info);
75            } else if let Some(expected_seq) = overlap(&mut segment, next_seq) {
76                // Segment starts before the next expected segment but has new data
77                self.consumed_flags |= segment.flags();
78                info.consume_pdu(segment, subscription, registry);
79                self.flush_ooo_buffer::<T>(expected_seq, info, subscription, registry);
80            } else {
81                // Segment contains old data
82                log::debug!(
83                    "Dropping old segment. cur: {} expect: {}",
84                    cur_seq,
85                    next_seq
86                );
87                drop(segment);
88            }
89        } else {
90            // expecting SYNACK in response to the originator's SYN
91            if segment.flags() & (SYN | ACK) != 0 {
92                let expected_seq = cur_seq.wrapping_add(1 + length);
93                self.next_seq = Some(expected_seq);
94                self.consumed_flags |= segment.flags();
95                info.consume_pdu(segment, subscription, registry);
96                self.flush_ooo_buffer::<T>(expected_seq, info, subscription, registry);
97            } else {
98                // Buffer out-of-order non-SYNACK packets
99                self.buffer_ooo_seg(segment, info);
100            }
101        }
102    }
103
104    /// Insert packet into ooo buffer and handle overflow
105    #[inline]
106    fn buffer_ooo_seg<T: Trackable>(&mut self, segment: L4Pdu, info: &mut ConnInfo<T>) {
107        if self.ooo_buf.insert_back(segment).is_err() {
108            log::warn!("Out-of-order buffer overflow");
109            info.state = ConnState::Remove;
110        }
111    }
112
113    /// Flushes the flow's out-of-order buffer given the next expected
114    /// sequence number and updates the flow's new next expected
115    /// sequence number and status after the flush.
116    #[inline]
117    pub(super) fn flush_ooo_buffer<T: Trackable>(
118        &mut self,
119        expected_seq: u32,
120        info: &mut ConnInfo<T>,
121        subscription: &Subscription<T::Subscribed>,
122        registry: &ParserRegistry,
123    ) {
124        if info.state == ConnState::Remove {
125            return;
126        }
127        let next_seq = self.ooo_buf.flush_ordered::<T>(
128            expected_seq,
129            &mut self.consumed_flags,
130            info,
131            subscription,
132            registry,
133        );
134        self.next_seq = Some(next_seq);
135    }
136}
137
138/// A buffer to hold reordered TCP segments
139#[derive(Debug)]
140pub(crate) struct OutOfOrderBuffer {
141    capacity: usize,
142    pub(crate) buf: VecDeque<L4Pdu>,
143}
144
145impl OutOfOrderBuffer {
146    /// Creates a new OutOfOrderBuffer with capacity
147    fn new(capacity: usize) -> Self {
148        OutOfOrderBuffer {
149            capacity,
150            buf: VecDeque::new(),
151        }
152    }
153
154    /// Returns the number of elements in the buffer
155    pub(crate) fn len(&self) -> usize {
156        self.buf.len()
157    }
158
159    /// Inserts segment at the end of the buffer.
160    fn insert_back(&mut self, segment: L4Pdu) -> Result<()> {
161        log::debug!("insert with seq : {:#?}", segment.seq_no());
162        if self.len() >= self.capacity {
163            // // must clear to drop buffered Mbufs
164            // self.buf.clear();
165            bail!("Out-of-order buffer overflow.");
166        }
167        self.buf.push_back(segment);
168        Ok(())
169    }
170
171    /// Consumes segments with expected data, retains segments with future data,
172    /// and drops segments with old data.
173    /// Returns the next expected sequence number and control flags of consumed segments.
174    #[inline]
175    fn flush_ordered<T: Trackable>(
176        &mut self,
177        expected_seq: u32,
178        consumed_flags: &mut u8,
179        info: &mut ConnInfo<T>,
180        subscription: &Subscription<T::Subscribed>,
181        registry: &ParserRegistry,
182    ) -> u32 {
183        let mut next_seq = expected_seq;
184        let mut index = 0;
185        while index < self.len() {
186            if info.state == ConnState::Remove {
187                return next_seq;
188            }
189
190            // unwraps ok because index < len
191            let cur_seq = self.buf.get_mut(index).unwrap().seq_no();
192            log::debug!("Flushing...current seq: {:#?}", cur_seq);
193
194            if next_seq == cur_seq {
195                let segment = self.buf.remove(index).unwrap();
196                *consumed_flags |= segment.flags();
197                if segment.flags() & RST != 0 {
198                    info.consume_pdu(segment, subscription, registry);
199                    return next_seq;
200                }
201                next_seq = next_seq.wrapping_add(segment.length() as u32);
202                if segment.flags() & FIN != 0 {
203                    next_seq = next_seq.wrapping_add(1);
204                }
205                info.consume_pdu(segment, subscription, registry);
206                index = 0;
207            } else if wrapping_lt(next_seq, cur_seq) {
208                index += 1;
209            } else {
210                let mut segment = self.buf.remove(index).unwrap();
211                if let Some(update_seq) = overlap(&mut segment, next_seq) {
212                    next_seq = update_seq;
213                    *consumed_flags |= segment.flags();
214                    info.consume_pdu(segment, subscription, registry);
215                    index = 0;
216                } else {
217                    log::debug!("Dropping old segment during flush.");
218                    drop(segment);
219                    index += 1;
220                }
221            }
222        }
223        next_seq
224    }
225}
226
227pub(crate) fn wrapping_lt(lhs: u32, rhs: u32) -> bool {
228    // From RFC1323:
229    //     TCP determines if a data segment is "old" or "new" by testing
230    //     whether its sequence number is within 2**31 bytes of the left edge
231    //     of the window, and if it is not, discarding the data as "old".  To
232    //     insure that new data is never mistakenly considered old and vice-
233    //     versa, the left edge of the sender's window has to be at most
234    //     2**31 away from the right edge of the receiver's window.
235    lhs.wrapping_sub(rhs) > (1 << 31)
236}
237
238/// Check if a segment has overlapping data with the received bytes.
239/// Returns the new expected sequence number if there is overlap
240fn overlap(segment: &mut L4Pdu, expected_seq: u32) -> Option<u32> {
241    let length = segment.length();
242    let cur_seq = segment.seq_no();
243    let end_seq = cur_seq.wrapping_add(length as u32);
244
245    if wrapping_lt(expected_seq, end_seq) {
246        // contains new data
247        let new_data_len = end_seq.wrapping_sub(expected_seq);
248        let overlap_data_len = expected_seq.wrapping_sub(cur_seq);
249
250        log::debug!("Overlap with new data size : {:#?}", new_data_len);
251        segment.ctxt.offset += overlap_data_len as usize;
252        segment.ctxt.length = new_data_len as usize;
253        Some(end_seq)
254    } else {
255        None
256    }
257}