Skip to main content

trillium_aws_lambda/
response.rs

1use 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}