Skip to main content

armature_lambda/
runtime.rs

1//! Lambda runtime for Armature applications.
2
3use lambda_http::{Body, Error, Request, Response, run, service_fn};
4use std::sync::Arc;
5use tracing::{debug, info};
6
7use crate::{LambdaRequest, LambdaResponse};
8
9/// Lambda runtime configuration.
10#[derive(Debug, Clone)]
11pub struct LambdaConfig {
12    /// Enable request logging.
13    pub log_requests: bool,
14    /// Enable response logging.
15    pub log_responses: bool,
16    /// Custom base path to strip (e.g., "/prod", "/dev").
17    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    /// Enable request logging.
32    pub fn log_requests(mut self, enabled: bool) -> Self {
33        self.log_requests = enabled;
34        self
35    }
36
37    /// Enable response logging.
38    pub fn log_responses(mut self, enabled: bool) -> Self {
39        self.log_responses = enabled;
40        self
41    }
42
43    /// Set a base path to strip from requests.
44    pub fn base_path(mut self, path: impl Into<String>) -> Self {
45        self.base_path = Some(path.into());
46        self
47    }
48}
49
50/// Lambda runtime driver.
51///
52/// Wraps any [`RequestHandler`] and runs it on the Lambda runtime, translating
53/// API Gateway / ALB / Function URL events into [`LambdaRequest`] and the
54/// handler's [`LambdaResponse`] back into a Lambda response. It does not know
55/// about `armature_core`'s request and response types; see
56/// `impl_lambda_handler!` for connecting an application to it.
57pub struct LambdaRuntime<App> {
58    app: Arc<App>,
59    config: LambdaConfig,
60}
61
62impl<App> LambdaRuntime<App>
63where
64    App: Send + Sync + 'static,
65{
66    /// Create a new Lambda runtime.
67    pub fn new(app: App) -> Self {
68        Self {
69            app: Arc::new(app),
70            config: LambdaConfig::default(),
71        }
72    }
73
74    /// Set the runtime configuration.
75    pub fn with_config(mut self, config: LambdaConfig) -> Self {
76        self.config = config;
77        self
78    }
79
80    /// Run the Lambda runtime.
81    ///
82    /// This function never returns under normal operation.
83    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/// The contract [`LambdaRuntime`] drives.
102///
103/// Implemented here for any `Fn(LambdaRequest) -> Future<Output = LambdaResponse>`.
104/// Other types implement it themselves, or via `impl_lambda_handler!`; there
105/// is no blanket implementation for an Armature `Application`.
106#[async_trait::async_trait]
107pub trait RequestHandler: Send + Sync {
108    /// Handle an HTTP request.
109    async fn handle(&self, request: LambdaRequest) -> LambdaResponse;
110}
111
112/// Handle a Lambda request.
113async fn handle_request<App: RequestHandler>(
114    app: Arc<App>,
115    config: LambdaConfig,
116    request: Request,
117) -> Result<Response<Body>, Error> {
118    // Convert Lambda request
119    let mut lambda_request = LambdaRequest::from_lambda_request(request);
120
121    // Strip base path if configured
122    if let Some(base_path) = &config.base_path {
123        lambda_request.path = strip_base_path(&lambda_request.path, base_path);
124    }
125
126    // Log request if enabled
127    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    // Handle request
137    let response = app.handle(lambda_request).await;
138
139    // Log response if enabled
140    if config.log_responses {
141        debug!(status = response.status, "Lambda response");
142    }
143
144    Ok(response.into_lambda_response())
145}
146
147/// Strip a configured base path prefix (e.g. an API Gateway stage like
148/// `/prod`) from a request path. When stripping empties the path it is
149/// normalized back to `/`. Paths that do not start with `base_path` are
150/// returned unchanged.
151pub(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/// Implement [`RequestHandler`] for a type that already exposes an inherent
160/// async `handle_request` method.
161///
162/// This macro knows nothing about `armature_core::Application`; there is no
163/// conversion in this crate between `armature_core`'s `HttpRequest`/
164/// `HttpResponse` and the Lambda event types. What it targets is a
165/// **user-supplied** shape:
166///
167/// ```rust,ignore
168/// impl MyApp {
169///     async fn handle_request(
170///         &self,
171///         request: armature_lambda::LambdaRequest,
172///     ) -> Result<MyResponse, MyError> { /* ... */ }
173/// }
174/// ```
175///
176/// where `MyResponse` has the fields
177/// - `status: u16`,
178/// - `body: impl Into<bytes::Bytes>`,
179/// - `headers: impl IntoIterator<Item = (impl Into<String>, impl Into<String>)>`,
180///
181/// and `MyError: std::fmt::Display`. Given that, expand the macro once per type:
182///
183/// ```rust,ignore
184/// use armature_lambda::impl_lambda_handler;
185///
186/// impl_lambda_handler!(MyApp);
187/// ```
188///
189/// If your application is an Armature `Application`, write that
190/// `handle_request` adapter yourself — the macro only removes the trait
191/// boilerplate around it.
192#[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                // Forward the full request so the application handler has
199                // access to headers, query string, path parameters, stage
200                // variables, and the request context (including authorizer
201                // claims) — not just the method/path/body.
202                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/// Example implementation for a simple handler function.
219#[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    // A minimal response/error/app trio mirroring the shape the
253    // `impl_lambda_handler!` macro documents.
254    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        // Response mapping is preserved.
342        assert_eq!(response.status, 201);
343        assert_eq!(&response.body[..], b"ok");
344        assert_eq!(response.header_value("x-app"), Some("yes"));
345
346        // The app received every part of the request, not just method/path/body.
347        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        // Repeated names survive the hand-off rather than collapsing.
360        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}