axdriver_net/
net_buf.rs

1use alloc::{boxed::Box, sync::Arc, vec, vec::Vec};
2use core::ptr::NonNull;
3
4use spin::Mutex;
5
6use crate::{DevError, DevResult};
7
8/// A raw buffer struct for network device.
9pub struct NetBufPtr {
10    // The raw pointer of the original object.
11    raw_ptr: NonNull<u8>,
12    // The pointer to the net buffer.
13    buf_ptr: NonNull<u8>,
14    len: usize,
15}
16
17impl NetBufPtr {
18    /// Create a new [`NetBufPtr`].
19    pub fn new(raw_ptr: NonNull<u8>, buf_ptr: NonNull<u8>, len: usize) -> Self {
20        Self {
21            raw_ptr,
22            buf_ptr,
23            len,
24        }
25    }
26
27    /// Return raw pointer of the original object.
28    pub fn raw_ptr<T>(&self) -> *mut T {
29        self.raw_ptr.as_ptr() as *mut T
30    }
31
32    /// Return [`NetBufPtr`] buffer len.
33    pub fn packet_len(&self) -> usize {
34        self.len
35    }
36
37    /// Return [`NetBufPtr`] buffer as &[u8].
38    pub fn packet(&self) -> &[u8] {
39        unsafe { core::slice::from_raw_parts(self.buf_ptr.as_ptr() as *const u8, self.len) }
40    }
41
42    /// Return [`NetBufPtr`] buffer as &mut [u8].
43    pub fn packet_mut(&mut self) -> &mut [u8] {
44        unsafe { core::slice::from_raw_parts_mut(self.buf_ptr.as_ptr(), self.len) }
45    }
46}
47
48const MIN_BUFFER_LEN: usize = 1526;
49const MAX_BUFFER_LEN: usize = 65535;
50
51/// A RAII network buffer wrapped in a [`Box`].
52pub type NetBufBox = Box<NetBuf>;
53
54/// A RAII network buffer.
55///
56/// It should be allocated from the [`NetBufPool`], and it will be
57/// deallocated into the pool automatically when dropped.
58///
59/// The layout of the buffer is:
60///
61/// ```text
62///   ______________________ capacity ______________________
63///  /                                                      \
64/// +------------------+------------------+------------------+
65/// |      Header      |      Packet      |      Unused      |
66/// +------------------+------------------+------------------+
67/// |\__ header_len __/ \__ packet_len __/
68/// |
69/// buf_ptr
70/// ```
71pub struct NetBuf {
72    header_len: usize,
73    packet_len: usize,
74    capacity: usize,
75    buf_ptr: NonNull<u8>,
76    pool_offset: usize,
77    pool: Arc<NetBufPool>,
78}
79
80unsafe impl Send for NetBuf {}
81unsafe impl Sync for NetBuf {}
82
83impl NetBuf {
84    const unsafe fn get_slice(&self, start: usize, len: usize) -> &[u8] {
85        unsafe { core::slice::from_raw_parts(self.buf_ptr.as_ptr().add(start), len) }
86    }
87
88    const unsafe fn get_slice_mut(&mut self, start: usize, len: usize) -> &mut [u8] {
89        unsafe { core::slice::from_raw_parts_mut(self.buf_ptr.as_ptr().add(start), len) }
90    }
91
92    /// Returns the capacity of the buffer.
93    pub const fn capacity(&self) -> usize {
94        self.capacity
95    }
96
97    /// Returns the length of the header part.
98    pub const fn header_len(&self) -> usize {
99        self.header_len
100    }
101
102    /// Returns the header part of the buffer.
103    pub const fn header(&self) -> &[u8] {
104        unsafe { self.get_slice(0, self.header_len) }
105    }
106
107    /// Returns the packet part of the buffer.
108    pub const fn packet(&self) -> &[u8] {
109        unsafe { self.get_slice(self.header_len, self.packet_len) }
110    }
111
112    /// Returns the mutable reference to the packet part.
113    pub const fn packet_mut(&mut self) -> &mut [u8] {
114        unsafe { self.get_slice_mut(self.header_len, self.packet_len) }
115    }
116
117    /// Returns both the header and the packet parts, as a contiguous slice.
118    pub const fn packet_with_header(&self) -> &[u8] {
119        unsafe { self.get_slice(0, self.header_len + self.packet_len) }
120    }
121
122    /// Returns the entire buffer.
123    pub const fn raw_buf(&self) -> &[u8] {
124        unsafe { self.get_slice(0, self.capacity) }
125    }
126
127    /// Returns the mutable reference to the entire buffer.
128    pub const fn raw_buf_mut(&mut self) -> &mut [u8] {
129        unsafe { self.get_slice_mut(0, self.capacity) }
130    }
131
132    /// Set the length of the header part.
133    pub fn set_header_len(&mut self, header_len: usize) {
134        debug_assert!(header_len + self.packet_len <= self.capacity);
135        self.header_len = header_len;
136    }
137
138    /// Set the length of the packet part.
139    pub fn set_packet_len(&mut self, packet_len: usize) {
140        debug_assert!(self.header_len + packet_len <= self.capacity);
141        self.packet_len = packet_len;
142    }
143
144    /// Converts the buffer into a [`NetBufPtr`].
145    pub fn into_buf_ptr(mut self: Box<Self>) -> NetBufPtr {
146        let buf_ptr = self.packet_mut().as_mut_ptr();
147        let len = self.packet_len;
148        NetBufPtr::new(
149            NonNull::new(Box::into_raw(self) as *mut u8).unwrap(),
150            NonNull::new(buf_ptr).unwrap(),
151            len,
152        )
153    }
154
155    /// Restore [`NetBuf`] struct from a raw pointer.
156    ///
157    /// # Safety
158    ///
159    /// This function is unsafe because it may cause some memory issues,
160    /// so we must ensure that it is called after calling `into_buf_ptr`.
161    pub unsafe fn from_buf_ptr(ptr: NetBufPtr) -> Box<Self> {
162        unsafe { Box::from_raw(ptr.raw_ptr::<Self>()) }
163    }
164}
165
166impl Drop for NetBuf {
167    /// Deallocates the buffer into the [`NetBufPool`].
168    fn drop(&mut self) {
169        self.pool.dealloc(self.pool_offset);
170    }
171}
172
173/// A pool of [`NetBuf`]s to speed up buffer allocation.
174///
175/// It divides a large memory into several equal parts for each buffer.
176pub struct NetBufPool {
177    capacity: usize,
178    buf_len: usize,
179    pool: Vec<u8>,
180    free_list: Mutex<Vec<usize>>,
181}
182
183impl NetBufPool {
184    /// Creates a new pool with the given `capacity`, and all buffer lengths are
185    /// set to `buf_len`.
186    pub fn new(capacity: usize, buf_len: usize) -> DevResult<Arc<Self>> {
187        if capacity == 0 {
188            return Err(DevError::InvalidParam);
189        }
190        if !(MIN_BUFFER_LEN..=MAX_BUFFER_LEN).contains(&buf_len) {
191            return Err(DevError::InvalidParam);
192        }
193
194        let pool = vec![0; capacity * buf_len];
195        let mut free_list = Vec::with_capacity(capacity);
196        for i in 0..capacity {
197            free_list.push(i * buf_len);
198        }
199        Ok(Arc::new(Self {
200            capacity,
201            buf_len,
202            pool,
203            free_list: Mutex::new(free_list),
204        }))
205    }
206
207    /// Returns the capacity of the pool.
208    pub const fn capacity(&self) -> usize {
209        self.capacity
210    }
211
212    /// Returns the length of each buffer.
213    pub const fn buffer_len(&self) -> usize {
214        self.buf_len
215    }
216
217    /// Allocates a buffer from the pool.
218    ///
219    /// Returns `None` if no buffer is available.
220    pub fn alloc(self: &Arc<Self>) -> Option<NetBuf> {
221        let pool_offset = self.free_list.lock().pop()?;
222        let buf_ptr =
223            unsafe { NonNull::new(self.pool.as_ptr().add(pool_offset) as *mut u8).unwrap() };
224        Some(NetBuf {
225            header_len: 0,
226            packet_len: 0,
227            capacity: self.buf_len,
228            buf_ptr,
229            pool_offset,
230            pool: Arc::clone(self),
231        })
232    }
233
234    /// Allocates a buffer wrapped in a [`Box`] from the pool.
235    ///
236    /// Returns `None` if no buffer is available.
237    pub fn alloc_boxed(self: &Arc<Self>) -> Option<NetBufBox> {
238        Some(Box::new(self.alloc()?))
239    }
240
241    /// Deallocates a buffer at the given offset.
242    ///
243    /// `pool_offset` must be a multiple of `buf_len`.
244    fn dealloc(&self, pool_offset: usize) {
245        debug_assert_eq!(pool_offset % self.buf_len, 0);
246        self.free_list.lock().push(pool_offset);
247    }
248}
axdriver_net/net_buf.rs

axdriver_net/
net_buf.rs
