s2n_quic_transport/endpoint/
packet_buffer.rs

1// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use bytes::{Bytes, BytesMut};
5use s2n_codec::{Encoder, EncoderBuffer};
6use s2n_quic_core::path::MINIMUM_MAX_DATAGRAM_SIZE;
7
8/// Allocates a large single buffer, rather than several small buffers
9///
10/// Used for sending packets that contain CONNECTION_CLOSE frames.
11#[derive(Debug)]
12pub struct Buffer {
13    buffer: BytesMut,
14    max_size: usize,
15    count: usize,
16}
17
18// This number shouldn't be _too_ small, otherwise we're performing a bunch
19// of allocations. It also shouldn't be _too_ big so we hold on to those allocations
20// for an extended period of time.
21const DEFAULT_PACKETS: usize = 64;
22
23impl Default for Buffer {
24    fn default() -> Self {
25        Self {
26            buffer: BytesMut::new(),
27            max_size: MINIMUM_MAX_DATAGRAM_SIZE as usize,
28            count: DEFAULT_PACKETS,
29        }
30    }
31}
32
33impl Buffer {
34    pub fn write<F: FnOnce(EncoderBuffer) -> EncoderBuffer>(
35        &mut self,
36        on_write: F,
37    ) -> Option<Bytes> {
38        let max_size = self.max_size;
39
40        if self.buffer.capacity() < max_size {
41            let len = max_size * self.count;
42            let mut buffer = BytesMut::with_capacity(len);
43            // extend the length of the buffer to the capacity so we can
44            // take a slice of it
45            //
46            // We could use `bytes::UninitSlice` but EncoderBuffer uses a
47            // concrete slice instead.
48            unsafe {
49                // Safety: the EncoderBuffer only allows writing (no reading) from
50                //         uninitialized memory
51                buffer.set_len(len);
52            }
53            self.buffer = buffer;
54        }
55
56        let buffer = EncoderBuffer::new(&mut self.buffer[..max_size]);
57
58        let new_buff = on_write(buffer);
59
60        let len = max_size - new_buff.remaining_capacity();
61
62        if len == 0 {
63            return None;
64        }
65
66        debug_assert!(len <= max_size);
67
68        Some(self.buffer.split_to(len).freeze())
69    }
70}
71
72#[cfg(test)]
73mod tests {
74    use super::*;
75
76    #[test]
77    fn empty_test() {
78        let mut buffer = Buffer::default();
79        assert_eq!(buffer.buffer.capacity(), 0);
80
81        assert!(
82            buffer.write(|buffer| buffer).is_none(),
83            "empty writes should return None"
84        );
85
86        assert!(buffer.buffer.capacity() > 0);
87    }
88
89    #[test]
90    fn non_empty_test() {
91        let mut buffer = Buffer::default();
92
93        let packet = buffer
94            .write(|mut buffer| {
95                assert_eq!(
96                    buffer.remaining_capacity(),
97                    MINIMUM_MAX_DATAGRAM_SIZE as usize,
98                    "the provider buffer should be the MINIMUM_MAX_DATAGRAM_SIZE"
99                );
100                buffer.encode(&1337u16);
101                buffer
102            })
103            .expect("non-empty writes should return a packet");
104
105        assert_eq!(packet, 1337u16.to_be_bytes()[..]);
106
107        assert!(buffer.buffer.capacity() > 0);
108        assert!(
109            buffer.buffer.capacity() < MINIMUM_MAX_DATAGRAM_SIZE as usize * DEFAULT_PACKETS,
110            "space should be trimmed off for the returned packet"
111        );
112    }
113}
s2n_quic_transport/endpoint/packet_buffer.rs

s2n_quic_transport/endpoint/
packet_buffer.rs
