packet/
serialize.rs

1// Copyright 2018 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5//! Serialization.
6
7use std::cmp;
8use std::convert::Infallible as Never;
9use std::fmt::{self, Debug, Formatter};
10use std::ops::{Range, RangeBounds};
11
12use arrayvec::ArrayVec;
13use zerocopy::SplitByteSlice;
14
15use crate::{
16    canonicalize_range, take_back, take_back_mut, take_front, take_front_mut,
17    AsFragmentedByteSlice, Buffer, BufferView, BufferViewMut, ContiguousBuffer, EmptyBuf,
18    FragmentedBuffer, FragmentedBufferMut, FragmentedBytes, FragmentedBytesMut, GrowBuffer,
19    GrowBufferMut, ParsablePacket, ParseBuffer, ParseBufferMut, ReusableBuffer, ShrinkBuffer,
20};
21
22/// Either of two buffers.
23///
24/// An `Either` wraps one of two different buffer types. It implements all of
25/// the relevant traits by calling the corresponding methods on the wrapped
26/// buffer.
27#[derive(Copy, Clone, Debug)]
28pub enum Either<A, B> {
29    A(A),
30    B(B),
31}
32
33impl<A, B> Either<A, B> {
34    /// Maps the `A` variant of an `Either`.
35    ///
36    /// Given an `Either<A, B>` and a function from `A` to `AA`, `map_a`
37    /// produces an `Either<AA, B>` by applying the function to the `A` variant
38    /// or passing on the `B` variant unmodified.
39    pub fn map_a<AA, F: FnOnce(A) -> AA>(self, f: F) -> Either<AA, B> {
40        match self {
41            Either::A(a) => Either::A(f(a)),
42            Either::B(b) => Either::B(b),
43        }
44    }
45
46    /// Maps the `B` variant of an `Either`.
47    ///
48    /// Given an `Either<A, B>` and a function from `B` to `BB`, `map_b`
49    /// produces an `Either<A, BB>` by applying the function to the `B` variant
50    /// or passing on the `A` variant unmodified.
51    pub fn map_b<BB, F: FnOnce(B) -> BB>(self, f: F) -> Either<A, BB> {
52        match self {
53            Either::A(a) => Either::A(a),
54            Either::B(b) => Either::B(f(b)),
55        }
56    }
57
58    /// Returns the `A` variant in an `Either<A, B>`.
59    ///
60    /// # Panics
61    ///
62    /// Panics if this `Either<A, B>` does not hold the `A` variant.
63    pub fn unwrap_a(self) -> A {
64        match self {
65            Either::A(x) => x,
66            Either::B(_) => panic!("This `Either<A, B>` does not hold the `A` variant"),
67        }
68    }
69
70    /// Returns the `B` variant in an `Either<A, B>`.
71    ///
72    /// # Panics
73    ///
74    /// Panics if this `Either<A, B>` does not hold the `B` variant.
75    pub fn unwrap_b(self) -> B {
76        match self {
77            Either::A(_) => panic!("This `Either<A, B>` does not hold the `B` variant"),
78            Either::B(x) => x,
79        }
80    }
81}
82
83impl<A> Either<A, A> {
84    /// Returns the inner value held by this `Either` when both possible values
85    /// `Either::A` and `Either::B` contain the same inner types.
86    pub fn into_inner(self) -> A {
87        match self {
88            Either::A(x) => x,
89            Either::B(x) => x,
90        }
91    }
92}
93
94impl<A> Either<A, Never> {
95    /// Returns the `A` value in an `Either<A, Never>`.
96    #[inline]
97    pub fn into_a(self) -> A {
98        match self {
99            Either::A(a) => a,
100        }
101    }
102}
103
104impl<B> Either<Never, B> {
105    /// Returns the `B` value in an `Either<Never, B>`.
106    #[inline]
107    pub fn into_b(self) -> B {
108        match self {
109            Either::B(b) => b,
110        }
111    }
112}
113
114macro_rules! call_method_on_either {
115    ($val:expr, $method:ident, $($args:expr),*) => {
116        match $val {
117            Either::A(a) => a.$method($($args),*),
118            Either::B(b) => b.$method($($args),*),
119        }
120    };
121    ($val:expr, $method:ident) => {
122        call_method_on_either!($val, $method,)
123    };
124}
125
126// NOTE(joshlf): We override the default implementations of all methods for
127// Either. Many of the default implementations make multiple calls to other
128// Buffer methods, each of which performs a match statement to figure out which
129// Either variant is present. We assume that doing this match once is more
130// performant than doing it multiple times.
131
132impl<A, B> FragmentedBuffer for Either<A, B>
133where
134    A: FragmentedBuffer,
135    B: FragmentedBuffer,
136{
137    fn len(&self) -> usize {
138        call_method_on_either!(self, len)
139    }
140
141    fn with_bytes<R, F>(&self, f: F) -> R
142    where
143        F: for<'a, 'b> FnOnce(FragmentedBytes<'a, 'b>) -> R,
144    {
145        call_method_on_either!(self, with_bytes, f)
146    }
147}
148
149impl<A, B> ContiguousBuffer for Either<A, B>
150where
151    A: ContiguousBuffer,
152    B: ContiguousBuffer,
153{
154}
155
156impl<A, B> ShrinkBuffer for Either<A, B>
157where
158    A: ShrinkBuffer,
159    B: ShrinkBuffer,
160{
161    fn shrink<R: RangeBounds<usize>>(&mut self, range: R) {
162        call_method_on_either!(self, shrink, range)
163    }
164    fn shrink_front(&mut self, n: usize) {
165        call_method_on_either!(self, shrink_front, n)
166    }
167    fn shrink_back(&mut self, n: usize) {
168        call_method_on_either!(self, shrink_back, n)
169    }
170}
171
172impl<A, B> ParseBuffer for Either<A, B>
173where
174    A: ParseBuffer,
175    B: ParseBuffer,
176{
177    fn parse<'a, P: ParsablePacket<&'a [u8], ()>>(&'a mut self) -> Result<P, P::Error> {
178        call_method_on_either!(self, parse)
179    }
180    fn parse_with<'a, ParseArgs, P: ParsablePacket<&'a [u8], ParseArgs>>(
181        &'a mut self,
182        args: ParseArgs,
183    ) -> Result<P, P::Error> {
184        call_method_on_either!(self, parse_with, args)
185    }
186}
187
188impl<A, B> FragmentedBufferMut for Either<A, B>
189where
190    A: FragmentedBufferMut,
191    B: FragmentedBufferMut,
192{
193    fn with_bytes_mut<R, F>(&mut self, f: F) -> R
194    where
195        F: for<'a, 'b> FnOnce(FragmentedBytesMut<'a, 'b>) -> R,
196    {
197        call_method_on_either!(self, with_bytes_mut, f)
198    }
199}
200
201impl<A, B> ParseBufferMut for Either<A, B>
202where
203    A: ParseBufferMut,
204    B: ParseBufferMut,
205{
206    fn parse_mut<'a, P: ParsablePacket<&'a mut [u8], ()>>(&'a mut self) -> Result<P, P::Error> {
207        call_method_on_either!(self, parse_mut)
208    }
209    fn parse_with_mut<'a, ParseArgs, P: ParsablePacket<&'a mut [u8], ParseArgs>>(
210        &'a mut self,
211        args: ParseArgs,
212    ) -> Result<P, P::Error> {
213        call_method_on_either!(self, parse_with_mut, args)
214    }
215}
216
217impl<A, B> GrowBuffer for Either<A, B>
218where
219    A: GrowBuffer,
220    B: GrowBuffer,
221{
222    #[inline]
223    fn with_parts<O, F>(&self, f: F) -> O
224    where
225        F: for<'a, 'b> FnOnce(&'a [u8], FragmentedBytes<'a, 'b>, &'a [u8]) -> O,
226    {
227        call_method_on_either!(self, with_parts, f)
228    }
229    fn capacity(&self) -> usize {
230        call_method_on_either!(self, capacity)
231    }
232    fn prefix_len(&self) -> usize {
233        call_method_on_either!(self, prefix_len)
234    }
235    fn suffix_len(&self) -> usize {
236        call_method_on_either!(self, suffix_len)
237    }
238    fn grow_front(&mut self, n: usize) {
239        call_method_on_either!(self, grow_front, n)
240    }
241    fn grow_back(&mut self, n: usize) {
242        call_method_on_either!(self, grow_back, n)
243    }
244    fn reset(&mut self) {
245        call_method_on_either!(self, reset)
246    }
247}
248
249impl<A, B> GrowBufferMut for Either<A, B>
250where
251    A: GrowBufferMut,
252    B: GrowBufferMut,
253{
254    fn with_parts_mut<O, F>(&mut self, f: F) -> O
255    where
256        F: for<'a, 'b> FnOnce(&'a mut [u8], FragmentedBytesMut<'a, 'b>, &'a mut [u8]) -> O,
257    {
258        call_method_on_either!(self, with_parts_mut, f)
259    }
260
261    fn serialize<BB: PacketBuilder>(&mut self, builder: BB) {
262        call_method_on_either!(self, serialize, builder)
263    }
264}
265
266impl<A, B> Buffer for Either<A, B>
267where
268    A: Buffer,
269    B: Buffer,
270{
271    fn parse_with_view<'a, ParseArgs, P: ParsablePacket<&'a [u8], ParseArgs>>(
272        &'a mut self,
273        args: ParseArgs,
274    ) -> Result<(P, &'a [u8]), P::Error> {
275        call_method_on_either!(self, parse_with_view, args)
276    }
277}
278
279impl<A: AsRef<[u8]>, B: AsRef<[u8]>> AsRef<[u8]> for Either<A, B> {
280    fn as_ref(&self) -> &[u8] {
281        call_method_on_either!(self, as_ref)
282    }
283}
284
285impl<A: AsMut<[u8]>, B: AsMut<[u8]>> AsMut<[u8]> for Either<A, B> {
286    fn as_mut(&mut self) -> &mut [u8] {
287        call_method_on_either!(self, as_mut)
288    }
289}
290
291/// A byte slice wrapper providing buffer functionality.
292///
293/// A `Buf` wraps a byte slice (a type which implements `AsRef<[u8]>` or
294/// `AsMut<[u8]>`) and implements various buffer traits by keeping track of
295/// prefix, body, and suffix offsets within the byte slice.
296#[derive(Clone, Debug)]
297pub struct Buf<B> {
298    buf: B,
299    body: Range<usize>,
300}
301
302impl<B: AsRef<[u8]>> PartialEq for Buf<B> {
303    fn eq(&self, other: &Self) -> bool {
304        let self_slice = AsRef::<[u8]>::as_ref(self);
305        let other_slice = AsRef::<[u8]>::as_ref(other);
306        PartialEq::eq(self_slice, other_slice)
307    }
308}
309
310impl<B: AsRef<[u8]>> Eq for Buf<B> {}
311
312impl Buf<Vec<u8>> {
313    /// Extracts the contained data trimmed to the buffer's range.
314    pub fn into_inner(self) -> Vec<u8> {
315        let Buf { mut buf, body } = self;
316        let len = body.end - body.start;
317        let _ = buf.drain(..body.start);
318        buf.truncate(len);
319        buf
320    }
321}
322
323impl<B: AsRef<[u8]>> Buf<B> {
324    /// Constructs a new `Buf`.
325    ///
326    /// `new` constructs a new `Buf` from a buffer and a body range. The bytes
327    /// within the range will be the body, the bytes before the range will be
328    /// the prefix, and the bytes after the range will be the suffix.
329    ///
330    /// # Panics
331    ///
332    /// Panics if `range` is out of bounds of `buf`, or if it is nonsensical
333    /// (the end precedes the start).
334    pub fn new<R: RangeBounds<usize>>(buf: B, body: R) -> Buf<B> {
335        let len = buf.as_ref().len();
336        Buf { buf, body: canonicalize_range(len, &body) }
337    }
338
339    /// Constructs a [`BufView`] which will be a [`BufferView`] into this `Buf`.
340    pub fn buffer_view(&mut self) -> BufView<'_> {
341        BufView { buf: &self.buf.as_ref()[self.body.clone()], body: &mut self.body }
342    }
343}
344
345impl<B: AsRef<[u8]> + AsMut<[u8]>> Buf<B> {
346    /// Constructs a [`BufViewMut`] which will be a [`BufferViewMut`] into this `Buf`.
347    pub fn buffer_view_mut(&mut self) -> BufViewMut<'_> {
348        BufViewMut { buf: &mut self.buf.as_mut()[self.body.clone()], body: &mut self.body }
349    }
350}
351
352impl<B: AsRef<[u8]>> FragmentedBuffer for Buf<B> {
353    fragmented_buffer_method_impls!();
354}
355impl<B: AsRef<[u8]>> ContiguousBuffer for Buf<B> {}
356impl<B: AsRef<[u8]>> ShrinkBuffer for Buf<B> {
357    fn shrink<R: RangeBounds<usize>>(&mut self, range: R) {
358        let len = self.len();
359        let mut range = canonicalize_range(len, &range);
360        range.start += self.body.start;
361        range.end += self.body.start;
362        self.body = range;
363    }
364
365    fn shrink_front(&mut self, n: usize) {
366        assert!(n <= self.len());
367        self.body.start += n;
368    }
369    fn shrink_back(&mut self, n: usize) {
370        assert!(n <= self.len());
371        self.body.end -= n;
372    }
373}
374impl<B: AsRef<[u8]>> ParseBuffer for Buf<B> {
375    fn parse_with<'a, ParseArgs, P: ParsablePacket<&'a [u8], ParseArgs>>(
376        &'a mut self,
377        args: ParseArgs,
378    ) -> Result<P, P::Error> {
379        P::parse(self.buffer_view(), args)
380    }
381}
382
383impl<B: AsRef<[u8]> + AsMut<[u8]>> FragmentedBufferMut for Buf<B> {
384    fragmented_buffer_mut_method_impls!();
385}
386
387impl<B: AsRef<[u8]> + AsMut<[u8]>> ParseBufferMut for Buf<B> {
388    fn parse_with_mut<'a, ParseArgs, P: ParsablePacket<&'a mut [u8], ParseArgs>>(
389        &'a mut self,
390        args: ParseArgs,
391    ) -> Result<P, P::Error> {
392        P::parse_mut(self.buffer_view_mut(), args)
393    }
394}
395
396impl<B: AsRef<[u8]>> GrowBuffer for Buf<B> {
397    fn with_parts<O, F>(&self, f: F) -> O
398    where
399        F: for<'a, 'b> FnOnce(&'a [u8], FragmentedBytes<'a, 'b>, &'a [u8]) -> O,
400    {
401        let (prefix, buf) = self.buf.as_ref().split_at(self.body.start);
402        let (body, suffix) = buf.split_at(self.body.end - self.body.start);
403        let mut body = [&body[..]];
404        f(prefix, body.as_fragmented_byte_slice(), suffix)
405    }
406    fn capacity(&self) -> usize {
407        self.buf.as_ref().len()
408    }
409    fn prefix_len(&self) -> usize {
410        self.body.start
411    }
412    fn suffix_len(&self) -> usize {
413        self.buf.as_ref().len() - self.body.end
414    }
415    fn grow_front(&mut self, n: usize) {
416        assert!(n <= self.body.start);
417        self.body.start -= n;
418    }
419    fn grow_back(&mut self, n: usize) {
420        assert!(n <= self.buf.as_ref().len() - self.body.end);
421        self.body.end += n;
422    }
423}
424
425impl<B: AsRef<[u8]> + AsMut<[u8]>> GrowBufferMut for Buf<B> {
426    fn with_parts_mut<O, F>(&mut self, f: F) -> O
427    where
428        F: for<'a, 'b> FnOnce(&'a mut [u8], FragmentedBytesMut<'a, 'b>, &'a mut [u8]) -> O,
429    {
430        let (prefix, buf) = self.buf.as_mut().split_at_mut(self.body.start);
431        let (body, suffix) = buf.split_at_mut(self.body.end - self.body.start);
432        let mut body = [&mut body[..]];
433        f(prefix, body.as_fragmented_byte_slice(), suffix)
434    }
435}
436
437impl<B: AsRef<[u8]>> AsRef<[u8]> for Buf<B> {
438    fn as_ref(&self) -> &[u8] {
439        &self.buf.as_ref()[self.body.clone()]
440    }
441}
442
443impl<B: AsMut<[u8]>> AsMut<[u8]> for Buf<B> {
444    fn as_mut(&mut self) -> &mut [u8] {
445        &mut self.buf.as_mut()[self.body.clone()]
446    }
447}
448
449impl<B: AsRef<[u8]>> Buffer for Buf<B> {
450    fn parse_with_view<'a, ParseArgs, P: ParsablePacket<&'a [u8], ParseArgs>>(
451        &'a mut self,
452        args: ParseArgs,
453    ) -> Result<(P, &'a [u8]), P::Error> {
454        let Self { body, ref buf } = self;
455        let body_before = body.clone();
456        let view = BufView { buf: &buf.as_ref()[body.clone()], body };
457        P::parse(view, args).map(|r| (r, &buf.as_ref()[body_before]))
458    }
459}
460
461/// A [`BufferView`] into a [`Buf`].
462///
463/// A `BufView` is constructed by [`Buf::buffer_view`], and implements
464/// `BufferView`, providing a view into the `Buf` from which it was constructed.
465pub struct BufView<'a> {
466    buf: &'a [u8],
467    body: &'a mut Range<usize>,
468}
469
470impl<'a> BufferView<&'a [u8]> for BufView<'a> {
471    fn take_front(&mut self, n: usize) -> Option<&'a [u8]> {
472        if self.len() < n {
473            return None;
474        }
475        self.body.start += n;
476        Some(take_front(&mut self.buf, n))
477    }
478
479    fn take_back(&mut self, n: usize) -> Option<&'a [u8]> {
480        if self.len() < n {
481            return None;
482        }
483        self.body.end -= n;
484        Some(take_back(&mut self.buf, n))
485    }
486
487    fn into_rest(self) -> &'a [u8] {
488        self.buf
489    }
490}
491
492impl<'a> AsRef<[u8]> for BufView<'a> {
493    fn as_ref(&self) -> &[u8] {
494        self.buf
495    }
496}
497
498/// A [`BufferViewMut`] into a [`Buf`].
499///
500/// A `BufViewMut` is constructed by [`Buf::buffer_view_mut`], and implements
501/// `BufferViewMut`, providing a mutable view into the `Buf` from which it was
502/// constructed.
503pub struct BufViewMut<'a> {
504    buf: &'a mut [u8],
505    body: &'a mut Range<usize>,
506}
507
508impl<'a> BufferView<&'a mut [u8]> for BufViewMut<'a> {
509    fn take_front(&mut self, n: usize) -> Option<&'a mut [u8]> {
510        if self.len() < n {
511            return None;
512        }
513        self.body.start += n;
514        Some(take_front_mut(&mut self.buf, n))
515    }
516
517    fn take_back(&mut self, n: usize) -> Option<&'a mut [u8]> {
518        if self.len() < n {
519            return None;
520        }
521        self.body.end -= n;
522        Some(take_back_mut(&mut self.buf, n))
523    }
524
525    fn into_rest(self) -> &'a mut [u8] {
526        self.buf
527    }
528}
529
530impl<'a> BufferViewMut<&'a mut [u8]> for BufViewMut<'a> {}
531
532impl<'a> AsRef<[u8]> for BufViewMut<'a> {
533    fn as_ref(&self) -> &[u8] {
534        self.buf
535    }
536}
537
538impl<'a> AsMut<[u8]> for BufViewMut<'a> {
539    fn as_mut(&mut self) -> &mut [u8] {
540        self.buf
541    }
542}
543
544/// The constraints required by a [`PacketBuilder`].
545///
546/// `PacketConstraints` represents the constraints that must be satisfied in
547/// order to serialize a `PacketBuilder`.
548///
549/// A `PacketConstraints`, `c`, guarantees two properties:
550/// - `c.max_body_len() >= c.min_body_len()`
551/// - `c.header_len() + c.min_body_len() + c.footer_len()` does not overflow
552///   `usize`
553///
554/// It is not possible (using safe code) to obtain a `PacketConstraints` which
555/// violates these properties, so code may rely for its correctness on the
556/// assumption that these properties hold.
557#[derive(Copy, Clone, Debug, Eq, PartialEq)]
558pub struct PacketConstraints {
559    header_len: usize,
560    footer_len: usize,
561    min_body_len: usize,
562    max_body_len: usize,
563}
564
565impl PacketConstraints {
566    /// A no-op `PacketConstraints` which does not add any constraints - there
567    /// is no header, footer, minimum body length requirement, or maximum body
568    /// length requirement.
569    pub const UNCONSTRAINED: Self =
570        Self { header_len: 0, footer_len: 0, min_body_len: 0, max_body_len: usize::MAX };
571
572    /// Constructs a new `PacketConstraints`.
573    ///
574    /// # Panics
575    ///
576    /// `new` panics if the arguments violate the validity properties of
577    /// `PacketConstraints` - if `max_body_len < min_body_len`, or if
578    /// `header_len + min_body_len + footer_len` overflows `usize`.
579    #[inline]
580    pub fn new(
581        header_len: usize,
582        footer_len: usize,
583        min_body_len: usize,
584        max_body_len: usize,
585    ) -> PacketConstraints {
586        PacketConstraints::try_new(header_len, footer_len, min_body_len, max_body_len).expect(
587            "max_body_len < min_body_len or header_len + min_body_len + footer_len overflows usize",
588        )
589    }
590
591    /// Tries to construct a new `PacketConstraints`.
592    ///
593    /// `new` returns `None` if the provided values violate the validity
594    /// properties of `PacketConstraints` - if `max_body_len < min_body_len`, or
595    /// if `header_len + min_body_len + footer_len` overflows `usize`.
596    #[inline]
597    pub fn try_new(
598        header_len: usize,
599        footer_len: usize,
600        min_body_len: usize,
601        max_body_len: usize,
602    ) -> Option<PacketConstraints> {
603        // Test case 3 in test_packet_constraints
604        let header_min_body_footer_overflows = header_len
605            .checked_add(min_body_len)
606            .and_then(|sum| sum.checked_add(footer_len))
607            .is_none();
608        // Test case 5 in test_packet_constraints
609        let max_less_than_min = max_body_len < min_body_len;
610        if max_less_than_min || header_min_body_footer_overflows {
611            return None;
612        }
613        Some(PacketConstraints { header_len, footer_len, min_body_len, max_body_len })
614    }
615
616    /// Constructs a new `PacketConstraints` with a given `max_body_len`.
617    ///
618    /// The `header_len`, `footer_len`, and `min_body_len` are all `0`.
619    #[inline]
620    pub fn with_max_body_len(max_body_len: usize) -> PacketConstraints {
621        // SAFETY:
622        // - `max_body_len >= min_body_len` by construction
623        // - `header_len + min_body_len + footer_len` is 0 and thus does not
624        //   overflow `usize`
625        PacketConstraints { header_len: 0, footer_len: 0, min_body_len: 0, max_body_len }
626    }
627
628    /// The number of bytes in this packet's header.
629    #[inline]
630    pub fn header_len(&self) -> usize {
631        self.header_len
632    }
633
634    /// The number of bytes in this packet's footer.
635    #[inline]
636    pub fn footer_len(&self) -> usize {
637        self.footer_len
638    }
639
640    /// The minimum body length (in bytes) required by this packet in order to
641    /// avoid adding padding.
642    ///
643    /// `min_body_len` returns the minimum number of body bytes required in
644    /// order to avoid adding padding. Note that, if padding bytes are required,
645    /// they may not necessarily belong immediately following the body,
646    /// depending on which packet layer imposes the minimum. In particular, in a
647    /// nested packet, padding goes after the body of the layer which imposes
648    /// the minimum. This means that, if the layer that imposes the minimum is
649    /// not the innermost one, then padding must be added not after the
650    /// innermost body, but instead in between footers.
651    /// [`NestedPacketBuilder::serialize_into`] is responsible for inserting
652    /// padding when serializing nested packets.
653    ///
654    /// If there is no minimum body length, this returns 0.
655    #[inline]
656    pub fn min_body_len(&self) -> usize {
657        self.min_body_len
658    }
659
660    /// The maximum length (in bytes) of a body allowed by this packet.
661    ///
662    /// If there is no maximum body length, this returns [`core::usize::MAX`].
663    #[inline]
664    pub fn max_body_len(&self) -> usize {
665        self.max_body_len
666    }
667
668    /// Attempts to encapsulate `self` in `outer`.
669    ///
670    /// Upon success, `try_encapsulate` returns a `PacketConstraints` which
671    /// represents the encapsulation of `self` in `outer`. Its header length,
672    /// footer length, minimum body length, and maximum body length are set
673    /// accordingly.
674    ///
675    /// This is probably not the method you want to use; consider
676    /// [`Serializer::encapsulate`] instead.
677    pub fn try_encapsulate(&self, outer: &Self) -> Option<PacketConstraints> {
678        let inner = self;
679        // Test case 1 in test_packet_constraints
680        let header_len = inner.header_len.checked_add(outer.header_len)?;
681        // Test case 2 in test_packet_constraints
682        let footer_len = inner.footer_len.checked_add(outer.footer_len)?;
683        // This is guaranteed not to overflow by the invariants on
684        // PacketConstraint.
685        let inner_header_footer_len = inner.header_len + inner.footer_len;
686        // Note the saturating_sub here - it's OK if the inner PacketBuilder
687        // more than satisfies the outer PacketBuilder's minimum body length
688        // requirement.
689        let min_body_len = cmp::max(
690            outer.min_body_len.saturating_sub(inner_header_footer_len),
691            inner.min_body_len,
692        );
693        // Note the checked_sub here - it's NOT OK if the inner PacketBuilder
694        // exceeds the outer PacketBuilder's maximum body length requirement.
695        //
696        // Test case 4 in test_packet_constraints
697        let max_body_len =
698            cmp::min(outer.max_body_len.checked_sub(inner_header_footer_len)?, inner.max_body_len);
699        // It's still possible that `min_body_len > max_body_len` or that
700        // `header_len + min_body_len + footer_len` overflows `usize`; `try_new`
701        // checks those constraints for us.
702        PacketConstraints::try_new(header_len, footer_len, min_body_len, max_body_len)
703    }
704}
705
706/// The target buffers into which [`PacketBuilder::serialize`] serializes its
707/// header and footer.
708pub struct SerializeTarget<'a> {
709    #[allow(missing_docs)]
710    pub header: &'a mut [u8],
711    #[allow(missing_docs)]
712    pub footer: &'a mut [u8],
713}
714
715/// A builder capable of serializing a packet's headers and footers.
716///
717/// A `PacketBuilder` describes a packet's headers and footers, and is capable
718/// of serializing the header and the footer into an existing buffer via the
719/// `serialize` method. A `PacketBuilder` never describes a body.
720/// [`PacketBuilder::wrap_body`] must be used to create a packet serializer
721/// for a whole packet.
722///
723/// `()` may be used as an "empty" `PacketBuilder` with no header, footer,
724/// minimum body length requirement, or maximum body length requirement.
725pub trait PacketBuilder: Sized {
726    /// Gets the constraints for this `PacketBuilder`.
727    fn constraints(&self) -> PacketConstraints;
728
729    /// Serializes this packet into an existing buffer.
730    ///
731    /// *This method is usually called by this crate during the serialization of
732    /// a [`Serializer`], not directly by the user.*
733    ///
734    /// # Preconditions
735    ///
736    /// The caller is responsible for initializing `body` with the body to be
737    /// encapsulated, and for ensuring that the body satisfies both the minimum
738    /// and maximum body length requirements, possibly by adding padding or by
739    /// truncating the body.
740    ///
741    /// # Postconditions
742    ///
743    /// `serialize` is responsible for serializing its header and footer into
744    /// `target.header` and `target.footer` respectively.
745    ///
746    /// # Security
747    ///
748    /// `serialize` must initialize the bytes of the header and footer, even if
749    /// only to zero, in order to avoid leaking the contents of packets
750    /// previously stored in the same buffer.
751    ///
752    /// # Panics
753    ///
754    /// May panic if the `target.header` or `target.footer` are not large enough
755    /// to fit the packet's header and footer respectively, or if the body does
756    /// not satisfy the minimum or maximum body length requirements.
757    fn serialize(&self, target: &mut SerializeTarget<'_>, body: FragmentedBytesMut<'_, '_>);
758
759    /// Wraps given packet `body` in this packet.
760    ///
761    /// Consumes the [`PacketBuilder`] and the `body`. If the `body` implements
762    /// `Serializer` then the result implement `Serializer` as well.
763    #[inline]
764    fn wrap_body<B>(self, body: B) -> Nested<B, Self> {
765        Nested { inner: body, outer: self }
766    }
767}
768
769impl<'a, B: PacketBuilder> PacketBuilder for &'a B {
770    #[inline]
771    fn constraints(&self) -> PacketConstraints {
772        B::constraints(self)
773    }
774    #[inline]
775    fn serialize(&self, target: &mut SerializeTarget<'_>, body: FragmentedBytesMut<'_, '_>) {
776        B::serialize(self, target, body)
777    }
778}
779
780impl<'a, B: PacketBuilder> PacketBuilder for &'a mut B {
781    #[inline]
782    fn constraints(&self) -> PacketConstraints {
783        B::constraints(self)
784    }
785    #[inline]
786    fn serialize(&self, target: &mut SerializeTarget<'_>, body: FragmentedBytesMut<'_, '_>) {
787        B::serialize(self, target, body)
788    }
789}
790
791impl PacketBuilder for () {
792    #[inline]
793    fn constraints(&self) -> PacketConstraints {
794        PacketConstraints::UNCONSTRAINED
795    }
796    #[inline]
797    fn serialize(&self, _target: &mut SerializeTarget<'_>, _body: FragmentedBytesMut<'_, '_>) {}
798}
799
800impl PacketBuilder for Never {
801    fn constraints(&self) -> PacketConstraints {
802        match *self {}
803    }
804    fn serialize(&self, _target: &mut SerializeTarget<'_>, _body: FragmentedBytesMut<'_, '_>) {}
805}
806
807/// One object encapsulated in another one.
808///
809/// `Nested`s are constructed using the [`PacketBuilder::wrap_body`] and
810/// [`Serializer::wrap_in`] methods.
811///
812/// When `I: Serializer` and `O: PacketBuilder`, `Nested<I, O>` implements
813/// [`Serializer`].
814#[derive(Copy, Clone, Debug, Eq, PartialEq)]
815pub struct Nested<I, O> {
816    inner: I,
817    outer: O,
818}
819
820impl<I, O> Nested<I, O> {
821    /// Creates a new `Nestable` instance with the `outer` wrapping the
822    /// `inner`. Normally `outer` should be a `PacketBuilder` and `inner`
823    /// should implement either `Serializer` or `PartialSerializer`.
824    pub fn new(inner: I, outer: O) -> Nested<I, O> {
825        Nested { inner, outer }
826    }
827
828    /// Consumes this `Nested` and returns the inner object, discarding the
829    /// outer one.
830    #[inline]
831    pub fn into_inner(self) -> I {
832        self.inner
833    }
834
835    /// Consumes this `Nested` and returns the outer object, discarding the
836    /// inner one.
837    #[inline]
838    pub fn into_outer(self) -> O {
839        self.outer
840    }
841
842    #[inline]
843    pub fn inner(&self) -> &I {
844        &self.inner
845    }
846
847    #[inline]
848    pub fn inner_mut(&mut self) -> &mut I {
849        &mut self.inner
850    }
851
852    #[inline]
853    pub fn outer(&self) -> &O {
854        &self.outer
855    }
856
857    #[inline]
858    pub fn outer_mut(&mut self) -> &mut O {
859        &mut self.outer
860    }
861}
862
863/// A [`PacketBuilder`] which has no header or footer, but which imposes a
864/// maximum body length constraint.
865///
866/// `LimitedSizePacketBuilder`s are constructed using the
867/// [`Serializer::with_size_limit`] method.
868#[derive(Copy, Clone, Debug)]
869#[cfg_attr(test, derive(Eq, PartialEq))]
870pub struct LimitedSizePacketBuilder {
871    /// The maximum body length.
872    pub limit: usize,
873}
874
875impl PacketBuilder for LimitedSizePacketBuilder {
876    fn constraints(&self) -> PacketConstraints {
877        PacketConstraints::with_max_body_len(self.limit)
878    }
879
880    fn serialize(&self, _target: &mut SerializeTarget<'_>, _body: FragmentedBytesMut<'_, '_>) {}
881}
882
883/// A builder capable of serializing packets - which do not encapsulate other
884/// packets - into an existing buffer.
885///
886/// An `InnerPacketBuilder` describes a packet, and is capable of serializing
887/// that packet into an existing buffer via the `serialize` method. Unlike the
888/// [`PacketBuilder`] trait, it describes a packet which does not encapsulate
889/// other packets.
890///
891/// # Notable implementations
892///
893/// `InnerPacketBuilder` is implemented for `&[u8]`, `&mut [u8]`, and `Vec<u8>`
894/// by treating the contents of the slice/`Vec` as the contents of the packet to
895/// be serialized.
896pub trait InnerPacketBuilder {
897    /// The number of bytes consumed by this packet.
898    fn bytes_len(&self) -> usize;
899
900    /// Serializes this packet into an existing buffer.
901    ///
902    /// `serialize` is called with a buffer of length `self.bytes_len()`, and is
903    /// responsible for serializing the packet into the buffer.
904    ///
905    /// # Security
906    ///
907    /// All of the bytes of the buffer should be initialized, even if only to
908    /// zero, in order to avoid leaking the contents of packets previously
909    /// stored in the same buffer.
910    ///
911    /// # Panics
912    ///
913    /// May panic if `buffer.len() != self.bytes_len()`.
914    fn serialize(&self, buffer: &mut [u8]);
915
916    /// Converts this `InnerPacketBuilder` into a [`Serializer`].
917    ///
918    /// `into_serializer` is like [`into_serializer_with`], except that no
919    /// buffer is provided for reuse in serialization.
920    ///
921    /// [`into_serializer_with`]: InnerPacketBuilder::into_serializer_with
922    #[inline]
923    fn into_serializer(self) -> InnerSerializer<Self, EmptyBuf>
924    where
925        Self: Sized,
926    {
927        self.into_serializer_with(EmptyBuf)
928    }
929
930    /// Converts this `InnerPacketBuilder` into a [`Serializer`] with a buffer
931    /// that can be used for serialization.
932    ///
933    /// `into_serializer_with` consumes a buffer and converts `self` into a type
934    /// which implements `Serialize` by treating it as the innermost body to be
935    /// contained within any encapsulating [`PacketBuilder`]s. During
936    /// serialization, `buffer` will be provided to the [`BufferProvider`],
937    /// allowing it to reuse the buffer for serialization and avoid allocating a
938    /// new one if possible.
939    ///
940    /// `buffer` will have its body shrunk to be zero bytes before the
941    /// `InnerSerializer` is constructed.
942    fn into_serializer_with<B: ShrinkBuffer>(self, mut buffer: B) -> InnerSerializer<Self, B>
943    where
944        Self: Sized,
945    {
946        buffer.shrink_back_to(0);
947        InnerSerializer { inner: self, buffer }
948    }
949}
950
951impl<'a, I: InnerPacketBuilder> InnerPacketBuilder for &'a I {
952    #[inline]
953    fn bytes_len(&self) -> usize {
954        I::bytes_len(self)
955    }
956    #[inline]
957    fn serialize(&self, buffer: &mut [u8]) {
958        I::serialize(self, buffer)
959    }
960}
961impl<'a, I: InnerPacketBuilder> InnerPacketBuilder for &'a mut I {
962    #[inline]
963    fn bytes_len(&self) -> usize {
964        I::bytes_len(self)
965    }
966    #[inline]
967    fn serialize(&self, buffer: &mut [u8]) {
968        I::serialize(self, buffer)
969    }
970}
971impl<'a> InnerPacketBuilder for &'a [u8] {
972    #[inline]
973    fn bytes_len(&self) -> usize {
974        self.len()
975    }
976    #[inline]
977    fn serialize(&self, buffer: &mut [u8]) {
978        buffer.copy_from_slice(self);
979    }
980}
981impl<'a> InnerPacketBuilder for &'a mut [u8] {
982    #[inline]
983    fn bytes_len(&self) -> usize {
984        self.len()
985    }
986    #[inline]
987    fn serialize(&self, buffer: &mut [u8]) {
988        buffer.copy_from_slice(self);
989    }
990}
991impl<'a> InnerPacketBuilder for Vec<u8> {
992    #[inline]
993    fn bytes_len(&self) -> usize {
994        self.len()
995    }
996    #[inline]
997    fn serialize(&self, buffer: &mut [u8]) {
998        buffer.copy_from_slice(self.as_slice());
999    }
1000}
1001impl<const N: usize> InnerPacketBuilder for ArrayVec<u8, N> {
1002    fn bytes_len(&self) -> usize {
1003        self.as_slice().bytes_len()
1004    }
1005    fn serialize(&self, buffer: &mut [u8]) {
1006        self.as_slice().serialize(buffer);
1007    }
1008}
1009
1010/// An [`InnerPacketBuilder`] created from any [`B: SplitByteSlice`].
1011///
1012/// `ByteSliceInnerPacketBuilder<B>` implements `InnerPacketBuilder` so long as
1013/// `B: SplitByteSlice`.
1014///
1015/// [`B: SplitByteSlice`]: zerocopy::SplitByteSlice
1016pub struct ByteSliceInnerPacketBuilder<B>(pub B);
1017
1018impl<B: SplitByteSlice> InnerPacketBuilder for ByteSliceInnerPacketBuilder<B> {
1019    fn bytes_len(&self) -> usize {
1020        self.0.deref().bytes_len()
1021    }
1022    fn serialize(&self, buffer: &mut [u8]) {
1023        self.0.deref().serialize(buffer)
1024    }
1025}
1026
1027impl<B: SplitByteSlice> Debug for ByteSliceInnerPacketBuilder<B> {
1028    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
1029        write!(f, "ByteSliceInnerPacketBuilder({:?})", self.0.as_ref())
1030    }
1031}
1032
1033/// An error in serializing a packet.
1034///
1035/// `SerializeError` is the type of errors returned from methods on the
1036/// [`Serializer`] trait. The `Alloc` variant indicates that a new buffer could
1037/// not be allocated, while the `SizeLimitExceeded` variant indicates that a
1038/// size limit constraint was exceeded.
1039#[derive(Copy, Clone, Debug, Eq, PartialEq)]
1040pub enum SerializeError<A> {
1041    /// A new buffer could not be allocated.
1042    Alloc(A),
1043    /// The size limit constraint was exceeded.
1044    SizeLimitExceeded,
1045}
1046
1047impl<A> SerializeError<A> {
1048    /// Is this `SerializeError::Alloc`?
1049    #[inline]
1050    pub fn is_alloc(&self) -> bool {
1051        match self {
1052            SerializeError::Alloc(_) => true,
1053            SerializeError::SizeLimitExceeded => false,
1054        }
1055    }
1056
1057    /// Is this `SerializeError::SizeLimitExceeded`?
1058    #[inline]
1059    pub fn is_size_limit_exceeded(&self) -> bool {
1060        match self {
1061            SerializeError::Alloc(_) => false,
1062            SerializeError::SizeLimitExceeded => true,
1063        }
1064    }
1065
1066    /// Maps the [`SerializeError::Alloc`] error type.
1067    pub fn map_alloc<T, F: FnOnce(A) -> T>(self, f: F) -> SerializeError<T> {
1068        match self {
1069            SerializeError::Alloc(a) => SerializeError::Alloc(f(a)),
1070            SerializeError::SizeLimitExceeded => SerializeError::SizeLimitExceeded,
1071        }
1072    }
1073}
1074
1075impl<A> From<A> for SerializeError<A> {
1076    fn from(a: A) -> SerializeError<A> {
1077        SerializeError::Alloc(a)
1078    }
1079}
1080
1081/// The error returned when a buffer is too short to hold a serialized packet,
1082/// and the [`BufferProvider`] is incapable of allocating a new one.
1083///
1084/// `BufferTooShortError` is returned by the [`Serializer`] methods
1085/// [`serialize_no_alloc`] and [`serialize_no_alloc_outer`].
1086///
1087/// [`serialize_no_alloc`]: Serializer::serialize_no_alloc
1088/// [`serialize_no_alloc_outer`]: Serializer::serialize_no_alloc_outer
1089#[derive(Copy, Clone, Debug, Eq, PartialEq)]
1090pub struct BufferTooShortError;
1091
1092/// An object capable of providing buffers which satisfy certain constraints.
1093///
1094/// A `BufferProvider<Input, Output>` is an object which is capable of consuming
1095/// a buffer of type `Input` and, either by reusing it or by allocating a new
1096/// one and copying the input buffer's body into it, producing a buffer of type
1097/// `Output` which meets certain prefix and suffix length constraints.
1098///
1099/// A `BufferProvider` must always be provided when serializing a
1100/// [`Serializer`].
1101///
1102/// Implementors may find the helper function [`try_reuse_buffer`] useful.
1103///
1104/// For clients who don't need the full expressive power of this trait, the
1105/// simpler [`BufferAlloc`] trait is provided. It only defines how to allocate
1106/// new buffers, and two blanket impls of `BufferProvider` are provided for all
1107/// `BufferAlloc` types.
1108pub trait BufferProvider<Input, Output> {
1109    /// The type of errors returned from [`reuse_or_realloc`].
1110    ///
1111    /// [`reuse_or_realloc`]: BufferProvider::reuse_or_realloc
1112    type Error;
1113
1114    /// Attempts to produce an output buffer with the given constraints by
1115    /// allocating a new one.
1116    ///
1117    /// `alloc_no_reuse` produces a new buffer with the following invariants:
1118    /// - The output buffer must have at least `prefix` bytes of prefix
1119    /// - The output buffer must have at least `suffix` bytes of suffix
1120    /// - The output buffer must have a body of length `body` bytes.
1121    ///
1122    /// If these requirements cannot be met, then an error is returned.
1123    fn alloc_no_reuse(
1124        self,
1125        prefix: usize,
1126        body: usize,
1127        suffix: usize,
1128    ) -> Result<Output, Self::Error>;
1129
1130    /// Consumes an input buffer and attempts to produce an output buffer with
1131    /// the given constraints, either by reusing the input buffer or by
1132    /// allocating a new one and copying the body into it.
1133    ///
1134    /// `reuse_or_realloc` consumes a buffer by value, and produces a new buffer
1135    /// with the following invariants:
1136    /// - The output buffer must have at least `prefix` bytes of prefix
1137    /// - The output buffer must have at least `suffix` bytes of suffix
1138    /// - The output buffer must have the same body as the input buffer
1139    ///
1140    /// If these requirements cannot be met, then an error is returned along
1141    /// with the input buffer, which is unmodified.
1142    fn reuse_or_realloc(
1143        self,
1144        buffer: Input,
1145        prefix: usize,
1146        suffix: usize,
1147    ) -> Result<Output, (Self::Error, Input)>;
1148}
1149
1150/// An object capable of allocating new buffers.
1151///
1152/// A `BufferAlloc<Output>` is an object which is capable of allocating new
1153/// buffers of type `Output`.
1154///
1155/// [Two blanket implementations] of [`BufferProvider`] are given for any type
1156/// which implements `BufferAlloc<O>`. One blanket implementation works for any
1157/// input buffer type, `I`, and produces buffers of type `Either<I, O>` as
1158/// output. One blanket implementation works only when the input and output
1159/// buffer types are the same, and produces buffers of that type. See the
1160/// documentation on those impls for more details.
1161///
1162/// The following implementations of `BufferAlloc` are provided:
1163/// - Any `FnOnce(usize) -> Result<O, E>` implements `BufferAlloc<O, Error = E>`
1164/// - `()` implements `BufferAlloc<Never, Error = ()>` (an allocator which
1165///   always fails)
1166/// - [`new_buf_vec`] implements `BufferAlloc<Buf<Vec<u8>>, Error = Never>` (an
1167///   allocator which infallibly heap-allocates `Vec`s)
1168///
1169/// [Two blanket implementations]: trait.BufferProvider.html#implementors
1170pub trait BufferAlloc<Output> {
1171    /// The type of errors returned from [`alloc`].
1172    ///
1173    /// [`alloc`]: BufferAlloc::alloc
1174    type Error;
1175
1176    /// Attempts to allocate a new buffer of size `len`.
1177    fn alloc(self, len: usize) -> Result<Output, Self::Error>;
1178}
1179
1180impl<O, E, F: FnOnce(usize) -> Result<O, E>> BufferAlloc<O> for F {
1181    type Error = E;
1182
1183    #[inline]
1184    fn alloc(self, len: usize) -> Result<O, E> {
1185        self(len)
1186    }
1187}
1188
1189impl BufferAlloc<Never> for () {
1190    type Error = ();
1191
1192    #[inline]
1193    fn alloc(self, _len: usize) -> Result<Never, ()> {
1194        Err(())
1195    }
1196}
1197
1198/// Allocates a new `Buf<Vec<u8>>`.
1199///
1200/// `new_buf_vec(len)` is shorthand for `Ok(Buf::new(vec![0; len], ..))`. It
1201/// implements [`BufferAlloc<Buf<Vec<u8>>, Error = Never>`], and, thanks to a
1202/// blanket impl, [`BufferProvider<I, Either<I, Buf<Vec<u8>>>, Error = Never>`]
1203/// for all `I: BufferMut`, and `BufferProvider<Buf<Vec<u8>>, Buf<Vec<u8>>,
1204/// Error = Never>`.
1205///
1206/// [`BufferAlloc<Buf<Vec<u8>>, Error = Never>`]: BufferAlloc
1207/// [`BufferProvider<I, Either<I, Buf<Vec<u8>>>, Error = Never>`]: BufferProvider
1208pub fn new_buf_vec(len: usize) -> Result<Buf<Vec<u8>>, Never> {
1209    Ok(Buf::new(vec![0; len], ..))
1210}
1211
1212/// Attempts to reuse a buffer for the purposes of implementing
1213/// [`BufferProvider::reuse_or_realloc`].
1214///
1215/// `try_reuse_buffer` attempts to reuse an existing buffer to satisfy the given
1216/// prefix and suffix constraints. If it succeeds, it returns `Ok` containing a
1217/// buffer with the same body as the input, and with at least `prefix` prefix
1218/// bytes and at least `suffix` suffix bytes. Otherwise, it returns `Err`
1219/// containing the original, unmodified input buffer.
1220///
1221/// Concretely, `try_reuse_buffer` has the following behavior:
1222/// - If the prefix and suffix constraints are already met, it returns `Ok` with
1223///   the input unmodified
1224/// - If the prefix and suffix constraints are not yet met, then...
1225///   - If there is enough capacity to meet the constraints and the body is not
1226///     larger than `max_copy_bytes`, the body will be moved within the buffer
1227///     in order to meet the constraints, and it will be returned
1228///   - Otherwise, if there is not enough capacity or the body is larger than
1229///     `max_copy_bytes`, it returns `Err` with the input unmodified
1230///
1231/// `max_copy_bytes` is meant to be an estimate of how many bytes can be copied
1232/// before allocating a new buffer will be cheaper than copying.
1233#[inline]
1234pub fn try_reuse_buffer<B: GrowBufferMut + ShrinkBuffer>(
1235    mut buffer: B,
1236    prefix: usize,
1237    suffix: usize,
1238    max_copy_bytes: usize,
1239) -> Result<B, B> {
1240    let need_prefix = prefix;
1241    let need_suffix = suffix;
1242    let have_prefix = buffer.prefix_len();
1243    let have_body = buffer.len();
1244    let have_suffix = buffer.suffix_len();
1245    let need_capacity = need_prefix + have_body + need_suffix;
1246
1247    if have_prefix >= need_prefix && have_suffix >= need_suffix {
1248        // We already satisfy the prefix and suffix requirements.
1249        Ok(buffer)
1250    } else if buffer.capacity() >= need_capacity && have_body <= max_copy_bytes {
1251        // The buffer is large enough, but the body is currently too far
1252        // forward or too far backwards to satisfy the prefix or suffix
1253        // requirements, so we need to move the body within the buffer.
1254        buffer.reset();
1255
1256        // Copy the original body range to a point starting immediatley
1257        // after `prefix`. This satisfies the `prefix` constraint by
1258        // definition, and satisfies the `suffix` constraint since we know
1259        // that the total buffer capacity is sufficient to hold the total
1260        // length of the prefix, body, and suffix.
1261        buffer.copy_within(have_prefix..(have_prefix + have_body), need_prefix);
1262        buffer.shrink(need_prefix..(need_prefix + have_body));
1263        debug_assert_eq!(buffer.prefix_len(), need_prefix);
1264        debug_assert!(buffer.suffix_len() >= need_suffix);
1265        debug_assert_eq!(buffer.len(), have_body);
1266        Ok(buffer)
1267    } else {
1268        Err(buffer)
1269    }
1270}
1271
1272/// Provides an implementation of [`BufferProvider`] from a [`BufferAlloc`] `A`
1273/// that attempts to reuse the input buffer and falls back to the allocator if
1274/// the input buffer can't be reused.
1275pub struct MaybeReuseBufferProvider<A>(pub A);
1276
1277impl<I: ReusableBuffer, O: ReusableBuffer, A: BufferAlloc<O>> BufferProvider<I, Either<I, O>>
1278    for MaybeReuseBufferProvider<A>
1279{
1280    type Error = A::Error;
1281
1282    fn alloc_no_reuse(
1283        self,
1284        prefix: usize,
1285        body: usize,
1286        suffix: usize,
1287    ) -> Result<Either<I, O>, Self::Error> {
1288        let Self(alloc) = self;
1289        let need_capacity = prefix + body + suffix;
1290        BufferAlloc::alloc(alloc, need_capacity).map(|mut buf| {
1291            buf.shrink(prefix..(prefix + body));
1292            Either::B(buf)
1293        })
1294    }
1295
1296    /// If `buffer` has enough capacity to store `need_prefix + need_suffix +
1297    /// buffer.len()` bytes, then reuse `buffer`. Otherwise, allocate a new
1298    /// buffer using `A`'s [`BufferAlloc`] implementation.
1299    ///
1300    /// If there is enough capacity, but the body is too far forwards or
1301    /// backwards in the buffer to satisfy the prefix and suffix constraints,
1302    /// the body will be moved within the buffer in order to satisfy the
1303    /// constraints. This operation is linear in the length of the body.
1304    #[inline]
1305    fn reuse_or_realloc(
1306        self,
1307        buffer: I,
1308        need_prefix: usize,
1309        need_suffix: usize,
1310    ) -> Result<Either<I, O>, (A::Error, I)> {
1311        // TODO(joshlf): Maybe it's worth coming up with a heuristic for when
1312        // moving the body is likely to be more expensive than allocating
1313        // (rather than just using `usize::MAX`)? This will be tough since we
1314        // don't know anything about the performance of `A::alloc`.
1315        match try_reuse_buffer(buffer, need_prefix, need_suffix, usize::MAX) {
1316            Ok(buffer) => Ok(Either::A(buffer)),
1317            Err(buffer) => {
1318                let have_body = buffer.len();
1319                let mut buf = match BufferProvider::<I, Either<I, O>>::alloc_no_reuse(
1320                    self,
1321                    need_prefix,
1322                    have_body,
1323                    need_suffix,
1324                ) {
1325                    Ok(buf) => buf,
1326                    Err(err) => return Err((err, buffer)),
1327                };
1328
1329                buf.copy_from(&buffer);
1330                debug_assert_eq!(buf.prefix_len(), need_prefix);
1331                debug_assert!(buf.suffix_len() >= need_suffix);
1332                debug_assert_eq!(buf.len(), have_body);
1333                Ok(buf)
1334            }
1335        }
1336    }
1337}
1338
1339impl<B: ReusableBuffer, A: BufferAlloc<B>> BufferProvider<B, B> for MaybeReuseBufferProvider<A> {
1340    type Error = A::Error;
1341
1342    fn alloc_no_reuse(self, prefix: usize, body: usize, suffix: usize) -> Result<B, Self::Error> {
1343        BufferProvider::<B, Either<B, B>>::alloc_no_reuse(self, prefix, body, suffix)
1344            .map(Either::into_inner)
1345    }
1346
1347    /// If `buffer` has enough capacity to store `need_prefix + need_suffix +
1348    /// buffer.len()` bytes, then reuse `buffer`. Otherwise, allocate a new
1349    /// buffer using `A`'s [`BufferAlloc`] implementation.
1350    ///
1351    /// If there is enough capacity, but the body is too far forwards or
1352    /// backwards in the buffer to satisfy the prefix and suffix constraints,
1353    /// the body will be moved within the buffer in order to satisfy the
1354    /// constraints. This operation is linear in the length of the body.
1355    #[inline]
1356    fn reuse_or_realloc(self, buffer: B, prefix: usize, suffix: usize) -> Result<B, (A::Error, B)> {
1357        BufferProvider::<B, Either<B, B>>::reuse_or_realloc(self, buffer, prefix, suffix)
1358            .map(Either::into_inner)
1359    }
1360}
1361
1362/// Provides an implementation of [`BufferProvider`] from a [`BufferAlloc`] `A`
1363/// that never attempts to reuse the input buffer, and always create a new
1364/// buffer from the allocator `A`.
1365pub struct NoReuseBufferProvider<A>(pub A);
1366
1367impl<I: FragmentedBuffer, O: ReusableBuffer, A: BufferAlloc<O>> BufferProvider<I, O>
1368    for NoReuseBufferProvider<A>
1369{
1370    type Error = A::Error;
1371
1372    fn alloc_no_reuse(self, prefix: usize, body: usize, suffix: usize) -> Result<O, A::Error> {
1373        let Self(alloc) = self;
1374        alloc.alloc(prefix + body + suffix).map(|mut b| {
1375            b.shrink(prefix..prefix + body);
1376            b
1377        })
1378    }
1379
1380    fn reuse_or_realloc(self, buffer: I, prefix: usize, suffix: usize) -> Result<O, (A::Error, I)> {
1381        BufferProvider::<I, O>::alloc_no_reuse(self, prefix, buffer.len(), suffix)
1382            .map(|mut b| {
1383                b.copy_from(&buffer);
1384                b
1385            })
1386            .map_err(|e| (e, buffer))
1387    }
1388}
1389
1390pub trait Serializer: Sized {
1391    /// The type of buffers returned from serialization methods on this trait.
1392    type Buffer;
1393
1394    /// Serializes this `Serializer`, producing a buffer.
1395    ///
1396    /// As `Serializer`s can be nested using the [`Nested`] type (constructed
1397    /// using [`PacketBuilder::wrap_body`] and [`Serializer::wrap_in`]), the
1398    /// `serialize` method is recursive - calling it on a `Nested` will recurse
1399    /// into the inner `Serializer`, which might itself be a `Nested`, and so
1400    /// on. When the innermost `Serializer` is reached, the contained buffer is
1401    /// passed to the `provider`, allowing it to decide how to produce a buffer
1402    /// which is large enough to fit the entire packet - either by reusing the
1403    /// existing buffer, or by discarding it and allocating a new one. `outer`
1404    /// specifies [`PacketConstraints`] for the outer parts of the packet
1405    /// (header and footer).
1406    fn serialize<B: GrowBufferMut, P: BufferProvider<Self::Buffer, B>>(
1407        self,
1408        outer: PacketConstraints,
1409        provider: P,
1410    ) -> Result<B, (SerializeError<P::Error>, Self)>;
1411
1412    /// Serializes the data into a new buffer without consuming `self`.
1413    ///
1414    /// Creates a new buffer using `alloc` and serializes the data into that
1415    /// that new buffer. Unlike all other serialize methods,
1416    /// `serialize_new_buf` takes `self` by reference. This allows to use the
1417    /// same `Serializer` to serialize the data more than once.
1418    fn serialize_new_buf<B: ReusableBuffer, A: BufferAlloc<B>>(
1419        &self,
1420        outer: PacketConstraints,
1421        alloc: A,
1422    ) -> Result<B, SerializeError<A::Error>>;
1423
1424    /// Serializes this `Serializer`, allocating a [`Buf<Vec<u8>>`] if the
1425    /// contained buffer isn't large enough.
1426    ///
1427    /// `serialize_vec` is like [`serialize`], except that, if the contained
1428    /// buffer isn't large enough to contain the packet, a new `Vec<u8>` is
1429    /// allocated and wrapped in a [`Buf`]. If the buffer is large enough, but
1430    /// the body is too far forwards or backwards to fit the encapsulating
1431    /// headers or footers, the body will be moved within the buffer (this
1432    /// operation's cost is linear in the size of the body).
1433    ///
1434    /// `serialize_vec` is equivalent to calling `serialize` with
1435    /// [`new_buf_vec`] as the [`BufferProvider`].
1436    ///
1437    /// [`Buf<Vec<u8>>`]: Buf
1438    /// [`serialize`]: Serializer::serialize
1439    #[inline]
1440    #[allow(clippy::type_complexity)]
1441    fn serialize_vec(
1442        self,
1443        outer: PacketConstraints,
1444    ) -> Result<Either<Self::Buffer, Buf<Vec<u8>>>, (SerializeError<Never>, Self)>
1445    where
1446        Self::Buffer: ReusableBuffer,
1447    {
1448        self.serialize(outer, MaybeReuseBufferProvider(new_buf_vec))
1449    }
1450
1451    /// Serializes this `Serializer`, failing if the existing buffer is not
1452    /// large enough.
1453    ///
1454    /// `serialize_no_alloc` is like [`serialize`], except that it will fail if
1455    /// the existing buffer isn't large enough. If the buffer is large enough,
1456    /// but the body is too far forwards or backwards to fit the encapsulating
1457    /// headers or footers, the body will be moved within the buffer (this
1458    /// operation's cost is linear in the size of the body).
1459    ///
1460    /// `serialize_no_alloc` is equivalent to calling `serialize` with a
1461    /// `BufferProvider` which cannot allocate a new buffer (such as `()`).
1462    ///
1463    /// [`serialize`]: Serializer::serialize
1464    #[inline]
1465    fn serialize_no_alloc(
1466        self,
1467        outer: PacketConstraints,
1468    ) -> Result<Self::Buffer, (SerializeError<BufferTooShortError>, Self)>
1469    where
1470        Self::Buffer: ReusableBuffer,
1471    {
1472        self.serialize(outer, MaybeReuseBufferProvider(())).map(Either::into_a).map_err(
1473            |(err, slf)| {
1474                (
1475                    match err {
1476                        SerializeError::Alloc(()) => BufferTooShortError.into(),
1477                        SerializeError::SizeLimitExceeded => SerializeError::SizeLimitExceeded,
1478                    },
1479                    slf,
1480                )
1481            },
1482        )
1483    }
1484
1485    /// Serializes this `Serializer` as the outermost packet.
1486    ///
1487    /// `serialize_outer` is like [`serialize`], except that it is called when
1488    /// this `Serializer` describes the outermost packet, not encapsulated in
1489    /// any other packets. It is equivalent to calling `serialize` with an empty
1490    /// [`PacketBuilder`] (such as `()`).
1491    ///
1492    /// [`serialize`]: Serializer::serialize
1493    #[inline]
1494    fn serialize_outer<B: GrowBufferMut, P: BufferProvider<Self::Buffer, B>>(
1495        self,
1496        provider: P,
1497    ) -> Result<B, (SerializeError<P::Error>, Self)> {
1498        self.serialize(PacketConstraints::UNCONSTRAINED, provider)
1499    }
1500
1501    /// Serializes this `Serializer` as the outermost packet, allocating a
1502    /// [`Buf<Vec<u8>>`] if the contained buffer isn't large enough.
1503    ///
1504    /// `serialize_vec_outer` is like [`serialize_vec`], except that it is
1505    /// called when this `Serializer` describes the outermost packet, not
1506    /// encapsulated in any other packets. It is equivalent to calling
1507    /// `serialize_vec` with an empty [`PacketBuilder`] (such as `()`).
1508    ///
1509    /// [`Buf<Vec<u8>>`]: Buf
1510    /// [`serialize_vec`]: Serializer::serialize_vec
1511    #[inline]
1512    #[allow(clippy::type_complexity)]
1513    fn serialize_vec_outer(
1514        self,
1515    ) -> Result<Either<Self::Buffer, Buf<Vec<u8>>>, (SerializeError<Never>, Self)>
1516    where
1517        Self::Buffer: ReusableBuffer,
1518    {
1519        self.serialize_vec(PacketConstraints::UNCONSTRAINED)
1520    }
1521
1522    /// Serializes this `Serializer` as the outermost packet, failing if the
1523    /// existing buffer is not large enough.
1524    ///
1525    /// `serialize_no_alloc_outer` is like [`serialize_no_alloc`], except that
1526    /// it is called when this `Serializer` describes the outermost packet, not
1527    /// encapsulated in any other packets. It is equivalent to calling
1528    /// `serialize_no_alloc` with an empty [`PacketBuilder`] (such as `()`).
1529    ///
1530    /// [`serialize_no_alloc`]: Serializer::serialize_no_alloc
1531    #[inline]
1532    fn serialize_no_alloc_outer(
1533        self,
1534    ) -> Result<Self::Buffer, (SerializeError<BufferTooShortError>, Self)>
1535    where
1536        Self::Buffer: ReusableBuffer,
1537    {
1538        self.serialize_no_alloc(PacketConstraints::UNCONSTRAINED)
1539    }
1540
1541    /// Encapsulates this `Serializer` in a packet, producing a new
1542    /// `Serializer`.
1543    ///
1544    /// `wrap_in()` consumes this `Serializer` and a [`PacketBuilder`], and
1545    /// produces a new `Serializer` which describes encapsulating this one in
1546    /// the packet described by `outer`.
1547    #[inline]
1548    fn wrap_in<B: PacketBuilder>(self, outer: B) -> Nested<Self, B> {
1549        outer.wrap_body(self)
1550    }
1551
1552    /// DEPRECATED. [`PacketBuilder::wrap_body()`] or [`Serializer::wrap_in()`]
1553    /// should be used instead of this method.
1554    #[inline]
1555    fn encapsulate<B>(self, outer: B) -> Nested<Self, B> {
1556        Nested { inner: self, outer }
1557    }
1558
1559    /// Creates a new `Serializer` which will enforce a size limit.
1560    ///
1561    /// `with_size_limit` consumes this `Serializer` and limit, and produces a
1562    /// new `Serializer` which will enforce the given limit on all serialization
1563    /// requests. Note that the given limit will be enforced at this layer -
1564    /// serialization requests will be rejected if the body produced by the
1565    /// request at this layer would exceed the limit. It has no effect on
1566    /// headers or footers added by encapsulating layers outside of this one.
1567    #[inline]
1568    fn with_size_limit(self, limit: usize) -> Nested<Self, LimitedSizePacketBuilder> {
1569        self.wrap_in(LimitedSizePacketBuilder { limit })
1570    }
1571}
1572
1573/// A [`Serializer`] constructed from an [`InnerPacketBuilder`].
1574///
1575/// An `InnerSerializer` wraps an `InnerPacketBuilder` and a buffer, and
1576/// implements the `Serializer` trait. When a serialization is requested, it
1577/// either reuses the stored buffer or allocates a new one large enough to hold
1578/// itself and all outer `PacketBuilder`s.
1579#[derive(Copy, Clone, Debug, Eq, PartialEq)]
1580pub struct InnerSerializer<I, B> {
1581    inner: I,
1582    // The buffer's length must be zero since we encapsulate the buffer in a
1583    // PacketBuilder. If the length were non-zero, that would have the effect of
1584    // retaining the contents of the buffer when serializing, and putting them
1585    // immediately after the bytes of `inner`.
1586    buffer: B,
1587}
1588
1589impl<I, B> InnerSerializer<I, B> {
1590    pub fn inner(&self) -> &I {
1591        &self.inner
1592    }
1593}
1594
1595/// A wrapper for `InnerPacketBuilders` which implements `PacketBuilder` by
1596/// treating the entire `InnerPacketBuilder` as the header of the
1597/// `PacketBuilder`. This allows us to compose our InnerPacketBuilder with
1598/// the outer `PacketBuilders` into a single, large `PacketBuilder`, and then
1599/// serialize it using `self.buffer`.
1600struct InnerPacketBuilderWrapper<I>(I);
1601
1602impl<I: InnerPacketBuilder> PacketBuilder for InnerPacketBuilderWrapper<I> {
1603    fn constraints(&self) -> PacketConstraints {
1604        let Self(wrapped) = self;
1605        PacketConstraints::new(wrapped.bytes_len(), 0, 0, usize::MAX)
1606    }
1607
1608    fn serialize(&self, target: &mut SerializeTarget<'_>, _body: FragmentedBytesMut<'_, '_>) {
1609        let Self(wrapped) = self;
1610
1611        // Note that the body might be non-empty if an outer
1612        // PacketBuilder added a minimum body length constraint that
1613        // required padding.
1614        debug_assert_eq!(target.header.len(), wrapped.bytes_len());
1615        debug_assert_eq!(target.footer.len(), 0);
1616
1617        InnerPacketBuilder::serialize(wrapped, target.header);
1618    }
1619}
1620
1621impl<I: InnerPacketBuilder, B: GrowBuffer + ShrinkBuffer> Serializer for InnerSerializer<I, B> {
1622    type Buffer = B;
1623
1624    #[inline]
1625    fn serialize<BB: GrowBufferMut, P: BufferProvider<B, BB>>(
1626        self,
1627        outer: PacketConstraints,
1628        provider: P,
1629    ) -> Result<BB, (SerializeError<P::Error>, InnerSerializer<I, B>)> {
1630        debug_assert_eq!(self.buffer.len(), 0);
1631        InnerPacketBuilderWrapper(self.inner)
1632            .wrap_body(self.buffer)
1633            .serialize(outer, provider)
1634            .map_err(|(err, Nested { inner: buffer, outer: pb })| {
1635                (err, InnerSerializer { inner: pb.0, buffer })
1636            })
1637    }
1638
1639    #[inline]
1640    fn serialize_new_buf<BB: ReusableBuffer, A: BufferAlloc<BB>>(
1641        &self,
1642        outer: PacketConstraints,
1643        alloc: A,
1644    ) -> Result<BB, SerializeError<A::Error>> {
1645        InnerPacketBuilderWrapper(&self.inner).wrap_body(EmptyBuf).serialize_new_buf(outer, alloc)
1646    }
1647}
1648
1649impl<B: GrowBuffer + ShrinkBuffer> Serializer for B {
1650    type Buffer = B;
1651
1652    #[inline]
1653    fn serialize<BB: GrowBufferMut, P: BufferProvider<Self::Buffer, BB>>(
1654        self,
1655        outer: PacketConstraints,
1656        provider: P,
1657    ) -> Result<BB, (SerializeError<P::Error>, Self)> {
1658        TruncatingSerializer::new(self, TruncateDirection::NoTruncating)
1659            .serialize(outer, provider)
1660            .map_err(|(err, ser)| (err, ser.buffer))
1661    }
1662
1663    fn serialize_new_buf<BB: ReusableBuffer, A: BufferAlloc<BB>>(
1664        &self,
1665        outer: PacketConstraints,
1666        alloc: A,
1667    ) -> Result<BB, SerializeError<A::Error>> {
1668        if self.len() > outer.max_body_len() {
1669            return Err(SerializeError::SizeLimitExceeded);
1670        }
1671
1672        let padding = outer.min_body_len().saturating_sub(self.len());
1673        let tail_size = padding + outer.footer_len();
1674        let buffer_size = outer.header_len() + self.len() + tail_size;
1675        let mut buffer = alloc.alloc(buffer_size)?;
1676        buffer.shrink_front(outer.header_len());
1677        buffer.shrink_back(tail_size);
1678        buffer.copy_from(self);
1679        buffer.grow_back(padding);
1680        Ok(buffer)
1681    }
1682}
1683
1684/// Either of two serializers.
1685///
1686/// An `EitherSerializer` wraps one of two different serializer types.
1687pub enum EitherSerializer<A, B> {
1688    A(A),
1689    B(B),
1690}
1691
1692impl<A: Serializer, B: Serializer<Buffer = A::Buffer>> Serializer for EitherSerializer<A, B> {
1693    type Buffer = A::Buffer;
1694
1695    fn serialize<TB: GrowBufferMut, P: BufferProvider<Self::Buffer, TB>>(
1696        self,
1697        outer: PacketConstraints,
1698        provider: P,
1699    ) -> Result<TB, (SerializeError<P::Error>, Self)> {
1700        match self {
1701            EitherSerializer::A(s) => {
1702                s.serialize(outer, provider).map_err(|(err, s)| (err, EitherSerializer::A(s)))
1703            }
1704            EitherSerializer::B(s) => {
1705                s.serialize(outer, provider).map_err(|(err, s)| (err, EitherSerializer::B(s)))
1706            }
1707        }
1708    }
1709
1710    fn serialize_new_buf<TB: ReusableBuffer, BA: BufferAlloc<TB>>(
1711        &self,
1712        outer: PacketConstraints,
1713        alloc: BA,
1714    ) -> Result<TB, SerializeError<BA::Error>> {
1715        match self {
1716            EitherSerializer::A(s) => s.serialize_new_buf(outer, alloc),
1717            EitherSerializer::B(s) => s.serialize_new_buf(outer, alloc),
1718        }
1719    }
1720}
1721
1722/// The direction a buffer's body should be truncated from to force
1723/// it to fit within a size limit.
1724#[derive(Copy, Clone, Debug, Eq, PartialEq)]
1725pub enum TruncateDirection {
1726    /// If a buffer cannot fit within a limit, discard bytes from the
1727    /// front of the body.
1728    DiscardFront,
1729    /// If a buffer cannot fit within a limit, discard bytes from the
1730    /// end of the body.
1731    DiscardBack,
1732    /// Do not attempt to truncate a buffer to make it fit within a limit.
1733    NoTruncating,
1734}
1735
1736/// A [`Serializer`] that truncates its body if it would exceed a size limit.
1737///
1738/// `TruncatingSerializer` wraps a buffer, and implements `Serializer`. Unlike
1739/// the blanket impl of `Serializer` for `B: GrowBuffer + ShrinkBuffer`, if the
1740/// buffer's body exceeds the size limit constraint passed to
1741/// `Serializer::serialize`, the body is truncated to fit.
1742///
1743/// Note that this does not guarantee that size limit exceeded errors will not
1744/// occur. The size limit may be small enough that the encapsulating headers
1745/// alone exceed the size limit.  There may also be a minimum body length
1746/// constraint which is larger than the size limit.
1747#[derive(Copy, Clone, Debug, Eq, PartialEq)]
1748pub struct TruncatingSerializer<B> {
1749    buffer: B,
1750    direction: TruncateDirection,
1751}
1752
1753impl<B> TruncatingSerializer<B> {
1754    /// Constructs a new `TruncatingSerializer`.
1755    pub fn new(buffer: B, direction: TruncateDirection) -> TruncatingSerializer<B> {
1756        TruncatingSerializer { buffer, direction }
1757    }
1758
1759    /// Provides shared access to the inner buffer.
1760    pub fn buffer(&self) -> &B {
1761        &self.buffer
1762    }
1763
1764    /// Provides mutable access to the inner buffer.
1765    pub fn buffer_mut(&mut self) -> &mut B {
1766        &mut self.buffer
1767    }
1768}
1769
1770impl<B: GrowBuffer + ShrinkBuffer> Serializer for TruncatingSerializer<B> {
1771    type Buffer = B;
1772
1773    fn serialize<BB: GrowBufferMut, P: BufferProvider<B, BB>>(
1774        mut self,
1775        outer: PacketConstraints,
1776        provider: P,
1777    ) -> Result<BB, (SerializeError<P::Error>, Self)> {
1778        let original_len = self.buffer.len();
1779        let excess_bytes = if original_len > outer.max_body_len {
1780            Some(original_len - outer.max_body_len)
1781        } else {
1782            None
1783        };
1784        if let Some(excess_bytes) = excess_bytes {
1785            match self.direction {
1786                TruncateDirection::DiscardFront => self.buffer.shrink_front(excess_bytes),
1787                TruncateDirection::DiscardBack => self.buffer.shrink_back(excess_bytes),
1788                TruncateDirection::NoTruncating => {
1789                    return Err((SerializeError::SizeLimitExceeded, self))
1790                }
1791            }
1792        }
1793
1794        let padding = outer.min_body_len().saturating_sub(self.buffer.len());
1795
1796        // At this point, the body and padding MUST fit within the limit. Note
1797        // that PacketConstraints guarantees that min_body_len <= max_body_len,
1798        // so the padding can't cause this assertion to fail.
1799        debug_assert!(self.buffer.len() + padding <= outer.max_body_len());
1800        match provider.reuse_or_realloc(
1801            self.buffer,
1802            outer.header_len(),
1803            padding + outer.footer_len(),
1804        ) {
1805            Ok(buffer) => Ok(buffer),
1806            Err((err, mut buffer)) => {
1807                // Undo the effects of shrinking the buffer so that the buffer
1808                // we return is unmodified from its original (which is required
1809                // by the contract of this method).
1810                if let Some(excess_bytes) = excess_bytes {
1811                    match self.direction {
1812                        TruncateDirection::DiscardFront => buffer.grow_front(excess_bytes),
1813                        TruncateDirection::DiscardBack => buffer.grow_back(excess_bytes),
1814                        TruncateDirection::NoTruncating => unreachable!(),
1815                    }
1816                }
1817
1818                Err((
1819                    SerializeError::Alloc(err),
1820                    TruncatingSerializer { buffer, direction: self.direction },
1821                ))
1822            }
1823        }
1824    }
1825
1826    fn serialize_new_buf<BB: ReusableBuffer, A: BufferAlloc<BB>>(
1827        &self,
1828        outer: PacketConstraints,
1829        alloc: A,
1830    ) -> Result<BB, SerializeError<A::Error>> {
1831        let truncated_size = cmp::min(self.buffer.len(), outer.max_body_len());
1832        let discarded_bytes = self.buffer.len() - truncated_size;
1833        let padding = outer.min_body_len().saturating_sub(truncated_size);
1834        let tail_size = padding + outer.footer_len();
1835        let buffer_size = outer.header_len() + truncated_size + tail_size;
1836        let mut buffer = alloc.alloc(buffer_size)?;
1837        buffer.shrink_front(outer.header_len());
1838        buffer.shrink_back(tail_size);
1839        buffer.with_bytes_mut(|mut dst| {
1840            self.buffer.with_bytes(|src| {
1841                let src = match (discarded_bytes > 0, self.direction) {
1842                    (false, _) => src,
1843                    (true, TruncateDirection::DiscardFront) => src.slice(discarded_bytes..),
1844                    (true, TruncateDirection::DiscardBack) => src.slice(..truncated_size),
1845                    (true, TruncateDirection::NoTruncating) => {
1846                        return Err(SerializeError::SizeLimitExceeded)
1847                    }
1848                };
1849                dst.copy_from(&src);
1850                Ok(())
1851            })
1852        })?;
1853        buffer.grow_back_zero(padding);
1854        Ok(buffer)
1855    }
1856}
1857
1858impl<I: Serializer, O: PacketBuilder> Serializer for Nested<I, O> {
1859    type Buffer = I::Buffer;
1860
1861    #[inline]
1862    fn serialize<B: GrowBufferMut, P: BufferProvider<I::Buffer, B>>(
1863        self,
1864        outer: PacketConstraints,
1865        provider: P,
1866    ) -> Result<B, (SerializeError<P::Error>, Self)> {
1867        let Some(outer) = self.outer.constraints().try_encapsulate(&outer) else {
1868            return Err((SerializeError::SizeLimitExceeded, self));
1869        };
1870
1871        match self.inner.serialize(outer, provider) {
1872            Ok(mut buf) => {
1873                buf.serialize(&self.outer);
1874                Ok(buf)
1875            }
1876            Err((err, inner)) => Err((err, self.outer.wrap_body(inner))),
1877        }
1878    }
1879
1880    #[inline]
1881    fn serialize_new_buf<B: ReusableBuffer, A: BufferAlloc<B>>(
1882        &self,
1883        outer: PacketConstraints,
1884        alloc: A,
1885    ) -> Result<B, SerializeError<A::Error>> {
1886        let Some(outer) = self.outer.constraints().try_encapsulate(&outer) else {
1887            return Err(SerializeError::SizeLimitExceeded);
1888        };
1889
1890        let mut buf = self.inner.serialize_new_buf(outer, alloc)?;
1891        GrowBufferMut::serialize(&mut buf, &self.outer);
1892        Ok(buf)
1893    }
1894}
1895
1896/// A packet builder used for partial packet serialization.
1897pub trait PartialPacketBuilder: PacketBuilder {
1898    /// Serializes the header to the specified `buffer`.
1899    ///
1900    /// Checksums (if any) should not calculated. The corresponding fields
1901    /// should be set to 0.
1902    ///
1903    /// `body_len` specifies size of the packet body wrapped by this
1904    /// `PacketBuilder`. It is supplied so the correct packet size can be
1905    /// written in the header.
1906    fn partial_serialize(&self, body_len: usize, buffer: &mut [u8]);
1907}
1908
1909impl PartialPacketBuilder for () {
1910    fn partial_serialize(&self, _body_len: usize, _buffer: &mut [u8]) {}
1911}
1912
1913/// Result returned by `PartialSerializer::partial_serialize`.
1914#[derive(Debug, Eq, PartialEq)]
1915pub struct PartialSerializeResult {
1916    // Number of bytes written to the output buffer.
1917    pub bytes_written: usize,
1918
1919    // Size of the whole packet.
1920    pub total_size: usize,
1921}
1922
1923/// A serializer that supports partial serialization.
1924///
1925/// Partial serialization allows to serialize only packet headers without
1926/// calculating packet checksums (if any).
1927pub trait PartialSerializer: Sized {
1928    /// Serializes the head of the packet to the specified `buffer`.
1929    ///
1930    /// If the packet contains network or transport level headers that fit in
1931    /// the provided buffer then they will be serialized. Complete or partial
1932    /// body may be copied to the output buffer as well, depending on the
1933    /// serializer type.
1934    ///
1935    /// `PartialSerializeResult.bytes_written` indicates how many bytes were
1936    /// actually serialized.
1937    fn partial_serialize(
1938        &self,
1939        outer: PacketConstraints,
1940        buffer: &mut [u8],
1941    ) -> Result<PartialSerializeResult, SerializeError<Never>>;
1942}
1943
1944impl<B: GrowBuffer + ShrinkBuffer> PartialSerializer for B {
1945    fn partial_serialize(
1946        &self,
1947        _outer: PacketConstraints,
1948        _buffer: &mut [u8],
1949    ) -> Result<PartialSerializeResult, SerializeError<Never>> {
1950        Ok(PartialSerializeResult { bytes_written: 0, total_size: self.len() })
1951    }
1952}
1953
1954impl<B: GrowBuffer + ShrinkBuffer> PartialSerializer for TruncatingSerializer<B> {
1955    fn partial_serialize(
1956        &self,
1957        outer: PacketConstraints,
1958        _buffer: &mut [u8],
1959    ) -> Result<PartialSerializeResult, SerializeError<Never>> {
1960        let total_size =
1961            cmp::max(outer.min_body_len(), cmp::min(self.buffer().len(), outer.max_body_len()));
1962        Ok(PartialSerializeResult { bytes_written: 0, total_size })
1963    }
1964}
1965
1966impl<I: InnerPacketBuilder, B: GrowBuffer + ShrinkBuffer> PartialSerializer
1967    for InnerSerializer<I, B>
1968{
1969    fn partial_serialize(
1970        &self,
1971        outer: PacketConstraints,
1972        _buffer: &mut [u8],
1973    ) -> Result<PartialSerializeResult, SerializeError<Never>> {
1974        Ok(PartialSerializeResult {
1975            bytes_written: 0,
1976            total_size: cmp::max(self.inner().bytes_len(), outer.min_body_len()),
1977        })
1978    }
1979}
1980
1981impl<A: Serializer + PartialSerializer, B: Serializer + PartialSerializer> PartialSerializer
1982    for EitherSerializer<A, B>
1983{
1984    fn partial_serialize(
1985        &self,
1986        outer: PacketConstraints,
1987        buffer: &mut [u8],
1988    ) -> Result<PartialSerializeResult, SerializeError<Never>> {
1989        match self {
1990            EitherSerializer::A(s) => s.partial_serialize(outer, buffer),
1991            EitherSerializer::B(s) => s.partial_serialize(outer, buffer),
1992        }
1993    }
1994}
1995
1996impl<I: PartialSerializer, O: PartialPacketBuilder> PartialSerializer for Nested<I, O> {
1997    fn partial_serialize(
1998        &self,
1999        outer: PacketConstraints,
2000        buffer: &mut [u8],
2001    ) -> Result<PartialSerializeResult, SerializeError<Never>> {
2002        let header_constraints = self.outer.constraints();
2003        let Some(constraints) = outer.try_encapsulate(&header_constraints) else {
2004            return Err(SerializeError::SizeLimitExceeded);
2005        };
2006
2007        let header_len = header_constraints.header_len();
2008        let inner_buf = buffer.get_mut(header_len..).unwrap_or(&mut []);
2009        let mut result = self.inner.partial_serialize(constraints, inner_buf)?;
2010        if header_len <= buffer.len() {
2011            self.outer.partial_serialize(result.total_size, &mut buffer[..header_len]);
2012            result.bytes_written += header_len;
2013        }
2014        result.total_size += header_len + header_constraints.footer_len();
2015        Ok(result)
2016    }
2017}
2018
2019#[cfg(test)]
2020mod tests {
2021    use super::*;
2022    use crate::BufferMut;
2023    use std::fmt::Debug;
2024    use test_case::test_case;
2025    use test_util::{assert_geq, assert_leq};
2026
2027    // DummyPacketBuilder:
2028    // - Implements PacketBuilder with the stored constraints; it fills the
2029    //   header with 0xFF and the footer with 0xFE
2030    // - Implements InnerPacketBuilder by consuming a `header_len`-bytes body,
2031    //   and filling it with 0xFF
2032    #[derive(Copy, Clone, Debug, Eq, PartialEq)]
2033    struct DummyPacketBuilder {
2034        header_len: usize,
2035        footer_len: usize,
2036        min_body_len: usize,
2037        max_body_len: usize,
2038    }
2039
2040    impl DummyPacketBuilder {
2041        fn new(
2042            header_len: usize,
2043            footer_len: usize,
2044            min_body_len: usize,
2045            max_body_len: usize,
2046        ) -> DummyPacketBuilder {
2047            DummyPacketBuilder { header_len, footer_len, min_body_len, max_body_len }
2048        }
2049    }
2050
2051    fn fill(bytes: &mut [u8], byte: u8) {
2052        for b in bytes {
2053            *b = byte;
2054        }
2055    }
2056
2057    impl PacketBuilder for DummyPacketBuilder {
2058        fn constraints(&self) -> PacketConstraints {
2059            PacketConstraints::new(
2060                self.header_len,
2061                self.footer_len,
2062                self.min_body_len,
2063                self.max_body_len,
2064            )
2065        }
2066
2067        fn serialize(&self, target: &mut SerializeTarget<'_>, body: FragmentedBytesMut<'_, '_>) {
2068            assert_eq!(target.header.len(), self.header_len);
2069            assert_eq!(target.footer.len(), self.footer_len);
2070            assert!(body.len() >= self.min_body_len);
2071            assert!(body.len() <= self.max_body_len);
2072            fill(target.header, 0xFF);
2073            fill(target.footer, 0xFE);
2074        }
2075    }
2076
2077    impl InnerPacketBuilder for DummyPacketBuilder {
2078        fn bytes_len(&self) -> usize {
2079            self.header_len
2080        }
2081
2082        fn serialize(&self, buffer: &mut [u8]) {
2083            assert_eq!(buffer.len(), self.header_len);
2084            fill(buffer, 0xFF);
2085        }
2086    }
2087
2088    // Helper for `VerifyingSerializer` used to verify the serialization result.
2089    #[derive(Copy, Clone, Debug, Eq, PartialEq)]
2090    struct SerializerVerifier {
2091        // Total size if the inner body if not truncated or `None` if
2092        // serialization is expected to fail due to size limit.
2093        inner_len: Option<usize>,
2094
2095        // Is the inner Serializer truncating (a TruncatingSerializer with
2096        // TruncateDirection::DiscardFront or DiscardBack)?
2097        truncating: bool,
2098    }
2099
2100    impl SerializerVerifier {
2101        fn new<S: Serializer>(serializer: &S, truncating: bool) -> Self {
2102            let inner_len = serializer
2103                .serialize_new_buf(PacketConstraints::UNCONSTRAINED, new_buf_vec)
2104                .map(|buf| buf.len())
2105                .inspect_err(|err| assert!(err.is_size_limit_exceeded()))
2106                .ok();
2107            Self { inner_len, truncating }
2108        }
2109
2110        fn verify_result<B: GrowBufferMut, A>(
2111            &self,
2112            result: Result<&B, &SerializeError<A>>,
2113            outer: PacketConstraints,
2114        ) {
2115            let should_exceed_size_limit = match self.inner_len {
2116                Some(inner_len) => outer.max_body_len() < inner_len && !self.truncating,
2117                None => true,
2118            };
2119
2120            match result {
2121                Ok(buf) => {
2122                    assert_geq!(buf.prefix_len(), outer.header_len());
2123                    assert_geq!(buf.suffix_len(), outer.footer_len());
2124                    assert_leq!(buf.len(), outer.max_body_len());
2125
2126                    // It is `Serialize::serialize()`'s responsibility to ensure that there
2127                    // is enough suffix room to fit any post-body padding and the footer,
2128                    // but it is the caller's responsibility to actually add that padding
2129                    // (ie, move it from the suffix to the body).
2130                    let padding = outer.min_body_len().saturating_sub(buf.len());
2131                    assert_leq!(padding + outer.footer_len(), buf.suffix_len());
2132
2133                    assert!(!should_exceed_size_limit);
2134                }
2135                Err(err) => {
2136                    // If we shouldn't fail as a result of a size limit exceeded
2137                    // error, we might still fail as a result of allocation.
2138                    if should_exceed_size_limit {
2139                        assert!(err.is_size_limit_exceeded());
2140                    } else {
2141                        assert!(err.is_alloc());
2142                    }
2143                }
2144            }
2145        }
2146    }
2147
2148    // A Serializer that verifies certain invariants while operating. In
2149    // particular:
2150    // - If serialization fails, the original Serializer is returned unmodified.
2151    // - If `outer.try_constraints()` returns `None`, serialization fails.
2152    // - If the size limit is exceeded and truncation is disabled, serialization
2153    //   fails.
2154    // - If serialization succeeds, the body has the correct length, including
2155    //   taking into account `outer`'s minimum body length requirement
2156    #[derive(Copy, Clone, Debug, Eq, PartialEq)]
2157    struct VerifyingSerializer<S> {
2158        ser: S,
2159        verifier: SerializerVerifier,
2160    }
2161
2162    impl<S: Serializer + Debug + Clone + Eq> Serializer for VerifyingSerializer<S>
2163    where
2164        S::Buffer: ReusableBuffer,
2165    {
2166        type Buffer = S::Buffer;
2167
2168        fn serialize<B: GrowBufferMut, P: BufferProvider<Self::Buffer, B>>(
2169            self,
2170            outer: PacketConstraints,
2171            provider: P,
2172        ) -> Result<B, (SerializeError<P::Error>, Self)> {
2173            let Self { ser, verifier } = self;
2174            let orig = ser.clone();
2175
2176            let result = ser.serialize(outer, provider).map_err(|(err, ser)| {
2177                // If serialization fails, the original Serializer should be
2178                // unmodified.
2179                assert_eq!(ser, orig);
2180                (err, Self { ser, verifier })
2181            });
2182
2183            verifier.verify_result(result.as_ref().map_err(|(err, _ser)| err), outer);
2184
2185            result
2186        }
2187
2188        fn serialize_new_buf<B: ReusableBuffer, A: BufferAlloc<B>>(
2189            &self,
2190            outer: PacketConstraints,
2191            alloc: A,
2192        ) -> Result<B, SerializeError<A::Error>> {
2193            let res = self.ser.serialize_new_buf(outer, alloc);
2194            self.verifier.verify_result(res.as_ref(), outer);
2195            res
2196        }
2197    }
2198
2199    trait SerializerExt: Serializer {
2200        fn into_verifying(self, truncating: bool) -> VerifyingSerializer<Self>
2201        where
2202            Self::Buffer: ReusableBuffer,
2203        {
2204            let verifier = SerializerVerifier::new(&self, truncating);
2205            VerifyingSerializer { ser: self, verifier }
2206        }
2207
2208        fn wrap_in_verifying<B: PacketBuilder>(
2209            self,
2210            outer: B,
2211            truncating: bool,
2212        ) -> VerifyingSerializer<Nested<Self, B>>
2213        where
2214            Self::Buffer: ReusableBuffer,
2215        {
2216            self.wrap_in(outer).into_verifying(truncating)
2217        }
2218
2219        fn with_size_limit_verifying(
2220            self,
2221            limit: usize,
2222            truncating: bool,
2223        ) -> VerifyingSerializer<Nested<Self, LimitedSizePacketBuilder>>
2224        where
2225            Self::Buffer: ReusableBuffer,
2226        {
2227            self.with_size_limit(limit).into_verifying(truncating)
2228        }
2229    }
2230
2231    impl<S: Serializer> SerializerExt for S {}
2232
2233    #[test]
2234    fn test_either_into_inner() {
2235        fn ret_either(a: u32, b: u32, c: bool) -> Either<u32, u32> {
2236            if c {
2237                Either::A(a)
2238            } else {
2239                Either::B(b)
2240            }
2241        }
2242
2243        assert_eq!(ret_either(1, 2, true).into_inner(), 1);
2244        assert_eq!(ret_either(1, 2, false).into_inner(), 2);
2245    }
2246
2247    #[test]
2248    fn test_either_unwrap_success() {
2249        assert_eq!(Either::<u16, u32>::A(5).unwrap_a(), 5);
2250        assert_eq!(Either::<u16, u32>::B(10).unwrap_b(), 10);
2251    }
2252
2253    #[test]
2254    #[should_panic]
2255    fn test_either_unwrap_a_panic() {
2256        let _: u16 = Either::<u16, u32>::B(10).unwrap_a();
2257    }
2258
2259    #[test]
2260    #[should_panic]
2261    fn test_either_unwrap_b_panic() {
2262        let _: u32 = Either::<u16, u32>::A(5).unwrap_b();
2263    }
2264
2265    #[test_case(Buf::new((0..100).collect(), ..); "entire buf")]
2266    #[test_case(Buf::new((0..100).collect(), 0..0); "empty range")]
2267    #[test_case(Buf::new((0..100).collect(), ..50); "prefix")]
2268    #[test_case(Buf::new((0..100).collect(), 50..); "suffix")]
2269    #[test_case(Buf::new((0..100).collect(), 25..75); "middle")]
2270    fn test_buf_into_inner(buf: Buf<Vec<u8>>) {
2271        assert_eq!(buf.clone().as_ref(), buf.into_inner());
2272    }
2273
2274    #[test]
2275    fn test_packet_constraints() {
2276        use PacketConstraints as PC;
2277
2278        // Test try_new
2279
2280        // Sanity check.
2281        assert!(PC::try_new(0, 0, 0, 0).is_some());
2282        // header_len + min_body_len + footer_len doesn't overflow usize
2283        assert!(PC::try_new(usize::MAX / 2, usize::MAX / 2, 0, 0).is_some());
2284        // header_len + min_body_len + footer_len overflows usize
2285        assert_eq!(PC::try_new(usize::MAX, 1, 0, 0), None);
2286        // min_body_len > max_body_len
2287        assert_eq!(PC::try_new(0, 0, 1, 0), None);
2288
2289        // Test PacketConstraints::try_encapsulate
2290
2291        // Sanity check.
2292        let pc = PC::new(10, 10, 0, usize::MAX);
2293        assert_eq!(pc.try_encapsulate(&pc).unwrap(), PC::new(20, 20, 0, usize::MAX - 20));
2294
2295        let pc = PC::new(10, 10, 0, usize::MAX);
2296        assert_eq!(pc.try_encapsulate(&pc).unwrap(), PC::new(20, 20, 0, usize::MAX - 20));
2297
2298        // Starting here, each failure test case corresponds to one check in
2299        // either PacketConstraints::try_encapsulate or PacketConstraints::new
2300        // (which is called from PacketConstraints::try_encapsulate). Each test
2301        // case is labeled "Test case N", and a corresponding comment in either
2302        // of those two functions identifies which line is being tested.
2303
2304        // The outer PC's minimum body length requirement of 10 is more than
2305        // satisfied by the inner PC's combined 20 bytes of header and footer.
2306        // The resulting PC has its minimum body length requirement saturated to
2307        // 0.
2308        let inner = PC::new(10, 10, 0, usize::MAX);
2309        let outer = PC::new(0, 0, 10, usize::MAX);
2310        assert_eq!(inner.try_encapsulate(&outer).unwrap(), PC::new(10, 10, 0, usize::MAX - 20));
2311
2312        // Test case 1
2313        //
2314        // The sum of the inner and outer header lengths overflows `usize`.
2315        let inner = PC::new(usize::MAX, 0, 0, usize::MAX);
2316        let outer = PC::new(1, 0, 0, usize::MAX);
2317        assert_eq!(inner.try_encapsulate(&outer), None);
2318
2319        // Test case 2
2320        //
2321        // The sum of the inner and outer footer lengths overflows `usize`.
2322        let inner = PC::new(0, usize::MAX, 0, usize::MAX);
2323        let outer = PC::new(0, 1, 0, usize::MAX);
2324        assert_eq!(inner.try_encapsulate(&outer), None);
2325
2326        // Test case 3
2327        //
2328        // The sum of the resulting header, footer, and minimum body lengths
2329        // overflows `usize`. We use usize::MAX / 5 + 1 as the constant so that
2330        // none of the intermediate additions overflow, so we make sure to test
2331        // that an overflow in the final addition will be caught.
2332        let one_fifth_max = (usize::MAX / 5) + 1;
2333        let inner = PC::new(one_fifth_max, one_fifth_max, one_fifth_max, usize::MAX);
2334        let outer = PC::new(one_fifth_max, one_fifth_max, 0, usize::MAX);
2335        assert_eq!(inner.try_encapsulate(&outer), None);
2336
2337        // Test case 4
2338        //
2339        // The header and footer of the inner PC exceed the maximum body length
2340        // requirement of the outer PC.
2341        let inner = PC::new(10, 10, 0, usize::MAX);
2342        let outer = PC::new(0, 0, 0, 10);
2343        assert_eq!(inner.try_encapsulate(&outer), None);
2344
2345        // Test case 5
2346        //
2347        // The resulting minimum body length (thanks to the inner
2348        // PacketBuilder's minimum body length) is larger than the resulting
2349        // maximum body length.
2350        let inner = PC::new(0, 0, 10, usize::MAX);
2351        let outer = PC::new(0, 0, 0, 5);
2352        assert_eq!(inner.try_encapsulate(&outer), None);
2353    }
2354
2355    #[test]
2356    fn test_inner_serializer() {
2357        const INNER: &[u8] = &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
2358
2359        fn concat<'a, I: IntoIterator<Item = &'a &'a [u8]>>(slices: I) -> Vec<u8> {
2360            let mut v = Vec::new();
2361            for slc in slices.into_iter() {
2362                v.extend_from_slice(slc);
2363            }
2364            v
2365        }
2366
2367        // Sanity check.
2368        let buf = INNER.into_serializer().serialize_vec_outer().unwrap();
2369        assert_eq!(buf.as_ref(), INNER);
2370
2371        // A larger minimum body length requirement will cause padding to be
2372        // added.
2373        let buf = INNER
2374            .into_serializer()
2375            .into_verifying(false)
2376            .wrap_in(DummyPacketBuilder::new(0, 0, 20, usize::MAX))
2377            .serialize_vec_outer()
2378            .unwrap();
2379        assert_eq!(buf.as_ref(), concat(&[INNER, vec![0; 10].as_ref()]).as_slice());
2380
2381        // Headers and footers are added as appropriate (note that
2382        // DummyPacketBuilder fills its header with 0xFF and its footer with
2383        // 0xFE).
2384        let buf = INNER
2385            .into_serializer()
2386            .into_verifying(false)
2387            .wrap_in(DummyPacketBuilder::new(10, 10, 0, usize::MAX))
2388            .serialize_vec_outer()
2389            .unwrap();
2390        assert_eq!(
2391            buf.as_ref(),
2392            concat(&[vec![0xFF; 10].as_ref(), INNER, vec![0xFE; 10].as_ref()]).as_slice()
2393        );
2394
2395        // An exceeded maximum body size is rejected.
2396        assert_eq!(
2397            INNER
2398                .into_serializer()
2399                .into_verifying(false)
2400                .wrap_in(DummyPacketBuilder::new(0, 0, 0, 9))
2401                .serialize_vec_outer()
2402                .unwrap_err()
2403                .0,
2404            SerializeError::SizeLimitExceeded
2405        );
2406
2407        // `into_serializer_with` truncates the buffer's body to zero before
2408        // returning, so those body bytes are not included in the serialized
2409        // output.
2410        assert_eq!(
2411            INNER
2412                .into_serializer_with(Buf::new(vec![0xFF], ..))
2413                .into_verifying(false)
2414                .serialize_vec_outer()
2415                .unwrap()
2416                .as_ref(),
2417            INNER
2418        );
2419    }
2420
2421    #[test]
2422    fn test_buffer_serializer_and_inner_serializer() {
2423        fn verify_buffer_serializer<B: BufferMut + Debug>(
2424            buffer: B,
2425            header_len: usize,
2426            footer_len: usize,
2427            min_body_len: usize,
2428        ) {
2429            let old_body = buffer.to_flattened_vec();
2430            let serializer =
2431                DummyPacketBuilder::new(header_len, footer_len, min_body_len, usize::MAX)
2432                    .wrap_body(buffer);
2433
2434            let buffer0 = serializer
2435                .serialize_new_buf(PacketConstraints::UNCONSTRAINED, new_buf_vec)
2436                .unwrap();
2437            verify(buffer0, &old_body, header_len, footer_len, min_body_len);
2438
2439            let buffer = serializer.serialize_vec_outer().unwrap();
2440            verify(buffer, &old_body, header_len, footer_len, min_body_len);
2441        }
2442
2443        fn verify_inner_packet_builder_serializer(
2444            body: &[u8],
2445            header_len: usize,
2446            footer_len: usize,
2447            min_body_len: usize,
2448        ) {
2449            let buffer = DummyPacketBuilder::new(header_len, footer_len, min_body_len, usize::MAX)
2450                .wrap_body(body.into_serializer())
2451                .serialize_vec_outer()
2452                .unwrap();
2453            verify(buffer, body, header_len, footer_len, min_body_len);
2454        }
2455
2456        fn verify<B: Buffer>(
2457            buffer: B,
2458            body: &[u8],
2459            header_len: usize,
2460            footer_len: usize,
2461            min_body_len: usize,
2462        ) {
2463            let flat = buffer.to_flattened_vec();
2464            let header_bytes = &flat[..header_len];
2465            let body_bytes = &flat[header_len..header_len + body.len()];
2466            let padding_len = min_body_len.saturating_sub(body.len());
2467            let padding_bytes =
2468                &flat[header_len + body.len()..header_len + body.len() + padding_len];
2469            let total_body_len = body.len() + padding_len;
2470            let footer_bytes = &flat[header_len + total_body_len..];
2471            assert_eq!(
2472                buffer.len() - total_body_len,
2473                header_len + footer_len,
2474                "buffer.len()({}) - total_body_len({}) != header_len({}) + footer_len({})",
2475                buffer.len(),
2476                header_len,
2477                footer_len,
2478                min_body_len,
2479            );
2480
2481            // DummyPacketBuilder fills its header with 0xFF
2482            assert!(
2483                header_bytes.iter().all(|b| *b == 0xFF),
2484                "header_bytes {:?} are not filled with 0xFF's",
2485                header_bytes,
2486            );
2487            assert_eq!(body_bytes, body);
2488            // Padding bytes must be initialized to zero
2489            assert!(
2490                padding_bytes.iter().all(|b| *b == 0),
2491                "padding_bytes {:?} are not filled with 0s",
2492                padding_bytes,
2493            );
2494            // DummyPacketBuilder fills its footer with 0xFE
2495            assert!(
2496                footer_bytes.iter().all(|b| *b == 0xFE),
2497                "footer_bytes {:?} are not filled with 0xFE's",
2498                footer_bytes,
2499            );
2500        }
2501
2502        // Test for every valid combination of buf_len, range_start, range_end,
2503        // prefix, suffix, and min_body within [0, 8).
2504        for buf_len in 0..8 {
2505            for range_start in 0..buf_len {
2506                for range_end in range_start..buf_len {
2507                    for prefix in 0..8 {
2508                        for suffix in 0..8 {
2509                            for min_body in 0..8 {
2510                                let mut vec = vec![0; buf_len];
2511                                // Initialize the vector with values 0, 1, 2,
2512                                // ... so that we can check to make sure that
2513                                // the range bytes have been properly copied if
2514                                // the buffer is reallocated.
2515                                #[allow(clippy::needless_range_loop)]
2516                                for i in 0..vec.len() {
2517                                    vec[i] = i as u8;
2518                                }
2519                                verify_buffer_serializer(
2520                                    Buf::new(vec.as_mut_slice(), range_start..range_end),
2521                                    prefix,
2522                                    suffix,
2523                                    min_body,
2524                                );
2525                                if range_start == 0 {
2526                                    // Unlike verify_buffer_serializer, this
2527                                    // test doesn't make use of the prefix or
2528                                    // suffix. In order to avoid running the
2529                                    // exact same test multiple times, we only
2530                                    // run this when `range_start == 0`, which
2531                                    // has the effect of reducing the number of
2532                                    // times that this test is run by roughly a
2533                                    // factor of 8.
2534                                    verify_inner_packet_builder_serializer(
2535                                        &vec.as_slice()[range_start..range_end],
2536                                        prefix,
2537                                        suffix,
2538                                        min_body,
2539                                    );
2540                                }
2541                            }
2542                        }
2543                    }
2544                }
2545            }
2546        }
2547    }
2548
2549    #[test]
2550    fn test_min_body_len() {
2551        // Test that padding is added after the body of the packet whose minimum
2552        // body length constraint requires it. A previous version of this code
2553        // had a bug where padding was always added after the innermost body.
2554
2555        let body = &[1, 2];
2556
2557        // 4 bytes of header and footer for a total of 6 bytes (including the
2558        // body).
2559        let inner = DummyPacketBuilder::new(2, 2, 0, usize::MAX);
2560        // Minimum body length of 8 will require 2 bytes of padding.
2561        let outer = DummyPacketBuilder::new(2, 2, 8, usize::MAX);
2562        let buf = body
2563            .into_serializer()
2564            .into_verifying(false)
2565            .wrap_in_verifying(inner, false)
2566            .wrap_in_verifying(outer, false)
2567            .serialize_vec_outer()
2568            .unwrap();
2569        assert_eq!(buf.prefix_len(), 0);
2570        assert_eq!(buf.suffix_len(), 0);
2571        assert_eq!(
2572            buf.as_ref(),
2573            &[
2574                0xFF, 0xFF, // Outer header
2575                0xFF, 0xFF, // Inner header
2576                1, 2, // Inner body
2577                0xFE, 0xFE, // Inner footer
2578                0, 0, // Padding to satisfy outer minimum body length requirement
2579                0xFE, 0xFE // Outer footer
2580            ]
2581        );
2582    }
2583
2584    #[test]
2585    fn test_size_limit() {
2586        // ser is a Serializer that will consume 1 byte of buffer space
2587        fn test<S: Serializer + Clone + Debug + Eq>(ser: S)
2588        where
2589            S::Buffer: ReusableBuffer,
2590        {
2591            // Each of these tests encapsulates ser in a DummyPacketBuilder
2592            // which consumes 1 byte for the header and one byte for the footer.
2593            // Thus, the inner serializer will consume 1 byte, while the
2594            // DummyPacketBuilder will consume 2 bytes, for a total of 3 bytes.
2595
2596            let pb = DummyPacketBuilder::new(1, 1, 0, usize::MAX);
2597
2598            // Test that a size limit of 3 is OK. Note that this is an important
2599            // test since it tests the case when the size limit is exactly
2600            // sufficient. A previous version of this code had a bug where a
2601            // packet which fit the size limit exactly would be rejected.
2602            assert!(ser
2603                .clone()
2604                .wrap_in_verifying(pb, false)
2605                .with_size_limit_verifying(3, false)
2606                .serialize_vec_outer()
2607                .is_ok());
2608            // Test that a more-than-large-enough size limit of 4 is OK.
2609            assert!(ser
2610                .clone()
2611                .wrap_in_verifying(pb, false)
2612                .with_size_limit_verifying(4, false)
2613                .serialize_vec_outer()
2614                .is_ok());
2615            // Test that the inner size limit of 1 only applies to the inner
2616            // serializer, and so is still OK even though the outer serializer
2617            // consumes 3 bytes total.
2618            assert!(ser
2619                .clone()
2620                .with_size_limit_verifying(1, false)
2621                .wrap_in_verifying(pb, false)
2622                .with_size_limit_verifying(3, false)
2623                .serialize_vec_outer()
2624                .is_ok());
2625            // Test that the inner size limit of 0 is exceeded by the inner
2626            // serializer's 1 byte length.
2627            assert!(ser
2628                .clone()
2629                .with_size_limit_verifying(0, false)
2630                .wrap_in_verifying(pb, false)
2631                .serialize_vec_outer()
2632                .is_err());
2633            // Test that a size limit which would be exceeded by the
2634            // encapsulating layer is rejected by Nested's implementation. If
2635            // this doesn't work properly, then the size limit should underflow,
2636            // resulting in a panic (see the Nested implementation of
2637            // Serialize).
2638            assert!(ser
2639                .clone()
2640                .wrap_in_verifying(pb, false)
2641                .with_size_limit_verifying(1, false)
2642                .serialize_vec_outer()
2643                .is_err());
2644        }
2645
2646        // We use this as an InnerPacketBuilder which consumes 1 byte of body.
2647        test(DummyPacketBuilder::new(1, 0, 0, usize::MAX).into_serializer().into_verifying(false));
2648        test(Buf::new(vec![0], ..).into_verifying(false));
2649    }
2650
2651    #[test]
2652    fn test_truncating_serializer() {
2653        fn verify_result<S: Serializer + Debug>(ser: S, expected: &[u8])
2654        where
2655            S::Buffer: ReusableBuffer + AsRef<[u8]>,
2656        {
2657            let buf = ser.serialize_new_buf(PacketConstraints::UNCONSTRAINED, new_buf_vec).unwrap();
2658            assert_eq!(buf.as_ref(), &expected[..]);
2659            let buf = ser.serialize_vec_outer().unwrap();
2660            assert_eq!(buf.as_ref(), &expected[..]);
2661        }
2662
2663        // Test truncate front.
2664        let body = vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
2665        let ser =
2666            TruncatingSerializer::new(Buf::new(body.clone(), ..), TruncateDirection::DiscardFront)
2667                .into_verifying(true)
2668                .with_size_limit_verifying(4, true);
2669        verify_result(ser, &[6, 7, 8, 9]);
2670
2671        // Test truncate back.
2672        let ser =
2673            TruncatingSerializer::new(Buf::new(body.clone(), ..), TruncateDirection::DiscardBack)
2674                .into_verifying(true)
2675                .with_size_limit_verifying(7, true);
2676        verify_result(ser, &[0, 1, 2, 3, 4, 5, 6]);
2677
2678        // Test no truncating (default/original case).
2679        let ser =
2680            TruncatingSerializer::new(Buf::new(body.clone(), ..), TruncateDirection::NoTruncating)
2681                .into_verifying(false)
2682                .with_size_limit_verifying(5, true);
2683        assert!(ser.clone().serialize_vec_outer().is_err());
2684        assert!(ser.serialize_new_buf(PacketConstraints::UNCONSTRAINED, new_buf_vec).is_err());
2685        assert!(ser.serialize_vec_outer().is_err());
2686
2687        // Test that, when serialization fails, any truncation is undone.
2688
2689        // `ser` has a body of `[1, 2]` and no prefix or suffix
2690        fn test_serialization_failure<S: Serializer + Clone + Eq + Debug>(
2691            ser: S,
2692            err: SerializeError<BufferTooShortError>,
2693        ) where
2694            S::Buffer: ReusableBuffer + Debug,
2695        {
2696            // Serialize with a PacketBuilder with a size limit of 1 so that the
2697            // body (of length 2) is too large. If `ser` is configured not to
2698            // truncate, it should result in a size limit error. If it is
2699            // configured to truncate, the 2 + 2 = 4 combined bytes of header
2700            // and footer will cause allocating a new buffer to fail, and it
2701            // should result in an allocation failure. Even if the body was
2702            // truncated, it should be returned to its original un-truncated
2703            // state before being returned from `serialize`.
2704            let (e, new_ser) = DummyPacketBuilder::new(2, 2, 0, 1)
2705                .wrap_body(ser.clone())
2706                .serialize_no_alloc_outer()
2707                .unwrap_err();
2708            assert_eq!(err, e);
2709            assert_eq!(new_ser.into_inner(), ser);
2710        }
2711
2712        let body = Buf::new(vec![1, 2], ..);
2713        test_serialization_failure(
2714            TruncatingSerializer::new(body.clone(), TruncateDirection::DiscardFront)
2715                .into_verifying(true),
2716            SerializeError::Alloc(BufferTooShortError),
2717        );
2718        test_serialization_failure(
2719            TruncatingSerializer::new(body.clone(), TruncateDirection::DiscardFront)
2720                .into_verifying(true),
2721            SerializeError::Alloc(BufferTooShortError),
2722        );
2723        test_serialization_failure(
2724            TruncatingSerializer::new(body.clone(), TruncateDirection::NoTruncating)
2725                .into_verifying(false),
2726            SerializeError::SizeLimitExceeded,
2727        );
2728    }
2729
2730    #[test]
2731    fn test_try_reuse_buffer() {
2732        fn test_expect_success(
2733            body_range: Range<usize>,
2734            prefix: usize,
2735            suffix: usize,
2736            max_copy_bytes: usize,
2737        ) {
2738            let mut bytes = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
2739            let buffer = Buf::new(&mut bytes[..], body_range);
2740            let body = buffer.as_ref().to_vec();
2741            let buffer = try_reuse_buffer(buffer, prefix, suffix, max_copy_bytes).unwrap();
2742            assert_eq!(buffer.as_ref(), body.as_slice());
2743            assert!(buffer.prefix_len() >= prefix);
2744            assert!(buffer.suffix_len() >= suffix);
2745        }
2746
2747        fn test_expect_failure(
2748            body_range: Range<usize>,
2749            prefix: usize,
2750            suffix: usize,
2751            max_copy_bytes: usize,
2752        ) {
2753            let mut bytes = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
2754            let buffer = Buf::new(&mut bytes[..], body_range.clone());
2755            let mut bytes = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
2756            let orig = Buf::new(&mut bytes[..], body_range.clone());
2757            let buffer = try_reuse_buffer(buffer, prefix, suffix, max_copy_bytes).unwrap_err();
2758            assert_eq!(buffer, orig);
2759        }
2760
2761        // No prefix or suffix trivially succeeds.
2762        test_expect_success(0..10, 0, 0, 0);
2763        // If we have enough prefix/suffix, it succeeds.
2764        test_expect_success(1..9, 1, 1, 0);
2765        // If we don't have enough prefix/suffix, but we have enough capacity to
2766        // move the buffer within the body, it succeeds...
2767        test_expect_success(0..9, 1, 0, 9);
2768        test_expect_success(1..10, 0, 1, 9);
2769        // ...but if we don't provide a large enough max_copy_bytes, it will fail.
2770        test_expect_failure(0..9, 1, 0, 8);
2771        test_expect_failure(1..10, 0, 1, 8);
2772    }
2773
2774    #[test]
2775    fn test_maybe_reuse_buffer_provider() {
2776        fn test_expect(body_range: Range<usize>, prefix: usize, suffix: usize, expect_a: bool) {
2777            let mut bytes = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
2778            let buffer = Buf::new(&mut bytes[..], body_range);
2779            let body = buffer.as_ref().to_vec();
2780            let buffer = BufferProvider::reuse_or_realloc(
2781                MaybeReuseBufferProvider(new_buf_vec),
2782                buffer,
2783                prefix,
2784                suffix,
2785            )
2786            .unwrap();
2787            match &buffer {
2788                Either::A(_) if expect_a => {}
2789                Either::B(_) if !expect_a => {}
2790                Either::A(_) => panic!("expected Eitehr::B variant"),
2791                Either::B(_) => panic!("expected Eitehr::A variant"),
2792            }
2793            let bytes: &[u8] = buffer.as_ref();
2794            assert_eq!(bytes, body.as_slice());
2795            assert!(buffer.prefix_len() >= prefix);
2796            assert!(buffer.suffix_len() >= suffix);
2797        }
2798
2799        // Expect that we'll be able to reuse the existing buffer.
2800        fn test_expect_reuse(body_range: Range<usize>, prefix: usize, suffix: usize) {
2801            test_expect(body_range, prefix, suffix, true);
2802        }
2803
2804        // Expect that we'll need to allocate a new buffer.
2805        fn test_expect_realloc(body_range: Range<usize>, prefix: usize, suffix: usize) {
2806            test_expect(body_range, prefix, suffix, false);
2807        }
2808
2809        // No prefix or suffix trivially succeeds.
2810        test_expect_reuse(0..10, 0, 0);
2811        // If we have enough prefix/suffix, it succeeds.
2812        test_expect_reuse(1..9, 1, 1);
2813        // If we don't have enough prefix/suffix, but we have enough capacity to
2814        // move the buffer within the body, it succeeds.
2815        test_expect_reuse(0..9, 1, 0);
2816        test_expect_reuse(1..10, 0, 1);
2817        // If we don't have enough capacity, it fails and must realloc.
2818        test_expect_realloc(0..9, 1, 1);
2819        test_expect_realloc(1..10, 1, 1);
2820    }
2821
2822    #[test]
2823    fn test_no_reuse_buffer_provider() {
2824        #[track_caller]
2825        fn test_expect(body_range: Range<usize>, prefix: usize, suffix: usize) {
2826            let mut bytes = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
2827            // The buffer that will not be reused.
2828            let internal_buffer: Buf<&mut [u8]> = Buf::new(&mut bytes[..], body_range);
2829            let body = internal_buffer.as_ref().to_vec();
2830            // The newly allocated buffer, note the type is different from
2831            // internal_buffer.
2832            let buffer: Buf<Vec<u8>> = BufferProvider::reuse_or_realloc(
2833                NoReuseBufferProvider(new_buf_vec),
2834                internal_buffer,
2835                prefix,
2836                suffix,
2837            )
2838            .unwrap();
2839            let bytes: &[u8] = buffer.as_ref();
2840            assert_eq!(bytes, body.as_slice());
2841            assert_eq!(buffer.prefix_len(), prefix);
2842            assert_eq!(buffer.suffix_len(), suffix);
2843        }
2844        // No prefix or suffix trivially succeeds, reuse opportunity is ignored.
2845        test_expect(0..10, 0, 0);
2846        // If we have enough prefix/suffix, reuse opportunity is ignored.
2847        test_expect(1..9, 1, 1);
2848        // Prefix and suffix and properly allocated and the body is copied.
2849        test_expect(0..9, 10, 10);
2850        test_expect(1..10, 15, 15);
2851    }
2852
2853    /// Simple Vec-backed buffer to test fragmented buffers implementation.
2854    ///
2855    /// `ScatterGatherBuf` keeps:
2856    /// - an inner buffer `inner`, which is always part of its body.
2857    /// - extra backing memory in `data`.
2858    ///
2859    /// `data` has two "root" regions, marked by the midpoint `mid`. Everything
2860    /// left of `mid` is this buffer's prefix, and after `mid` is this buffer's
2861    /// suffix.
2862    ///
2863    /// The `range` field keeps the range in `data` that contains *filled*
2864    /// prefix and suffix information. `range.start` is always less than or
2865    /// equal to `mid` and `range.end` is always greater than or equal to `mid`,
2866    /// such that growing the front of the buffer means decrementing
2867    /// `range.start` and growing the back of the buffer means incrementing
2868    /// `range.end`.
2869    ///
2870    ///  At any time this buffer's parts are:
2871    /// - Free prefix data in range `0..range.start`.
2872    /// - Used prefix data (now part of body) in range `range.start..mid`.
2873    /// - Inner buffer body in `inner`.
2874    /// - Used suffix data (now part of body) in range `mid..range.end`.
2875    /// - Free suffix data in range `range.end..`
2876    struct ScatterGatherBuf<B> {
2877        data: Vec<u8>,
2878        mid: usize,
2879        range: Range<usize>,
2880        inner: B,
2881    }
2882
2883    impl<B: BufferMut> FragmentedBuffer for ScatterGatherBuf<B> {
2884        fn len(&self) -> usize {
2885            self.inner.len() + (self.range.end - self.range.start)
2886        }
2887
2888        fn with_bytes<R, F>(&self, f: F) -> R
2889        where
2890            F: for<'a, 'b> FnOnce(FragmentedBytes<'a, 'b>) -> R,
2891        {
2892            let (_, rest) = self.data.split_at(self.range.start);
2893            let (prefix_b, rest) = rest.split_at(self.mid - self.range.start);
2894            let (suffix_b, _) = rest.split_at(self.range.end - self.mid);
2895            let mut bytes = [prefix_b, self.inner.as_ref(), suffix_b];
2896            f(FragmentedBytes::new(&mut bytes[..]))
2897        }
2898    }
2899
2900    impl<B: BufferMut> FragmentedBufferMut for ScatterGatherBuf<B> {
2901        fn with_bytes_mut<R, F>(&mut self, f: F) -> R
2902        where
2903            F: for<'a, 'b> FnOnce(FragmentedBytesMut<'a, 'b>) -> R,
2904        {
2905            let (_, rest) = self.data.split_at_mut(self.range.start);
2906            let (prefix_b, rest) = rest.split_at_mut(self.mid - self.range.start);
2907            let (suffix_b, _) = rest.split_at_mut(self.range.end - self.mid);
2908            let mut bytes = [prefix_b, self.inner.as_mut(), suffix_b];
2909            f(FragmentedBytesMut::new(&mut bytes[..]))
2910        }
2911    }
2912
2913    impl<B: BufferMut> GrowBuffer for ScatterGatherBuf<B> {
2914        fn with_parts<O, F>(&self, f: F) -> O
2915        where
2916            F: for<'a, 'b> FnOnce(&'a [u8], FragmentedBytes<'a, 'b>, &'a [u8]) -> O,
2917        {
2918            let (prefix, rest) = self.data.split_at(self.range.start);
2919            let (prefix_b, rest) = rest.split_at(self.mid - self.range.start);
2920            let (suffix_b, suffix) = rest.split_at(self.range.end - self.mid);
2921            let mut bytes = [prefix_b, self.inner.as_ref(), suffix_b];
2922            f(prefix, bytes.as_fragmented_byte_slice(), suffix)
2923        }
2924        fn prefix_len(&self) -> usize {
2925            self.range.start
2926        }
2927
2928        fn suffix_len(&self) -> usize {
2929            self.data.len() - self.range.end
2930        }
2931
2932        fn grow_front(&mut self, n: usize) {
2933            self.range.start -= n;
2934        }
2935
2936        fn grow_back(&mut self, n: usize) {
2937            self.range.end += n;
2938            assert!(self.range.end <= self.data.len());
2939        }
2940    }
2941
2942    impl<B: BufferMut> GrowBufferMut for ScatterGatherBuf<B> {
2943        fn with_parts_mut<O, F>(&mut self, f: F) -> O
2944        where
2945            F: for<'a, 'b> FnOnce(&'a mut [u8], FragmentedBytesMut<'a, 'b>, &'a mut [u8]) -> O,
2946        {
2947            let (prefix, rest) = self.data.split_at_mut(self.range.start);
2948            let (prefix_b, rest) = rest.split_at_mut(self.mid - self.range.start);
2949            let (suffix_b, suffix) = rest.split_at_mut(self.range.end - self.mid);
2950            let mut bytes = [prefix_b, self.inner.as_mut(), suffix_b];
2951            f(prefix, bytes.as_fragmented_byte_slice(), suffix)
2952        }
2953    }
2954
2955    struct ScatterGatherProvider;
2956
2957    impl<B: BufferMut> BufferProvider<B, ScatterGatherBuf<B>> for ScatterGatherProvider {
2958        type Error = Never;
2959
2960        fn alloc_no_reuse(
2961            self,
2962            _prefix: usize,
2963            _body: usize,
2964            _suffix: usize,
2965        ) -> Result<ScatterGatherBuf<B>, Self::Error> {
2966            unimplemented!("not used in tests")
2967        }
2968
2969        fn reuse_or_realloc(
2970            self,
2971            buffer: B,
2972            prefix: usize,
2973            suffix: usize,
2974        ) -> Result<ScatterGatherBuf<B>, (Self::Error, B)> {
2975            let inner = buffer;
2976            let data = vec![0; prefix + suffix];
2977            let range = Range { start: prefix, end: prefix };
2978            let mid = prefix;
2979            Ok(ScatterGatherBuf { inner, data, range, mid })
2980        }
2981    }
2982
2983    #[test]
2984    fn test_scatter_gather_serialize() {
2985        // Assert that a buffer composed of different allocations can be used as
2986        // a serialization target, while reusing an internal body buffer.
2987        let buf = Buf::new(vec![10, 20, 30, 40, 50], ..);
2988        let pb = DummyPacketBuilder::new(3, 2, 0, usize::MAX);
2989        let ser = pb.wrap_body(buf);
2990        let result = ser.serialize_outer(ScatterGatherProvider {}).unwrap();
2991        let flattened = result.to_flattened_vec();
2992        assert_eq!(&flattened[..], &[0xFF, 0xFF, 0xFF, 10, 20, 30, 40, 50, 0xFE, 0xFE]);
2993    }
2994}