Implement OpenID Connect custom claims using AdditionalClaims (#36167)

GitOrigin-RevId: 2ae1db008c82cc96c99959252b68c086cb08c222

Author: goffrie

Cargo.lock

@@ -18,14 +18,16 @@ use chrono::TimeZone;
 use common::auth::AuthInfo;
 use errors::ErrorMetadata;
 use futures::Future;
-use keybroker::UserIdentity;
+use keybroker::{
+    CoreIdTokenWithCustomClaims,
+    UserIdentity,
+};
 use oauth2::{
     HttpRequest,
     HttpResponse,
 };
 use openidconnect::{
     core::{
-        CoreIdToken,
         CoreIdTokenVerifier,
         CoreProviderMetadata,
     },
@@ -108,10 +110,9 @@ where
     F: Future<Output = Result<HttpResponse, E>>,
     E: std::error::Error + 'static + Send + Sync,
 {
-    let token = CoreIdToken::from_str(&token_str.0).context(ErrorMetadata::unauthenticated(
-        "InvalidAuthHeader",
-        "Could not parse as id token",
-    ))?;
+    let token = CoreIdTokenWithCustomClaims::from_str(&token_str.0).context(
+        ErrorMetadata::unauthenticated("InvalidAuthHeader", "Could not parse as id token"),
+    )?;
     let (audiences, issuer) = {
         let verifier = CoreIdTokenVerifier::new_insecure_without_verification();
         let claims = match token.claims(&verifier, |_: Option<&openidconnect::Nonce>| Ok(())) {

crates/authentication/src/lib.rs

@@ -27,11 +27,13 @@ prost = { workspace = true }
 rand = { workspace = true }
 rsa = { workspace = true, optional = true }
 serde = { workspace = true }
+serde_json = { workspace = true }
 sodium_secretbox = { path = "../sodium_secretbox" }
 sync_types = { package = "convex_sync_types", path = "../convex/sync_types" }
 tracing = { workspace = true }
 
 [dev-dependencies]
+base64 = { workspace = true }
 common = { path = "../common", features = ["testing"] }
 errors = { path = "../errors", features = ["testing"] }
 metrics = { path = "../metrics", features = ["testing"] }

crates/keybroker/Cargo.toml

@@ -1,6 +1,9 @@
 use core::panic;
 use std::{
-    collections::BTreeMap,
+    collections::{
+        BTreeMap,
+        HashMap,
+    },
     fmt,
     time::{
         Duration,
@@ -43,9 +46,14 @@ use errors::ErrorMetadata;
 use metrics::StaticMetricLabel;
 use openidconnect::{
     core::{
-        CoreIdToken,
+        CoreGenderClaim,
         CoreIdTokenVerifier,
+        CoreJsonWebKeyType,
+        CoreJweContentEncryptionAlgorithm,
+        CoreJwsSigningAlgorithm,
     },
+    AdditionalClaims,
+    IdToken,
     Nonce,
 };
 use pb::{
@@ -370,7 +378,7 @@ pub struct UserIdentity {
     pub expiration: SystemTime,
     pub attributes: UserIdentityAttributes,
     // The original token this user identity was created from.
-    pub original_token: CoreIdToken,
+    pub original_token: CoreIdTokenWithCustomClaims,
 }
 
 #[cfg(any(test, feature = "testing"))]
@@ -422,9 +430,21 @@ macro_rules! get_localized_string {
     };
 }
 
+pub type CoreIdTokenWithCustomClaims = IdToken<
+    CustomClaims,
+    CoreGenderClaim,
+    CoreJweContentEncryptionAlgorithm,
+    CoreJwsSigningAlgorithm,
+    CoreJsonWebKeyType,
+>;
+
+#[derive(Deserialize, Serialize, Debug, Clone, Default, PartialEq, Eq)]
+pub struct CustomClaims(HashMap<String, serde_json::Value>);
+impl AdditionalClaims for CustomClaims {}
+
 impl UserIdentity {
     pub fn from_token(
-        token: CoreIdToken,
+        token: CoreIdTokenWithCustomClaims,
         verifier: CoreIdTokenVerifier,
     ) -> Result<Self, anyhow::Error> {
         // NB: Nonce verification is optional, and we'd need the developer to create and
@@ -436,7 +456,7 @@ impl UserIdentity {
         let subject = claims.subject().to_string();
         let issuer = claims.issuer().to_string();
         let mut custom_claims = BTreeMap::new();
-        for claim in claims.custom_claims() {
+        for claim in &claims.additional_claims().0 {
             // Filter out standard claims and claims set by auth providers
             match claim.0.as_str() {
                 // Standard claims that we support: see https://docs.convex.dev/api/interfaces/server.UserIdentity

crates/keybroker/src/broker.rs

@@ -9,12 +9,17 @@ mod secret;
 #[cfg(any(test, feature = "testing"))]
 pub mod testing;
 
+#[cfg(test)]
+mod tests;
+
 pub use sync_types::UserIdentityAttributes;
 
 pub use self::{
     broker::{
         AdminIdentity,
         AdminIdentityPrincipal,
+        CoreIdTokenWithCustomClaims,
+        CustomClaims,
         GetFileAuthorization,
         Identity,
         KeyBroker,

crates/keybroker/src/lib.rs

@@ -6,16 +6,14 @@ use chrono::{
 };
 use openidconnect::{
     core::{
-        CoreIdToken,
-        CoreIdTokenClaims,
         CoreIdTokenVerifier,
         CoreJwsSigningAlgorithm,
         CoreRsaPrivateSigningKey,
     },
     Audience,
-    EmptyAdditionalClaims,
     EndUserEmail,
     EndUserName,
+    IdTokenClaims,
     IssuerUrl,
     JsonWebKeyId,
     StandardClaims,
@@ -24,7 +22,13 @@ use openidconnect::{
 use rsa::pkcs1::EncodeRsaPrivateKey;
 use sync_types::UserIdentityAttributes;
 
-use crate::UserIdentity;
+use crate::{
+    broker::{
+        CoreIdTokenWithCustomClaims,
+        CustomClaims,
+    },
+    UserIdentity,
+};
 
 pub static TEST_SIGNING_KEY: LazyLock<CoreRsaPrivateSigningKey> = LazyLock::new(|| {
     let key = rsa::RsaPrivateKey::new(&mut rsa::rand_core::OsRng, 2048).unwrap();
@@ -42,16 +46,16 @@ impl TestUserIdentity for UserIdentity {
         let issuer = "https://testauth.fake.domain".to_owned();
         let audience = Audience::new("client-id-123".to_string());
 
-        let token = CoreIdToken::new(
-            CoreIdTokenClaims::new(
+        let token = CoreIdTokenWithCustomClaims::new(
+            IdTokenClaims::new(
                 IssuerUrl::new(issuer).unwrap(),
                 vec![audience],
                 Utc::now() + Duration::seconds(600),
                 Utc::now(),
                 StandardClaims::new(SubjectIdentifier::new(subject))
                     .set_email(Some(EndUserEmail::new("foo@bar.com".to_string())))
                     .set_name(Some(EndUserName::new("Al Pastor".to_string()).into())),
-                EmptyAdditionalClaims {},
+                CustomClaims::default(),
             ),
             &*TEST_SIGNING_KEY,
             CoreJwsSigningAlgorithm::RsaSsaPkcs1V15Sha256,

crates/keybroker/src/testing.rs

@@ -0,0 +1,37 @@
+use std::str::FromStr;
+
+use chrono::{
+    TimeZone,
+    Utc,
+};
+use openidconnect::core::CoreIdTokenVerifier;
+
+use crate::{
+    CoreIdTokenWithCustomClaims,
+    UserIdentity,
+};
+
+#[test]
+fn test_custom_claims() {
+    let header_json = "{\"alg\":\"none\"}";
+    let payload_json = "{
+        \"iss\": \"https://server.example.com\",
+        \"sub\": \"24400320\",
+        \"aud\": [\"s6BhdRkqt3\"],
+        \"exp\": 1311281970,
+        \"iat\": 1311280970,
+        \"tfa_method\": \"u2f\",
+        \"bool_me\": true
+    }";
+    let signature = "foo";
+    let token_str = [header_json, payload_json, signature]
+        .map(|s| base64::encode_config(s, base64::URL_SAFE_NO_PAD))
+        .join(".");
+    let token = CoreIdTokenWithCustomClaims::from_str(&token_str).expect("failed to parse");
+    let verifier = CoreIdTokenVerifier::new_insecure_without_verification()
+        .set_time_fn(|| Utc.timestamp_opt(1311281969, 0).unwrap());
+    let identity = UserIdentity::from_token(token, verifier).unwrap();
+    let custom_claims = identity.attributes.custom_claims;
+    assert_eq!(custom_claims.get("tfa_method").unwrap(), "\"u2f\"");
+    assert_eq!(custom_claims.get("bool_me").unwrap(), "true");
+}