tokio_smoltcp/
socket_allocator.rs

1use parking_lot::Mutex;
2use smoltcp::{
3    iface::{SocketHandle as InnerSocketHandle, SocketSet},
4    socket::{raw, tcp, udp},
5    wire::{IpProtocol, IpVersion},
6};
7use std::{
8    ops::{Deref, DerefMut},
9    sync::Arc,
10};
11
12/// `BufferSize` is used to configure the size of the socket buffer.
13#[derive(Debug, Clone, Copy)]
14pub struct BufferSize {
15    pub tcp_rx_size: usize,
16    pub tcp_tx_size: usize,
17    pub udp_rx_size: usize,
18    pub udp_tx_size: usize,
19    pub udp_rx_meta_size: usize,
20    pub udp_tx_meta_size: usize,
21    pub raw_rx_size: usize,
22    pub raw_tx_size: usize,
23    pub raw_rx_meta_size: usize,
24    pub raw_tx_meta_size: usize,
25}
26
27impl Default for BufferSize {
28    fn default() -> Self {
29        BufferSize {
30            tcp_rx_size: 8192,
31            tcp_tx_size: 8192,
32            udp_rx_size: 8192,
33            udp_tx_size: 8192,
34            udp_rx_meta_size: 32,
35            udp_tx_meta_size: 32,
36            raw_rx_size: 8192,
37            raw_tx_size: 8192,
38            raw_rx_meta_size: 32,
39            raw_tx_meta_size: 32,
40        }
41    }
42}
43
44type SharedSocketSet = Arc<Mutex<SocketSet<'static>>>;
45
46#[derive(Clone)]
47pub struct SocketAlloctor {
48    sockets: SharedSocketSet,
49    buffer_size: BufferSize,
50}
51
52impl SocketAlloctor {
53    pub(crate) fn new(buffer_size: BufferSize) -> SocketAlloctor {
54        let sockets = Arc::new(Mutex::new(SocketSet::new(Vec::new())));
55        SocketAlloctor {
56            sockets,
57            buffer_size,
58        }
59    }
60    pub(crate) fn sockets(&self) -> &SharedSocketSet {
61        &self.sockets
62    }
63    pub fn new_tcp_socket(&self) -> SocketHandle {
64        let mut set = self.sockets.lock();
65        let handle = set.add(self.alloc_tcp_socket());
66        SocketHandle::new(handle, self.sockets.clone())
67    }
68    pub fn new_udp_socket(&self) -> SocketHandle {
69        let mut set = self.sockets.lock();
70        let handle = set.add(self.alloc_udp_socket());
71        SocketHandle::new(handle, self.sockets.clone())
72    }
73    pub fn new_raw_socket(&self, ip_version: IpVersion, ip_protocol: IpProtocol) -> SocketHandle {
74        let mut set = self.sockets.lock();
75        let handle = set.add(self.alloc_raw_socket(ip_version, ip_protocol));
76        SocketHandle::new(handle, self.sockets.clone())
77    }
78    fn alloc_tcp_socket(&self) -> tcp::Socket<'static> {
79        let rx_buffer = tcp::SocketBuffer::new(vec![0; self.buffer_size.tcp_rx_size]);
80        let tx_buffer = tcp::SocketBuffer::new(vec![0; self.buffer_size.tcp_tx_size]);
81        let tcp = tcp::Socket::new(rx_buffer, tx_buffer);
82
83        tcp
84    }
85    fn alloc_udp_socket(&self) -> udp::Socket<'static> {
86        let rx_buffer = udp::PacketBuffer::new(
87            vec![udp::PacketMetadata::EMPTY; self.buffer_size.udp_rx_meta_size],
88            vec![0; self.buffer_size.udp_rx_size],
89        );
90        let tx_buffer = udp::PacketBuffer::new(
91            vec![udp::PacketMetadata::EMPTY; self.buffer_size.udp_tx_meta_size],
92            vec![0; self.buffer_size.udp_tx_size],
93        );
94        let udp = udp::Socket::new(rx_buffer, tx_buffer);
95
96        udp
97    }
98    fn alloc_raw_socket(
99        &self,
100        ip_version: IpVersion,
101        ip_protocol: IpProtocol,
102    ) -> raw::Socket<'static> {
103        let rx_buffer = raw::PacketBuffer::new(
104            vec![raw::PacketMetadata::EMPTY; self.buffer_size.raw_rx_meta_size],
105            vec![0; self.buffer_size.raw_rx_size],
106        );
107        let tx_buffer = raw::PacketBuffer::new(
108            vec![raw::PacketMetadata::EMPTY; self.buffer_size.raw_tx_meta_size],
109            vec![0; self.buffer_size.raw_tx_size],
110        );
111        let raw = raw::Socket::new(ip_version, ip_protocol, rx_buffer, tx_buffer);
112
113        raw
114    }
115}
116
117pub struct SocketHandle(InnerSocketHandle, SharedSocketSet);
118
119impl SocketHandle {
120    fn new(inner: InnerSocketHandle, set: SharedSocketSet) -> SocketHandle {
121        SocketHandle(inner, set)
122    }
123}
124
125impl Drop for SocketHandle {
126    fn drop(&mut self) {
127        let mut iface = self.1.lock();
128        iface.remove(self.0);
129    }
130}
131
132impl Deref for SocketHandle {
133    type Target = InnerSocketHandle;
134
135    fn deref(&self) -> &Self::Target {
136        &self.0
137    }
138}
139
140impl DerefMut for SocketHandle {
141    fn deref_mut(&mut self) -> &mut Self::Target {
142        &mut self.0
143    }
144}
tokio_smoltcp/socket_allocator.rs

tokio_smoltcp/
socket_allocator.rs
