Skip to main content

libtw2_net/
protocol7.rs

1use arrayvec::ArrayVec;
2use libtw2_buffer as buffer;
3use libtw2_buffer::with_buffer;
4use libtw2_buffer::Buffer;
5use libtw2_buffer::BufferRef;
6use libtw2_common::boilerplate_packed;
7use libtw2_common::bytes::FromBytesExt as _;
8use libtw2_common::num::Cast;
9use libtw2_common::pretty;
10use libtw2_common::unwrap_or_return;
11use libtw2_huffman::instances::TEEWORLDS as HUFFMAN;
12use libtw2_warn::Ignore;
13use libtw2_warn::Warn;
14use std::cmp;
15use std::fmt;
16use std::mem;
17use zerocopy::AsBytes;
18use zerocopy_derive::AsBytes;
19use zerocopy_derive::FromBytes;
20use zerocopy_derive::FromZeroes;
21
22pub const CHUNK_HEADER_SIZE: usize = 2;
23pub const CHUNK_HEADER_SIZE_VITAL: usize = 3;
24pub const HEADER_SIZE: usize = 7;
25pub const HEADER_SIZE_CONNLESS: usize = 9;
26pub const MAX_PACKETSIZE: usize = 1400;
27
28// For connectionless packets, this is obvious (MAX_PACKETSIZE -
29// HEADER_SIZE_CONNLESS). For packets sent in a connection context, you also get
30// a chunk header which, which makes the maximum payload size (MAX_PACKETSIZE -
31// HEADER_SIZE - CHUNK_HEADER_SIZE_VITAL).
32pub const MAX_PAYLOAD: usize = 1390;
33
34pub const PACKETFLAG_CONTROL: u8 = 1 << 0;
35pub const PACKETFLAG_REQUEST_RESEND: u8 = 1 << 1;
36pub const PACKETFLAG_COMPRESSION: u8 = 1 << 2;
37pub const PACKETFLAG_CONNLESS: u8 = 1 << 3;
38
39pub const CHUNKFLAG_VITAL: u8 = 1 << 0;
40pub const CHUNKFLAG_RESEND: u8 = 1 << 1;
41
42pub const CTRLMSG_KEEPALIVE: u8 = 0;
43pub const CTRLMSG_CONNECT: u8 = 1;
44pub const CTRLMSG_ACCEPT: u8 = 2;
45pub const CTRLMSG_CLOSE: u8 = 4;
46pub const CTRLMSG_TOKEN: u8 = 5;
47
48pub const CONNLESS_VERSION: u8 = 1;
49pub const CTRLMSG_CLOSE_REASON_LENGTH: usize = 127;
50pub const TOKEN_REQUEST_PACKET_SIZE: usize = 519;
51pub const TOKEN_NONE: Token = Token([0xff, 0xff, 0xff, 0xff]);
52
53pub const CHUNK_FLAGS_BITS: u32 = 2;
54pub const CHUNK_SIZE_BITS: u32 = 12;
55pub const PACKET_FLAGS_BITS: u32 = 4;
56pub const SEQUENCE_BITS: u32 = 10;
57pub const SEQUENCE_MODULUS: u16 = 1 << SEQUENCE_BITS;
58pub const VERSION_BITS: u32 = 2;
59
60#[derive(Debug)]
61pub enum Error {
62    Capacity(buffer::CapacityError),
63    TooLongData,
64}
65
66impl From<buffer::CapacityError> for Error {
67    fn from(e: buffer::CapacityError) -> Error {
68        Error::Capacity(e)
69    }
70}
71
72pub fn chunk_header_size(vital: bool) -> usize {
73    if vital {
74        CHUNK_HEADER_SIZE_VITAL
75    } else {
76        CHUNK_HEADER_SIZE
77    }
78}
79
80#[derive(Debug, Eq, PartialEq)]
81pub enum Warning {
82    ChunkHeaderPadding,
83    ChunksNoChunks,
84    ChunksNumChunks,
85    ChunksUnknownData,
86    ConnlessFlags,
87    ControlExcessData,
88    ControlFlags,
89    ControlNulTermination,
90    ControlNumChunks,
91    PacketHeaderPadding,
92}
93
94#[derive(Debug, Eq, PartialEq)]
95pub enum PacketReadError {
96    Compression,
97    ControlMissing,
98    ControlResponseTokenMissing,
99    ControlTokenRequestTooShort,
100    TooLong,
101    TooShort,
102    UnknownConnlessVersion,
103    UnknownControl,
104}
105
106#[derive(Clone, Copy)]
107pub enum ControlPacket<'a> {
108    KeepAlive,
109    /// Connect(response_token)
110    Connect(Token),
111    Accept,
112    /// Close(reason)
113    Close(&'a [u8]),
114    /// Token(response_token)
115    Token(Token),
116}
117
118impl<'a> fmt::Debug for ControlPacket<'a> {
119    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
120        match self {
121            ControlPacket::KeepAlive => f.debug_tuple("KeepAlive").finish(),
122            ControlPacket::Connect(rt) => f.debug_tuple("Connect").field(rt).finish(),
123            ControlPacket::Accept => f.debug_tuple("Accept").finish(),
124            ControlPacket::Close(reason) => f
125                .debug_tuple("Close")
126                .field(&pretty::AlmostString::new(reason))
127                .finish(),
128            ControlPacket::Token(rt) => f.debug_tuple("Token").field(rt).finish(),
129        }
130    }
131}
132
133#[derive(Clone, Copy, Eq, Ord, PartialEq, PartialOrd)]
134pub struct Token(pub [u8; 4]);
135
136impl fmt::Debug for Token {
137    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
138        write!(f, "{:08x}", u32::from_be_bytes(self.0))
139    }
140}
141
142impl fmt::Display for Token {
143    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
144        fmt::Debug::fmt(self, f)
145    }
146}
147
148impl Token {
149    pub fn random<F: FnMut(&mut [u8])>(mut f: F) -> Token {
150        loop {
151            let mut token = TOKEN_NONE;
152            f(&mut token.0);
153            if token != TOKEN_NONE {
154                return token;
155            }
156        }
157    }
158}
159
160#[derive(Clone, Copy, Debug)]
161pub struct ConnlessPacket<'a> {
162    pub payload: &'a [u8],
163    pub token: Token,
164    pub response_token: Token,
165}
166
167#[derive(Clone, Copy, Debug)]
168pub struct ConnectedPacket<'a> {
169    pub ack: u16, // u10
170    pub token: Token,
171    pub type_: ConnectedPacketType<'a>,
172}
173
174#[derive(Clone, Copy, Debug)]
175pub enum ConnectedPacketType<'a> {
176    // Chunks(request_resend, num_chunks, payload)
177    Chunks(bool, u8, &'a [u8]),
178    Control(ControlPacket<'a>),
179}
180
181#[derive(Clone, Copy, Debug)]
182pub enum Packet<'a> {
183    Connless(ConnlessPacket<'a>),
184    Connected(ConnectedPacket<'a>),
185}
186
187#[derive(Clone, Copy, Debug)]
188pub struct Chunk<'a> {
189    pub data: &'a [u8],
190    // vital: Some((sequence, resend))
191    pub vital: Option<(u16, bool)>,
192}
193
194#[derive(Clone, Debug)]
195pub struct ChunksIter<'a> {
196    data: &'a [u8],
197    initial_len: usize,
198    num_remaining_chunks: i32,
199    checked_num_chunks_warning: bool,
200}
201
202impl<'a> ChunksIter<'a> {
203    pub fn new(data: &'a [u8], num_chunks: u8) -> ChunksIter<'a> {
204        ChunksIter {
205            data: data,
206            initial_len: data.len(),
207            num_remaining_chunks: num_chunks.i32(),
208            checked_num_chunks_warning: false,
209        }
210    }
211    fn excess_data<W: Warn<Warning>>(&mut self, warn: &mut W) -> Option<Chunk<'static>> {
212        warn.warn(Warning::ChunksUnknownData);
213        self.data = &[];
214        None
215    }
216    pub fn pos(&self) -> usize {
217        self.initial_len - self.data.len()
218    }
219    pub fn next_warn<W>(&mut self, warn: &mut W) -> Option<Chunk<'a>>
220    where
221        W: Warn<Warning>,
222    {
223        if self.data.len() == 0 {
224            if !self.checked_num_chunks_warning {
225                self.checked_num_chunks_warning = true;
226                if self.num_remaining_chunks != 0 {
227                    warn.warn(Warning::ChunksNumChunks);
228                }
229            }
230            return None;
231        }
232        let (header, sequence, chunk_data_and_rest) =
233            unwrap_or_return!(read_chunk_header(warn, self.data), self.excess_data(warn));
234        let vital = sequence.map(|s| (s, header.flags & CHUNKFLAG_RESEND != 0));
235        let size = header.size.usize();
236        if chunk_data_and_rest.len() < size {
237            return self.excess_data(warn);
238        }
239        let (chunk_data, rest) = chunk_data_and_rest.split_at(size);
240        self.data = rest;
241        self.num_remaining_chunks -= 1;
242        Some(Chunk {
243            data: chunk_data,
244            vital: vital,
245        })
246    }
247}
248
249impl<'a> Iterator for ChunksIter<'a> {
250    type Item = Chunk<'a>;
251    fn next(&mut self) -> Option<Chunk<'a>> {
252        self.next_warn(&mut Ignore)
253    }
254    fn size_hint(&self) -> (usize, Option<usize>) {
255        let len = self.clone().count();
256        (len, Some(len))
257    }
258}
259
260impl<'a> ExactSizeIterator for ChunksIter<'a> {}
261
262// TODO: Make this a member function of `Chunk`
263// vital: Some((sequence, resend))
264pub fn write_chunk<'a, B: Buffer<'a>>(
265    bytes: &[u8],
266    vital: Option<(u16, bool)>,
267    buffer: B,
268) -> Result<&'a [u8], buffer::CapacityError> {
269    with_buffer(buffer, |b| write_chunk_impl(bytes, vital, b))
270}
271
272pub fn write_chunk_impl<'d, 's>(
273    bytes: &[u8],
274    vital: Option<(u16, bool)>,
275    mut buffer: BufferRef<'d, 's>,
276) -> Result<&'d [u8], buffer::CapacityError> {
277    assert!(bytes.len() >> CHUNK_SIZE_BITS == 0);
278    let size = bytes.len().assert_u16();
279
280    let (sequence, resend) = vital.unwrap_or((0, false));
281    let resend_flag = if resend { CHUNKFLAG_RESEND } else { 0 };
282    let vital_flag = if vital.is_some() { CHUNKFLAG_VITAL } else { 0 };
283    let flags = vital_flag | resend_flag;
284
285    let header_nonvital = ChunkHeader {
286        flags: flags,
287        size: size,
288    };
289
290    let header1;
291    let header2;
292    let header: &[u8] = if vital.is_some() {
293        header1 = ChunkHeaderVital {
294            h: header_nonvital,
295            sequence: sequence,
296        }
297        .pack();
298        header1.as_bytes()
299    } else {
300        header2 = header_nonvital.pack();
301        header2.as_bytes()
302    };
303    buffer.write(header)?;
304    buffer.write(bytes)?;
305    Ok(buffer.initialized())
306}
307
308fn write_connless_packet<'a, B: Buffer<'a>>(
309    packet: &ConnlessPacket<'_>,
310    buffer: B,
311) -> Result<&'a [u8], Error> {
312    fn inner<'d, 's>(
313        packet: &ConnlessPacket<'_>,
314        mut buffer: BufferRef<'d, 's>,
315    ) -> Result<&'d [u8], Error> {
316        let ConnlessPacket {
317            payload,
318            token,
319            response_token,
320        } = *packet;
321        if payload.len() > MAX_PAYLOAD {
322            return Err(Error::TooLongData);
323        }
324        let header = PacketHeaderConnless {
325            flags: PACKETFLAG_CONNLESS,
326            version: CONNLESS_VERSION,
327            token,
328            response_token,
329        };
330        buffer.write(header.pack().as_bytes())?;
331        buffer.write(payload)?;
332        Ok(buffer.initialized())
333    }
334
335    with_buffer(buffer, |b| inner(packet, b))
336}
337
338impl<'a> Packet<'a> {
339    fn needs_decompression(packet: &[u8]) -> bool {
340        if packet.len() > MAX_PACKETSIZE {
341            return false;
342        }
343        let (header, _) = unwrap_or_return!(PacketHeaderPacked::ref_and_rest_from(packet), false);
344        let header = header.unpack_warn(&mut Ignore);
345        header.flags & PACKETFLAG_CONNLESS == 0 && header.flags & PACKETFLAG_COMPRESSION != 0
346    }
347    /// Parse a packet.
348    ///
349    /// `buffer` needs to have at least size `MAX_PAYLOAD`.
350    pub fn read<'b, B, W>(
351        warn: &mut W,
352        bytes: &'b [u8],
353        buffer: B,
354    ) -> Result<Packet<'b>, PacketReadError>
355    where
356        B: Buffer<'b>,
357        W: Warn<Warning>,
358    {
359        with_buffer(buffer, |b| Packet::read_impl(warn, bytes, Some(b)))
360    }
361    pub fn read_panic_on_decompression<'b, W>(
362        warn: &mut W,
363        bytes: &'b [u8],
364    ) -> Result<Packet<'b>, PacketReadError>
365    where
366        W: Warn<Warning>,
367    {
368        Packet::read_impl(warn, bytes, None)
369    }
370    fn read_impl<'d, 's, W>(
371        warn: &mut W,
372        bytes: &'d [u8],
373        buffer: Option<BufferRef<'d, 's>>,
374    ) -> Result<Packet<'d>, PacketReadError>
375    where
376        W: Warn<Warning>,
377    {
378        use self::PacketReadError::*;
379
380        assert!(buffer
381            .as_ref()
382            .map(|b| b.remaining() >= MAX_PACKETSIZE)
383            .unwrap_or(true));
384        if bytes.len() > MAX_PACKETSIZE {
385            return Err(TooLong);
386        }
387        let (header, payload) =
388            unwrap_or_return!(PacketHeaderPacked::ref_and_rest_from(bytes), Err(TooShort));
389        let header = header.unpack_warn(warn);
390        if header.flags & PACKETFLAG_CONNLESS != 0 {
391            let (header, payload) = unwrap_or_return!(
392                PacketHeaderConnlessPacked::ref_and_rest_from(bytes),
393                Err(TooShort)
394            );
395            let header = header.unpack_warn(warn);
396            if header.version != CONNLESS_VERSION {
397                return Err(UnknownConnlessVersion);
398            }
399            if header.flags & PACKETFLAG_COMPRESSION != 0
400                || header.flags & PACKETFLAG_REQUEST_RESEND != 0
401                || header.flags & PACKETFLAG_CONTROL != 0
402            {
403                // TODO: Should we handle these flags? Vanilla does that too.
404                warn.warn(Warning::ConnlessFlags);
405            }
406
407            return Ok(Packet::Connless(ConnlessPacket {
408                payload,
409                token: header.token,
410                response_token: header.response_token,
411            }));
412        }
413
414        let payload = if header.flags & PACKETFLAG_COMPRESSION != 0 {
415            let mut buffer =
416                buffer.expect("read_panic_on_decompression called on compressed packet");
417            let decompressed = Packet::decompress(bytes, &mut buffer).map_err(|_| Compression)?;
418            let (_, payload) = PacketHeaderPacked::ref_and_rest_from(decompressed).unwrap();
419            payload
420        } else {
421            payload
422        };
423
424        if payload.len() > MAX_PACKETSIZE - HEADER_SIZE {
425            return Err(Compression);
426        }
427
428        let ack = header.ack;
429        let type_ = if header.flags & PACKETFLAG_CONTROL != 0 {
430            if header.num_chunks != 0 {
431                warn.warn(Warning::ControlNumChunks);
432            }
433            if header.flags & PACKETFLAG_COMPRESSION != 0
434                || header.flags & PACKETFLAG_REQUEST_RESEND != 0
435            {
436                // TODO: Should we handle these flags? Vanilla does that too.
437                warn.warn(Warning::ControlFlags);
438            }
439
440            let (&control, payload) = unwrap_or_return!(payload.split_first(), Err(ControlMissing));
441            let empty = |warn: &mut W| {
442                if !payload.is_empty() {
443                    warn.warn(Warning::ControlExcessData);
444                }
445            };
446            let token = |warn: &mut W, warn_more: bool| {
447                if payload.len() < 4 {
448                    return Err(ControlResponseTokenMissing);
449                }
450                let (token, rest) = payload.split_at(4);
451                if warn_more && !rest.is_empty() {
452                    warn.warn(Warning::ControlExcessData);
453                }
454                Ok(Token([token[0], token[1], token[2], token[3]]))
455            };
456            let control = match control {
457                CTRLMSG_KEEPALIVE => {
458                    empty(warn);
459                    ControlPacket::KeepAlive
460                }
461                CTRLMSG_CONNECT => ControlPacket::Connect(token(warn, true)?),
462                CTRLMSG_ACCEPT => {
463                    empty(warn);
464                    ControlPacket::Accept
465                }
466                CTRLMSG_CLOSE => {
467                    let nul = payload
468                        .iter()
469                        .position(|&b| b == 0)
470                        .unwrap_or(payload.len());
471                    let nul = cmp::min(nul, CTRLMSG_CLOSE_REASON_LENGTH);
472                    if payload.len() != 0 && nul + 1 != payload.len() {
473                        if nul + 1 < payload.len() {
474                            warn.warn(Warning::ControlExcessData);
475                        } else {
476                            warn.warn(Warning::ControlNulTermination);
477                        }
478                    }
479                    ControlPacket::Close(&payload[..nul])
480                }
481                CTRLMSG_TOKEN => {
482                    if header.token == TOKEN_NONE && bytes.len() < TOKEN_REQUEST_PACKET_SIZE {
483                        return Err(ControlTokenRequestTooShort);
484                    }
485                    ControlPacket::Token(token(warn, header.token != TOKEN_NONE)?)
486                }
487                _ => {
488                    // Unrecognized control packet.
489                    return Err(UnknownControl);
490                }
491            };
492
493            ConnectedPacketType::Control(control)
494        } else {
495            let request_resend = header.flags & PACKETFLAG_REQUEST_RESEND != 0;
496            if header.num_chunks == 0 && !request_resend {
497                warn.warn(Warning::ChunksNoChunks);
498            }
499            ConnectedPacketType::Chunks(request_resend, header.num_chunks, payload)
500        };
501
502        Ok(Packet::Connected(ConnectedPacket {
503            ack: ack,
504            token: header.token,
505            type_: type_,
506        }))
507    }
508    /// `buffer` needs to have at least size `MAX_PACKETSIZE`.
509    pub fn decompress_if_needed<B: Buffer<'a>>(
510        packet: &[u8],
511        buffer: B,
512    ) -> Result<bool, libtw2_huffman::DecompressionError> {
513        with_buffer(buffer, |b| Packet::decompress_if_needed_impl(packet, b))
514    }
515    fn decompress_if_needed_impl<'d, 's>(
516        packet: &[u8],
517        mut buffer: BufferRef<'d, 's>,
518    ) -> Result<bool, libtw2_huffman::DecompressionError> {
519        assert!(buffer.remaining() >= MAX_PACKETSIZE);
520        if !Packet::needs_decompression(packet) {
521            return Ok(false);
522        }
523        Packet::decompress(packet, &mut buffer)?;
524        Ok(true)
525    }
526
527    fn decompress<B: Buffer<'a>>(
528        packet: &[u8],
529        buffer: B,
530    ) -> Result<&'a [u8], libtw2_huffman::DecompressionError> {
531        with_buffer(buffer, |b| Packet::decompress_impl(packet, b))
532    }
533    fn decompress_impl<'d, 's>(
534        packet: &[u8],
535        mut buffer: BufferRef<'d, 's>,
536    ) -> Result<&'d [u8], libtw2_huffman::DecompressionError> {
537        assert!(buffer.remaining() >= MAX_PACKETSIZE);
538        assert!(Packet::needs_decompression(packet));
539        let (header, payload) = PacketHeaderPacked::ref_and_rest_from(packet)
540            .expect("packet passed to decompress too short for header");
541        let header = header.unpack_warn(&mut Ignore);
542        assert!(header.flags & PACKETFLAG_CONNLESS == 0);
543        assert!(header.flags & PACKETFLAG_COMPRESSION != 0);
544
545        let fake_header = PacketHeader {
546            flags: header.flags & !PACKETFLAG_COMPRESSION,
547            ack: header.ack,
548            num_chunks: header.num_chunks,
549            token: header.token,
550        };
551        buffer.write(fake_header.pack().as_bytes()).unwrap();
552        HUFFMAN.decompress(payload, &mut buffer)?;
553
554        Ok(buffer.initialized())
555    }
556
557    pub fn write<'b, B: Buffer<'b>>(&self, buffer: B) -> Result<&'b [u8], Error> {
558        match *self {
559            Packet::Connected(ref p) => with_buffer(buffer, |b| p.write_impl(b)),
560            Packet::Connless(ref d) => write_connless_packet(d, buffer),
561        }
562    }
563}
564
565impl<'a> ConnectedPacket<'a> {
566    pub fn write<'b, B: Buffer<'b>>(&self, buffer: B) -> Result<&'b [u8], Error> {
567        with_buffer(buffer, |b| self.write_impl(b))
568    }
569
570    fn write_impl<'d, 's>(&self, mut buffer: BufferRef<'d, 's>) -> Result<&'d [u8], Error> {
571        match self.type_ {
572            ConnectedPacketType::Chunks(request_resend, num_chunks, payload) => {
573                let mut compression_buffer: ArrayVec<[u8; 2048]> = ArrayVec::new();
574                let mut compression = 0;
575                let comp_result = HUFFMAN.compress(payload, &mut compression_buffer);
576                if comp_result
577                    .map(|s| s.len() < payload.len())
578                    .unwrap_or(false)
579                {
580                    compression = PACKETFLAG_COMPRESSION;
581                }
582                let request_resend = if request_resend {
583                    PACKETFLAG_REQUEST_RESEND
584                } else {
585                    0
586                };
587                buffer.write(
588                    PacketHeader {
589                        flags: request_resend | compression,
590                        ack: self.ack,
591                        num_chunks: num_chunks,
592                        token: self.token,
593                    }
594                    .pack()
595                    .as_bytes(),
596                )?;
597                buffer.write(if compression != 0 {
598                    &compression_buffer
599                } else {
600                    payload
601                })?;
602                Ok(buffer.initialized())
603            }
604            ConnectedPacketType::Control(c) => c.write(self.token, self.ack, buffer),
605        }
606    }
607}
608
609impl<'a> ControlPacket<'a> {
610    fn write<'d, 's>(
611        &self,
612        token: Token,
613        ack: u16,
614        mut buffer: BufferRef<'d, 's>,
615    ) -> Result<&'d [u8], Error> {
616        buffer.write(
617            PacketHeader {
618                flags: PACKETFLAG_CONTROL,
619                ack: ack,
620                num_chunks: 0,
621                token: token,
622            }
623            .pack()
624            .as_bytes(),
625        )?;
626        let magic = match *self {
627            ControlPacket::KeepAlive => CTRLMSG_KEEPALIVE,
628            ControlPacket::Connect(..) => CTRLMSG_CONNECT,
629            ControlPacket::Accept => CTRLMSG_ACCEPT,
630            ControlPacket::Close(..) => CTRLMSG_CLOSE,
631            ControlPacket::Token(..) => CTRLMSG_TOKEN,
632        };
633        buffer.write(&[magic])?;
634        match *self {
635            ControlPacket::KeepAlive => {}
636            ControlPacket::Connect(response_token) => {
637                assert!(response_token != TOKEN_NONE);
638                buffer.write(&response_token.0)?;
639            }
640            ControlPacket::Accept => {}
641            ControlPacket::Close(m) => {
642                assert!(m.iter().all(|&b| b != 0));
643                buffer.write(m)?;
644                buffer.write(&[0])?;
645            }
646            ControlPacket::Token(response_token) => {
647                assert!(response_token != TOKEN_NONE);
648                buffer.write(&response_token.0)?;
649                if token == TOKEN_NONE {
650                    const ADDITIONAL: usize = TOKEN_REQUEST_PACKET_SIZE
651                        - mem::size_of::<PacketHeaderPacked>()
652                        - 1
653                        - mem::size_of::<Token>();
654                    buffer.write(&[0; ADDITIONAL])?;
655                }
656            }
657        }
658        let result = buffer.initialized();
659        assert!(result.len() <= MAX_PACKETSIZE);
660        Ok(result)
661    }
662}
663
664#[repr(C, packed)]
665#[derive(AsBytes, Clone, Copy, FromBytes, FromZeroes)]
666pub struct PacketHeaderPacked {
667    padding_flags_ack: u8, // u2 u4 u2
668    ack: u8,
669    num_chunks: u8,
670    token: [u8; 4],
671}
672
673#[derive(Copy, Clone, Debug, Eq, PartialEq)]
674pub struct PacketHeader {
675    pub flags: u8, // u4
676    pub ack: u16,  // u10
677    pub num_chunks: u8,
678    pub token: Token,
679}
680
681#[repr(C, packed)]
682#[derive(AsBytes, Clone, Copy, FromBytes, FromZeroes)]
683pub struct PacketHeaderConnlessPacked {
684    padding_flags_version: u8, // u2 u4 u2
685    token: [u8; 4],
686    response_token: [u8; 4],
687}
688
689#[derive(Copy, Clone, Debug, Eq, PartialEq)]
690pub struct PacketHeaderConnless {
691    pub flags: u8,   // u4
692    pub version: u8, // u2
693    pub token: Token,
694    pub response_token: Token,
695}
696
697impl PacketHeaderPacked {
698    pub fn unpack_warn<W: Warn<Warning>>(self, warn: &mut W) -> PacketHeader {
699        let PacketHeaderPacked {
700            padding_flags_ack,
701            ack,
702            num_chunks,
703            token,
704        } = self;
705        if padding_flags_ack & 0b1100_0000 != 0 {
706            warn.warn(Warning::PacketHeaderPadding);
707        }
708        PacketHeader {
709            flags: (padding_flags_ack & 0b0011_1100) >> 2,
710            ack: (((padding_flags_ack & 0b0000_0011) as u16) << 8) | (ack as u16),
711            num_chunks,
712            token: Token(token),
713        }
714    }
715    pub fn unpack(self) -> PacketHeader {
716        self.unpack_warn(&mut Ignore)
717    }
718}
719
720impl PacketHeader {
721    pub fn pack(self) -> PacketHeaderPacked {
722        let PacketHeader {
723            flags,
724            ack,
725            num_chunks,
726            token,
727        } = self;
728        // Check that the fields do not exceed their maximal size.
729        assert!(flags >> PACKET_FLAGS_BITS == 0);
730        assert!(ack >> SEQUENCE_BITS == 0);
731        PacketHeaderPacked {
732            padding_flags_ack: flags << 2 | (ack >> 8) as u8,
733            ack: ack as u8,
734            num_chunks,
735            token: token.0,
736        }
737    }
738}
739
740impl PacketHeaderConnlessPacked {
741    pub fn unpack_warn<W: Warn<Warning>>(self, warn: &mut W) -> PacketHeaderConnless {
742        let PacketHeaderConnlessPacked {
743            padding_flags_version,
744            token,
745            response_token,
746        } = self;
747        if padding_flags_version & 0b1100_0000 != 0 {
748            warn.warn(Warning::PacketHeaderPadding);
749        }
750        PacketHeaderConnless {
751            flags: (padding_flags_version & 0b0011_1100) >> 2,
752            version: padding_flags_version & 0b0000_0011,
753            token: Token(token),
754            response_token: Token(response_token),
755        }
756    }
757    pub fn unpack(self) -> PacketHeaderConnless {
758        self.unpack_warn(&mut Ignore)
759    }
760}
761
762impl PacketHeaderConnless {
763    pub fn pack(self) -> PacketHeaderConnlessPacked {
764        let PacketHeaderConnless {
765            flags,
766            version,
767            token,
768            response_token,
769        } = self;
770        // Check that the fields do not exceed their maximal size.
771        assert!(flags >> PACKET_FLAGS_BITS == 0);
772        assert!(version >> VERSION_BITS == 0);
773        PacketHeaderConnlessPacked {
774            padding_flags_version: flags << 2 | version as u8,
775            token: token.0,
776            response_token: response_token.0,
777        }
778    }
779}
780
781#[derive(Copy, Clone, Debug, Eq, PartialEq)]
782pub struct ChunkHeader {
783    pub flags: u8, // u2
784    pub size: u16, // u10
785}
786
787#[derive(Copy, Clone, Debug, Eq, PartialEq)]
788pub struct ChunkHeaderVital {
789    pub h: ChunkHeader,
790    pub sequence: u16, // u16
791}
792
793#[repr(C, packed)]
794#[derive(AsBytes, Clone, Copy, FromBytes, FromZeroes)]
795pub struct ChunkHeaderPacked {
796    flags_size: u8,   // u2 u6
797    padding_size: u8, // u2 u6
798}
799
800#[repr(C, packed)]
801#[derive(AsBytes, Clone, Copy, FromBytes, FromZeroes)]
802pub struct ChunkHeaderVitalPacked {
803    flags_size: u8,    // u2 u6
804    sequence_size: u8, // u2 u6
805    sequence: u8,
806}
807
808/// -> Some((chunk_header, sequence, rest))
809pub fn read_chunk_header<'a, W>(
810    warn: &mut W,
811    data: &'a [u8],
812) -> Option<(ChunkHeader, Option<u16>, &'a [u8])>
813where
814    W: Warn<Warning>,
815{
816    let (raw_header, chunk_data_and_rest) =
817        unwrap_or_return!(ChunkHeaderPacked::ref_and_rest_from(data));
818
819    let header = raw_header.unpack_warn(&mut Ignore);
820    Some(if header.flags & CHUNKFLAG_VITAL != 0 {
821        let (header, chunk_data_and_rest_vital) =
822            unwrap_or_return!(ChunkHeaderVitalPacked::ref_and_rest_from(data));
823        let header = header.unpack_warn(warn);
824        (header.h, Some(header.sequence), chunk_data_and_rest_vital)
825    } else {
826        raw_header.unpack_warn(warn);
827        (header, None, chunk_data_and_rest)
828    })
829}
830
831impl ChunkHeaderPacked {
832    pub fn unpack_warn<W: Warn<Warning>>(self, warn: &mut W) -> ChunkHeader {
833        let ChunkHeaderPacked {
834            flags_size,
835            padding_size,
836        } = self;
837        if padding_size & 0b1111_0000 != 0 {
838            warn.warn(Warning::ChunkHeaderPadding);
839        }
840        ChunkHeader {
841            flags: (flags_size & 0b1100_0000) >> 6,
842            size: ((((flags_size & 0b0011_1111) as u16) << 6)
843                | (padding_size & 0b0011_1111) as u16),
844        }
845    }
846    pub fn unpack(self) -> ChunkHeader {
847        self.unpack_warn(&mut Ignore)
848    }
849}
850
851impl ChunkHeader {
852    pub fn pack(self) -> ChunkHeaderPacked {
853        let ChunkHeader { flags, size } = self;
854        // Check that the fields do not exceed their maximal size.
855        assert!(flags >> CHUNK_FLAGS_BITS == 0);
856        assert!(size >> CHUNK_SIZE_BITS == 0);
857        ChunkHeaderPacked {
858            flags_size: (flags & 0b11) << 6 | ((size & 0b1111_1100_0000) >> 6) as u8,
859            padding_size: (size & 0b0000_0011_1111) as u8,
860        }
861    }
862}
863
864impl ChunkHeaderVitalPacked {
865    pub fn unpack_warn<W: Warn<Warning>>(self, warn: &mut W) -> ChunkHeaderVital {
866        let ChunkHeaderVitalPacked {
867            flags_size,
868            sequence_size,
869            sequence,
870        } = self;
871        ChunkHeaderVital {
872            h: ChunkHeaderPacked {
873                flags_size: flags_size,
874                padding_size: sequence_size & 0b0011_1111,
875            }
876            .unpack_warn(warn),
877            sequence: ((sequence_size & 0b1100_0000) as u16) << 2
878                | ((sequence & 0b1111_1111) as u16),
879        }
880    }
881    pub fn unpack(self) -> ChunkHeaderVital {
882        self.unpack_warn(&mut Ignore)
883    }
884}
885
886impl ChunkHeaderVital {
887    pub fn pack(self) -> ChunkHeaderVitalPacked {
888        let ChunkHeaderVital { h, sequence } = self;
889        assert!(sequence >> SEQUENCE_BITS == 0);
890        let ChunkHeaderPacked {
891            flags_size,
892            padding_size,
893        } = h.pack();
894        ChunkHeaderVitalPacked {
895            flags_size: flags_size,
896            sequence_size: (padding_size & 0b0011_1111) | ((sequence & 0b11_0000_0000) >> 2) as u8,
897            sequence: (sequence & 0b00_1111_1111) as u8,
898        }
899    }
900}
901
902boilerplate_packed!(PacketHeaderPacked, HEADER_SIZE, test_ph_size);
903boilerplate_packed!(
904    PacketHeaderConnlessPacked,
905    HEADER_SIZE_CONNLESS,
906    test_phc_size
907);
908boilerplate_packed!(ChunkHeaderPacked, CHUNK_HEADER_SIZE, test_ch_size);
909boilerplate_packed!(
910    ChunkHeaderVitalPacked,
911    CHUNK_HEADER_SIZE_VITAL,
912    test_chv_size
913);
914
915#[cfg(test)]
916#[rustfmt::skip]
917mod test {
918    use libtw2_common::bytes::FromBytesExt as _;
919    use libtw2_warn::Panic;
920    use libtw2_warn::Warn;
921    use quickcheck::quickcheck;
922    use super::CHUNK_FLAGS_BITS;
923    use super::CHUNK_SIZE_BITS;
924    use super::ChunkHeader;
925    use super::ChunkHeaderPacked;
926    use super::ChunkHeaderVital;
927    use super::ChunkHeaderVitalPacked;
928    use super::ChunksIter;
929    use super::ConnectedPacket;
930    use super::ConnectedPacketType;
931    use super::MAX_PACKETSIZE;
932    use super::PACKET_FLAGS_BITS;
933    use super::Packet;
934    use super::PacketHeader;
935    use super::PacketHeaderPacked;
936    use super::PacketReadError::*;
937    use super::PacketReadError;
938    use super::SEQUENCE_BITS;
939    use super::Token;
940    use super::Warning::*;
941    use super::Warning;
942    use zerocopy::AsBytes as _;
943
944    struct WarnVec<'a>(&'a mut Vec<Warning>);
945
946    impl<'a> Warn<Warning> for WarnVec<'a> {
947        fn warn(&mut self, warning: Warning) {
948            self.0.push(warning);
949        }
950    }
951
952    fn assert_warnings(input: &[u8], warnings: &[Warning]) {
953        let mut vec = vec![];
954        let mut buffer = Vec::with_capacity(4096);
955        let packet = Packet::read(&mut WarnVec(&mut vec), input, &mut buffer).unwrap();
956        if let Packet::Connected(ConnectedPacket {
957            type_: ConnectedPacketType::Chunks(_, num_chunks, chunk_data),
958            ..
959        }) = packet {
960            let mut chunks = ChunksIter::new(chunk_data, num_chunks);
961            while let Some(_) = chunks.next_warn(&mut WarnVec(&mut vec)) { }
962        }
963        if warnings != &[ChunksNoChunks] {
964            vec.retain(|w| *w != ChunksNoChunks);
965        }
966        assert_eq!(vec, warnings);
967    }
968
969    fn assert_warn(input: &[u8], warning: Warning) {
970        assert_warnings(input, &[warning]);
971    }
972
973    fn assert_no_warn(input: &[u8]) {
974        assert_warnings(input, &[]);
975    }
976
977    fn assert_err(input: &[u8], error: PacketReadError) {
978        let mut buffer = Vec::with_capacity(4096);
979        assert_eq!(Packet::read(&mut Panic, input, &mut buffer).unwrap_err(), error);
980    }
981
982    #[test] fn w_chp() { assert_warn(b"\x00\x00\x01\x00\x00\x00\x00\x00\xc0", ChunkHeaderPadding) }
983    #[test] fn w_cud1() { assert_warn(b"\x00\x00\x00\x00\x00\x00\x00\xff", ChunksUnknownData) }
984    #[test] fn w_cud2() { assert_warn(b"\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00", ChunksUnknownData) }
985    #[test] fn w_cud3() { assert_no_warn(b"\x00\x00\x01\x00\x00\x00\x00\x00\x00") }
986    #[test] fn w_cud4_ddnet_6() { assert_warn(b"\x00\x00\x01\x00\x00\x12\x34\x45\x67", ChunksUnknownData) }
987    #[test] fn w_cnc1() { assert_warn(b"\x00\x00\x01\x00\x00\x00\x00", ChunksNumChunks) }
988    #[test] fn w_cnc2() { assert_warn(b"\x00\x00\x00\x00\x00\x00\x00\x00\x00", ChunksNumChunks) }
989    #[test] fn w_cnc_() { assert_warn(b"\x00\x00\x00\x00\x00\x00\x00", ChunksNoChunks) }
990    #[test] fn w_cf_1() { assert_warn(b"\x31\x00\x00\x00\x00\x00\x00\x00\x00", ConnlessFlags) }
991    #[test] fn w_cf_2() { assert_warn(b"\x29\x00\x00\x00\x00\x00\x00\x00\x00", ConnlessFlags) }
992    #[test] fn w_cf_3() { assert_warn(b"\x25\x00\x00\x00\x00\x00\x00\x00\x00", ConnlessFlags) }
993    #[test] fn w_cf_4() { assert_no_warn(b"\x21\x00\x00\x00\x00\x00\x00\x00\x00") }
994    #[test] fn w_ced1() { assert_warn(b"\x04\x00\x00\x00\x00\x00\x00\x00\x00", ControlExcessData) }
995    #[test] fn w_ced2() { assert_warn(b"\x04\x00\x00\x00\x00\x00\x00\x04\x00\x00", ControlExcessData) }
996    #[test] fn w_cf1() { assert_warn(b"\x14\x00\x00\x00\x00\x00\x00\x15\x37", ControlFlags) }
997    #[test] fn w_cf2() { assert_warn(b"\x0c\x00\x00\x00\x00\x00\x00\x00", ControlFlags) }
998    #[test] fn w_cnt1() { assert_warn(b"\x04\x00\x00\x00\x00\x00\x00\x04\x01", ControlNulTermination) }
999    #[test] fn w_cnt2() { assert_no_warn(b"\x04\x00\x00\x00\x00\x00\x00\x04") }
1000    #[test] fn w_cnc() { assert_warn(b"\x04\x00\xff\x00\x00\x00\x00\x00", ControlNumChunks) }
1001    #[test] fn w_php1() { assert_warn(b"\x80\x00\x00\x00\x00\x00\x00", PacketHeaderPadding) }
1002    #[test] fn w_php2() { assert_warn(b"\x40\x00\x00\x00\x00\x00\x00", PacketHeaderPadding) }
1003
1004    #[test] fn e_cm() { assert_err(b"\x04\x00\x00\x00\x00\x00\x00", ControlMissing) }
1005    #[test] fn e_crtm() { assert_err(b"\x04\x00\x00\x00\x00\x00\x00\x05", ControlResponseTokenMissing) }
1006    #[test] fn e_ctrts() { assert_err(b"\x04\x00\x00\xff\xff\xff\xff\x05\x00\x00\x00\x00", ControlTokenRequestTooShort) }
1007    #[test] fn e_sc() { assert_err(b"\xff\xff\xff", TooShort) }
1008    #[test] fn e_tl() { assert_err(&[0; MAX_PACKETSIZE+1], TooLong) }
1009    #[test] fn e_ts1() { assert_err(b"\x00\x00", TooShort) }
1010    #[test] fn e_ts2() { assert_err(b"", TooShort) }
1011    #[test] fn e_uc1() { assert_err(b"\x04\x00\x00\x00\x00\x00\x00\x06", UnknownControl) }
1012    #[test] fn e_uc2() { assert_err(b"\x04\x00\x00\x00\x00\x00\x00\xff", UnknownControl) }
1013    #[test] fn e_c() { assert_err(b"\x10\x00\x00\x00\x00\x00\x00", Compression) }
1014    #[test] fn e_ucv() { assert_err(b"\x22\x00\x00\x00\x00\x00\x00\x00\x00", UnknownConnlessVersion) }
1015
1016    quickcheck! {
1017        fn packet_header_roundtrip(flags: u8, ack: u16, num_chunks: u8, token: (u8, u8, u8, u8)) -> bool {
1018            let flags = flags ^ (flags >> PACKET_FLAGS_BITS << PACKET_FLAGS_BITS);
1019            let ack = ack ^ (ack >> SEQUENCE_BITS << SEQUENCE_BITS);
1020            let token = Token([token.0, token.1, token.2, token.3]);
1021            let packet_header = PacketHeader {
1022                flags,
1023                ack,
1024                num_chunks,
1025                token,
1026            };
1027            packet_header == packet_header.pack().unpack()
1028        }
1029
1030        fn packet_header_unpack(v: (u8, u8, u8, u8, u8, u8, u8)) -> bool {
1031            // Two bits must be zeroed (see doc/packet7.md).
1032            let v0 = v.0 & 0b0011_1111;
1033            let bytes = &[v0, v.1, v.2, v.3, v.4, v.5, v.6];
1034            PacketHeaderPacked::ref_from_array(bytes).unpack().pack().as_bytes() == bytes
1035        }
1036
1037        fn chunk_header_roundtrip(flags: u8, size: u16, sequence: u16) -> bool {
1038            let flags = flags ^ (flags >> CHUNK_FLAGS_BITS << CHUNK_FLAGS_BITS);
1039            let size = size ^ (size >> CHUNK_SIZE_BITS << CHUNK_SIZE_BITS);
1040            let sequence = sequence ^ (sequence >> SEQUENCE_BITS << SEQUENCE_BITS);
1041            let chunk_header = ChunkHeader {
1042                flags: flags,
1043                size: size,
1044            };
1045            let chunk_header_vital = ChunkHeaderVital {
1046                h: chunk_header,
1047                sequence: sequence,
1048            };
1049            chunk_header == chunk_header.pack().unpack()
1050                && chunk_header_vital == chunk_header_vital.pack().unpack()
1051        }
1052
1053        fn chunk_header_unpack(v: (u8, u8, u8)) -> bool {
1054            let bytes2 = &[v.0, v.1];
1055            let bytes3 = &[v.0, v.1, v.2];
1056            let bytes2_result = &[v.0, v.1 & 0b0011_1111];
1057            let bytes3_result = &[v.0, v.1, v.2];
1058            ChunkHeaderPacked::ref_from_array(bytes2).unpack().pack().as_bytes() == bytes2_result
1059                && ChunkHeaderVitalPacked::ref_from_array(bytes3).unpack().pack().as_bytes() == bytes3_result
1060        }
1061    }
1062}