fluke/
responder.rs

1use http::header;
2
3use crate::{h1::body::BodyWriteMode, Body, BodyChunk, Headers, HeadersExt, Response};
4use fluke_buffet::PieceCore;
5
6pub trait ResponseState {}
7
8pub struct ExpectResponseHeaders;
9impl ResponseState for ExpectResponseHeaders {}
10
11pub struct ExpectResponseBody {
12    mode: BodyWriteMode,
13}
14impl ResponseState for ExpectResponseBody {}
15
16pub struct ResponseDone;
17impl ResponseState for ResponseDone {}
18
19pub struct Responder<E, S>
20where
21    E: Encoder,
22    S: ResponseState,
23{
24    pub(crate) encoder: E,
25    pub(crate) state: S,
26}
27
28impl<E> Responder<E, ExpectResponseHeaders>
29where
30    E: Encoder,
31{
32    /// Send an informational status code, cf. <https://httpwg.org/specs/rfc9110.html#status.1xx>
33    /// Errors out if the response status is not 1xx
34    pub async fn write_interim_response(&mut self, res: Response) -> eyre::Result<()> {
35        if !res.status.is_informational() {
36            return Err(eyre::eyre!("interim response must have status code 1xx"));
37        }
38
39        self.encoder.write_response(res).await?;
40        Ok(())
41    }
42
43    /// Send the final response headers
44    /// Errors out if the response status is < 200.
45    /// Errors out if the client sent `expect: 100-continue`
46    pub async fn write_final_response(
47        mut self,
48        mut res: Response,
49    ) -> eyre::Result<Responder<E, ExpectResponseBody>> {
50        if res.status.is_informational() {
51            return Err(eyre::eyre!("final response must have status code >= 200"));
52        }
53
54        let mode = if res.means_empty_body() {
55            // do nothing
56            BodyWriteMode::Empty
57        } else {
58            match res.headers.content_length() {
59                Some(0) => BodyWriteMode::Empty,
60                Some(len) => {
61                    // TODO: can probably save that heap allocation
62                    res.headers
63                        .insert(header::CONTENT_LENGTH, format!("{len}").into_bytes().into());
64                    BodyWriteMode::ContentLength
65                }
66                None => {
67                    res.headers
68                        .insert(header::TRANSFER_ENCODING, "chunked".into());
69                    BodyWriteMode::Chunked
70                }
71            }
72        };
73        self.encoder.write_response(res).await?;
74
75        Ok(Responder {
76            state: ExpectResponseBody { mode },
77            encoder: self.encoder,
78        })
79    }
80
81    /// Writes a response with the given body. Sets `content-length` or
82    /// `transfer-encoding` as needed.
83    pub async fn write_final_response_with_body(
84        self,
85        mut res: Response,
86        body: &mut impl Body,
87    ) -> eyre::Result<Responder<E, ResponseDone>> {
88        if let Some(clen) = body.content_len() {
89            res.headers
90                .entry(header::CONTENT_LENGTH)
91                .or_insert_with(|| {
92                    // TODO: can probably get rid of this heap allocation, also
93                    // use `itoa`
94                    format!("{clen}").into_bytes().into()
95                });
96        }
97
98        let mut this = self.write_final_response(res).await?;
99
100        loop {
101            match body.next_chunk().await? {
102                BodyChunk::Chunk(chunk) => {
103                    this.write_chunk(chunk).await?;
104                }
105                BodyChunk::Done { trailers } => {
106                    // TODO: should we do something here in case of
107                    // content-length mismatches?
108                    return this.finish_body(trailers).await;
109                }
110            }
111        }
112    }
113}
114
115impl<E> Responder<E, ExpectResponseBody>
116where
117    E: Encoder,
118{
119    /// Send a response body chunk. Errors out if sending more than the
120    /// announced content-length.
121    pub async fn write_chunk(&mut self, chunk: PieceCore) -> eyre::Result<()> {
122        self.encoder.write_body_chunk(chunk, self.state.mode).await
123    }
124
125    /// Finish the body, with optional trailers, cf. <https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/TE>
126    /// Errors out if the sent body doesn't match the announced content-length.
127    /// Errors out if trailers that weren't announced are being sent, or if the client
128    /// didn't explicitly announce it accepted trailers, or if the response is a 204,
129    /// 205 or 304, or if the body wasn't sent with chunked transfer encoding.
130    pub async fn finish_body(
131        mut self,
132        trailers: Option<Box<Headers>>,
133    ) -> eyre::Result<Responder<E, ResponseDone>> {
134        self.encoder.write_body_end(self.state.mode).await?;
135
136        if let Some(trailers) = trailers {
137            self.encoder.write_trailers(trailers).await?;
138        }
139
140        // TODO: check announced content-length size vs actual, etc.
141
142        Ok(Responder {
143            state: ResponseDone,
144            encoder: self.encoder,
145        })
146    }
147}
148
149impl<E> Responder<E, ResponseDone>
150where
151    E: Encoder,
152{
153    pub fn into_inner(self) -> E {
154        self.encoder
155    }
156}
157
158#[allow(async_fn_in_trait)] // we never require Send
159pub trait Encoder {
160    async fn write_response(&mut self, res: Response) -> eyre::Result<()>;
161    async fn write_body_chunk(&mut self, chunk: PieceCore, mode: BodyWriteMode)
162        -> eyre::Result<()>;
163    async fn write_body_end(&mut self, mode: BodyWriteMode) -> eyre::Result<()>;
164    async fn write_trailers(&mut self, trailers: Box<Headers>) -> eyre::Result<()>;
165}