s2n_quic_core/packet/
interceptor.rs

1// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::{
5    event::api::{SocketAddress, Subject},
6    havoc,
7    packet::number::{PacketNumber, PacketNumberSpace},
8    path::{LocalAddress, RemoteAddress},
9    time::Timestamp,
10    varint::VarInt,
11};
12use core::ops::RangeInclusive;
13use s2n_codec::encoder::scatter;
14
15pub use s2n_codec::{DecoderBufferMut, EncoderBuffer};
16pub mod loss;
17pub use loss::Loss;
18
19/// TODO add `non_exhaustive` once/if this feature is stable
20#[derive(Debug)]
21pub struct Packet {
22    pub number: PacketNumber,
23    pub timestamp: Timestamp,
24}
25
26/// TODO add `non_exhaustive` once/if this feature is stable
27#[derive(Debug)]
28pub struct Datagram<'a> {
29    pub remote_address: SocketAddress<'a>,
30    pub local_address: SocketAddress<'a>,
31    pub timestamp: Timestamp,
32}
33
34pub trait Ack {
35    fn space(&self) -> PacketNumberSpace;
36
37    fn insert_range(&mut self, range: RangeInclusive<VarInt>);
38}
39
40/// Trait which enables an application to intercept packets that are transmitted and received
41pub trait Interceptor: 'static + Send {
42    #[inline(always)]
43    fn intercept_rx_ack<A: Ack>(&mut self, subject: &Subject, ack: &mut A) {
44        let _ = subject;
45        let _ = ack;
46    }
47
48    #[inline(always)]
49    fn intercept_rx_local_address(&mut self, subject: &Subject, addr: &mut LocalAddress) {
50        let _ = subject;
51        let _ = addr;
52    }
53
54    #[inline(always)]
55    fn intercept_rx_remote_address(&mut self, subject: &Subject, addr: &mut RemoteAddress) {
56        let _ = subject;
57        let _ = addr;
58    }
59
60    #[inline(always)]
61    fn intercept_rx_datagram<'a>(
62        &mut self,
63        subject: &Subject,
64        datagram: &Datagram,
65        payload: DecoderBufferMut<'a>,
66    ) -> DecoderBufferMut<'a> {
67        let _ = subject;
68        let _ = datagram;
69        payload
70    }
71
72    #[inline(always)]
73    fn intercept_rx_payload<'a>(
74        &mut self,
75        subject: &Subject,
76        packet: &Packet,
77        payload: DecoderBufferMut<'a>,
78    ) -> DecoderBufferMut<'a> {
79        let _ = subject;
80        let _ = packet;
81        payload
82    }
83
84    #[inline(always)]
85    fn intercept_tx_datagram(
86        &mut self,
87        subject: &Subject,
88        datagram: &Datagram,
89        payload: &mut EncoderBuffer,
90    ) {
91        let _ = subject;
92        let _ = datagram;
93        let _ = payload;
94    }
95
96    #[inline(always)]
97    fn intercept_tx_payload(
98        &mut self,
99        subject: &Subject,
100        packet: &Packet,
101        payload: &mut scatter::Buffer,
102    ) {
103        let _ = subject;
104        let _ = packet;
105        let _ = payload;
106    }
107}
108
109#[derive(Debug, Default)]
110pub struct Disabled(());
111
112impl Interceptor for Disabled {}
113
114impl<X, Y> Interceptor for (X, Y)
115where
116    X: Interceptor,
117    Y: Interceptor,
118{
119    #[inline(always)]
120    fn intercept_rx_ack<A: Ack>(&mut self, subject: &Subject, ack: &mut A) {
121        self.0.intercept_rx_ack(subject, ack);
122        self.1.intercept_rx_ack(subject, ack);
123    }
124
125    #[inline(always)]
126    fn intercept_rx_local_address(&mut self, subject: &Subject, addr: &mut LocalAddress) {
127        self.0.intercept_rx_local_address(subject, addr);
128        self.1.intercept_rx_local_address(subject, addr);
129    }
130
131    #[inline(always)]
132    fn intercept_rx_remote_address(&mut self, subject: &Subject, addr: &mut RemoteAddress) {
133        self.0.intercept_rx_remote_address(subject, addr);
134        self.1.intercept_rx_remote_address(subject, addr);
135    }
136
137    #[inline(always)]
138    fn intercept_rx_datagram<'a>(
139        &mut self,
140        subject: &Subject,
141        datagram: &Datagram,
142        payload: DecoderBufferMut<'a>,
143    ) -> DecoderBufferMut<'a> {
144        let payload = self.0.intercept_rx_datagram(subject, datagram, payload);
145        self.1.intercept_rx_datagram(subject, datagram, payload)
146    }
147
148    #[inline(always)]
149    fn intercept_rx_payload<'a>(
150        &mut self,
151        subject: &Subject,
152        packet: &Packet,
153        payload: DecoderBufferMut<'a>,
154    ) -> DecoderBufferMut<'a> {
155        let payload = self.0.intercept_rx_payload(subject, packet, payload);
156        self.1.intercept_rx_payload(subject, packet, payload)
157    }
158
159    #[inline(always)]
160    fn intercept_tx_datagram(
161        &mut self,
162        subject: &Subject,
163        datagram: &Datagram,
164        payload: &mut EncoderBuffer,
165    ) {
166        self.0.intercept_tx_datagram(subject, datagram, payload);
167        self.1.intercept_tx_datagram(subject, datagram, payload);
168    }
169
170    #[inline(always)]
171    fn intercept_tx_payload(
172        &mut self,
173        subject: &Subject,
174        packet: &Packet,
175        payload: &mut scatter::Buffer,
176    ) {
177        self.0.intercept_tx_payload(subject, packet, payload);
178        self.1.intercept_tx_payload(subject, packet, payload);
179    }
180}
181
182#[derive(Debug, Default)]
183pub struct Havoc<Rx, Tx, P, R>
184where
185    Rx: 'static + Send + havoc::Strategy,
186    Tx: 'static + Send + havoc::Strategy,
187    P: 'static + Send + havoc::Strategy,
188    R: 'static + Send + havoc::Random,
189{
190    pub rx: Rx,
191    pub tx: Tx,
192    pub port: P,
193    pub random: R,
194}
195
196impl<Rx, Tx, P, R> Interceptor for Havoc<Rx, Tx, P, R>
197where
198    Rx: 'static + Send + havoc::Strategy,
199    Tx: 'static + Send + havoc::Strategy,
200    P: 'static + Send + havoc::Strategy,
201    R: 'static + Send + havoc::Random,
202{
203    #[inline]
204    fn intercept_rx_remote_address(&mut self, _subject: &Subject, addr: &mut RemoteAddress) {
205        let port = addr.port();
206        let port = self.port.havoc_u16(&mut self.random, port);
207        addr.set_port(port);
208    }
209
210    #[inline]
211    fn intercept_rx_payload<'a>(
212        &mut self,
213        _subject: &Subject,
214        _packet: &Packet,
215        payload: DecoderBufferMut<'a>,
216    ) -> DecoderBufferMut<'a> {
217        let payload = payload.into_less_safe_slice();
218        let len = payload.len();
219
220        let len = {
221            use s2n_codec::Encoder;
222            let mut payload = EncoderBuffer::new(payload);
223            payload.set_position(len);
224            self.rx.havoc(&mut self.random, &mut payload);
225            payload.len()
226        };
227
228        let payload = &mut payload[..len];
229
230        DecoderBufferMut::new(payload)
231    }
232
233    #[inline]
234    fn intercept_tx_payload(
235        &mut self,
236        _subject: &Subject,
237        _packet: &Packet,
238        payload: &mut scatter::Buffer,
239    ) {
240        let payload = payload.flatten();
241        self.tx.havoc(&mut self.random, payload);
242    }
243}
244
245impl<T: Interceptor> Interceptor for Option<T> {
246    #[inline]
247    fn intercept_rx_local_address(&mut self, subject: &Subject, addr: &mut LocalAddress) {
248        if let Some(inner) = self.as_mut() {
249            inner.intercept_rx_local_address(subject, addr)
250        }
251    }
252
253    #[inline]
254    fn intercept_rx_remote_address(&mut self, subject: &Subject, addr: &mut RemoteAddress) {
255        if let Some(inner) = self.as_mut() {
256            inner.intercept_rx_remote_address(subject, addr)
257        }
258    }
259
260    #[inline]
261    fn intercept_rx_payload<'a>(
262        &mut self,
263        subject: &Subject,
264        packet: &Packet,
265        payload: DecoderBufferMut<'a>,
266    ) -> DecoderBufferMut<'a> {
267        if let Some(inner) = self.as_mut() {
268            inner.intercept_rx_payload(subject, packet, payload)
269        } else {
270            payload
271        }
272    }
273
274    #[inline]
275    fn intercept_tx_payload(
276        &mut self,
277        subject: &Subject,
278        packet: &Packet,
279        payload: &mut scatter::Buffer,
280    ) {
281        if let Some(inner) = self.as_mut() {
282            inner.intercept_tx_payload(subject, packet, payload)
283        }
284    }
285}