connect_rpc/response/
builder.rs

1use http::{header, HeaderMap, HeaderName, StatusCode};
2
3use crate::{
4    common::{is_valid_http_token, CONNECT_CONTENT_ENCODING, CONTENT_TYPE_PREFIX},
5    metadata::Metadata,
6    Error,
7};
8
9use super::{StreamingResponse, UnaryResponse};
10
11#[derive(Debug, Default)]
12pub struct ResponseBuilder {
13    status: StatusCode,
14    metadata: HeaderMap,
15    message_codec: Option<String>,
16    content_encoding: Option<String>,
17}
18
19impl ResponseBuilder {
20    /// Sets the response status code.
21    pub fn status(mut self, status: StatusCode) -> Self {
22        self.status = status;
23        self
24    }
25
26    /// Appends ASCII metadata to the response.
27    pub fn ascii_metadata(
28        mut self,
29        key: impl TryInto<HeaderName, Error: Into<Error>>,
30        val: impl Into<String>,
31    ) -> Result<Self, Error> {
32        self.metadata.append_ascii(key, val)?;
33        Ok(self)
34    }
35
36    /// Appends binary metadata to the response.
37    pub fn binary_metadata(
38        mut self,
39        key: impl TryInto<HeaderName, Error: Into<Error>>,
40        val: impl AsRef<[u8]>,
41    ) -> Result<Self, Error> {
42        self.metadata.append_binary(key, val)?;
43        Ok(self)
44    }
45
46    /// Sets the message codec for this response.
47    ///
48    /// Typical codecs are 'json' and 'proto', corresponding to the
49    /// `content-type`s `application/json` and `application/proto`.
50    ///
51    /// The caller is responsible for making sure the response payload matches
52    /// this message codec.
53    pub fn message_codec(mut self, message_codec: impl Into<String>) -> Result<Self, Error> {
54        let mut message_codec: String = message_codec.into();
55        message_codec.make_ascii_lowercase();
56        if !is_valid_http_token(&message_codec) {
57            return Err(Error::invalid_request("invalid message codec"));
58        }
59        self.message_codec = Some(message_codec);
60        Ok(self)
61    }
62
63    /// Sets the response content encoding (e.g. compression).
64    pub fn content_encoding(mut self, content_encoding: impl Into<String>) -> Result<Self, Error> {
65        let content_encoding = content_encoding.into();
66        if !is_valid_http_token(&content_encoding) {
67            return Err(Error::invalid_request("invalid content encoding"));
68        }
69        self.content_encoding = Some(content_encoding);
70        Ok(self)
71    }
72
73    /// Build logic common to all responses.
74    fn common_response<T>(&mut self, body: T) -> http::Response<T> {
75        let mut resp = http::Response::new(body);
76        *resp.status_mut() = self.status;
77        *resp.headers_mut() = std::mem::take(&mut self.metadata);
78        resp
79    }
80
81    /// Builds a [`UnaryResponse`].
82    pub fn unary<T>(mut self, body: T) -> Result<UnaryResponse<T>, Error> {
83        let mut resp = self.common_response(body);
84        // Unary-Content-Type → "content-type" "application/" Message-Codec
85        if let Some(message_codec) = &self.message_codec {
86            resp.headers_mut().insert(
87                header::CONTENT_TYPE,
88                (format!("{CONTENT_TYPE_PREFIX}{message_codec}")).try_into()?,
89            );
90        }
91        // Content-Encoding → "content-encoding" Content-Coding
92        if let Some(content_encoding) = self.content_encoding.take() {
93            resp.headers_mut()
94                .insert(header::CONTENT_ENCODING, content_encoding.try_into()?);
95        }
96        Ok(resp.into())
97    }
98
99    /// Builds a [`StreamingResponse`].
100    pub fn streaming<T>(mut self, body: T) -> Result<StreamingResponse<T>, Error> {
101        let mut resp = self.common_response(body);
102        // Streaming-Content-Type → "content-type" "application/connect+" [...]
103        if let Some(message_codec) = &self.message_codec {
104            resp.headers_mut().insert(
105                header::CONTENT_TYPE,
106                (format!("{CONTENT_TYPE_PREFIX}{message_codec}")).try_into()?,
107            );
108        }
109        // Streaming-Content-Encoding → "connect-content-encoding" Content-Coding
110        if let Some(content_encoding) = self.content_encoding.take() {
111            resp.headers_mut()
112                .insert(CONNECT_CONTENT_ENCODING, content_encoding.try_into()?);
113        }
114        Ok(resp.into())
115    }
116}