1use packet::{
14 BufferView, FragmentedBytesMut, NestablePacketBuilder, PacketBuilder, PacketConstraints,
15 ParsablePacket, ParseMetadata, SerializationContext, SerializeTarget,
16};
17use std::num::NonZeroU16;
18use zerocopy::byteorder::little_endian::U32;
19use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout, Ref, SplitByteSlice, Unaligned};
20
21pub use witness::ErrorValue;
23
24pub const SERVER_PORT: NonZeroU16 = NonZeroU16::new(33330).unwrap();
26pub const ADVERT_PORT: NonZeroU16 = NonZeroU16::new(33331).unwrap();
28
29const MAGIC: u32 = 0xAA774217;
30
31mod witness {
32 #[derive(Debug, Copy, Clone, Eq, PartialEq)]
37 pub struct ErrorValue(u32);
38
39 impl ErrorValue {
40 const ERROR_MASK: u32 = 0x80000000;
41
42 pub const fn new(v: u32) -> Option<Self> {
45 if v & Self::ERROR_MASK != 0 { Some(Self(v)) } else { None }
46 }
47 }
48
49 impl From<ErrorValue> for u32 {
50 fn from(v: ErrorValue) -> Self {
51 let ErrorValue(v) = v;
52 v
53 }
54 }
55}
56
57#[derive(Debug, Copy, Clone, Eq, PartialEq)]
61pub enum Opcode {
62 Command,
63 SendFile,
64 Data,
65 Boot,
66 Query,
67 ShellCmd,
68 Open,
69 Read,
70 Write,
71 Close,
72 LastData,
73 Reboot,
74 GetAdvert,
75 Ack,
76 FileReceived,
77 Advertise,
78}
79
80impl From<Opcode> for u32 {
81 fn from(value: Opcode) -> u32 {
82 match value {
83 Opcode::Command => 1,
84 Opcode::SendFile => 2,
85 Opcode::Data => 3,
86 Opcode::Boot => 4,
87 Opcode::Query => 5,
88 Opcode::ShellCmd => 6,
89 Opcode::Open => 7,
90 Opcode::Read => 8,
91 Opcode::Write => 9,
92 Opcode::Close => 10,
93 Opcode::LastData => 11,
94 Opcode::Reboot => 12,
95 Opcode::GetAdvert => 13,
96 Opcode::Ack => 0,
97 Opcode::FileReceived => 0x70000001,
98 Opcode::Advertise => 0x77777777,
99 }
100 }
101}
102
103#[derive(Debug, Copy, Clone, Eq, PartialEq)]
105pub enum ErrorCode {
106 BadCommand,
107 BadParam,
108 TooLarge,
109 BadFile,
110 Unknown(ErrorValue),
111}
112
113impl From<ErrorCode> for u32 {
114 fn from(value: ErrorCode) -> u32 {
115 match value {
116 ErrorCode::BadCommand => 0x80000001,
117 ErrorCode::BadParam => 0x80000002,
118 ErrorCode::TooLarge => 0x80000003,
119 ErrorCode::BadFile => 0x80000004,
120 ErrorCode::Unknown(v) => v.into(),
121 }
122 }
123}
124
125#[derive(Debug, Copy, Clone, Eq, PartialEq)]
129pub enum OpcodeOrErr {
130 Op(Opcode),
131 Err(ErrorCode),
132}
133
134impl From<Opcode> for OpcodeOrErr {
135 fn from(op: Opcode) -> Self {
136 OpcodeOrErr::Op(op)
137 }
138}
139
140impl From<ErrorCode> for OpcodeOrErr {
141 fn from(err: ErrorCode) -> Self {
142 OpcodeOrErr::Err(err)
143 }
144}
145
146impl TryFrom<u32> for OpcodeOrErr {
147 type Error = u32;
148
149 fn try_from(value: u32) -> Result<Self, Self::Error> {
150 match value {
151 1 => Ok(Opcode::Command.into()),
152 2 => Ok(Opcode::SendFile.into()),
153 3 => Ok(Opcode::Data.into()),
154 4 => Ok(Opcode::Boot.into()),
155 5 => Ok(Opcode::Query.into()),
156 6 => Ok(Opcode::ShellCmd.into()),
157 7 => Ok(Opcode::Open.into()),
158 8 => Ok(Opcode::Read.into()),
159 9 => Ok(Opcode::Write.into()),
160 10 => Ok(Opcode::Close.into()),
161 11 => Ok(Opcode::LastData.into()),
162 12 => Ok(Opcode::Reboot.into()),
163 13 => Ok(Opcode::GetAdvert.into()),
164 0 => Ok(Opcode::Ack.into()),
165 0x70000001 => Ok(Opcode::FileReceived.into()),
166 0x77777777 => Ok(Opcode::Advertise.into()),
167 0x80000001 => Ok(ErrorCode::BadCommand.into()),
168 0x80000002 => Ok(ErrorCode::BadParam.into()),
169 0x80000003 => Ok(ErrorCode::TooLarge.into()),
170 0x80000004 => Ok(ErrorCode::BadFile.into()),
171 v => match ErrorValue::new(v) {
172 Some(e) => Ok(ErrorCode::Unknown(e).into()),
173 None => Err(v),
174 },
175 }
176 }
177}
178
179impl From<OpcodeOrErr> for u32 {
180 fn from(cmd: OpcodeOrErr) -> Self {
181 match cmd {
182 OpcodeOrErr::Op(op) => op.into(),
183 OpcodeOrErr::Err(e) => e.into(),
184 }
185 }
186}
187
188#[derive(Debug)]
190pub enum ParseError {
191 Malformed,
192 UnknownOpcode(u32),
193 BadMagic,
194}
195
196#[repr(C)]
197#[derive(KnownLayout, FromBytes, IntoBytes, Immutable, Unaligned, Debug)]
198struct MessageHead {
199 magic: U32,
200 cookie: U32,
201 cmd: U32,
202 arg: U32,
203}
204
205#[derive(Debug)]
207pub struct NetbootPacket<B: SplitByteSlice> {
208 command: OpcodeOrErr,
209 message: Ref<B, MessageHead>,
210 payload: B,
211}
212
213impl<B: SplitByteSlice> NetbootPacket<B> {
214 pub fn command(&self) -> OpcodeOrErr {
215 self.command
216 }
217
218 pub fn cookie(&self) -> u32 {
219 self.message.cookie.get()
220 }
221
222 pub fn arg(&self) -> u32 {
223 self.message.arg.get()
224 }
225
226 pub fn payload(&self) -> &[u8] {
227 self.payload.as_ref()
228 }
229}
230
231impl<B: SplitByteSlice> ParsablePacket<B, ()> for NetbootPacket<B> {
232 type Error = ParseError;
233
234 fn parse<BV: BufferView<B>>(mut buffer: BV, _args: ()) -> Result<Self, Self::Error> {
235 let message = buffer.take_obj_front::<MessageHead>().ok_or(ParseError::Malformed)?;
236 if message.magic.get() != MAGIC {
237 return Err(ParseError::BadMagic);
238 }
239 let opcode = message.cmd.get().try_into().map_err(ParseError::UnknownOpcode)?;
240 let payload = buffer.into_rest();
241 Ok(Self { command: opcode, message, payload })
242 }
243
244 fn parse_metadata(&self) -> ParseMetadata {
245 unimplemented!()
248 }
249}
250
251#[derive(Debug)]
253pub struct NetbootPacketBuilder {
254 cmd: OpcodeOrErr,
255 cookie: u32,
256 arg: u32,
257}
258
259impl NetbootPacketBuilder {
260 pub fn new(cmd: OpcodeOrErr, cookie: u32, arg: u32) -> Self {
261 Self { cmd, cookie, arg }
262 }
263}
264
265impl NestablePacketBuilder for NetbootPacketBuilder {
266 fn constraints(&self) -> PacketConstraints {
267 PacketConstraints::new(std::mem::size_of::<MessageHead>(), 0, 0, std::usize::MAX)
268 }
269}
270
271impl<C: SerializationContext> PacketBuilder<C> for NetbootPacketBuilder {
272 fn serialize(
273 &self,
274 _context: &mut C,
275 target: &mut SerializeTarget<'_>,
276 _body: FragmentedBytesMut<'_, '_>,
277 ) {
278 let mut bv = crate::as_buffer_view_mut(&mut target.header);
279 let mut message = bv.take_obj_front::<MessageHead>().expect("not enough space in buffer");
280 let MessageHead { magic, cookie, cmd, arg } = &mut *message;
281 magic.set(MAGIC);
282 cookie.set(self.cookie);
283 arg.set(self.arg);
284 cmd.set(self.cmd.into());
285 }
286}
287
288#[cfg(test)]
289mod tests {
290
291 use super::*;
292
293 use assert_matches::assert_matches;
294 use packet::{
295 InnerPacketBuilder as _, NoOpSerializationContext, ParseBuffer as _, Serializer as _,
296 };
297
298 #[test]
299 fn test_parse_serialize() {
300 const PAYLOAD: [u8; 10] = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
301 let mut pkt = NetbootPacketBuilder::new(Opcode::Ack.into(), 3, 4)
302 .wrap_body((&PAYLOAD[..]).into_serializer())
303 .serialize_vec_outer(&mut NoOpSerializationContext)
304 .expect("failed to serialize");
305 let parsed = pkt.parse::<NetbootPacket<_>>().expect("failed to parse");
306 assert_eq!(parsed.command(), OpcodeOrErr::Op(Opcode::Ack));
307 assert_eq!(parsed.cookie(), 3);
308 assert_eq!(parsed.arg(), 4);
309 assert_eq!(parsed.payload(), &PAYLOAD[..]);
310 }
311
312 #[test]
313 fn test_parse_serialize_opcodes() {
314 const TEST_OPCODES: [Opcode; 16] = [
315 Opcode::Command,
316 Opcode::SendFile,
317 Opcode::Data,
318 Opcode::Boot,
319 Opcode::Query,
320 Opcode::ShellCmd,
321 Opcode::Open,
322 Opcode::Read,
323 Opcode::Write,
324 Opcode::Close,
325 Opcode::LastData,
326 Opcode::Reboot,
327 Opcode::GetAdvert,
328 Opcode::Ack,
329 Opcode::FileReceived,
330 Opcode::Advertise,
331 ];
332
333 for opcode in TEST_OPCODES.iter() {
334 match opcode {
335 Opcode::Command
336 | Opcode::SendFile
337 | Opcode::Data
338 | Opcode::Boot
339 | Opcode::Query
340 | Opcode::ShellCmd
341 | Opcode::Open
342 | Opcode::Read
343 | Opcode::Write
344 | Opcode::Close
345 | Opcode::LastData
346 | Opcode::Reboot
347 | Opcode::GetAdvert
348 | Opcode::Ack
349 | Opcode::FileReceived
350 | Opcode::Advertise => {
351 }
354 }
355 let opcode_or_err = OpcodeOrErr::try_from(u32::from(*opcode)).expect("failed to parse");
356 assert_matches!(opcode_or_err, OpcodeOrErr::Op(op) if op == *opcode);
357 }
358 }
359
360 #[test]
361 fn test_parse_serialize_error_codes() {
362 let test_error_codes = [
363 ErrorCode::BadCommand,
364 ErrorCode::BadParam,
365 ErrorCode::TooLarge,
366 ErrorCode::BadFile,
367 ErrorCode::Unknown(ErrorValue::new(0x80001234).unwrap()),
368 ];
369 for error in test_error_codes.iter() {
370 match error {
371 ErrorCode::BadCommand
372 | ErrorCode::BadParam
373 | ErrorCode::TooLarge
374 | ErrorCode::BadFile
375 | ErrorCode::Unknown(_) => {
376 }
379 }
380 let opcode_or_err = OpcodeOrErr::try_from(u32::from(*error)).expect("failed to parse");
381 assert_matches!(opcode_or_err, OpcodeOrErr::Err(e) if e == *error);
382 }
383 }
384}