1use std::fmt::Formatter;
2use zerocopy::{AsBytes, FromBytes, FromZeroes, LittleEndian, U32};
3
4pub type BytesMut = bytes::BytesMut;
5pub type Bytes = bytes::Bytes;
6
7bitflags::bitflags! {
8 #[derive(Debug)]
9 struct KcpPacketHeaderFlags: u8 {
10 const SYN = 0b0000_0001;
11 const ACK = 0b0000_0010;
12 const FIN = 0b0000_0100;
13 const DATA = 0b0000_1000;
14 const RST = 0b0001_0000;
15
16 const PING = 0b0010_0000;
17 const PONG = 0b0100_0000;
18
19 const _ = !0;
20 }
21}
22
23#[repr(C, packed)]
24#[derive(AsBytes, FromBytes, FromZeroes, Clone, Default)]
25pub struct KcpPacketHeader {
26 conv: U32<LittleEndian>,
27 src_session_id: U32<LittleEndian>,
28 dst_session_id: U32<LittleEndian>,
29 flag: u8,
30 rsv: u8,
31}
32
33impl KcpPacketHeader {
34 pub fn conv(&self) -> u32 {
35 self.conv.into()
36 }
37
38 pub fn src_session_id(&self) -> u32 {
39 self.src_session_id.into()
40 }
41
42 pub fn dst_session_id(&self) -> u32 {
43 self.dst_session_id.into()
44 }
45
46 pub fn is_syn(&self) -> bool {
47 KcpPacketHeaderFlags::from_bits(self.flag)
48 .unwrap()
49 .contains(KcpPacketHeaderFlags::SYN)
50 }
51
52 pub fn is_ack(&self) -> bool {
53 KcpPacketHeaderFlags::from_bits(self.flag)
54 .unwrap()
55 .contains(KcpPacketHeaderFlags::ACK)
56 }
57
58 pub fn is_fin(&self) -> bool {
59 KcpPacketHeaderFlags::from_bits(self.flag)
60 .unwrap()
61 .contains(KcpPacketHeaderFlags::FIN)
62 }
63
64 pub fn is_data(&self) -> bool {
65 KcpPacketHeaderFlags::from_bits(self.flag)
66 .unwrap()
67 .contains(KcpPacketHeaderFlags::DATA)
68 }
69
70 pub fn is_rst(&self) -> bool {
71 KcpPacketHeaderFlags::from_bits(self.flag)
72 .unwrap()
73 .contains(KcpPacketHeaderFlags::RST)
74 }
75
76 pub fn is_ping(&self) -> bool {
77 KcpPacketHeaderFlags::from_bits(self.flag)
78 .unwrap()
79 .contains(KcpPacketHeaderFlags::PING)
80 }
81
82 pub fn is_pong(&self) -> bool {
83 KcpPacketHeaderFlags::from_bits(self.flag)
84 .unwrap()
85 .contains(KcpPacketHeaderFlags::PONG)
86 }
87
88 pub fn set_conv(&mut self, conv: u32) -> &mut Self {
89 self.conv = conv.into();
90 self
91 }
92
93 pub fn set_src_session_id(&mut self, session_id: u32) -> &mut Self {
94 self.src_session_id = session_id.into();
95 self
96 }
97
98 pub fn set_dst_session_id(&mut self, session_id: u32) -> &mut Self {
99 self.dst_session_id = session_id.into();
100 self
101 }
102
103 pub fn set_syn(&mut self, syn: bool) -> &mut Self {
104 let mut flags = KcpPacketHeaderFlags::from_bits(self.flag).unwrap();
105 if syn {
106 flags.insert(KcpPacketHeaderFlags::SYN);
107 } else {
108 flags.remove(KcpPacketHeaderFlags::SYN);
109 }
110 self.flag = flags.bits();
111 self
112 }
113
114 pub fn set_ack(&mut self, ack: bool) -> &mut Self {
115 let mut flags = KcpPacketHeaderFlags::from_bits(self.flag).unwrap();
116 if ack {
117 flags.insert(KcpPacketHeaderFlags::ACK);
118 } else {
119 flags.remove(KcpPacketHeaderFlags::ACK);
120 }
121 self.flag = flags.bits();
122 self
123 }
124
125 pub fn set_fin(&mut self, fin: bool) -> &mut Self {
126 let mut flags = KcpPacketHeaderFlags::from_bits(self.flag).unwrap();
127 if fin {
128 flags.insert(KcpPacketHeaderFlags::FIN);
129 } else {
130 flags.remove(KcpPacketHeaderFlags::FIN);
131 }
132 self.flag = flags.bits();
133 self
134 }
135
136 pub fn set_data(&mut self, data: bool) -> &mut Self {
137 let mut flags = KcpPacketHeaderFlags::from_bits(self.flag).unwrap();
138 if data {
139 flags.insert(KcpPacketHeaderFlags::DATA);
140 } else {
141 flags.remove(KcpPacketHeaderFlags::DATA);
142 }
143 self.flag = flags.bits();
144 self
145 }
146
147 pub fn set_rst(&mut self, rst: bool) -> &mut Self {
148 let mut flags = KcpPacketHeaderFlags::from_bits(self.flag).unwrap();
149 if rst {
150 flags.insert(KcpPacketHeaderFlags::RST);
151 } else {
152 flags.remove(KcpPacketHeaderFlags::RST);
153 }
154 self.flag = flags.bits();
155 self
156 }
157
158 pub fn set_ping(&mut self, ping: bool) -> &mut Self {
159 let mut flags = KcpPacketHeaderFlags::from_bits(self.flag).unwrap();
160 if ping {
161 flags.insert(KcpPacketHeaderFlags::PING);
162 } else {
163 flags.remove(KcpPacketHeaderFlags::PING);
164 }
165 self.flag = flags.bits();
166 self
167 }
168
169 pub fn set_pong(&mut self, pong: bool) -> &mut Self {
170 let mut flags = KcpPacketHeaderFlags::from_bits(self.flag).unwrap();
171 if pong {
172 flags.insert(KcpPacketHeaderFlags::PONG);
173 } else {
174 flags.remove(KcpPacketHeaderFlags::PONG);
175 }
176 self.flag = flags.bits();
177 self
178 }
179}
180
181impl std::fmt::Debug for KcpPacketHeader {
182 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
183 f.debug_struct("KcpPacketHeader")
184 .field("conv", &self.conv())
185 .field("src_session_id", &self.src_session_id())
186 .field("dst_session_id", &self.dst_session_id())
187 .field("flag", &KcpPacketHeaderFlags::from_bits(self.flag).unwrap())
188 .finish()
189 }
190}
191
192
193#[derive(Clone)]
194pub struct KcpPacket {
195 inner: BytesMut,
196}
197
198impl Default for KcpPacket {
199 fn default() -> Self {
200 Self::new(0)
201 }
202}
203
204impl From<BytesMut> for KcpPacket {
205 fn from(inner: BytesMut) -> Self {
206 Self { inner }
207 }
208}
209
210impl std::fmt::Debug for KcpPacket {
211 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
212 f.debug_struct("KcpPacket")
213 .field("header", &self.header())
214 .field("payload", &self.payload())
215 .finish()
216 }
217}
218
219impl Into<BytesMut> for KcpPacket {
220 fn into(self) -> BytesMut {
221 self.inner
222 }
223}
224
225impl Into<Bytes> for KcpPacket {
226 fn into(self) -> Bytes {
227 self.inner.freeze()
228 }
229}
230
231impl KcpPacket {
232 pub fn new(body_size: usize) -> Self {
233 let mut inner = BytesMut::with_capacity(std::mem::size_of::<KcpPacketHeader>() + body_size);
234 inner.resize(inner.capacity(), 0);
235 Self { inner }
236 }
237
238 pub fn new_with_payload(payload: &[u8]) -> Self {
239 let mut inner =
240 BytesMut::with_capacity(std::mem::size_of::<KcpPacketHeader>() + payload.len());
241 inner.resize(std::mem::size_of::<KcpPacketHeader>(), 0);
242 inner.extend_from_slice(payload);
243 Self { inner }
244 }
245
246 pub fn mut_header(&mut self) -> &mut KcpPacketHeader {
247 KcpPacketHeader::mut_from_prefix(&mut self.inner).unwrap()
248 }
249
250 pub fn header(&self) -> &KcpPacketHeader {
251 KcpPacketHeader::ref_from_prefix(&self.inner).unwrap()
252 }
253
254 pub fn payload(&self) -> &[u8] {
255 &self.inner[std::mem::size_of::<KcpPacketHeader>()..]
256 }
257
258 pub fn inner(self) -> BytesMut {
259 self.inner
260 }
261
262 pub fn len(&self) -> usize {
263 self.inner.len()
264 }
265}