multistore-sts 0.7.1

//! Route handler for STS `AssumeRoleWithWebIdentity` requests.
//!
//! Intercepts STS queries before they reach the proxy dispatch pipeline
//! and delegates to [`try_handle_sts`].

use crate::{try_handle_sts, JwksCache, TokenKey};
use multistore::registry::CredentialRegistry;
use multistore::route_handler::{ProxyResult, RequestInfo, RouteHandler, RouteHandlerFuture};
use multistore::router::Router;

/// Handler that intercepts `AssumeRoleWithWebIdentity` STS requests.
struct StsHandler<C> {
    config: C,
    cache: JwksCache,
    key: Option<TokenKey>,
}

impl<C: CredentialRegistry> RouteHandler for StsHandler<C> {
    fn handle<'a>(&'a self, req: &'a RequestInfo<'a>) -> RouteHandlerFuture<'a> {
        Box::pin(async move {
            let (status, xml) = try_handle_sts(
                req.query,
                req.form_body,
                &self.config,
                &self.cache,
                self.key.as_ref(),
            )
            .await?;
            Some(ProxyResult::xml(status, xml))
        })
    }
}

/// Extension trait for registering STS routes on a [`Router`].
pub trait StsRouterExt {
    /// Register the STS handler on the given `path`.
    ///
    /// STS requests are identified by their parameters
    /// (`Action=AssumeRoleWithWebIdentity`) — in the query string or, as AWS
    /// SDKs send them, in a form-encoded `POST` body surfaced via
    /// [`RequestInfo::form_body`] — not by path, so any path can be used
    /// (e.g. `"/"` or `"/.sts"`). Runtimes must populate `form_body` for
    /// form-encoded `POST`s or SDK clients will fall through unhandled.
    fn with_sts<C: CredentialRegistry + 'static>(
        self,
        path: &str,
        config: C,
        cache: JwksCache,
        key: Option<TokenKey>,
    ) -> Self;
}

impl StsRouterExt for Router {
    fn with_sts<C: CredentialRegistry + 'static>(
        self,
        path: &str,
        config: C,
        cache: JwksCache,
        key: Option<TokenKey>,
    ) -> Self {
        self.route(path, StsHandler { config, cache, key })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use multistore::error::ProxyError;
    use multistore::types::{RoleConfig, StoredCredential};

    /// Minimal stub that satisfies `CredentialRegistry` without real data.
    #[derive(Clone)]
    struct EmptyRegistry;

    impl CredentialRegistry for EmptyRegistry {
        async fn get_credential(
            &self,
            _access_key_id: &str,
        ) -> Result<Option<StoredCredential>, ProxyError> {
            Ok(None)
        }
        async fn get_role(&self, _role_id: &str) -> Result<Option<RoleConfig>, ProxyError> {
            Ok(None)
        }
    }

    fn test_router() -> Router {
        let cache = JwksCache::new(reqwest::Client::new(), std::time::Duration::from_secs(60));
        Router::new().with_sts("/", EmptyRegistry, cache, None)
    }

    #[tokio::test]
    async fn sts_query_on_root_path_is_handled() {
        let router = test_router();
        let headers = http::HeaderMap::new();
        let req = RequestInfo::new(
            &http::Method::GET,
            "/",
            Some("Action=AssumeRoleWithWebIdentity&RoleArn=test&WebIdentityToken=tok"),
            &headers,
            None,
        );
        assert!(
            router.dispatch(&req).await.is_some(),
            "STS request to / must be intercepted by the router"
        );
    }

    #[tokio::test]
    async fn non_sts_query_on_root_path_falls_through() {
        let router = test_router();
        let headers = http::HeaderMap::new();
        let req = RequestInfo::new(&http::Method::GET, "/", Some("prefix=foo/"), &headers, None);
        assert!(
            router.dispatch(&req).await.is_none(),
            "non-STS request to / must fall through"
        );
    }

    #[tokio::test]
    async fn sts_form_body_post_is_handled() {
        // AWS SDKs send AssumeRoleWithWebIdentity as a form-encoded POST body
        // with no query string at all.
        let router = test_router();
        let headers = http::HeaderMap::new();
        let req = RequestInfo::new(&http::Method::POST, "/", None, &headers, None)
            .with_form_body(Some(
            "Action=AssumeRoleWithWebIdentity&Version=2011-06-15&RoleArn=test&WebIdentityToken=tok",
        ));
        assert!(
            router.dispatch(&req).await.is_some(),
            "STS form-body POST must be intercepted by the router"
        );
    }

    #[tokio::test]
    async fn non_sts_form_body_falls_through() {
        let router = test_router();
        let headers = http::HeaderMap::new();
        let req = RequestInfo::new(&http::Method::POST, "/", None, &headers, None)
            .with_form_body(Some("grant_type=client_credentials"));
        assert!(
            router.dispatch(&req).await.is_none(),
            "non-STS form body must fall through"
        );
    }
}