Skip to main content

sark_core/http/response/
stream.rs

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