nex_packet/
packet.rs

1use bytes::{Bytes, BytesMut};
2use std::marker::PhantomData;
3
4/// Represents a generic network packet.
5pub trait Packet: Sized {
6    type Header;
7
8    /// Parse from a byte slice.
9    fn from_buf(buf: &[u8]) -> Option<Self>;
10
11    /// Parse from raw bytes. (with ownership)
12    fn from_bytes(bytes: Bytes) -> Option<Self>;
13
14    /// Serialize into raw bytes.
15    fn to_bytes(&self) -> Bytes;
16
17    /// Get the header of the packet.
18    fn header(&self) -> Bytes;
19
20    /// Get the payload of the packet.
21    fn payload(&self) -> Bytes;
22
23    /// Get the length of the header.
24    fn header_len(&self) -> usize;
25
26    /// Get the length of the payload.
27    fn payload_len(&self) -> usize;
28    /// Get the total length of the packet (header + payload).
29    fn total_len(&self) -> usize;
30    /// Convert the packet to a mutable byte buffer.
31    fn to_bytes_mut(&self) -> BytesMut {
32        let mut buf = BytesMut::with_capacity(self.total_len());
33        buf.extend_from_slice(&self.to_bytes());
34        buf
35    }
36    /// Get a mutable byte buffer for the header.
37    fn header_mut(&self) -> BytesMut {
38        let mut buf = BytesMut::with_capacity(self.header_len());
39        buf.extend_from_slice(&self.header());
40        buf
41    }
42    /// Get a mutable byte buffer for the payload.
43    fn payload_mut(&self) -> BytesMut {
44        let mut buf = BytesMut::with_capacity(self.payload_len());
45        buf.extend_from_slice(&self.payload());
46        buf
47    }
48
49    fn into_parts(self) -> (Self::Header, Bytes);
50}
51
52/// Represents a mutable network packet that can be parsed and modified in place.
53///
54/// Types implementing this trait work on top of the same backing buffer and allow
55/// layered packet parsing to be chained without additional allocations.
56pub trait MutablePacket<'a>: Sized {
57    /// The immutable packet type associated with this mutable view.
58    type Packet: Packet;
59
60    /// Construct a mutable packet from the provided buffer.
61    fn new(buffer: &'a mut [u8]) -> Option<Self>;
62
63    /// Get a shared view over the entire packet buffer.
64    fn packet(&self) -> &[u8];
65
66    /// Get a mutable view over the entire packet buffer.
67    fn packet_mut(&mut self) -> &mut [u8];
68
69    /// Get the serialized header bytes of the packet.
70    fn header(&self) -> &[u8];
71
72    /// Get a mutable view over the serialized header bytes of the packet.
73    fn header_mut(&mut self) -> &mut [u8];
74
75    /// Get the payload bytes of the packet.
76    fn payload(&self) -> &[u8];
77
78    /// Get a mutable view over the payload bytes of the packet.
79    fn payload_mut(&mut self) -> &mut [u8];
80
81    /// Convert the mutable packet into its immutable counterpart.
82    fn freeze(&self) -> Option<Self::Packet> {
83        Self::Packet::from_buf(self.packet())
84    }
85}
86
87/// A generic mutable packet wrapper that validates using the immutable packet
88/// parser and exposes the raw buffer for in-place mutation.
89pub struct GenericMutablePacket<'a, P: Packet> {
90    buffer: &'a mut [u8],
91    _marker: PhantomData<P>,
92}
93
94impl<'a, P: Packet> MutablePacket<'a> for GenericMutablePacket<'a, P> {
95    type Packet = P;
96
97    fn new(buffer: &'a mut [u8]) -> Option<Self> {
98        P::from_buf(buffer)?;
99        Some(Self {
100            buffer,
101            _marker: PhantomData,
102        })
103    }
104
105    fn packet(&self) -> &[u8] {
106        &*self.buffer
107    }
108
109    fn packet_mut(&mut self) -> &mut [u8] {
110        &mut *self.buffer
111    }
112
113    fn header(&self) -> &[u8] {
114        let (header_len, _) = self.lengths();
115        &self.packet()[..header_len]
116    }
117
118    fn header_mut(&mut self) -> &mut [u8] {
119        let (header_len, _) = self.lengths();
120        let (header, _) = (&mut *self.buffer).split_at_mut(header_len);
121        header
122    }
123
124    fn payload(&self) -> &[u8] {
125        let (header_len, payload_len) = self.lengths();
126        &self.packet()[header_len..header_len + payload_len]
127    }
128
129    fn payload_mut(&mut self) -> &mut [u8] {
130        let (header_len, payload_len) = self.lengths();
131        let (_, payload) = (&mut *self.buffer).split_at_mut(header_len);
132        &mut payload[..payload_len]
133    }
134}
135
136impl<'a, P: Packet> GenericMutablePacket<'a, P> {
137    /// Construct a mutable packet without running additional validation.
138    pub fn new_unchecked(buffer: &'a mut [u8]) -> Self {
139        Self {
140            buffer,
141            _marker: PhantomData,
142        }
143    }
144
145    fn lengths(&self) -> (usize, usize) {
146        match P::from_buf(self.packet()) {
147            Some(packet) => {
148                let header_len = packet.header_len();
149                let payload_len = packet.payload_len();
150                (header_len, payload_len)
151            }
152            _ => (self.buffer.len(), 0),
153        }
154    }
155}