Skip to main content

sark_core/http/response/
stream.rs

1use std::pin::Pin;
2use std::task::Poll;
3
4use dope_fiber::abi::Fiber;
5use dope_fiber::raw::task::Context;
6use http::StatusCode;
7use o3::buffer::Shared;
8
9use super::wire_emit::{HeadWrite, ProvenWireWriter, TransferEncodingChunked};
10
11pub const CHUNK_TERMINATOR: &[u8; 5] = b"0\r\n\r\n";
12
13pub struct Stream<S> {
14    status: StatusCode,
15    wire_headers: Vec<Shared>,
16    stream: S,
17}
18
19pub struct IterStream<I> {
20    iter: I,
21}
22
23impl<'d, I> Fiber<'d> for IterStream<I>
24where
25    I: Iterator<Item = Shared> + Unpin + 'static,
26{
27    type Output = Option<Shared>;
28    fn poll(self: Pin<&mut Self>, _cx: Pin<&mut Context<'_, 'd>>) -> Poll<Option<Shared>> {
29        Poll::Ready(self.get_mut().iter.next())
30    }
31}
32
33impl<S> Stream<S> {
34    pub fn new(stream: S) -> Self {
35        Self {
36            status: StatusCode::OK,
37            wire_headers: Vec::new(),
38            stream,
39        }
40    }
41
42    pub fn header(mut self, name: &[u8], value: &[u8]) -> Self {
43        let mut buf = Vec::new();
44        buf.extend_from_slice(name);
45        buf.extend_from_slice(b": ");
46        buf.extend_from_slice(value);
47        buf.extend_from_slice(b"\r\n");
48        self.wire_headers.push(Shared::from(buf));
49        self
50    }
51
52    pub fn write_head_stream(self, out: &mut [u8], date: &[u8; 29]) -> Option<(usize, S)> {
53        let head = HeadWrite {
54            status: self.status,
55            headers: self.wire_headers.as_slice(),
56            framing: TransferEncodingChunked,
57        };
58        let head_len = head.wire_len();
59        let mut out = ProvenWireWriter::new(out, head_len)?;
60        head.emit(&mut out, date);
61        Some((out.finish(head_len), self.stream))
62    }
63}
64
65impl<II> Stream<IterStream<II>>
66where
67    II: Iterator<Item = Shared> + Unpin + 'static,
68{
69    pub fn from_chunks<I>(iter: I) -> Self
70    where
71        I: IntoIterator<Item = Shared, IntoIter = II>,
72    {
73        Self::new(IterStream {
74            iter: iter.into_iter(),
75        })
76    }
77}