1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82

/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0
 */

use crate::endpoint;
use crate::endpoint::{apply_endpoint, EndpointPrefix, ResolveEndpointError};
use crate::middleware::MapRequest;
use crate::operation::Request;
use http::header::HeaderName;
use http::{HeaderValue, Uri};
use std::str::FromStr;

/// Middleware to apply an HTTP endpoint to the request
///
/// This middleware reads [`aws_smithy_types::endpoint::Endpoint`] out of the request properties and applies
/// it to the HTTP request.
#[non_exhaustive]
#[derive(Default, Debug, Clone)]
pub struct SmithyEndpointStage;
impl SmithyEndpointStage {
    pub fn new() -> Self {
        Self::default()
    }
}

impl MapRequest for SmithyEndpointStage {
    type Error = ResolveEndpointError;

    fn name(&self) -> &'static str {
        "resolve_endpoint"
    }

    fn apply(&self, request: Request) -> Result<Request, Self::Error> {
        request.augment(|mut http_req, props| {
            // we need to do a little dance so that this works with retries.
            // the first pass through, we convert the result into just an endpoint, early returning
            // the error. Put the endpoint back in the bag in case this request gets retried.
            //
            // the next pass through, there is no result, so in that case, we'll look for the
            // endpoint directly.
            //
            // In an ideal world, we would do this in make_operation, but it's much easier for
            // certain protocol tests if we allow requests with invalid endpoint to be constructed.
            if let Some(endpoint) = props.remove::<endpoint::Result>().transpose()? {
                props.insert(endpoint);
            };
            let endpoint = props.get::<aws_smithy_types::endpoint::Endpoint>();
            let endpoint =
                endpoint.ok_or_else(|| ResolveEndpointError::message("no endpoint present"))?;

            let uri: Uri = endpoint.url().parse().map_err(|err| {
                ResolveEndpointError::from_source("endpoint did not have a valid uri", err)
            })?;
            apply_endpoint(http_req.uri_mut(), &uri, props.get::<EndpointPrefix>()).map_err(
                |err| {
                    ResolveEndpointError::message(format!(
                        "failed to apply endpoint `{:?}` to request `{:?}`",
                        uri, http_req
                    ))
                    .with_source(Some(err.into()))
                },
            )?;
            for (header_name, header_values) in endpoint.headers() {
                http_req.headers_mut().remove(header_name);
                for value in header_values {
                    http_req.headers_mut().insert(
                        HeaderName::from_str(header_name).map_err(|err| {
                            ResolveEndpointError::message("invalid header name")
                                .with_source(Some(err.into()))
                        })?,
                        HeaderValue::from_str(value).map_err(|err| {
                            ResolveEndpointError::message("invalid header value")
                                .with_source(Some(err.into()))
                        })?,
                    );
                }
            }
            Ok(http_req)
        })
    }
}