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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0
 */

use aws_credential_types::provider::{
    error::CredentialsError, ProvideCredentials, SharedCredentialsProvider,
};
use aws_smithy_http::middleware::AsyncMapRequest;
use aws_smithy_http::operation::Request;
use aws_smithy_http::property_bag::PropertyBag;
use std::future::Future;
use std::pin::Pin;

/// Sets the credentials provider in the given property bag.
pub fn set_provider(bag: &mut PropertyBag, provider: SharedCredentialsProvider) {
    bag.insert(provider);
}

/// Middleware stage that loads credentials from a [CredentialsProvider](aws_credential_types::provider::ProvideCredentials)
/// and places them in the property bag of the request.
///
/// [CredentialsStage] implements [`AsyncMapRequest`](aws_smithy_http::middleware::AsyncMapRequest), and:
/// 1. Retrieves a `CredentialsProvider` from the property bag.
/// 2. Calls the credential provider's `provide_credentials` and awaits its result.
/// 3. Places returned `Credentials` into the property bad to drive downstream signing middleware.
#[derive(Clone, Debug, Default)]
#[non_exhaustive]
pub struct CredentialsStage;

impl CredentialsStage {
    /// Creates a new credentials stage.
    pub fn new() -> Self {
        CredentialsStage
    }

    async fn load_creds(mut request: Request) -> Result<Request, CredentialsStageError> {
        let provider = request
            .properties()
            .get::<SharedCredentialsProvider>()
            .cloned();
        let provider = match provider {
            Some(provider) => provider,
            None => {
                tracing::info!("no credentials provider for request");
                return Ok(request);
            }
        };
        match provider.provide_credentials().await {
            Ok(creds) => {
                request.properties_mut().insert(creds);
            }
            // ignore the case where there is no provider wired up
            Err(CredentialsError::CredentialsNotLoaded { .. }) => {
                tracing::info!("provider returned CredentialsNotLoaded, ignoring")
            }
            // if we get another error class, there is probably something actually wrong that the user will
            // want to know about
            Err(other) => return Err(other.into()),
        }
        Ok(request)
    }
}

mod error {
    use aws_credential_types::provider::error::CredentialsError;
    use std::error::Error as StdError;
    use std::fmt;

    /// Failures that can occur in the credentials middleware.
    #[derive(Debug)]
    pub struct CredentialsStageError {
        source: CredentialsError,
    }

    impl StdError for CredentialsStageError {
        fn source(&self) -> Option<&(dyn StdError + 'static)> {
            Some(&self.source as _)
        }
    }

    impl fmt::Display for CredentialsStageError {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            write!(
                f,
                "failed to load credentials from the credentials provider"
            )
        }
    }

    impl From<CredentialsError> for CredentialsStageError {
        fn from(source: CredentialsError) -> Self {
            CredentialsStageError { source }
        }
    }
}

pub use error::*;

type BoxFuture<T> = Pin<Box<dyn Future<Output = T> + Send>>;

impl AsyncMapRequest for CredentialsStage {
    type Error = CredentialsStageError;
    type Future = Pin<Box<dyn Future<Output = Result<Request, Self::Error>> + Send + 'static>>;

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

    fn apply(&self, request: Request) -> BoxFuture<Result<Request, Self::Error>> {
        Box::pin(Self::load_creds(request))
    }
}

#[cfg(test)]
mod tests {
    use super::set_provider;
    use super::CredentialsStage;
    use aws_credential_types::{
        provider::{
            error::CredentialsError, future, ProvideCredentials, SharedCredentialsProvider,
        },
        Credentials,
    };
    use aws_smithy_http::body::SdkBody;
    use aws_smithy_http::middleware::AsyncMapRequest;
    use aws_smithy_http::operation;

    #[derive(Debug)]
    struct Unhandled;
    impl ProvideCredentials for Unhandled {
        fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
        where
            Self: 'a,
        {
            future::ProvideCredentials::ready(Err(CredentialsError::unhandled("whoops")))
        }
    }

    #[derive(Debug)]
    struct NoCreds;
    impl ProvideCredentials for NoCreds {
        fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
        where
            Self: 'a,
        {
            future::ProvideCredentials::ready(Err(CredentialsError::not_loaded("no creds")))
        }
    }

    #[tokio::test]
    async fn no_cred_provider_is_ok() {
        let req = operation::Request::new(http::Request::new(SdkBody::from("some body")));
        CredentialsStage::new()
            .apply(req)
            .await
            .expect("no credential provider should not populate credentials");
    }

    #[tokio::test]
    async fn provider_failure_is_failure() {
        let mut req = operation::Request::new(http::Request::new(SdkBody::from("some body")));
        set_provider(
            &mut req.properties_mut(),
            SharedCredentialsProvider::new(Unhandled),
        );
        CredentialsStage::new()
            .apply(req)
            .await
            .expect_err("no credential provider should not populate credentials");
    }

    #[tokio::test]
    async fn credentials_not_loaded_is_ok() {
        let mut req = operation::Request::new(http::Request::new(SdkBody::from("some body")));
        set_provider(
            &mut req.properties_mut(),
            SharedCredentialsProvider::new(NoCreds),
        );
        CredentialsStage::new()
            .apply(req)
            .await
            .expect("credentials not loaded is OK");
    }

    #[tokio::test]
    async fn async_map_request_apply_populates_credentials() {
        let mut req = operation::Request::new(http::Request::new(SdkBody::from("some body")));
        set_provider(
            &mut req.properties_mut(),
            SharedCredentialsProvider::new(Credentials::for_tests()),
        );
        let req = CredentialsStage::new()
            .apply(req)
            .await
            .expect("credential provider is in the bag; should succeed");
        assert!(
            req.properties().get::<Credentials>().is_some(),
            "it should set credentials on the request config"
        );
    }
}