1use lambda_http::{Body, Error, Request, Response, run, service_fn};
4use std::sync::Arc;
5use tracing::{debug, info};
6
7use crate::{LambdaRequest, LambdaResponse};
8
9#[derive(Debug, Clone)]
11pub struct LambdaConfig {
12 pub log_requests: bool,
14 pub log_responses: bool,
16 pub base_path: Option<String>,
18}
19
20impl Default for LambdaConfig {
21 fn default() -> Self {
22 Self {
23 log_requests: true,
24 log_responses: false,
25 base_path: None,
26 }
27 }
28}
29
30impl LambdaConfig {
31 pub fn log_requests(mut self, enabled: bool) -> Self {
33 self.log_requests = enabled;
34 self
35 }
36
37 pub fn log_responses(mut self, enabled: bool) -> Self {
39 self.log_responses = enabled;
40 self
41 }
42
43 pub fn base_path(mut self, path: impl Into<String>) -> Self {
45 self.base_path = Some(path.into());
46 self
47 }
48}
49
50pub struct LambdaRuntime<App> {
58 app: Arc<App>,
59 config: LambdaConfig,
60}
61
62impl<App> LambdaRuntime<App>
63where
64 App: Send + Sync + 'static,
65{
66 pub fn new(app: App) -> Self {
68 Self {
69 app: Arc::new(app),
70 config: LambdaConfig::default(),
71 }
72 }
73
74 pub fn with_config(mut self, config: LambdaConfig) -> Self {
76 self.config = config;
77 self
78 }
79
80 pub async fn run(self) -> Result<(), Error>
84 where
85 App: RequestHandler,
86 {
87 info!("Starting Armature Lambda runtime");
88
89 let app = self.app.clone();
90 let config = self.config.clone();
91
92 run(service_fn(move |request: Request| {
93 let app = app.clone();
94 let config = config.clone();
95 async move { handle_request(app, config, request).await }
96 }))
97 .await
98 }
99}
100
101#[async_trait::async_trait]
107pub trait RequestHandler: Send + Sync {
108 async fn handle(&self, request: LambdaRequest) -> LambdaResponse;
110}
111
112async fn handle_request<App: RequestHandler>(
114 app: Arc<App>,
115 config: LambdaConfig,
116 request: Request,
117) -> Result<Response<Body>, Error> {
118 let mut lambda_request = LambdaRequest::from_lambda_request(request);
120
121 if let Some(base_path) = &config.base_path {
123 lambda_request.path = strip_base_path(&lambda_request.path, base_path);
124 }
125
126 if config.log_requests {
128 debug!(
129 method = %lambda_request.method,
130 path = %lambda_request.path,
131 request_id = ?lambda_request.request_context.request_id,
132 "Handling Lambda request"
133 );
134 }
135
136 let response = app.handle(lambda_request).await;
138
139 if config.log_responses {
141 debug!(status = response.status, "Lambda response");
142 }
143
144 Ok(response.into_lambda_response())
145}
146
147pub(crate) fn strip_base_path(path: &str, base_path: &str) -> String {
152 match path.strip_prefix(base_path) {
153 Some("") => "/".to_string(),
154 Some(stripped) => stripped.to_string(),
155 None => path.to_string(),
156 }
157}
158
159#[macro_export]
193macro_rules! impl_lambda_handler {
194 ($app_type:ty) => {
195 #[$crate::async_trait::async_trait]
196 impl $crate::RequestHandler for $app_type {
197 async fn handle(&self, request: $crate::LambdaRequest) -> $crate::LambdaResponse {
198 match self.handle_request(request).await {
203 Ok(response) => {
204 let mut lambda_response =
205 $crate::LambdaResponse::new(response.status, response.body);
206 for (name, value) in response.headers {
207 lambda_response = lambda_response.header(name, value);
208 }
209 lambda_response
210 }
211 Err(e) => $crate::LambdaResponse::internal_error(e.to_string()),
212 }
213 }
214 }
215 };
216}
217
218#[async_trait::async_trait]
220impl<F, Fut> RequestHandler for F
221where
222 F: Fn(LambdaRequest) -> Fut + Send + Sync,
223 Fut: std::future::Future<Output = LambdaResponse> + Send,
224{
225 async fn handle(&self, request: LambdaRequest) -> LambdaResponse {
226 self(request).await
227 }
228}
229
230#[cfg(test)]
231mod tests {
232 use super::*;
233 use crate::LambdaRequest;
234 use std::collections::HashMap;
235 use std::sync::Mutex;
236
237 #[test]
238 fn strip_base_path_removes_stage_prefix() {
239 assert_eq!(strip_base_path("/prod/users", "/prod"), "/users");
240 }
241
242 #[test]
243 fn strip_base_path_normalizes_empty_to_root() {
244 assert_eq!(strip_base_path("/prod", "/prod"), "/");
245 }
246
247 #[test]
248 fn strip_base_path_leaves_non_matching_paths() {
249 assert_eq!(strip_base_path("/other/users", "/prod"), "/other/users");
250 }
251
252 struct MockResponse {
255 status: u16,
256 body: Vec<u8>,
257 headers: Vec<(String, String)>,
258 }
259
260 #[derive(Debug)]
261 struct MockError(String);
262
263 impl std::fmt::Display for MockError {
264 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
265 write!(f, "{}", self.0)
266 }
267 }
268
269 #[derive(Default)]
270 struct Captured {
271 method: Option<String>,
272 path: Option<String>,
273 headers: Vec<(String, String)>,
274 query_string: Option<String>,
275 path_parameters: HashMap<String, String>,
276 claims: HashMap<String, String>,
277 body: Vec<u8>,
278 }
279
280 struct MockApp {
281 captured: Mutex<Captured>,
282 }
283
284 impl MockApp {
285 async fn handle_request(
286 &self,
287 request: LambdaRequest,
288 ) -> std::result::Result<MockResponse, MockError> {
289 let mut captured = self.captured.lock().unwrap();
290 captured.method = Some(request.method.to_string());
291 captured.path = Some(request.path.clone());
292 captured.headers = request.headers.clone();
293 captured.query_string = request.query_string.clone();
294 captured.path_parameters = request.path_parameters.clone();
295 captured.claims = request.request_context.authorizer_claims.clone();
296 captured.body = request.body.to_vec();
297 Ok(MockResponse {
298 status: 201,
299 body: b"ok".to_vec(),
300 headers: vec![("x-app".to_string(), "yes".to_string())],
301 })
302 }
303 }
304
305 impl_lambda_handler!(MockApp);
306
307 fn sample_request() -> LambdaRequest {
308 let headers = vec![
309 ("x-custom".to_string(), "value".to_string()),
310 ("cookie".to_string(), "a=1".to_string()),
311 ("cookie".to_string(), "b=2".to_string()),
312 ];
313 let mut path_parameters = HashMap::new();
314 path_parameters.insert("id".to_string(), "42".to_string());
315 let mut claims = HashMap::new();
316 claims.insert("sub".to_string(), "user-1".to_string());
317
318 LambdaRequest {
319 method: http::Method::POST,
320 path: "/users/42".to_string(),
321 query_string: Some("page=2".to_string()),
322 headers,
323 body: bytes::Bytes::from_static(b"payload"),
324 path_parameters,
325 stage_variables: HashMap::new(),
326 request_context: crate::request::RequestContext {
327 authorizer_claims: claims,
328 ..Default::default()
329 },
330 }
331 }
332
333 #[tokio::test]
334 async fn macro_forwards_full_request_to_app() {
335 let app = MockApp {
336 captured: Mutex::new(Captured::default()),
337 };
338
339 let response = RequestHandler::handle(&app, sample_request()).await;
340
341 assert_eq!(response.status, 201);
343 assert_eq!(&response.body[..], b"ok");
344 assert_eq!(response.header_value("x-app"), Some("yes"));
345
346 let captured = app.captured.lock().unwrap();
348 assert_eq!(captured.method.as_deref(), Some("POST"));
349 assert_eq!(captured.path.as_deref(), Some("/users/42"));
350 assert_eq!(captured.query_string.as_deref(), Some("page=2"));
351 assert_eq!(
352 captured
353 .headers
354 .iter()
355 .find(|(name, _)| name == "x-custom")
356 .map(|(_, value)| value.as_str()),
357 Some("value")
358 );
359 assert_eq!(
361 captured
362 .headers
363 .iter()
364 .filter(|(name, _)| name == "cookie")
365 .map(|(_, value)| value.as_str())
366 .collect::<Vec<_>>(),
367 vec!["a=1", "b=2"]
368 );
369 assert_eq!(
370 captured.path_parameters.get("id").map(String::as_str),
371 Some("42")
372 );
373 assert_eq!(
374 captured.claims.get("sub").map(String::as_str),
375 Some("user-1")
376 );
377 assert_eq!(captured.body, b"payload".to_vec());
378 }
379}