Skip to main content

multistore_sts/
route_handler.rs

1//! Route handler for STS `AssumeRoleWithWebIdentity` requests.
2//!
3//! Intercepts STS queries before they reach the proxy dispatch pipeline
4//! and delegates to [`try_handle_sts`].
5
6use crate::{
7    handle_get_caller_identity, is_get_caller_identity, try_handle_sts, JwksCache, TokenKey,
8};
9use multistore::registry::CredentialRegistry;
10use multistore::route_handler::{ProxyResult, RequestInfo, RouteHandler, RouteHandlerFuture};
11use multistore::router::Router;
12
13/// Handler that intercepts STS `AssumeRoleWithWebIdentity` and
14/// `GetCallerIdentity` requests.
15struct StsHandler<C> {
16    config: C,
17    cache: JwksCache,
18    key: Option<TokenKey>,
19}
20
21impl<C: CredentialRegistry> RouteHandler for StsHandler<C> {
22    fn handle<'a>(&'a self, req: &'a RequestInfo<'a>) -> RouteHandlerFuture<'a> {
23        Box::pin(async move {
24            // GetCallerIdentity is authenticated (SigV4 over the temporary
25            // credentials) and needs the full request, so it is dispatched
26            // before the unauthenticated AssumeRoleWithWebIdentity exchange.
27            if is_get_caller_identity(req.query) || is_get_caller_identity(req.form_body) {
28                let (status, xml) = handle_get_caller_identity(req, self.key.as_ref());
29                return Some(ProxyResult::xml(status, xml));
30            }
31            let (status, xml) = try_handle_sts(
32                req.query,
33                req.form_body,
34                &self.config,
35                &self.cache,
36                self.key.as_ref(),
37            )
38            .await?;
39            Some(ProxyResult::xml(status, xml))
40        })
41    }
42}
43
44/// Extension trait for registering STS routes on a [`Router`].
45pub trait StsRouterExt {
46    /// Register the STS handler on the given `path`.
47    ///
48    /// STS requests are identified by their parameters (`Action=...`) — in the
49    /// query string or, as AWS SDKs send them, in a form-encoded `POST` body
50    /// surfaced via [`RequestInfo::form_body`] — not by path, so any path can be
51    /// used (e.g. `"/"` or `"/.sts"`). Runtimes must populate `form_body` for
52    /// form-encoded `POST`s or SDK clients will fall through unhandled.
53    ///
54    /// The handler is registered on both `path` and its trailing-slash variant
55    /// (`/.sts` and `/.sts/`). AWS SDK JS v3 — the STS client inside
56    /// `aws-actions/configure-aws-credentials` — appends a trailing slash to the
57    /// configured endpoint path, so its requests arrive at `/.sts/`; without the
58    /// second route they would miss the handler entirely.
59    fn with_sts<C: CredentialRegistry + Clone + 'static>(
60        self,
61        path: &str,
62        config: C,
63        cache: JwksCache,
64        key: Option<TokenKey>,
65    ) -> Self;
66}
67
68impl StsRouterExt for Router {
69    fn with_sts<C: CredentialRegistry + Clone + 'static>(
70        self,
71        path: &str,
72        config: C,
73        cache: JwksCache,
74        key: Option<TokenKey>,
75    ) -> Self {
76        let router = self.route(
77            path,
78            StsHandler {
79                config: config.clone(),
80                cache: cache.clone(),
81                key: key.clone(),
82            },
83        );
84        // Also accept the trailing-slash form (see doc comment). Skip it when
85        // `path` already ends in `/` to avoid registering the same route twice
86        // (matchit panics on a duplicate).
87        let with_slash = format!("{}/", path.trim_end_matches('/'));
88        if with_slash == path {
89            router
90        } else {
91            router.route(&with_slash, StsHandler { config, cache, key })
92        }
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99    use multistore::error::ProxyError;
100    use multistore::types::{RoleConfig, StoredCredential};
101
102    /// Minimal stub that satisfies `CredentialRegistry` without real data.
103    #[derive(Clone)]
104    struct EmptyRegistry;
105
106    impl CredentialRegistry for EmptyRegistry {
107        async fn get_credential(
108            &self,
109            _access_key_id: &str,
110        ) -> Result<Option<StoredCredential>, ProxyError> {
111            Ok(None)
112        }
113        async fn get_role(&self, _role_id: &str) -> Result<Option<RoleConfig>, ProxyError> {
114            Ok(None)
115        }
116    }
117
118    fn test_router() -> Router {
119        let cache = JwksCache::new(reqwest::Client::new(), std::time::Duration::from_secs(60));
120        Router::new().with_sts("/", EmptyRegistry, cache, None)
121    }
122
123    #[tokio::test]
124    async fn sts_query_on_root_path_is_handled() {
125        let router = test_router();
126        let headers = http::HeaderMap::new();
127        let req = RequestInfo::new(
128            &http::Method::GET,
129            "/",
130            Some("Action=AssumeRoleWithWebIdentity&RoleArn=test&WebIdentityToken=tok"),
131            &headers,
132            None,
133        );
134        assert!(
135            router.dispatch(&req).await.is_some(),
136            "STS request to / must be intercepted by the router"
137        );
138    }
139
140    #[tokio::test]
141    async fn non_sts_query_on_root_path_falls_through() {
142        let router = test_router();
143        let headers = http::HeaderMap::new();
144        let req = RequestInfo::new(&http::Method::GET, "/", Some("prefix=foo/"), &headers, None);
145        assert!(
146            router.dispatch(&req).await.is_none(),
147            "non-STS request to / must fall through"
148        );
149    }
150
151    #[tokio::test]
152    async fn sts_form_body_post_is_handled() {
153        // AWS SDKs send AssumeRoleWithWebIdentity as a form-encoded POST body
154        // with no query string at all.
155        let router = test_router();
156        let headers = http::HeaderMap::new();
157        let req = RequestInfo::new(&http::Method::POST, "/", None, &headers, None)
158            .with_form_body(Some(
159            "Action=AssumeRoleWithWebIdentity&Version=2011-06-15&RoleArn=test&WebIdentityToken=tok",
160        ));
161        assert!(
162            router.dispatch(&req).await.is_some(),
163            "STS form-body POST must be intercepted by the router"
164        );
165    }
166
167    #[tokio::test]
168    async fn non_sts_form_body_falls_through() {
169        let router = test_router();
170        let headers = http::HeaderMap::new();
171        let req = RequestInfo::new(&http::Method::POST, "/", None, &headers, None)
172            .with_form_body(Some("grant_type=client_credentials"));
173        assert!(
174            router.dispatch(&req).await.is_none(),
175            "non-STS form body must fall through"
176        );
177    }
178}