feat(auth): support OIDC ID-token subjects in enterprise exchange

crates/rmcp/src/transport/auth/enterprise.rs

@@ -1,6 +1,6 @@
 //! Non-interactive enterprise-managed authorization (EMA/XAA) token exchanges.
 //!
-//! Exchanges an enterprise refresh token for an ID-JAG, then for an MCP access
+//! Exchanges an enterprise ID token or refresh token for an ID-JAG, then for an MCP access
 //! token. Callers must discover and approve both servers and their registrations
 //! before supplying a credential. This module does not discover servers, log in,
 //! persist credentials, or decide when to reauthenticate.
@@ -56,6 +56,8 @@ use super::{
 };
 
 const ID_JAG_TOKEN_TYPE: &str = "urn:ietf:params:oauth:token-type:id-jag";
+const ID_TOKEN_TYPE: &str = "urn:ietf:params:oauth:token-type:id_token";
+const REFRESH_TOKEN_TYPE: &str = "urn:ietf:params:oauth:token-type:refresh_token";
 
 /// Authentication approved for a pre-registered client at one authorization server.
 ///
@@ -164,12 +166,33 @@ impl EmaAuthorizationServer {
     }
 }
 
-/// A refresh-token exchange bound to one MCP resource and two registered clients.
+enum EmaSubjectToken<'a> {
+    RefreshToken(&'a RefreshToken),
+    IdToken(&'a str),
+}
+
+impl EmaSubjectToken<'_> {
+    fn secret(&self) -> &str {
+        match self {
+            Self::RefreshToken(token) => token.secret(),
+            Self::IdToken(token) => token,
+        }
+    }
+
+    fn token_type(&self) -> &'static str {
+        match self {
+            Self::RefreshToken(_) => REFRESH_TOKEN_TYPE,
+            Self::IdToken(_) => ID_TOKEN_TYPE,
+        }
+    }
+}
+
+/// An identity token exchange bound to one MCP resource and two registered clients.
 pub struct EmaExchangeRequest<'a> {
     idp: EmaAuthorizationServer,
     resource_as: EmaAuthorizationServer,
     resource: &'a str,
-    refresh_token: &'a RefreshToken,
+    subject_token: EmaSubjectToken<'a>,
     scopes: Vec<String>,
 }
 
@@ -185,7 +208,24 @@ impl<'a> EmaExchangeRequest<'a> {
             idp,
             resource_as,
             resource,
-            refresh_token,
+            subject_token: EmaSubjectToken::RefreshToken(refresh_token),
+            scopes: Vec::new(),
+        }
+    }
+
+    /// Exchange an OIDC ID token issued by the approved IdP to this client.
+    /// The caller owns ID-token freshness and renewal; there is no refresh-token fallback.
+    pub fn new_with_id_token(
+        idp: EmaAuthorizationServer,
+        resource_as: EmaAuthorizationServer,
+        resource: &'a str,
+        id_token: &'a str,
+    ) -> Self {
+        Self {
+            idp,
+            resource_as,
+            resource,
+            subject_token: EmaSubjectToken::IdToken(id_token),
             scopes: Vec::new(),
         }
     }
@@ -202,7 +242,7 @@ impl<'a> EmaExchangeRequest<'a> {
     }
 
     /// Obtain an ID-JAG without redirects or retries, checking its resource/client bindings.
-    /// The returned assertion does not retain the enterprise refresh token.
+    /// The returned assertion does not retain the enterprise subject token.
     pub async fn exchange_id_jag(
         self,
         idp_http: &dyn OAuthHttpClient,
@@ -223,10 +263,10 @@ impl<'a> EmaExchangeRequest<'a> {
         }
         if self.idp.client_id.trim().is_empty()
             || self.resource_as.client_id.trim().is_empty()
-            || self.refresh_token.secret().trim().is_empty()
+            || self.subject_token.secret().trim().is_empty()
         {
             return Err(EmaError::InvalidRequest(
-                "client IDs and refresh token must not be empty",
+                "client IDs and subject token must not be empty",
             ));
         }
         self.idp.client_authentication.validate()?;
@@ -243,11 +283,8 @@ impl<'a> EmaExchangeRequest<'a> {
                 "urn:ietf:params:oauth:grant-type:token-exchange",
             ),
             ("requested_token_type", ID_JAG_TOKEN_TYPE),
-            (
-                "subject_token_type",
-                "urn:ietf:params:oauth:token-type:refresh_token",
-            ),
-            ("subject_token", self.refresh_token.secret()),
+            ("subject_token_type", self.subject_token.token_type()),
+            ("subject_token", self.subject_token.secret()),
             ("audience", self.resource_as.issuer.as_str()),
             ("resource", self.resource),
         ];

crates/rmcp/src/transport/auth/enterprise_tests.rs

@@ -122,35 +122,56 @@ fn authenticated_form(request: &OAuthHttpRequest) -> BTreeMap<String, String> {
 }
 
 #[tokio::test]
-async fn refresh_exchange_preserves_exact_forms_and_signed_narrowing() {
-    for scopes in ["files.read files.write", ""] {
-        let response = jag(&claims());
-        let idp = MockHttp::new(200, response.clone());
-        let result = exchange(&idp, scopes).await.unwrap();
-        assert_eq!(
-            result.assertion().secret(),
-            response["access_token"].as_str().unwrap()
-        );
-        assert_eq!(result.scopes(), &HashSet::from(["files.read".to_owned()]));
-        assert!(!format!("{result:?}").contains(result.assertion().secret()));
-        assert!(!format!("{result:?}").contains("private-query"));
-        let requests = idp.requests.lock().unwrap();
-        assert_eq!(requests.len(), 1);
-        assert_eq!(requests[0].request.uri(), IDP_TOKEN);
-        let mut fields = form(&requests[0]);
-        assert_eq!(
-            fields.remove("scope"),
-            (!scopes.is_empty()).then(|| scopes.to_owned())
-        );
-        assert_eq!(
-            serde_json::to_value(fields).unwrap(),
-            json!({
-                "grant_type":"urn:ietf:params:oauth:grant-type:token-exchange",
-                "requested_token_type":ID_JAG_TOKEN_TYPE,"subject_token":"refresh-token",
-                "subject_token_type":"urn:ietf:params:oauth:token-type:refresh_token",
-                "audience":AS,"resource":RESOURCE,"client_id":"idp"
-            })
-        );
+async fn subject_exchange_preserves_exact_forms_and_signed_narrowing() {
+    for id_token in [None, Some("oidc-id-token")] {
+        for scopes in ["files.read files.write", ""] {
+            let response = jag(&claims());
+            let idp = MockHttp::new(200, response.clone());
+            let refresh = RefreshToken::new("refresh-token".into());
+            let idp_registration = EmaAuthorizationServer::new(IDP, IDP_TOKEN, "idp");
+            let resource_registration = EmaAuthorizationServer::new(AS, AS_TOKEN, "mcp");
+            let request = match id_token {
+                Some(id_token) => EmaExchangeRequest::new_with_id_token(
+                    idp_registration,
+                    resource_registration,
+                    RESOURCE,
+                    id_token,
+                ),
+                None => EmaExchangeRequest::new(
+                    idp_registration,
+                    resource_registration,
+                    RESOURCE,
+                    &refresh,
+                ),
+            }
+            .with_scopes(scopes.split_ascii_whitespace());
+            assert!(!format!("{request:?}").contains(id_token.unwrap_or("refresh-token")));
+            let result = request.exchange_id_jag(&idp).await.unwrap();
+            assert_eq!(
+                result.assertion().secret(),
+                response["access_token"].as_str().unwrap()
+            );
+            assert_eq!(result.scopes(), &HashSet::from(["files.read".to_owned()]));
+            assert!(!format!("{result:?}").contains(result.assertion().secret()));
+            assert!(!format!("{result:?}").contains("private-query"));
+            let requests = idp.requests.lock().unwrap();
+            assert_eq!(requests.len(), 1);
+            assert_eq!(requests[0].request.uri(), IDP_TOKEN);
+            let mut fields = form(&requests[0]);
+            assert_eq!(
+                fields.remove("scope"),
+                (!scopes.is_empty()).then(|| scopes.to_owned())
+            );
+            assert_eq!(
+                serde_json::to_value(fields).unwrap(),
+                json!({
+                    "grant_type":"urn:ietf:params:oauth:grant-type:token-exchange",
+                    "requested_token_type":ID_JAG_TOKEN_TYPE,"subject_token":id_token.unwrap_or("refresh-token"),
+                    "subject_token_type":if id_token.is_some() { "urn:ietf:params:oauth:token-type:id_token" } else { "urn:ietf:params:oauth:token-type:refresh_token" },
+                    "audience":AS,"resource":RESOURCE,"client_id":"idp"
+                })
+            );
+        }
     }
 }
 
@@ -342,18 +363,23 @@ async fn invalid_inputs_fail_before_http() {
     }
     cases.push((original, vec!["files.read", "files.read"]));
     for (fields, scopes) in cases {
-        let http = MockHttp::default();
-        let result = EmaExchangeRequest::new(
-            EmaAuthorizationServer::new(fields[0], fields[1], fields[2]),
-            EmaAuthorizationServer::new(fields[3], fields[4], fields[5]),
-            fields[6],
-            &RefreshToken::new(fields[7].into()),
-        )
-        .with_scopes(scopes)
-        .exchange_id_jag(&http)
-        .await;
-        assert!(matches!(result, Err(EmaError::InvalidRequest(_))));
-        assert!(http.requests.lock().unwrap().is_empty());
+        for use_id_token in [false, true] {
+            let http = MockHttp::default();
+            let refresh = RefreshToken::new(fields[7].into());
+            let idp = EmaAuthorizationServer::new(fields[0], fields[1], fields[2]);
+            let resource_as = EmaAuthorizationServer::new(fields[3], fields[4], fields[5]);
+            let request = if use_id_token {
+                EmaExchangeRequest::new_with_id_token(idp, resource_as, fields[6], fields[7])
+            } else {
+                EmaExchangeRequest::new(idp, resource_as, fields[6], &refresh)
+            };
+            let result = request
+                .with_scopes(scopes.iter().copied())
+                .exchange_id_jag(&http)
+                .await;
+            assert!(matches!(result, Err(EmaError::InvalidRequest(_))));
+            assert!(http.requests.lock().unwrap().is_empty());
+        }
     }
 }