awsim_core/protocol/
json.rs

1use axum::http::{HeaderMap, StatusCode};
2use bytes::Bytes;
3use serde_json::Value;
4
5use crate::error::AwsError;
6
7use super::ParsedRequest;
8
9/// Parse an awsJson (1.0/1.1) request.
10///
11/// Operation is extracted from the `X-Amz-Target` header:
12/// Format: `ServicePrefix_Version.OperationName` (e.g., `DynamoDB_20120810.GetItem`)
13pub fn parse_request(headers: &HeaderMap, body: &Bytes) -> Result<ParsedRequest, AwsError> {
14    let target = headers
15        .get("x-amz-target")
16        .and_then(|v| v.to_str().ok())
17        .ok_or_else(|| AwsError::bad_request("MissingHeader", "Missing X-Amz-Target header"))?;
18
19    let operation = target
20        .split('.')
21        .next_back()
22        .ok_or_else(|| {
23            AwsError::bad_request(
24                "InvalidHeader",
25                format!("Invalid X-Amz-Target format: {target}"),
26            )
27        })?
28        .to_string();
29
30    let input = if body.is_empty() {
31        Value::Object(serde_json::Map::new())
32    } else {
33        serde_json::from_slice(body).map_err(|e| {
34            AwsError::bad_request("SerializationException", format!("Invalid JSON body: {e}"))
35        })?
36    };
37
38    Ok(ParsedRequest { operation, input })
39}
40
41/// Serialize a successful JSON response.
42///
43/// Recognised "magic" keys on the output value:
44///   * `__raw_body` — base64-encoded bytes that become the response body verbatim.
45///   * `__content_type` — overrides the `content-type` header.
46///   * `__status_code` — overrides the HTTP status (defaults to 200).
47///   * `__headers` — extra response headers `{ "Header-Name": "value", ... }`.
48///
49/// When `__raw_body` is absent, the entire output is JSON-encoded as the body
50/// (after stripping the magic keys above).
51pub fn serialize_response(output: &Value, request_id: &str) -> (StatusCode, HeaderMap, Bytes) {
52    let status = extract_status(output);
53    let mut headers = HeaderMap::new();
54    headers.insert("x-amzn-requestid", request_id.parse().unwrap());
55
56    if let Some(raw_b64) = output.get("__raw_body").and_then(Value::as_str) {
57        use base64::Engine;
58        let data = base64::engine::general_purpose::STANDARD
59            .decode(raw_b64)
60            .unwrap_or_default();
61        let content_type = output
62            .get("__content_type")
63            .and_then(Value::as_str)
64            .unwrap_or("application/octet-stream");
65        headers.insert("content-type", content_type.parse().unwrap());
66        apply_extra_headers(&mut headers, output);
67        return (status, headers, Bytes::from(data));
68    }
69
70    headers.insert(
71        "content-type",
72        "application/x-amz-json-1.0".parse().unwrap(),
73    );
74    apply_extra_headers(&mut headers, output);
75
76    let body_value = strip_magic_keys(output);
77    let body = serde_json::to_vec(&body_value).unwrap_or_default();
78    (status, headers, Bytes::from(body))
79}
80
81fn extract_status(output: &Value) -> StatusCode {
82    output
83        .get("__status_code")
84        .and_then(Value::as_u64)
85        .and_then(|n| u16::try_from(n).ok())
86        .and_then(|n| StatusCode::from_u16(n).ok())
87        .unwrap_or(StatusCode::OK)
88}
89
90fn apply_extra_headers(headers: &mut HeaderMap, output: &Value) {
91    let Some(extra) = output.get("__headers").and_then(Value::as_object) else {
92        return;
93    };
94    for (name, value) in extra {
95        let Some(s) = value.as_str() else { continue };
96        if let (Ok(k), Ok(v)) = (
97            axum::http::header::HeaderName::from_bytes(name.as_bytes()),
98            axum::http::HeaderValue::from_str(s),
99        ) {
100            headers.insert(k, v);
101        }
102    }
103}
104
105fn strip_magic_keys(output: &Value) -> Value {
106    let Some(map) = output.as_object() else {
107        return output.clone();
108    };
109    let mut cleaned = serde_json::Map::with_capacity(map.len());
110    for (k, v) in map {
111        if matches!(
112            k.as_str(),
113            "__raw_body" | "__content_type" | "__status_code" | "__headers"
114        ) {
115            continue;
116        }
117        cleaned.insert(k.clone(), v.clone());
118    }
119    Value::Object(cleaned)
120}
121
122/// Serialize a JSON error response.
123pub fn serialize_error(error: &AwsError, request_id: &str) -> (StatusCode, HeaderMap, Bytes) {
124    let mut body = serde_json::Map::new();
125    body.insert("__type".to_string(), Value::String(error.code.clone()));
126    body.insert("message".to_string(), Value::String(error.message.clone()));
127    if let Some(extras) = &error.extras {
128        for (k, v) in extras.as_ref() {
129            body.insert(k.clone(), v.clone());
130        }
131    }
132    let body = serde_json::to_vec(&Value::Object(body)).unwrap_or_default();
133    let mut headers = HeaderMap::new();
134    headers.insert(
135        "content-type",
136        "application/x-amz-json-1.0".parse().unwrap(),
137    );
138    headers.insert("x-amzn-requestid", request_id.parse().unwrap());
139    headers.insert("x-amzn-errortype", error.code.parse().unwrap());
140    (error.status, headers, Bytes::from(body))
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146    use serde_json::json;
147
148    #[test]
149    fn applies_extra_headers_and_status() {
150        let output = json!({
151            "Foo": "bar",
152            "__status_code": 202u64,
153            "__headers": { "X-Amz-Function-Error": "Handled" },
154        });
155        let (status, headers, body) = serialize_response(&output, "req-1");
156        assert_eq!(status, StatusCode::ACCEPTED);
157        assert_eq!(
158            headers
159                .get("x-amz-function-error")
160                .and_then(|v| v.to_str().ok()),
161            Some("Handled")
162        );
163        let parsed: Value = serde_json::from_slice(&body).unwrap();
164        assert_eq!(parsed["Foo"], json!("bar"));
165        assert!(parsed.get("__status_code").is_none());
166        assert!(parsed.get("__headers").is_none());
167    }
168
169    #[test]
170    fn defaults_to_ok_when_no_status_provided() {
171        let (status, _headers, body) = serialize_response(&json!({"Hello": "world"}), "req-2");
172        assert_eq!(status, StatusCode::OK);
173        let parsed: Value = serde_json::from_slice(&body).unwrap();
174        assert_eq!(parsed, json!({"Hello": "world"}));
175    }
176
177    #[test]
178    fn raw_body_path_respects_status_and_extra_headers() {
179        use base64::Engine;
180        let payload = b"hello";
181        let encoded = base64::engine::general_purpose::STANDARD.encode(payload);
182        let output = json!({
183            "__raw_body": encoded,
184            "__content_type": "text/plain",
185            "__status_code": 201u64,
186            "__headers": { "X-Custom": "yes" },
187        });
188        let (status, headers, body) = serialize_response(&output, "req-3");
189        assert_eq!(status, StatusCode::CREATED);
190        assert_eq!(body.as_ref(), payload);
191        assert_eq!(
192            headers.get("content-type").and_then(|v| v.to_str().ok()),
193            Some("text/plain")
194        );
195        assert_eq!(
196            headers.get("x-custom").and_then(|v| v.to_str().ok()),
197            Some("yes")
198        );
199    }
200}
awsim_core/protocol/json.rs

awsim_core/protocol/
json.rs
