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::{try_handle_sts, JwksCache, TokenKey};
7use multistore::registry::CredentialRegistry;
8use multistore::route_handler::{ProxyResult, RequestInfo, RouteHandler, RouteHandlerFuture};
9use multistore::router::Router;
10
11/// Handler that intercepts `AssumeRoleWithWebIdentity` STS requests.
12struct StsHandler<C> {
13    config: C,
14    cache: JwksCache,
15    key: Option<TokenKey>,
16}
17
18impl<C: CredentialRegistry> RouteHandler for StsHandler<C> {
19    fn handle<'a>(&'a self, req: &'a RequestInfo<'a>) -> RouteHandlerFuture<'a> {
20        Box::pin(async move {
21            let (status, xml) = try_handle_sts(
22                req.query,
23                req.form_body,
24                &self.config,
25                &self.cache,
26                self.key.as_ref(),
27            )
28            .await?;
29            Some(ProxyResult::xml(status, xml))
30        })
31    }
32}
33
34/// Extension trait for registering STS routes on a [`Router`].
35pub trait StsRouterExt {
36    /// Register the STS handler on the given `path`.
37    ///
38    /// STS requests are identified by their parameters
39    /// (`Action=AssumeRoleWithWebIdentity`) — in the query string or, as AWS
40    /// SDKs send them, in a form-encoded `POST` body surfaced via
41    /// [`RequestInfo::form_body`] — not by path, so any path can be used
42    /// (e.g. `"/"` or `"/.sts"`). Runtimes must populate `form_body` for
43    /// form-encoded `POST`s or SDK clients will fall through unhandled.
44    fn with_sts<C: CredentialRegistry + 'static>(
45        self,
46        path: &str,
47        config: C,
48        cache: JwksCache,
49        key: Option<TokenKey>,
50    ) -> Self;
51}
52
53impl StsRouterExt for Router {
54    fn with_sts<C: CredentialRegistry + 'static>(
55        self,
56        path: &str,
57        config: C,
58        cache: JwksCache,
59        key: Option<TokenKey>,
60    ) -> Self {
61        self.route(path, StsHandler { config, cache, key })
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    use super::*;
68    use multistore::error::ProxyError;
69    use multistore::types::{RoleConfig, StoredCredential};
70
71    /// Minimal stub that satisfies `CredentialRegistry` without real data.
72    #[derive(Clone)]
73    struct EmptyRegistry;
74
75    impl CredentialRegistry for EmptyRegistry {
76        async fn get_credential(
77            &self,
78            _access_key_id: &str,
79        ) -> Result<Option<StoredCredential>, ProxyError> {
80            Ok(None)
81        }
82        async fn get_role(&self, _role_id: &str) -> Result<Option<RoleConfig>, ProxyError> {
83            Ok(None)
84        }
85    }
86
87    fn test_router() -> Router {
88        let cache = JwksCache::new(reqwest::Client::new(), std::time::Duration::from_secs(60));
89        Router::new().with_sts("/", EmptyRegistry, cache, None)
90    }
91
92    #[tokio::test]
93    async fn sts_query_on_root_path_is_handled() {
94        let router = test_router();
95        let headers = http::HeaderMap::new();
96        let req = RequestInfo::new(
97            &http::Method::GET,
98            "/",
99            Some("Action=AssumeRoleWithWebIdentity&RoleArn=test&WebIdentityToken=tok"),
100            &headers,
101            None,
102        );
103        assert!(
104            router.dispatch(&req).await.is_some(),
105            "STS request to / must be intercepted by the router"
106        );
107    }
108
109    #[tokio::test]
110    async fn non_sts_query_on_root_path_falls_through() {
111        let router = test_router();
112        let headers = http::HeaderMap::new();
113        let req = RequestInfo::new(&http::Method::GET, "/", Some("prefix=foo/"), &headers, None);
114        assert!(
115            router.dispatch(&req).await.is_none(),
116            "non-STS request to / must fall through"
117        );
118    }
119
120    #[tokio::test]
121    async fn sts_form_body_post_is_handled() {
122        // AWS SDKs send AssumeRoleWithWebIdentity as a form-encoded POST body
123        // with no query string at all.
124        let router = test_router();
125        let headers = http::HeaderMap::new();
126        let req = RequestInfo::new(&http::Method::POST, "/", None, &headers, None)
127            .with_form_body(Some(
128            "Action=AssumeRoleWithWebIdentity&Version=2011-06-15&RoleArn=test&WebIdentityToken=tok",
129        ));
130        assert!(
131            router.dispatch(&req).await.is_some(),
132            "STS form-body POST must be intercepted by the router"
133        );
134    }
135
136    #[tokio::test]
137    async fn non_sts_form_body_falls_through() {
138        let router = test_router();
139        let headers = http::HeaderMap::new();
140        let req = RequestInfo::new(&http::Method::POST, "/", None, &headers, None)
141            .with_form_body(Some("grant_type=client_credentials"));
142        assert!(
143            router.dispatch(&req).await.is_none(),
144            "non-STS form body must fall through"
145        );
146    }
147}
multistore_sts/route_handler.rs

multistore_sts/
route_handler.rs
