Skip to main content

sark_core/http/response/direct/
head.rs

1use std::slice;
2
3use o3::buffer::{ByteSink, Owned, Shared};
4
5use super::headers::{DEFAULT_HEADER_CAPACITY, Headers};
6use crate::http::Field;
7use crate::http::response::wire_emit::{HeaderSection, ProvenWireWriter};
8
9#[derive(Clone, Debug)]
10pub struct HeadInner<'req, const N: usize = DEFAULT_HEADER_CAPACITY, S = StaticHeaderBytes> {
11    static_headers: S,
12    headers: Headers<'req, N>,
13}
14
15#[derive(Clone, Copy, Debug)]
16pub struct StaticHeaderBytes(&'static [u8]);
17
18impl StaticHeaderBytes {
19    pub const fn new(wire: &'static [u8]) -> Self {
20        Self(wire)
21    }
22}
23
24#[derive(Debug)]
25pub struct HeaderTemplate {
26    wire: &'static [u8],
27    fields: &'static [Field<'static>],
28}
29
30impl HeaderTemplate {
31    pub const fn new(wire: &'static [u8], fields: &'static [Field<'static>]) -> Self {
32        Self { wire, fields }
33    }
34}
35
36#[derive(Clone, Copy, Debug)]
37pub struct StaticHeaderFields(&'static HeaderTemplate);
38
39impl StaticHeaderFields {
40    pub const fn new(template: &'static HeaderTemplate) -> Self {
41        Self(template)
42    }
43}
44
45#[doc(hidden)]
46pub trait StaticHeaders {
47    type Fields<'a>: Iterator<Item = Field<'a>>
48    where
49        Self: 'a;
50
51    fn wire(&self) -> &'static [u8];
52
53    fn fields(&self) -> Self::Fields<'_>;
54}
55
56pub struct WireHeaderFields<'a> {
57    remaining: &'a [u8],
58}
59
60impl<'a> WireHeaderFields<'a> {
61    pub fn new(wire: &'a [u8]) -> Self {
62        Self { remaining: wire }
63    }
64}
65
66impl<'a> Iterator for WireHeaderFields<'a> {
67    type Item = Field<'a>;
68
69    fn next(&mut self) -> Option<Self::Item> {
70        loop {
71            let (line, remaining) = match self.remaining.iter().position(|&b| b == b'\n') {
72                Some(end) => (&self.remaining[..end], &self.remaining[end + 1..]),
73                None => (self.remaining, &[][..]),
74            };
75            self.remaining = remaining;
76            let line = line.strip_suffix(b"\r").unwrap_or(line);
77            if line.is_empty() {
78                if self.remaining.is_empty() {
79                    return None;
80                }
81                continue;
82            }
83            let Some(colon) = line.iter().position(|&b| b == b':') else {
84                if self.remaining.is_empty() {
85                    return None;
86                }
87                continue;
88            };
89            let name = &line[..colon];
90            let mut value = &line[colon + 1..];
91            while let Some((&b' ', rest)) = value.split_first() {
92                value = rest;
93            }
94            return Some(Field::new(name, value));
95        }
96    }
97}
98
99pub struct TemplateHeaderFields<'a> {
100    fields: slice::Iter<'a, Field<'static>>,
101}
102
103impl<'a> Iterator for TemplateHeaderFields<'a> {
104    type Item = Field<'a>;
105
106    fn next(&mut self) -> Option<Self::Item> {
107        self.fields
108            .next()
109            .map(|field| Field::new(field.name, field.value))
110    }
111
112    fn size_hint(&self) -> (usize, Option<usize>) {
113        self.fields.size_hint()
114    }
115}
116
117impl ExactSizeIterator for TemplateHeaderFields<'_> {}
118
119impl StaticHeaders for StaticHeaderBytes {
120    type Fields<'a> = WireHeaderFields<'a>;
121
122    fn wire(&self) -> &'static [u8] {
123        self.0
124    }
125
126    fn fields(&self) -> Self::Fields<'_> {
127        WireHeaderFields::new(self.0)
128    }
129}
130
131impl StaticHeaders for StaticHeaderFields {
132    type Fields<'a> = TemplateHeaderFields<'a>;
133
134    fn wire(&self) -> &'static [u8] {
135        self.0.wire
136    }
137
138    fn fields(&self) -> Self::Fields<'_> {
139        TemplateHeaderFields {
140            fields: self.0.fields.iter(),
141        }
142    }
143}
144
145impl<'req, const N: usize> HeadInner<'req, N> {
146    pub fn new(static_headers: &'static [u8], headers: Headers<'req, N>) -> Self {
147        Self {
148            static_headers: StaticHeaderBytes::new(static_headers),
149            headers,
150        }
151    }
152}
153
154impl<'req, const N: usize> HeadInner<'req, N, StaticHeaderFields> {
155    pub fn structured(static_headers: &'static HeaderTemplate, headers: Headers<'req, N>) -> Self {
156        Self {
157            static_headers: StaticHeaderFields::new(static_headers),
158            headers,
159        }
160    }
161}
162
163impl<'req, const N: usize, S> HeadInner<'req, N, S>
164where
165    S: StaticHeaders,
166{
167    pub fn static_headers(&self) -> &'static [u8] {
168        self.static_headers.wire()
169    }
170
171    pub fn headers(&self) -> &Headers<'req, N> {
172        &self.headers
173    }
174
175    pub(super) fn headers_mut(&mut self) -> &mut Headers<'req, N> {
176        &mut self.headers
177    }
178
179    pub fn wire_len(&self) -> usize {
180        self.static_headers
181            .wire()
182            .len()
183            .saturating_add(self.headers.wire_len())
184    }
185
186    pub(crate) fn wire_headers(&self) -> Shared {
187        if self.headers.is_empty() {
188            return Shared::from_static(self.static_headers.wire());
189        }
190        let capacity = self.wire_len();
191        match Owned::try_build_exact(capacity, |out| self.write_into(out)) {
192            Ok(out) => out.freeze(),
193            Err(_) => {
194                let mut out = Vec::with_capacity(capacity);
195                match self.write_into(&mut out) {
196                    Ok(()) => Shared::from(out),
197                    Err(error) => match error {},
198                }
199            }
200        }
201    }
202
203    fn write_into<W: ByteSink>(&self, out: &mut W) -> Result<(), W::Error> {
204        out.write_slice(self.static_headers.wire())?;
205        self.headers.write_into(out)
206    }
207
208    pub(crate) fn fields(&self) -> impl Iterator<Item = Field<'_>> + '_ {
209        self.static_headers.fields().chain(self.headers.fields())
210    }
211
212    pub fn write_slice(&self, out: &mut [u8]) -> Option<usize> {
213        let total = self.wire_len();
214        if out.len() < total {
215            return None;
216        }
217        let static_headers = self.static_headers.wire();
218        let n = static_headers.len();
219        out[..n].copy_from_slice(static_headers);
220        let m = self.headers.write(&mut out[n..]);
221        Some(n + m)
222    }
223}
224
225impl<const N: usize, S> HeaderSection for HeadInner<'_, N, S>
226where
227    S: StaticHeaders,
228{
229    fn wire_len(&self) -> usize {
230        Self::wire_len(self)
231    }
232    fn write_headers(&self, out: &mut ProvenWireWriter<'_>) {
233        out.put(self.static_headers.wire());
234        self.headers.write_wire(out);
235    }
236}