Skip to main content

turul_a2a_aws_lambda/
adapter.rs

1//! Request/response conversion between Lambda and axum.
2
3use axum::body::Body;
4use http_body_util::BodyExt;
5
6/// Header prefix for authorizer context. Client-supplied headers with this
7/// prefix are stripped (anti-spoofing).
8pub const AUTHORIZER_HEADER_PREFIX: &str = "x-authorizer-";
9
10/// Convert a Lambda HTTP request to an axum request.
11///
12/// - Strips any client-supplied x-authorizer-* headers (anti-spoofing)
13/// - Injects authorizer context from Lambda requestContext as x-authorizer-* headers
14/// - Preserves all other headers, method, URI, body
15pub fn lambda_to_axum_request(
16    event: lambda_http::Request,
17) -> Result<http::Request<Body>, lambda_http::Error> {
18    let (mut parts, body) = event.into_parts();
19
20    // Anti-spoofing: strip ALL client-supplied x-authorizer-* headers
21    let keys_to_remove: Vec<http::header::HeaderName> = parts
22        .headers
23        .keys()
24        .filter(|k| k.as_str().starts_with(AUTHORIZER_HEADER_PREFIX))
25        .cloned()
26        .collect();
27    for key in keys_to_remove {
28        parts.headers.remove(&key);
29    }
30
31    // Extract authorizer context from Lambda extensions and inject as headers.
32    // The exact structure depends on the API Gateway type (v1, v2, ALB).
33    // We use a generic approach: look for the RequestContext extension and
34    // extract authorizer fields from it.
35    if let Some(ctx) = parts
36        .extensions
37        .get::<lambda_http::request::RequestContext>()
38    {
39        match ctx {
40            lambda_http::request::RequestContext::ApiGatewayV2(apigw) => {
41                if let Some(ref authorizer) = apigw.authorizer {
42                    // JWT authorizer: claims are in jwt.claims (HashMap<String, String>)
43                    if let Some(ref jwt) = authorizer.jwt {
44                        for (key, value) in &jwt.claims {
45                            inject_authorizer_header(&mut parts.headers, key, value);
46                        }
47                    }
48                    // Lambda authorizer: context fields
49                    // In lambda_http v1, these may be in authorizer.fields or similar
50                    // For now, JWT claims are the primary path
51                }
52            }
53            lambda_http::request::RequestContext::ApiGatewayV1(apigw) => {
54                // V1 authorizer context fields
55                for (key, value) in &apigw.authorizer.fields {
56                    if let serde_json::Value::String(s) = value {
57                        inject_authorizer_header(&mut parts.headers, key, s);
58                    }
59                }
60            }
61            _ => {} // ALB, etc. — no authorizer extraction
62        }
63    }
64
65    // Convert body
66    let body_bytes: bytes::Bytes = match body {
67        lambda_http::Body::Empty => bytes::Bytes::new(),
68        lambda_http::Body::Text(s) => bytes::Bytes::from(s),
69        lambda_http::Body::Binary(b) => bytes::Bytes::from(b),
70        _ => bytes::Bytes::new(),
71    };
72
73    let req = http::Request::from_parts(parts, Body::from(body_bytes));
74    Ok(req)
75}
76
77/// Rewrite `req.uri()` to drop a leading `prefix` from the path, if
78/// present. Preserves query string, scheme, authority. If `prefix` is
79/// `/` or the path does not start with `prefix`, the request is
80/// returned unchanged (callers downstream will 404 on a genuinely
81/// unknown path — that's the correct failure mode).
82///
83/// A path that equals `prefix` exactly becomes `/` so root-routed
84/// handlers reach the router.
85pub(crate) fn strip_request_path_prefix(
86    mut req: http::Request<Body>,
87    prefix: &str,
88) -> http::Request<Body> {
89    if prefix.is_empty() || prefix == "/" {
90        return req;
91    }
92    let uri = req.uri().clone();
93    let path = uri.path();
94    if !path.starts_with(prefix) {
95        return req;
96    }
97    let rest = &path[prefix.len()..];
98    let new_path = if rest.is_empty() || !rest.starts_with('/') {
99        // path == prefix → "/"; or prefix matches mid-segment
100        // (`/devs` vs `/dev`) → do not strip.
101        if rest.is_empty() {
102            "/"
103        } else {
104            return req;
105        }
106    } else {
107        rest
108    };
109
110    let path_and_query = match uri.query() {
111        Some(q) => format!("{new_path}?{q}"),
112        None => new_path.to_string(),
113    };
114
115    let mut parts = uri.into_parts();
116    parts.path_and_query = match http::uri::PathAndQuery::from_maybe_shared(path_and_query) {
117        Ok(pq) => Some(pq),
118        Err(_) => return req,
119    };
120    if let Ok(new_uri) = http::Uri::from_parts(parts) {
121        *req.uri_mut() = new_uri;
122    }
123    req
124}
125
126fn inject_authorizer_header(headers: &mut http::HeaderMap, key: &str, value: &str) {
127    let header_name = format!("{AUTHORIZER_HEADER_PREFIX}{}", key.to_lowercase());
128    if let Ok(name) = http::header::HeaderName::from_bytes(header_name.as_bytes()) {
129        if let Ok(val) = http::header::HeaderValue::from_str(value) {
130            headers.insert(name, val);
131        }
132    }
133}
134
135/// Convert an axum response to a Lambda HTTP response.
136pub async fn axum_to_lambda_response(
137    resp: http::Response<Body>,
138) -> Result<lambda_http::Response<lambda_http::Body>, lambda_http::Error> {
139    let (parts, body) = resp.into_parts();
140    let body_bytes = body
141        .collect()
142        .await
143        .map_err(|e| lambda_http::Error::from(format!("Body collect error: {e}")))?
144        .to_bytes();
145
146    let lambda_body = if body_bytes.is_empty() {
147        lambda_http::Body::Empty
148    } else {
149        lambda_http::Body::Text(String::from_utf8_lossy(&body_bytes).into_owned())
150    };
151
152    let resp = http::Response::from_parts(parts, lambda_body);
153    Ok(resp)
154}
155
156#[cfg(test)]
157mod tests {
158    use super::*;
159
160    #[test]
161    fn strips_client_supplied_authorizer_headers() {
162        let req: lambda_http::Request = http::Request::builder()
163            .method("POST")
164            .uri("/message:send")
165            .header("content-type", "application/json")
166            .header("a2a-version", "1.0")
167            .header("x-authorizer-userid", "forged-admin")
168            .header("x-authorizer-role", "superuser")
169            .header("x-real-header", "keep-this")
170            .body(lambda_http::Body::Text("{}".into()))
171            .unwrap();
172
173        let axum_req = lambda_to_axum_request(req).unwrap();
174
175        // Forged headers must be stripped
176        assert!(
177            axum_req.headers().get("x-authorizer-userid").is_none(),
178            "Forged x-authorizer-userid must be stripped"
179        );
180        assert!(
181            axum_req.headers().get("x-authorizer-role").is_none(),
182            "Forged x-authorizer-role must be stripped"
183        );
184
185        // Legitimate headers preserved
186        assert_eq!(
187            axum_req.headers().get("x-real-header").unwrap(),
188            "keep-this"
189        );
190        assert_eq!(
191            axum_req.headers().get("content-type").unwrap(),
192            "application/json"
193        );
194    }
195
196    #[test]
197    fn inject_authorizer_header_works() {
198        let mut headers = http::HeaderMap::new();
199        inject_authorizer_header(&mut headers, "userId", "user-123");
200        assert_eq!(headers.get("x-authorizer-userid").unwrap(), "user-123");
201    }
202
203    fn uri_only(path_and_query: &str) -> http::Request<Body> {
204        http::Request::builder()
205            .method("GET")
206            .uri(path_and_query)
207            .body(Body::empty())
208            .unwrap()
209    }
210
211    #[test]
212    fn strip_path_prefix_removes_api_gateway_stage_and_resource_prefix() {
213        let req = uri_only("/stage/agent/message:send");
214        let out = strip_request_path_prefix(req, "/stage/agent");
215        assert_eq!(out.uri().path(), "/message:send");
216    }
217
218    #[test]
219    fn strip_path_prefix_preserves_query_string() {
220        let req = uri_only("/stage/agent/tasks/abc?historyLength=5");
221        let out = strip_request_path_prefix(req, "/stage/agent");
222        assert_eq!(out.uri().path(), "/tasks/abc");
223        assert_eq!(out.uri().query(), Some("historyLength=5"));
224    }
225
226    #[test]
227    fn strip_path_prefix_handles_root_after_strip() {
228        // Path exactly equals prefix (e.g. agent-card at the prefix root).
229        let req = uri_only("/stage/agent");
230        let out = strip_request_path_prefix(req, "/stage/agent");
231        assert_eq!(out.uri().path(), "/");
232    }
233
234    #[test]
235    fn strip_path_prefix_passes_through_when_prefix_absent() {
236        let req = uri_only("/message:send");
237        let out = strip_request_path_prefix(req, "/stage/agent");
238        assert_eq!(
239            out.uri().path(),
240            "/message:send",
241            "non-matching prefix must not mutate path"
242        );
243    }
244
245    #[test]
246    fn strip_path_prefix_rejects_mid_segment_match() {
247        // "/dev" is a prefix of "/devs/foo" as a string, but not as a
248        // path segment — must not strip.
249        let req = uri_only("/devs/foo");
250        let out = strip_request_path_prefix(req, "/dev");
251        assert_eq!(out.uri().path(), "/devs/foo");
252    }
253
254    #[test]
255    fn strip_path_prefix_noop_on_empty_or_slash_config() {
256        let req = uri_only("/message:send");
257        let out = strip_request_path_prefix(req, "/");
258        assert_eq!(out.uri().path(), "/message:send");
259        let req = uri_only("/message:send");
260        let out = strip_request_path_prefix(req, "");
261        assert_eq!(out.uri().path(), "/message:send");
262    }
263
264    #[test]
265    fn strip_path_prefix_well_known_and_jsonrpc() {
266        let prefix = "/stage/agent";
267        for (input, expected) in [
268            (
269                "/stage/agent/.well-known/agent-card.json",
270                "/.well-known/agent-card.json",
271            ),
272            ("/stage/agent/jsonrpc", "/jsonrpc"),
273            ("/stage/agent/tasks/abc:cancel", "/tasks/abc:cancel"),
274            ("/stage/agent/extendedAgentCard", "/extendedAgentCard"),
275        ] {
276            let req = uri_only(input);
277            let out = strip_request_path_prefix(req, prefix);
278            assert_eq!(out.uri().path(), expected, "input: {input}");
279        }
280    }
281}