armature_lambda/
response.rs

1//! Lambda response conversion.
2
3use bytes::Bytes;
4use lambda_http::{Body, Response};
5
6/// Lambda HTTP response.
7pub struct LambdaResponse {
8    /// Status code.
9    pub status: u16,
10    /// Response headers, in emission order.
11    ///
12    /// A list rather than a map because HTTP allows the same field name to
13    /// appear more than once and a handler must be able to use that — most
14    /// importantly to emit several `Set-Cookie` lines, which cannot legally be
15    /// folded into one. A map would silently keep only the last one.
16    pub headers: Vec<(String, String)>,
17    /// Response body.
18    pub body: Bytes,
19    /// Whether body is base64 encoded.
20    pub is_base64: bool,
21}
22
23impl LambdaResponse {
24    /// Create a new response.
25    pub fn new(status: u16, body: impl Into<Bytes>) -> Self {
26        Self {
27            status,
28            headers: Vec::new(),
29            body: body.into(),
30            is_base64: false,
31        }
32    }
33
34    /// Create an OK response.
35    pub fn ok(body: impl Into<Bytes>) -> Self {
36        Self::new(200, body)
37    }
38
39    /// Create a JSON response.
40    pub fn json<T: serde::Serialize>(data: &T) -> Result<Self, serde_json::Error> {
41        let body = serde_json::to_vec(data)?;
42        Ok(Self::new(200, body).header("content-type", "application/json"))
43    }
44
45    /// Create an error response.
46    pub fn error(status: u16, message: impl Into<String>) -> Self {
47        let body = serde_json::json!({
48            "error": message.into()
49        });
50        Self::new(status, serde_json::to_vec(&body).unwrap_or_default())
51            .header("content-type", "application/json")
52    }
53
54    /// Create a not found response.
55    pub fn not_found() -> Self {
56        Self::error(404, "Not Found")
57    }
58
59    /// Create an internal server error response.
60    pub fn internal_error(message: impl Into<String>) -> Self {
61        Self::error(500, message)
62    }
63
64    /// Append a header.
65    ///
66    /// Appends rather than replaces, so calling this twice with the same name
67    /// emits two header lines (e.g. two `Set-Cookie`s). Use [`set_header`] to
68    /// replace instead.
69    ///
70    /// [`set_header`]: LambdaResponse::set_header
71    pub fn header(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
72        self.headers.push((name.into(), value.into()));
73        self
74    }
75
76    /// Set a header, removing any existing lines with the same name.
77    pub fn set_header(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
78        let name = name.into();
79        self.headers.retain(|(n, _)| !n.eq_ignore_ascii_case(&name));
80        self.headers.push((name, value.into()));
81        self
82    }
83
84    /// The first value for `name`, matched case-insensitively.
85    pub fn header_value(&self, name: &str) -> Option<&str> {
86        self.header_values(name).next()
87    }
88
89    /// Every value for `name`, in emission order, matched case-insensitively.
90    pub fn header_values<'a, 'n>(
91        &'a self,
92        name: &'n str,
93    ) -> impl Iterator<Item = &'a str> + use<'a, 'n> {
94        self.headers
95            .iter()
96            .filter(move |(n, _)| n.eq_ignore_ascii_case(name))
97            .map(|(_, v)| v.as_str())
98    }
99
100    /// Set content type.
101    pub fn content_type(self, content_type: impl Into<String>) -> Self {
102        self.set_header("content-type", content_type)
103    }
104
105    /// Mark body as base64 encoded.
106    pub fn base64(mut self) -> Self {
107        self.is_base64 = true;
108        self
109    }
110
111    /// Convert to lambda_http::Response.
112    pub fn into_lambda_response(self) -> Response<Body> {
113        let mut builder = Response::builder().status(self.status);
114
115        for (name, value) in &self.headers {
116            builder = builder.header(name, value);
117        }
118
119        let body = if self.is_base64 {
120            Body::Binary(self.body.to_vec())
121        } else if let Ok(s) = String::from_utf8(self.body.to_vec()) {
122            Body::Text(s)
123        } else {
124            Body::Binary(self.body.to_vec())
125        };
126
127        builder.body(body).unwrap_or_else(|_| {
128            Response::builder()
129                .status(500)
130                .body(Body::Text("Internal Server Error".to_string()))
131                .unwrap()
132        })
133    }
134}
135
136impl Default for LambdaResponse {
137    fn default() -> Self {
138        Self::new(200, Bytes::new())
139    }
140}
141
142#[cfg(test)]
143mod tests {
144    use super::*;
145
146    #[test]
147    fn utf8_body_maps_to_text() {
148        let resp = LambdaResponse::ok("hello world");
149        let lambda = resp.into_lambda_response();
150        assert_eq!(lambda.status(), 200);
151        match lambda.body() {
152            Body::Text(s) => assert_eq!(s, "hello world"),
153            other => panic!("expected text body, got {other:?}"),
154        }
155    }
156
157    #[test]
158    fn base64_flag_forces_binary_body() {
159        let resp = LambdaResponse::ok("hello").base64();
160        let lambda = resp.into_lambda_response();
161        match lambda.body() {
162            Body::Binary(b) => assert_eq!(b, b"hello"),
163            other => panic!("expected binary body, got {other:?}"),
164        }
165    }
166
167    #[test]
168    fn non_utf8_body_maps_to_binary() {
169        let resp = LambdaResponse::new(200, Bytes::from_static(&[0xff, 0xfe, 0x00]));
170        let lambda = resp.into_lambda_response();
171        match lambda.body() {
172            Body::Binary(b) => assert_eq!(b, &[0xff, 0xfe, 0x00]),
173            other => panic!("expected binary body, got {other:?}"),
174        }
175    }
176
177    #[test]
178    fn headers_are_forwarded() {
179        let resp = LambdaResponse::json(&serde_json::json!({ "ok": true })).unwrap();
180        let lambda = resp.into_lambda_response();
181        assert_eq!(
182            lambda
183                .headers()
184                .get("content-type")
185                .and_then(|v| v.to_str().ok()),
186            Some("application/json")
187        );
188    }
189
190    #[test]
191    fn duplicate_headers_are_all_emitted() {
192        // Session/auth flows need more than one Set-Cookie line, and these
193        // cannot be folded into a single comma-separated value.
194        let resp = LambdaResponse::ok("x")
195            .header("set-cookie", "a=1")
196            .header("set-cookie", "b=2");
197        let lambda = resp.into_lambda_response();
198        let cookies: Vec<_> = lambda
199            .headers()
200            .get_all("set-cookie")
201            .iter()
202            .map(|v| v.to_str().unwrap())
203            .collect();
204        assert_eq!(cookies, vec!["a=1", "b=2"]);
205    }
206
207    #[test]
208    fn set_header_replaces_existing_lines() {
209        let resp = LambdaResponse::ok("x")
210            .header("content-type", "text/plain")
211            .set_header("content-type", "application/json");
212        assert_eq!(resp.header_values("content-type").count(), 1);
213        assert_eq!(resp.header_value("content-type"), Some("application/json"));
214    }
215
216    #[test]
217    fn error_response_sets_status_and_json() {
218        let resp = LambdaResponse::not_found();
219        assert_eq!(resp.status, 404);
220        let lambda = resp.into_lambda_response();
221        assert_eq!(lambda.status(), 404);
222        match lambda.body() {
223            Body::Text(s) => assert!(s.contains("Not Found")),
224            other => panic!("expected text body, got {other:?}"),
225        }
226    }
227}
armature_lambda/response.rs

armature_lambda/
response.rs
