Skip to main content

aws_smithy_query/
protocol.rs

1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6use aws_smithy_runtime_api::client::orchestrator::Metadata;
7use aws_smithy_runtime_api::http::{Request, Response};
8use aws_smithy_schema::protocol::{apply_http_endpoint, ClientProtocolInner};
9use aws_smithy_schema::serde::{
10    SerdeError, SerializableStruct, ShapeDeserializer, ShapeSerializer,
11};
12use aws_smithy_schema::{shape_id, Schema, ShapeId};
13use aws_smithy_types::body::SdkBody;
14use aws_smithy_types::config_bag::ConfigBag;
15use aws_smithy_xml::codec::find_depth2_element_slice_by;
16
17use crate::codec::serializer::QueryShapeSerializer;
18
19#[derive(Debug)]
20pub struct AwsQueryProtocol {
21    protocol_id: ShapeId,
22    service_version: String,
23}
24
25impl AwsQueryProtocol {
26    pub fn new(version: impl Into<String>) -> Self {
27        Self {
28            protocol_id: shape_id!("aws.protocols", "awsQuery"),
29            service_version: version.into(),
30        }
31    }
32}
33
34impl ClientProtocolInner for AwsQueryProtocol {
35    type Request = Request;
36    type Response = Response;
37
38    fn protocol_id(&self) -> &ShapeId {
39        &self.protocol_id
40    }
41
42    fn serialize_request(
43        &self,
44        input: &dyn SerializableStruct,
45        input_schema: &Schema,
46        endpoint: &str,
47        cfg: &ConfigBag,
48    ) -> Result<Request, SerdeError> {
49        let op_name = cfg
50            .load::<Metadata>()
51            .map(|m| m.name().to_string())
52            .ok_or_else(|| {
53                SerdeError::custom(
54                    "operation Metadata is required to serialize an awsQuery request (Action=)",
55                )
56            })?;
57
58        let mut serializer = QueryShapeSerializer::new(&op_name, &self.service_version);
59        serializer.write_struct(input_schema, input)?;
60        let body = aws_smithy_schema::codec::FinishSerializer::finish(serializer);
61
62        let uri = if endpoint.is_empty() { "/" } else { endpoint };
63        let mut request = Request::new(SdkBody::from(body));
64        request
65            .set_method("POST")
66            .map_err(|e| SerdeError::custom(format!("{e}")))?;
67        request
68            .set_uri(uri)
69            .map_err(|e| SerdeError::custom(format!("{e}")))?;
70        request
71            .headers_mut()
72            .insert("Content-Type", "application/x-www-form-urlencoded");
73        if let Some(len) = request.body().content_length() {
74            request
75                .headers_mut()
76                .insert("Content-Length", len.to_string());
77        }
78        Ok(request)
79    }
80
81    fn deserialize_response<'a>(
82        &self,
83        response: &'a Response,
84        _output_schema: &Schema,
85        _cfg: &ConfigBag,
86    ) -> Result<Box<dyn ShapeDeserializer + 'a>, SerdeError> {
87        use aws_smithy_schema::codec::Codec;
88        use aws_smithy_xml::codec::{XmlCodec, XmlCodecSettings};
89
90        let body = response
91            .body()
92            .bytes()
93            .ok_or_else(|| SerdeError::custom("response body not available"))?;
94        let body_str = std::str::from_utf8(body).map_err(|e| SerdeError::InvalidInput {
95            message: e.to_string(),
96        })?;
97
98        // Strip the AWS Query response envelope down to the `<...Result>` (or
99        // `<Error>`) element, inclusive of its tags, so the XML deserializer
100        // can treat it as the output struct's root wrapper element (the merged
101        // `XmlDeserializer::read_struct` reads members from the root element's
102        // children).
103        let inner = strip_aws_query_envelope(body_str);
104
105        // AWS Query deserializes responses as XML with a default timestamp
106        // format of `date-time` (per the protocol spec). `XmlCodec` is
107        // stateless; the returned deserializer borrows `inner` (and therefore
108        // `response`), not the local codec — `create_deserializer` clones the
109        // shared settings `Arc` into the deserializer.
110        let codec = XmlCodec::new(
111            XmlCodecSettings::builder()
112                .default_timestamp_format(aws_smithy_types::date_time::Format::DateTime)
113                .build(),
114        );
115        Ok(Box::new(codec.create_deserializer(inner.as_bytes())))
116    }
117
118    fn update_endpoint(
119        &self,
120        request: &mut Request,
121        endpoint: &aws_smithy_types::endpoint::Endpoint,
122        cfg: &ConfigBag,
123    ) -> Result<(), SerdeError> {
124        apply_http_endpoint(request, endpoint, cfg)
125    }
126}
127
128/// Extracts the result (or error) element from an AWS Query XML response
129/// envelope, returning the element *inclusive of its tags*.
130///
131/// AWS Query responses are shaped like:
132/// ```xml
133/// <OperationNameResponse>
134///   <OperationNameResult> ... </OperationNameResult>
135///   <ResponseMetadata>...</ResponseMetadata>
136/// </OperationNameResponse>
137/// ```
138/// or, for errors:
139/// ```xml
140/// <ErrorResponse><Error> ... </Error></ErrorResponse>
141/// ```
142/// We locate the depth-2 element whose local name ends with `Result` or equals
143/// `Error` and return its full `<El>...</El>` slice so it can be handed to an
144/// `XmlDeserializer` as the output struct's root wrapper element.
145///
146/// Delegates the depth-2 lookup to `aws_smithy_xml`'s shared
147/// [`find_depth2_element_slice_by`] (the same utility the REST XML error path
148/// uses). If no such element is found — or the body isn't valid XML — we fall
149/// back to the whole body and let the downstream `XmlDeserializer` surface any
150/// error, mirroring the REST XML fallback.
151fn strip_aws_query_envelope(xml: &str) -> &str {
152    match find_depth2_element_slice_by(xml.as_bytes(), |name| {
153        name.ends_with("Result") || name == "Error"
154    }) {
155        // The returned slice is a sub-slice of `xml` bounded by ASCII `<`/`>`,
156        // so it is always valid UTF-8.
157        Some(slice) => std::str::from_utf8(slice).unwrap_or(xml),
158        None => xml,
159    }
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165    use aws_smithy_schema::protocol::ClientProtocolInner;
166    use aws_smithy_schema::serde::ShapeSerializer;
167    use aws_smithy_schema::ShapeType;
168    use aws_smithy_types::config_bag::Layer;
169
170    struct EmptyInput;
171    impl SerializableStruct for EmptyInput {
172        fn serialize_members(&self, _: &mut dyn ShapeSerializer) -> Result<(), SerdeError> {
173            Ok(())
174        }
175    }
176
177    static SCHEMA: Schema = Schema::new(shape_id!("test", "Input"), ShapeType::Structure);
178
179    fn cfg_with_metadata() -> ConfigBag {
180        let mut layer = Layer::new("test");
181        layer.store_put(Metadata::new("GetUser", "MyService"));
182        ConfigBag::of_layers(vec![layer])
183    }
184
185    #[test]
186    fn request_has_correct_content_type() {
187        let cfg = cfg_with_metadata();
188        let request = AwsQueryProtocol::new("2012-11-05")
189            .serialize_request(&EmptyInput, &SCHEMA, "https://example.com", &cfg)
190            .unwrap();
191        assert_eq!(
192            request.headers().get("Content-Type").unwrap(),
193            "application/x-www-form-urlencoded"
194        );
195    }
196
197    #[test]
198    fn request_has_action_and_version() {
199        let cfg = cfg_with_metadata();
200        let request = AwsQueryProtocol::new("2012-11-05")
201            .serialize_request(&EmptyInput, &SCHEMA, "https://example.com", &cfg)
202            .unwrap();
203        let body = std::str::from_utf8(request.body().bytes().unwrap()).unwrap();
204        assert!(body.contains("Action=GetUser"));
205        assert!(body.contains("Version=2012-11-05"));
206    }
207
208    #[test]
209    fn request_posts_to_endpoint() {
210        let cfg = cfg_with_metadata();
211        let request = AwsQueryProtocol::new("1.0")
212            .serialize_request(
213                &EmptyInput,
214                &SCHEMA,
215                "https://sqs.us-east-1.amazonaws.com",
216                &cfg,
217            )
218            .unwrap();
219        assert_eq!(request.uri(), "https://sqs.us-east-1.amazonaws.com");
220    }
221
222    #[test]
223    fn request_defaults_to_slash() {
224        let cfg = cfg_with_metadata();
225        let request = AwsQueryProtocol::new("1.0")
226            .serialize_request(&EmptyInput, &SCHEMA, "", &cfg)
227            .unwrap();
228        assert_eq!(request.uri(), "/");
229    }
230
231    #[test]
232    fn deserialize_response_strips_envelope() {
233        let xml = "<GetUserResponse><GetUserResult><Name>Alice</Name><Age>30</Age></GetUserResult></GetUserResponse>";
234        let response = Response::new(200u16.try_into().unwrap(), SdkBody::from(xml));
235
236        static NAME: Schema = Schema::new_member(shape_id!("t", "S"), ShapeType::String, "Name", 0);
237        static AGE: Schema = Schema::new_member(shape_id!("t", "S"), ShapeType::Integer, "Age", 1);
238        static OUT_SCHEMA: Schema =
239            Schema::new_struct(shape_id!("t", "S"), ShapeType::Structure, &[&NAME, &AGE]);
240
241        let mut deser = AwsQueryProtocol::new("1.0")
242            .deserialize_response(&response, &OUT_SCHEMA, &ConfigBag::base())
243            .unwrap();
244        let mut name = String::new();
245        let mut age = 0i32;
246        deser
247            .read_struct(&OUT_SCHEMA, &mut |member, d| {
248                match member.member_name() {
249                    Some("Name") => name = d.read_string(member)?,
250                    Some("Age") => age = d.read_integer(member)?,
251                    _ => {}
252                }
253                Ok(())
254            })
255            .unwrap();
256        assert_eq!(name, "Alice");
257        assert_eq!(age, 30);
258    }
259
260    #[test]
261    fn strip_envelope_returns_self_closing_result_element() {
262        // A self-closing result element (empty output) must still be returned
263        // inclusive of its tags, not fall through to the whole-document root.
264        let xml = "<GetUserResponse><GetUserResult/><ResponseMetadata><RequestId>r</RequestId></ResponseMetadata></GetUserResponse>";
265        assert_eq!(strip_aws_query_envelope(xml), "<GetUserResult/>");
266    }
267
268    #[test]
269    fn strip_envelope_returns_error_element() {
270        let xml = "<ErrorResponse><Error><Code>Boom</Code></Error></ErrorResponse>";
271        assert_eq!(
272            strip_aws_query_envelope(xml),
273            "<Error><Code>Boom</Code></Error>"
274        );
275    }
276
277    #[test]
278    fn protocol_id() {
279        assert_eq!(
280            AwsQueryProtocol::new("1.0").protocol_id().as_str(),
281            "aws.protocols#awsQuery"
282        );
283    }
284}