apimock_server/
response_handler.rs

1use http_body_util::{BodyExt, Empty, Full};
2use hyper::{
3    HeaderMap, StatusCode,
4    body::{Body, Bytes},
5    header::{
6        ACCESS_CONTROL_ALLOW_CREDENTIALS, ACCESS_CONTROL_ALLOW_ORIGIN, CONTENT_LENGTH, HeaderName,
7        HeaderValue, ORIGIN, VARY,
8    },
9    http::response::Builder,
10};
11
12use std::{collections::HashMap, str::FromStr};
13
14use super::{
15    constant::DEFAULT_RESPONSE_HEADERS, response::error_response::internal_server_error_response,
16};
17use crate::types::BoxBody;
18
19#[derive(Clone, Default)]
20pub enum BodyKind {
21    #[default]
22    Empty,
23    Text(String),
24    Binary(Vec<u8>),
25}
26
27#[derive(Default)]
28pub struct ResponseHandler {
29    response_builder: Builder,
30    status: Option<StatusCode>,
31    headers: HashMap<String, Option<String>>,
32    body_kind: BodyKind,
33}
34
35impl ResponseHandler {
36    /// build response
37    pub fn into_response(
38        self,
39        request_headers: &HeaderMap,
40    ) -> Result<hyper::Response<BoxBody>, hyper::http::Error> {
41        // - body + content-length
42        let response = match self.body_kind {
43            BodyKind::Text(s) => self
44                .response_builder
45                .body(Full::new(Bytes::from(s.to_owned())).boxed()),
46            BodyKind::Binary(b) => self
47                .response_builder
48                .body(Full::new(Bytes::from(b)).boxed()),
49            BodyKind::Empty => self.response_builder.body(Empty::new().boxed()),
50        };
51
52        let mut response = match response {
53            Ok(x) => x,
54            Err(err) => {
55                return internal_server_error_response(
56                    &format!("failed to create response: {}", err),
57                    request_headers,
58                );
59            }
60        };
61
62        // - http status code
63        *response.status_mut() = if let Some(status) = self.status {
64            status
65        } else {
66            StatusCode::OK
67        };
68
69        // - content-length
70        let content_length = response.body().size_hint().exact().unwrap_or_default();
71
72        let headers = response.headers_mut();
73
74        headers.insert(CONTENT_LENGTH, HeaderValue::from(content_length));
75
76        // - the other default headers
77        for (header_key, header_value) in default_response_headers(request_headers).iter() {
78            headers.insert(header_key, header_value.to_owned());
79        }
80
81        // - additional custom headers passed from caller
82        for (header_key, header_value) in self.headers {
83            match HeaderName::from_str(header_key.as_str()) {
84                Ok(header_key) => {
85                    match HeaderValue::from_str(header_value.unwrap_or_default().as_str()) {
86                        Ok(header_value) => {
87                            headers.insert(header_key, header_value);
88                        }
89                        Err(err) => {
90                            log::warn!(
91                                "failed to create header with the header value (header key = {}) ({})",
92                                header_key,
93                                err
94                            );
95                            headers.insert(header_key, HeaderValue::from_static(""));
96                        }
97                    }
98                }
99                Err(err) => log::warn!(
100                    "failed to create header with the header key: {} ({})",
101                    header_key,
102                    err
103                ),
104            };
105        }
106
107        Ok(response)
108    }
109
110    /// set http status code
111    pub fn with_status(mut self, status: &StatusCode) -> Self {
112        self.status = Some(status.to_owned());
113        self
114    }
115
116    /// add custom header
117    pub fn with_header(mut self, key: impl Into<String>, value: Option<impl Into<String>>) -> Self {
118        self.headers.insert(key.into(), value.map(|x| x.into()));
119        self
120    }
121
122    /// add custom headers
123    pub fn with_headers<K, V, I>(mut self, headers: I) -> Self
124    where
125        K: Into<String>,
126        V: Into<String>,
127        I: IntoIterator<Item = (K, Option<V>)>,
128    {
129        for (key, value) in headers {
130            self.headers.insert(key.into(), value.map(|x| x.into()));
131        }
132        self
133    }
134
135    /// Apply an operator's custom `respond.headers`, always **last** —
136    /// after whichever `with_text`/`with_json_body`/`with_binary_body`
137    /// call already set a default `content-type` (RFC 065's override
138    /// rule: an explicit `content-type` always wins over the derived
139    /// default, on every body source, uniformly).
140    ///
141    /// # Why one method instead of each call site's own `if let`
142    ///
143    /// Before this, every response-building function repeated
144    /// `if let Some(custom_headers) = custom_headers { response_handler
145    /// = response_handler.with_headers(custom_headers.to_owned()); }`
146    /// itself, and the two places that got the ordering wrong
147    /// (`json_response`, and `FileResponse`'s own binary-file path) each
148    /// silently let the derived content-type win instead — because
149    /// `self.headers` is a plain `HashMap`, whichever call happens last
150    /// wins for that key, and there was nothing forcing "last" to always
151    /// mean "the custom headers." Routing every call site through this
152    /// one method, called only after the body is set, makes that
153    /// ordering the only way to call it — not a convention that can
154    /// drift a third time.
155    pub fn with_custom_headers(
156        self,
157        custom_headers: Option<&HashMap<String, Option<String>>>,
158    ) -> Self {
159        match custom_headers {
160            Some(custom_headers) => self.with_headers(custom_headers.to_owned()),
161            None => self,
162        }
163    }
164
165    /// add text to body
166    pub fn with_text(mut self, text: impl Into<String>, content_type: Option<&str>) -> Self {
167        let content_type = if let Some(content_type) = content_type {
168            content_type.into()
169        } else {
170            "text/plain; charset=utf-8".to_owned()
171        };
172        self.headers
173            .insert("content-type".into(), Some(content_type));
174
175        self.body_kind = BodyKind::Text(text.into());
176        self
177    }
178
179    /// treat response as json
180    pub fn with_json_body(mut self, body: impl Into<String>) -> Self {
181        self.headers
182            .insert("content-type".into(), Some("application/json".into()));
183        self.body_kind = BodyKind::Text(body.into());
184        self
185    }
186
187    /// treat response as json
188    pub fn with_binary_body(
189        mut self,
190        body: Vec<u8>,
191        content_type: Option<impl Into<String>>,
192    ) -> Self {
193        let content_type = if let Some(content_type) = content_type {
194            content_type.into()
195        } else {
196            "application/octet-stream".to_owned()
197        };
198        self.headers
199            .insert("content-type".into(), Some(content_type));
200
201        self.body_kind = BodyKind::Binary(body);
202
203        self
204    }
205}
206
207/// default response headers key-value pairs
208pub fn default_response_headers(request_headers: &HeaderMap) -> HeaderMap {
209    let mut header_map_src = Vec::with_capacity(DEFAULT_RESPONSE_HEADERS.len() + 1);
210
211    // resource
212    // - the other default headers but access-control-allow-origin, vary
213    header_map_src.extend(
214        DEFAULT_RESPONSE_HEADERS
215            .iter()
216            .map(|(k, v)| (k.to_string(), v.to_string())),
217    );
218
219    // - access-control-allow-origin, vary
220    let origin = if is_likely_authenticated_request(request_headers) {
221        request_headers.get(ORIGIN).map(|x| x.to_owned())
222    } else {
223        None
224    };
225    let (origin, vary) = if let Some(origin) = origin {
226        header_map_src.push((
227            ACCESS_CONTROL_ALLOW_CREDENTIALS.to_string(),
228            "true".to_owned(),
229        ));
230
231        (origin, HeaderValue::from_static("Origin"))
232    } else {
233        (HeaderValue::from_static("*"), HeaderValue::from_static("*"))
234    };
235    header_map_src.push((
236        ACCESS_CONTROL_ALLOW_ORIGIN.to_string(),
237        origin.to_str().unwrap_or_default().to_owned(),
238    ));
239    header_map_src.push((
240        VARY.to_string(),
241        vary.to_str().unwrap_or_default().to_owned(),
242    ));
243
244    // header map
245    header_map_src
246        .iter()
247        .fold(HeaderMap::new(), |mut ret, (header_key, header_value)| {
248            match HeaderName::from_str(header_key) {
249                Ok(header_key) => match HeaderValue::from_str(header_value.as_str()) {
250                    Ok(header_value) => {
251                        ret.insert(header_key, header_value);
252                        ret
253                    }
254                    Err(err) => {
255                        log::warn!(
256                            "only header key set because failed to get header value: {} [key = {}] ({})",
257                            header_value.as_str(),
258                            header_key,
259                            err
260                        );
261                        ret.insert(header_key, HeaderValue::from_static(""));
262                        ret
263                    }
264                },
265                Err(err) => {
266                    log::warn!("failed to set header key: {} ({})", header_key, err);
267                    ret
268                }
269            }
270        })
271}
272
273/// guess if the request is likely related to authentication
274fn is_likely_authenticated_request(request_headers: &HeaderMap) -> bool {
275    request_headers.contains_key("cookie") || request_headers.contains_key("authorization")
276}
apimock_server/response_handler.rs

apimock_server/
response_handler.rs
