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 for Armature applications.
51///
52/// This wraps an Armature application and runs it on the Lambda runtime,
53/// translating API Gateway/ALB requests to Armature requests.
54pub struct LambdaRuntime<App> {
55    app: Arc<App>,
56    config: LambdaConfig,
57}
58
59impl<App> LambdaRuntime<App>
60where
61    App: Send + Sync + 'static,
62{
63    /// Create a new Lambda runtime.
64    pub fn new(app: App) -> Self {
65        Self {
66            app: Arc::new(app),
67            config: LambdaConfig::default(),
68        }
69    }
70
71    /// Set the runtime configuration.
72    pub fn with_config(mut self, config: LambdaConfig) -> Self {
73        self.config = config;
74        self
75    }
76
77    /// Run the Lambda runtime.
78    ///
79    /// This function never returns under normal operation.
80    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/// Request handler trait for Armature applications.
99///
100/// This is automatically implemented for Armature Application types.
101#[async_trait::async_trait]
102pub trait RequestHandler: Send + Sync {
103    /// Handle an HTTP request.
104    async fn handle(&self, request: LambdaRequest) -> LambdaResponse;
105}
106
107/// Handle a Lambda request.
108async fn handle_request<App: RequestHandler>(
109    app: Arc<App>,
110    config: LambdaConfig,
111    request: Request,
112) -> Result<Response<Body>, Error> {
113    // Convert Lambda request
114    let mut lambda_request = LambdaRequest::from_lambda_request(request);
115
116    // Strip base path if configured
117    if let Some(base_path) = &config.base_path {
118        lambda_request.path = strip_base_path(&lambda_request.path, base_path);
119    }
120
121    // Log request if enabled
122    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    // Handle request
132    let response = app.handle(lambda_request).await;
133
134    // Log response if enabled
135    if config.log_responses {
136        debug!(status = response.status, "Lambda response");
137    }
138
139    Ok(response.into_lambda_response())
140}
141
142/// Strip a configured base path prefix (e.g. an API Gateway stage like
143/// `/prod`) from a request path. When stripping empties the path it is
144/// normalized back to `/`. Paths that do not start with `base_path` are
145/// returned unchanged.
146pub(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 to implement RequestHandler for Armature applications.
155///
156/// Usage:
157/// ```rust,ignore
158/// use armature_lambda::impl_request_handler;
159///
160/// impl_request_handler!(MyApplication);
161/// ```
162#[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                // Forward the full request so the application handler has
169                // access to headers, query string, path parameters, stage
170                // variables, and the request context (including authorizer
171                // claims) — not just the method/path/body.
172                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/// Example implementation for a simple handler function.
189#[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    // A minimal response/error/app trio mirroring the shape the
223    // `impl_request_handler!` macro expects from an Armature application.
224    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        // Response mapping is preserved.
309        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        // The app received every part of the request, not just method/path/body.
317        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}