1use 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#[doc(hidden)]
15#[derive(Serialize, Debug)]
16#[serde(untagged)]
17pub enum LambdaResponse {
18 ApiGatewayV2(ApiGatewayV2Response),
19 Alb(AlbResponse),
20 ApiGateway(ApiGatewayResponse),
21}
22
23#[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#[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#[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
67fn 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
83fn 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
96fn 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
108impl 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 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
161pub trait IntoResponse {
178 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}