trillium_aws_lambda/
response.rs1use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
2use std::collections::HashMap;
3use trillium::Conn;
4use trillium_http::Status;
5
6#[derive(serde::Serialize, Debug)]
7#[serde(rename_all = "camelCase")]
8pub(crate) struct AlbMultiHeadersResponse {
9 pub is_base64_encoded: bool,
10 pub status_code: u16,
11 pub status_description: String,
12 pub multi_value_headers: HashMap<String, Vec<String>>,
13 pub body: Option<String>,
14}
15
16impl AlbMultiHeadersResponse {
17 pub async fn from_conn(mut conn: Conn) -> Self {
18 let status = conn.status().unwrap_or(Status::NotFound);
19 let (body, is_base64_encoded) = response_body(&mut conn).await;
20
21 let multi_value_headers = conn
22 .inner()
23 .response_headers()
24 .iter()
25 .map(|(n, v)| (n.to_string(), v.iter().map(|v| v.to_string()).collect()))
26 .collect();
27
28 Self {
29 is_base64_encoded,
30 status_code: status as u16,
31 status_description: format!("{} {}", status as u16, status.canonical_reason()),
32 multi_value_headers,
33 body,
34 }
35 }
36}
37
38#[derive(serde::Serialize, Debug)]
39#[serde(rename_all = "camelCase")]
40pub(crate) struct AlbResponse {
41 pub is_base64_encoded: bool,
42 pub status_code: u16,
43 pub status_description: String,
44 pub headers: HashMap<String, String>,
45 pub body: Option<String>,
46}
47
48async fn response_body(conn: &mut Conn) -> (Option<String>, bool) {
49 match conn.inner_mut().take_response_body() {
50 Some(body) => {
51 let bytes = body.into_bytes().await.unwrap();
52 match String::from_utf8(bytes.to_vec()) {
53 Ok(string) => (Some(string), false),
54 Err(e) => (Some(BASE64.encode(e.into_bytes())), true),
55 }
56 }
57 None => (None, false),
58 }
59}
60
61impl AlbResponse {
62 pub async fn from_conn(mut conn: Conn) -> Self {
63 let status = conn.status().unwrap_or(Status::NotFound);
64 let (body, is_base64_encoded) = response_body(&mut conn).await;
65 let headers =
66 conn.inner()
67 .response_headers()
68 .iter()
69 .fold(HashMap::new(), |mut h, (n, v)| {
70 if let Some(one) = v.one() {
71 h.insert(n.to_string(), one.to_string());
72 }
73 h
74 });
75
76 Self {
77 is_base64_encoded,
78 status_code: status as u16,
79 status_description: format!("{} {}", status as u16, status.canonical_reason()),
80 headers,
81 body,
82 }
83 }
84}
85
86#[derive(Debug, serde::Serialize)]
87#[serde(untagged)]
88pub(crate) enum LambdaResponse {
89 Alb(AlbResponse),
90 AlbMultiHeaders(AlbMultiHeadersResponse),
91}