fishrock_lambda_http/
response.rs

1//! Response types
2
3use crate::{body::Body, request::RequestOrigin};
4use http::{
5    header::{HeaderMap, HeaderValue, CONTENT_TYPE, SET_COOKIE},
6    Response,
7};
8use serde::{
9    ser::{Error as SerError, SerializeMap, SerializeSeq},
10    Serialize, Serializer,
11};
12
13/// Representation of Lambda response
14#[doc(hidden)]
15#[derive(Serialize, Debug)]
16#[serde(untagged)]
17pub enum LambdaResponse {
18    ApiGatewayV2(ApiGatewayV2Response),
19    Alb(AlbResponse),
20    ApiGateway(ApiGatewayResponse),
21}
22
23/// Representation of API Gateway v2 lambda response
24#[doc(hidden)]
25#[derive(Serialize, Debug)]
26#[serde(rename_all = "camelCase")]
27pub struct ApiGatewayV2Response {
28    status_code: u16,
29    #[serde(serialize_with = "serialize_headers")]
30    headers: HeaderMap<HeaderValue>,
31    #[serde(serialize_with = "serialize_headers_slice")]
32    cookies: Vec<HeaderValue>,
33    #[serde(skip_serializing_if = "Option::is_none")]
34    body: Option<Body>,
35    is_base64_encoded: bool,
36}
37
38/// Representation of ALB lambda response
39#[doc(hidden)]
40#[derive(Serialize, Debug)]
41#[serde(rename_all = "camelCase")]
42pub struct AlbResponse {
43    status_code: u16,
44    status_description: String,
45    #[serde(serialize_with = "serialize_headers")]
46    headers: HeaderMap<HeaderValue>,
47    #[serde(skip_serializing_if = "Option::is_none")]
48    body: Option<Body>,
49    is_base64_encoded: bool,
50}
51
52/// Representation of API Gateway lambda response
53#[doc(hidden)]
54#[derive(Serialize, Debug)]
55#[serde(rename_all = "camelCase")]
56pub struct ApiGatewayResponse {
57    status_code: u16,
58    #[serde(serialize_with = "serialize_headers")]
59    headers: HeaderMap<HeaderValue>,
60    #[serde(serialize_with = "serialize_multi_value_headers")]
61    multi_value_headers: HeaderMap<HeaderValue>,
62    #[serde(skip_serializing_if = "Option::is_none")]
63    body: Option<Body>,
64    is_base64_encoded: bool,
65}
66
67/// Serialize a http::HeaderMap into a serde str => str map
68fn serialize_multi_value_headers<S>(headers: &HeaderMap<HeaderValue>, serializer: S) -> Result<S::Ok, S::Error>
69where
70    S: Serializer,
71{
72    let mut map = serializer.serialize_map(Some(headers.keys_len()))?;
73    for key in headers.keys() {
74        let mut map_values = Vec::new();
75        for value in headers.get_all(key) {
76            map_values.push(value.to_str().map_err(S::Error::custom)?)
77        }
78        map.serialize_entry(key.as_str(), &map_values)?;
79    }
80    map.end()
81}
82
83/// Serialize a http::HeaderMap into a serde str => Vec<str> map
84fn serialize_headers<S>(headers: &HeaderMap<HeaderValue>, serializer: S) -> Result<S::Ok, S::Error>
85where
86    S: Serializer,
87{
88    let mut map = serializer.serialize_map(Some(headers.keys_len()))?;
89    for key in headers.keys() {
90        let map_value = headers[key].to_str().map_err(S::Error::custom)?;
91        map.serialize_entry(key.as_str(), map_value)?;
92    }
93    map.end()
94}
95
96/// Serialize a &[HeaderValue] into a Vec<str>
97fn serialize_headers_slice<S>(headers: &[HeaderValue], serializer: S) -> Result<S::Ok, S::Error>
98where
99    S: Serializer,
100{
101    let mut seq = serializer.serialize_seq(Some(headers.len()))?;
102    for header in headers {
103        seq.serialize_element(header.to_str().map_err(S::Error::custom)?)?;
104    }
105    seq.end()
106}
107
108/// tranformation from http type to internal type
109impl LambdaResponse {
110    pub fn from_response<T>(request_origin: &RequestOrigin, value: Response<T>) -> Self
111    where
112        T: Into<Body>,
113    {
114        let (parts, bod) = value.into_parts();
115        let (is_base64_encoded, body) = match bod.into() {
116            Body::Empty => (false, None),
117            b @ Body::Text(_) => (false, Some(b)),
118            b @ Body::Binary(_) => (true, Some(b)),
119        };
120
121        let mut headers = parts.headers;
122        let status_code = parts.status.as_u16();
123
124        match request_origin {
125            RequestOrigin::ApiGatewayV2 => {
126                // ApiGatewayV2 expects the set-cookies headers to be in the "cookies" attribute,
127                // so remove them from the headers.
128                let cookies: Vec<HeaderValue> = headers.get_all(SET_COOKIE).iter().cloned().collect();
129                headers.remove(SET_COOKIE);
130
131                LambdaResponse::ApiGatewayV2(ApiGatewayV2Response {
132                    body,
133                    status_code,
134                    is_base64_encoded,
135                    cookies,
136                    headers,
137                })
138            }
139            RequestOrigin::ApiGateway => LambdaResponse::ApiGateway(ApiGatewayResponse {
140                body,
141                status_code,
142                is_base64_encoded,
143                headers: headers.clone(),
144                multi_value_headers: headers,
145            }),
146            RequestOrigin::Alb => LambdaResponse::Alb(AlbResponse {
147                body,
148                status_code,
149                is_base64_encoded,
150                headers,
151                status_description: format!(
152                    "{} {}",
153                    status_code,
154                    parts.status.canonical_reason().unwrap_or_default()
155                ),
156            }),
157        }
158    }
159}
160
161/// A conversion of self into a `Response<Body>` for various types.
162///
163/// Implementations for `Response<B> where B: Into<Body>`,
164/// `B where B: Into<Body>` and `serde_json::Value` are provided
165/// by default.
166///
167/// # Example
168///
169/// ```rust
170/// use lambda_http::{Body, IntoResponse, Response};
171///
172/// assert_eq!(
173///   "hello".into_response().body(),
174///   Response::new(Body::from("hello")).body()
175/// );
176/// ```
177pub trait IntoResponse {
178    /// Return a translation of `self` into a `Response<Body>`
179    fn into_response(self) -> Response<Body>;
180}
181
182impl<B> IntoResponse for Response<B>
183where
184    B: Into<Body>,
185{
186    fn into_response(self) -> Response<Body> {
187        let (parts, body) = self.into_parts();
188        Response::from_parts(parts, body.into())
189    }
190}
191
192impl<B> IntoResponse for B
193where
194    B: Into<Body>,
195{
196    fn into_response(self) -> Response<Body> {
197        Response::new(self.into())
198    }
199}
200
201impl IntoResponse for serde_json::Value {
202    fn into_response(self) -> Response<Body> {
203        Response::builder()
204            .header(CONTENT_TYPE, "application/json")
205            .body(
206                serde_json::to_string(&self)
207                    .expect("unable to serialize serde_json::Value")
208                    .into(),
209            )
210            .expect("unable to build http::Response")
211    }
212}
213
214#[cfg(test)]
215mod tests {
216    use super::{
217        AlbResponse, ApiGatewayResponse, ApiGatewayV2Response, Body, IntoResponse, LambdaResponse, RequestOrigin,
218    };
219    use http::{header::CONTENT_TYPE, Response};
220    use serde_json::{self, json};
221
222    fn api_gateway_response() -> ApiGatewayResponse {
223        ApiGatewayResponse {
224            status_code: 200,
225            headers: Default::default(),
226            multi_value_headers: Default::default(),
227            body: Default::default(),
228            is_base64_encoded: Default::default(),
229        }
230    }
231
232    fn alb_response() -> AlbResponse {
233        AlbResponse {
234            status_code: 200,
235            status_description: "200 OK".to_string(),
236            headers: Default::default(),
237            body: Default::default(),
238            is_base64_encoded: Default::default(),
239        }
240    }
241
242    fn api_gateway_v2_response() -> ApiGatewayV2Response {
243        ApiGatewayV2Response {
244            status_code: 200,
245            headers: Default::default(),
246            body: Default::default(),
247            cookies: Default::default(),
248            is_base64_encoded: Default::default(),
249        }
250    }
251
252    #[test]
253    fn json_into_response() {
254        let response = json!({ "hello": "lambda"}).into_response();
255        match response.body() {
256            Body::Text(json) => assert_eq!(json, r#"{"hello":"lambda"}"#),
257            _ => panic!("invalid body"),
258        }
259        assert_eq!(
260            response
261                .headers()
262                .get(CONTENT_TYPE)
263                .map(|h| h.to_str().expect("invalid header")),
264            Some("application/json")
265        )
266    }
267
268    #[test]
269    fn text_into_response() {
270        let response = "text".into_response();
271        match response.body() {
272            Body::Text(text) => assert_eq!(text, "text"),
273            _ => panic!("invalid body"),
274        }
275    }
276
277    #[test]
278    fn serialize_body_for_api_gateway() {
279        let mut resp = api_gateway_response();
280        resp.body = Some("foo".into());
281        assert_eq!(
282            serde_json::to_string(&resp).expect("failed to serialize response"),
283            r#"{"statusCode":200,"headers":{},"multiValueHeaders":{},"body":"foo","isBase64Encoded":false}"#
284        );
285    }
286
287    #[test]
288    fn serialize_body_for_alb() {
289        let mut resp = alb_response();
290        resp.body = Some("foo".into());
291        assert_eq!(
292            serde_json::to_string(&resp).expect("failed to serialize response"),
293            r#"{"statusCode":200,"statusDescription":"200 OK","headers":{},"body":"foo","isBase64Encoded":false}"#
294        );
295    }
296
297    #[test]
298    fn serialize_body_for_api_gateway_v2() {
299        let mut resp = api_gateway_v2_response();
300        resp.body = Some("foo".into());
301        assert_eq!(
302            serde_json::to_string(&resp).expect("failed to serialize response"),
303            r#"{"statusCode":200,"headers":{},"cookies":[],"body":"foo","isBase64Encoded":false}"#
304        );
305    }
306
307    #[test]
308    fn serialize_multi_value_headers() {
309        let res = LambdaResponse::from_response(
310            &RequestOrigin::ApiGateway,
311            Response::builder()
312                .header("multi", "a")
313                .header("multi", "b")
314                .body(Body::from(()))
315                .expect("failed to create response"),
316        );
317        let json = serde_json::to_string(&res).expect("failed to serialize to json");
318        assert_eq!(
319            json,
320            r#"{"statusCode":200,"headers":{"multi":"a"},"multiValueHeaders":{"multi":["a","b"]},"isBase64Encoded":false}"#
321        )
322    }
323
324    #[test]
325    fn serialize_cookies() {
326        let res = LambdaResponse::from_response(
327            &RequestOrigin::ApiGatewayV2,
328            Response::builder()
329                .header("set-cookie", "cookie1=a")
330                .header("set-cookie", "cookie2=b")
331                .body(Body::from(()))
332                .expect("failed to create response"),
333        );
334        let json = serde_json::to_string(&res).expect("failed to serialize to json");
335        assert_eq!(
336            json,
337            r#"{"statusCode":200,"headers":{},"cookies":["cookie1=a","cookie2=b"],"isBase64Encoded":false}"#
338        )
339    }
340}