multistore_sts/
route_handler.rs1use 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
13struct 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 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
44pub trait StsRouterExt {
46 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 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 #[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 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}