Skip to main content

sark_http/h1/codec/decode/
chunks.rs

1use std::num;
2
3use chunk::{scan, trailer::Partial};
4use o3::{alloc, buffer::write::Sink as _};
5
6use crate::{
7    h1::{
8        chunk,
9        codec::decode::{fields, fields::Sink as _, trailer},
10    },
11    payload::response::headers::List,
12};
13
14const MAX_BODY_SIZE: usize = 100 * 1024 * 1024;
15
16#[derive(Debug)]
17pub enum Error {
18    Framing(chunk::Error),
19    Field(Conversion),
20    Allocation(alloc::Error),
21    Incomplete,
22}
23
24#[derive(Debug)]
25pub struct Conversion {
26    pub field: usize,
27    pub cause: fields::Error,
28}
29
30pub(crate) struct Output<'field> {
31    pub(crate) body: Vec<u8>,
32    pub(crate) trailers: List<'field>,
33}
34
35impl<'field> Output<'field> {
36    pub(crate) fn drain(buf: &[u8], max_body: usize) -> Result<Option<Self>, Error> {
37        let mut decoder = Decoder::with_limit(max_body);
38        let mut body = Vec::new();
39        let mut input = buf;
40
41        loop {
42            let plan = decoder.decode(input)?;
43            match plan.commit(&mut decoder) {
44                Event::NeedMore(_) => return Ok(None),
45                Event::Chunk {
46                    bytes, consumed, ..
47                } => {
48                    // SAFETY: decode just admitted this consumed prefix on
49                    // this unchanged input; the event retains that exact borrow.
50                    input = unsafe { input.get_unchecked(consumed.get()..) };
51                    body.write_slice(bytes).map_err(Error::Allocation)?;
52                }
53                Event::Done { fields, .. } => {
54                    let trailers = Decoder::parse(fields).map_err(Error::Field)?;
55                    return Ok(Some(Self { body, trailers }));
56                }
57            }
58        }
59    }
60}
61
62pub enum Event<'a> {
63    NeedMore(Option<num::NonZero<usize>>),
64    Chunk {
65        consumed: num::NonZero<usize>,
66        bytes: &'a [u8],
67        boundary: chunk::Boundary,
68    },
69    Done {
70        consumed: num::NonZero<usize>,
71        fields: trailer::Fields<'a>,
72    },
73}
74
75pub struct Plan<'a> {
76    event: Prepared<'a>,
77}
78
79enum Prepared<'a> {
80    Event { next: Decoder, event: Event<'a> },
81    Trailer(Partial<'a>),
82}
83
84impl<'a> Plan<'a> {
85    pub fn commit(self, decoder: &mut Decoder) -> Event<'a> {
86        match self.event {
87            Prepared::Event { next, event } => {
88                *decoder = next;
89                event
90            }
91            Prepared::Trailer(_) => Event::NeedMore(None),
92        }
93    }
94
95    /// Transfers partial trailer admission together with its exact input borrow.
96    /// Retaining that prefix permits later input to be scanned only once.
97    pub fn into_trailer(self) -> Result<trailer::Pending<'a>, Self> {
98        match self.event {
99            Prepared::Trailer(partial) => Ok(trailer::Pending::new(partial)),
100            event @ Prepared::Event { .. } => Err(Self { event }),
101        }
102    }
103}
104
105#[derive(Clone, Copy)]
106pub struct Decoder {
107    framer: chunk::Framer,
108    limits: chunk::Limit,
109}
110
111impl Default for Decoder {
112    fn default() -> Self {
113        Self::new()
114    }
115}
116
117impl Decoder {
118    pub const DEFAULT_MAX_BODY: usize = MAX_BODY_SIZE;
119
120    pub fn new() -> Self {
121        Self::with_limit(MAX_BODY_SIZE)
122    }
123
124    pub fn with_limit(max_body: usize) -> Self {
125        Self {
126            framer: chunk::Framer::new(),
127            limits: chunk::Limit::with_body(max_body),
128        }
129    }
130
131    /// Bounds an emitted body fragment without copying or limiting framing
132    /// progress.
133    pub fn decode_max<'a>(&self, buf: &'a [u8], max: usize) -> Result<Plan<'a>, Error> {
134        let (framer, event) = self
135            .framer
136            .prepare::<scan::Fragments>(buf, &self.limits, max)
137            .map_err(Error::Framing)?;
138        let event = match event {
139            scan::Decoded::Data {
140                consumed,
141                bytes,
142                boundary,
143            } => Event::Chunk {
144                consumed,
145                bytes,
146                boundary,
147            },
148            scan::Decoded::Trailer(partial) => {
149                return Ok(Plan {
150                    event: Prepared::Trailer(partial),
151                });
152            }
153            scan::Decoded::Complete {
154                consumed,
155                bytes: lines,
156            } => Event::Done {
157                consumed,
158                fields: trailer::Fields::new(lines),
159            },
160            scan::Decoded::More(consumed) => Event::NeedMore(num::NonZero::new(consumed)),
161        };
162        let mut next = *self;
163        next.framer = framer;
164        Ok(Plan {
165            event: Prepared::Event { next, event },
166        })
167    }
168
169    pub fn decode<'a>(&self, buf: &'a [u8]) -> Result<Plan<'a>, Error> {
170        self.decode_max(buf, usize::MAX)
171    }
172
173    pub(crate) fn require<'field>(buf: &[u8], max_body: usize) -> Result<Output<'field>, Error> {
174        Output::drain(buf, max_body)?.ok_or(Error::Incomplete)
175    }
176
177    pub fn body(buf: &[u8]) -> Result<Option<Vec<u8>>, Error> {
178        Ok(Output::drain(buf, MAX_BODY_SIZE)?.map(|result| result.body))
179    }
180
181    fn parse<'field>(fields: trailer::Fields<'_>) -> Result<List<'field>, Conversion> {
182        let mut trailers = List::new();
183        for (field, (name, value)) in fields.enumerate() {
184            trailers
185                .push(name, value)
186                .map_err(|cause| Conversion { field, cause })?;
187        }
188        Ok(trailers)
189    }
190}
191
192const _: () = {
193    assert!(core::mem::size_of::<Error>() <= 5 * core::mem::size_of::<usize>());
194    assert!(!core::mem::needs_drop::<Error>());
195};