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    /// add text to body
136    pub fn with_text(mut self, text: impl Into<String>, content_type: Option<&str>) -> Self {
137        let content_type = if let Some(content_type) = content_type {
138            content_type.into()
139        } else {
140            "text/plain; charset=utf-8".to_owned()
141        };
142        self.headers
143            .insert("content-type".into(), Some(content_type));
144
145        self.body_kind = BodyKind::Text(text.into());
146        self
147    }
148
149    /// treat response as json
150    pub fn with_json_body(mut self, body: impl Into<String>) -> Self {
151        self.headers
152            .insert("content-type".into(), Some("application/json".into()));
153        self.body_kind = BodyKind::Text(body.into());
154        self
155    }
156
157    /// treat response as json
158    pub fn with_binary_body(
159        mut self,
160        body: Vec<u8>,
161        content_type: Option<impl Into<String>>,
162    ) -> Self {
163        let content_type = if let Some(content_type) = content_type {
164            content_type.into()
165        } else {
166            "application/octet-stream".to_owned()
167        };
168        self.headers
169            .insert("content-type".into(), Some(content_type));
170
171        self.body_kind = BodyKind::Binary(body);
172
173        self
174    }
175}
176
177/// default response headers key-value pairs
178pub fn default_response_headers(request_headers: &HeaderMap) -> HeaderMap {
179    let mut header_map_src = Vec::with_capacity(DEFAULT_RESPONSE_HEADERS.len() + 1);
180
181    // resource
182    // - the other default headers but access-control-allow-origin, vary
183    header_map_src.extend(
184        DEFAULT_RESPONSE_HEADERS
185            .iter()
186            .map(|(k, v)| (k.to_string(), v.to_string())),
187    );
188
189    // - access-control-allow-origin, vary
190    let origin = if is_likely_authenticated_request(request_headers) {
191        request_headers.get(ORIGIN).map(|x| x.to_owned())
192    } else {
193        None
194    };
195    let (origin, vary) = if let Some(origin) = origin {
196        header_map_src.push((
197            ACCESS_CONTROL_ALLOW_CREDENTIALS.to_string(),
198            "true".to_owned(),
199        ));
200
201        (origin, HeaderValue::from_static("Origin"))
202    } else {
203        (HeaderValue::from_static("*"), HeaderValue::from_static("*"))
204    };
205    header_map_src.push((
206        ACCESS_CONTROL_ALLOW_ORIGIN.to_string(),
207        origin.to_str().unwrap_or_default().to_owned(),
208    ));
209    header_map_src.push((
210        VARY.to_string(),
211        vary.to_str().unwrap_or_default().to_owned(),
212    ));
213
214    // header map
215    header_map_src
216        .iter()
217        .fold(HeaderMap::new(), |mut ret, (header_key, header_value)| {
218            match HeaderName::from_str(header_key) {
219                Ok(header_key) => match HeaderValue::from_str(header_value.as_str()) {
220                    Ok(header_value) => {
221                        ret.insert(header_key, header_value);
222                        ret
223                    }
224                    Err(err) => {
225                        log::warn!(
226                            "only header key set because failed to get header value: {} [key = {}] ({})",
227                            header_value.as_str(),
228                            header_key,
229                            err
230                        );
231                        ret.insert(header_key, HeaderValue::from_static(""));
232                        ret
233                    }
234                },
235                Err(err) => {
236                    log::warn!("failed to set header key: {} ({})", header_key, err);
237                    ret
238                }
239            }
240        })
241}
242
243/// guess if the request is likely related to authentication
244fn is_likely_authenticated_request(request_headers: &HeaderMap) -> bool {
245    request_headers.contains_key("cookie") || request_headers.contains_key("authorization")
246}
apimock_server/response_handler.rs

apimock_server/
response_handler.rs
