Skip to main content

sark_core/http/response/
stream.rs

1use std::pin::Pin;
2use std::task::Poll;
3
4use dope_fiber::{Context, Fiber};
5use http::StatusCode;
6use o3::buffer::{Owned, Shared};
7
8use super::wire_emit::{HeadWrite, TransferEncodingChunked};
9
10pub const CHUNK_TERMINATOR: &[u8; 5] = b"0\r\n\r\n";
11
12pub struct Stream<S> {
13    status: StatusCode,
14    wire_headers: Shared,
15    stream: S,
16}
17
18pub struct IterStream<I> {
19    iter: I,
20}
21
22impl<'d, I> Fiber<'d> for IterStream<I>
23where
24    I: Iterator<Item = Shared> + Unpin + 'static,
25{
26    type Output = Option<Shared>;
27    fn poll(self: Pin<&mut Self>, _cx: Pin<&mut Context<'_, 'd>>) -> Poll<Option<Shared>> {
28        Poll::Ready(self.get_mut().iter.next())
29    }
30}
31
32impl<S> Stream<S> {
33    pub fn new(stream: S) -> Self {
34        Self {
35            status: StatusCode::OK,
36            wire_headers: Shared::new(),
37            stream,
38        }
39    }
40
41    pub fn header(mut self, name: &[u8], value: &[u8]) -> Self {
42        let capacity = self
43            .wire_headers
44            .len()
45            .checked_add(name.len())
46            .and_then(|len| len.checked_add(value.len()))
47            .and_then(|len| len.checked_add(4))
48            .expect("stream header length overflow");
49        let mut buf = Owned::with_capacity(capacity);
50        buf.extend_from_slice(self.wire_headers.as_ref());
51        buf.extend_from_slice(name);
52        buf.extend_from_slice(b": ");
53        buf.extend_from_slice(value);
54        buf.extend_from_slice(b"\r\n");
55        self.wire_headers = buf.freeze();
56        self
57    }
58
59    pub fn write_head_stream(self, out: &mut [u8], date: &[u8; 29]) -> Option<(usize, S)> {
60        let status_str = self.status.as_str().as_bytes();
61        let reason = self
62            .status
63            .canonical_reason()
64            .map(str::as_bytes)
65            .unwrap_or(b"");
66
67        let head = HeadWrite {
68            status_str,
69            reason,
70            headers: self.wire_headers.as_ref(),
71            framing: TransferEncodingChunked,
72        };
73        if out.len() < head.wire_len() {
74            return None;
75        }
76
77        let mut off = 0usize;
78        head.write(out, &mut off, date);
79        Some((off, self.stream))
80    }
81}
82
83impl<II> Stream<IterStream<II>>
84where
85    II: Iterator<Item = Shared> + Unpin + 'static,
86{
87    pub fn from_chunks<I>(iter: I) -> Self
88    where
89        I: IntoIterator<Item = Shared, IntoIter = II>,
90    {
91        Self::new(IterStream {
92            iter: iter.into_iter(),
93        })
94    }
95}
96
97impl<S> super::IntoServeResponse<'static> for Stream<S> {
98    fn into_serve_response(self) -> super::ServeInner<'static> {
99        unreachable!(
100            "Stream::into_serve_response — Stream routes \
101             route through Stream::write_head_stream"
102        )
103    }
104}