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> {
55 app: Arc<App>,
56 config: LambdaConfig,
57}
58
59impl<App> LambdaRuntime<App>
60where
61 App: Send + Sync + 'static,
62{
63 pub fn new(app: App) -> Self {
65 Self {
66 app: Arc::new(app),
67 config: LambdaConfig::default(),
68 }
69 }
70
71 pub fn with_config(mut self, config: LambdaConfig) -> Self {
73 self.config = config;
74 self
75 }
76
77 pub async fn run(self) -> Result<(), Error>
81 where
82 App: RequestHandler,
83 {
84 info!("Starting Armature Lambda runtime");
85
86 let app = self.app.clone();
87 let config = self.config.clone();
88
89 run(service_fn(move |request: Request| {
90 let app = app.clone();
91 let config = config.clone();
92 async move { handle_request(app, config, request).await }
93 }))
94 .await
95 }
96}
97
98#[async_trait::async_trait]
102pub trait RequestHandler: Send + Sync {
103 async fn handle(&self, request: LambdaRequest) -> LambdaResponse;
105}
106
107async fn handle_request<App: RequestHandler>(
109 app: Arc<App>,
110 config: LambdaConfig,
111 request: Request,
112) -> Result<Response<Body>, Error> {
113 let mut lambda_request = LambdaRequest::from_lambda_request(request);
115
116 if let Some(base_path) = &config.base_path {
118 lambda_request.path = strip_base_path(&lambda_request.path, base_path);
119 }
120
121 if config.log_requests {
123 debug!(
124 method = %lambda_request.method,
125 path = %lambda_request.path,
126 request_id = ?lambda_request.request_context.request_id,
127 "Handling Lambda request"
128 );
129 }
130
131 let response = app.handle(lambda_request).await;
133
134 if config.log_responses {
136 debug!(status = response.status, "Lambda response");
137 }
138
139 Ok(response.into_lambda_response())
140}
141
142pub(crate) fn strip_base_path(path: &str, base_path: &str) -> String {
147 match path.strip_prefix(base_path) {
148 Some("") => "/".to_string(),
149 Some(stripped) => stripped.to_string(),
150 None => path.to_string(),
151 }
152}
153
154#[macro_export]
163macro_rules! impl_request_handler {
164 ($app_type:ty) => {
165 #[async_trait::async_trait]
166 impl $crate::runtime::RequestHandler for $app_type {
167 async fn handle(&self, request: $crate::LambdaRequest) -> $crate::LambdaResponse {
168 match self.handle_request(request).await {
173 Ok(response) => {
174 let mut lambda_response =
175 $crate::LambdaResponse::new(response.status, response.body);
176 for (name, value) in response.headers {
177 lambda_response = lambda_response.header(name, value);
178 }
179 lambda_response
180 }
181 Err(e) => $crate::LambdaResponse::internal_error(e.to_string()),
182 }
183 }
184 }
185 };
186}
187
188#[async_trait::async_trait]
190impl<F, Fut> RequestHandler for F
191where
192 F: Fn(LambdaRequest) -> Fut + Send + Sync,
193 Fut: std::future::Future<Output = LambdaResponse> + Send,
194{
195 async fn handle(&self, request: LambdaRequest) -> LambdaResponse {
196 self(request).await
197 }
198}
199
200#[cfg(test)]
201mod tests {
202 use super::*;
203 use crate::LambdaRequest;
204 use std::collections::HashMap;
205 use std::sync::Mutex;
206
207 #[test]
208 fn strip_base_path_removes_stage_prefix() {
209 assert_eq!(strip_base_path("/prod/users", "/prod"), "/users");
210 }
211
212 #[test]
213 fn strip_base_path_normalizes_empty_to_root() {
214 assert_eq!(strip_base_path("/prod", "/prod"), "/");
215 }
216
217 #[test]
218 fn strip_base_path_leaves_non_matching_paths() {
219 assert_eq!(strip_base_path("/other/users", "/prod"), "/other/users");
220 }
221
222 struct MockResponse {
225 status: u16,
226 body: Vec<u8>,
227 headers: Vec<(String, String)>,
228 }
229
230 #[derive(Debug)]
231 struct MockError(String);
232
233 impl std::fmt::Display for MockError {
234 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
235 write!(f, "{}", self.0)
236 }
237 }
238
239 #[derive(Default)]
240 struct Captured {
241 method: Option<String>,
242 path: Option<String>,
243 headers: HashMap<String, String>,
244 query_string: Option<String>,
245 path_parameters: HashMap<String, String>,
246 claims: HashMap<String, String>,
247 body: Vec<u8>,
248 }
249
250 struct MockApp {
251 captured: Mutex<Captured>,
252 }
253
254 impl MockApp {
255 async fn handle_request(
256 &self,
257 request: LambdaRequest,
258 ) -> std::result::Result<MockResponse, MockError> {
259 let mut captured = self.captured.lock().unwrap();
260 captured.method = Some(request.method.to_string());
261 captured.path = Some(request.path.clone());
262 captured.headers = request.headers.clone();
263 captured.query_string = request.query_string.clone();
264 captured.path_parameters = request.path_parameters.clone();
265 captured.claims = request.request_context.authorizer_claims.clone();
266 captured.body = request.body.to_vec();
267 Ok(MockResponse {
268 status: 201,
269 body: b"ok".to_vec(),
270 headers: vec![("x-app".to_string(), "yes".to_string())],
271 })
272 }
273 }
274
275 impl_request_handler!(MockApp);
276
277 fn sample_request() -> LambdaRequest {
278 let mut headers = HashMap::new();
279 headers.insert("x-custom".to_string(), "value".to_string());
280 let mut path_parameters = HashMap::new();
281 path_parameters.insert("id".to_string(), "42".to_string());
282 let mut claims = HashMap::new();
283 claims.insert("sub".to_string(), "user-1".to_string());
284
285 LambdaRequest {
286 method: http::Method::POST,
287 path: "/users/42".to_string(),
288 query_string: Some("page=2".to_string()),
289 headers,
290 body: bytes::Bytes::from_static(b"payload"),
291 path_parameters,
292 stage_variables: HashMap::new(),
293 request_context: crate::request::RequestContext {
294 authorizer_claims: claims,
295 ..Default::default()
296 },
297 }
298 }
299
300 #[tokio::test]
301 async fn macro_forwards_full_request_to_app() {
302 let app = MockApp {
303 captured: Mutex::new(Captured::default()),
304 };
305
306 let response = RequestHandler::handle(&app, sample_request()).await;
307
308 assert_eq!(response.status, 201);
310 assert_eq!(&response.body[..], b"ok");
311 assert_eq!(
312 response.headers.get("x-app").map(String::as_str),
313 Some("yes")
314 );
315
316 let captured = app.captured.lock().unwrap();
318 assert_eq!(captured.method.as_deref(), Some("POST"));
319 assert_eq!(captured.path.as_deref(), Some("/users/42"));
320 assert_eq!(captured.query_string.as_deref(), Some("page=2"));
321 assert_eq!(
322 captured.headers.get("x-custom").map(String::as_str),
323 Some("value")
324 );
325 assert_eq!(
326 captured.path_parameters.get("id").map(String::as_str),
327 Some("42")
328 );
329 assert_eq!(
330 captured.claims.get("sub").map(String::as_str),
331 Some("user-1")
332 );
333 assert_eq!(captured.body, b"payload");
334 }
335}