support any valid json values in custom claims (#28563)

GitOrigin-RevId: c9425165bc45f3f5d40ce008ac15cf9842433f17

Author: rrwang7

sync_types/src/json.rs

@@ -790,7 +790,7 @@ struct UserIdentityAttributesJson {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub updated_at: Option<String>,
     #[serde(flatten)]
-    pub custom_claims: Option<BTreeMap<String, String>>,
+    pub custom_claims: Option<BTreeMap<String, JsonValue>>,
 }
 
 impl TryFrom<JsonValue> for UserIdentityAttributes {
@@ -808,6 +808,13 @@ impl TryFrom<JsonValue> for UserIdentityAttributes {
         let custom_claims = raw
             .custom_claims
             .context("expected custom claims to be set")?;
+        let custom_claims_string = custom_claims
+            .into_iter()
+            .map(|(key, value)| {
+                let value_string = serde_json::to_string(&value)?;
+                Ok((key, value_string))
+            })
+            .collect::<anyhow::Result<_>>()?;
 
         Ok(UserIdentityAttributes {
             token_identifier,
@@ -831,7 +838,7 @@ impl TryFrom<JsonValue> for UserIdentityAttributes {
             phone_number_verified: raw.phone_number_verified,
             address: raw.address,
             updated_at: raw.updated_at,
-            custom_claims,
+            custom_claims: custom_claims_string,
         })
     }
 }
@@ -840,6 +847,14 @@ impl TryFrom<UserIdentityAttributes> for JsonValue {
     type Error = anyhow::Error;
 
     fn try_from(value: UserIdentityAttributes) -> Result<Self, Self::Error> {
+        let custom_claims_json = value
+            .custom_claims
+            .into_iter()
+            .map(|(key, value)| {
+                let value_json = serde_json::from_str(&value)?;
+                Ok((key, value_json))
+            })
+            .collect::<anyhow::Result<_>>()?;
         let raw = UserIdentityAttributesJson {
             token_identifier: Some(value.token_identifier),
             issuer: value.issuer,
@@ -862,7 +877,7 @@ impl TryFrom<UserIdentityAttributes> for JsonValue {
             phone_number_verified: value.phone_number_verified,
             address: value.address,
             updated_at: value.updated_at,
-            custom_claims: Some(value.custom_claims),
+            custom_claims: Some(custom_claims_json),
         };
         Ok(serde_json::to_value(raw)?)
     }

sync_types/src/testing.rs

@@ -22,14 +22,14 @@ pub fn arb_json() -> impl Strategy<Value = JsonValue> {
         ".*".prop_map(JsonValue::String),
     ];
     leaf.prop_recursive(
-        8,   // 8 levels deep
-        256, // Shoot for maximum size of 256 nodes
-        10,  // We put up to 10 items per collection
+        4,   // 4 levels deep
+        128, // Shoot for maximum size of 128 nodes
+        5,   // We put up to 5 items per collection
         |inner| {
             prop_oneof![
                 // Take the inner strategy and make the two recursive cases.
-                prop::collection::vec(inner.clone(), 0..10).prop_map(JsonValue::Array),
-                prop::collection::hash_map(".*", inner, 0..10)
+                prop::collection::vec(inner.clone(), 0..5).prop_map(JsonValue::Array),
+                prop::collection::hash_map(".*", inner, 0..5)
                     .prop_map(|m| JsonValue::Object(m.into_iter().collect())),
             ]
         },

sync_types/src/types.rs

@@ -17,6 +17,8 @@ use serde::{
 use serde_json::Value as JsonValue;
 use uuid::Uuid;
 
+#[cfg(any(test, feature = "testing"))]
+use crate::testing::arb_json;
 use crate::{
     Timestamp,
     UdfPath,
@@ -51,7 +53,17 @@ pub type IdentityVersion = u32;
 /// it's good enough for our tests here.
 #[cfg(any(test, feature = "testing"))]
 fn string_json_args_strategy() -> impl proptest::strategy::Strategy<Value = Vec<JsonValue>> {
-    Vec::<String>::arbitrary().prop_map(|v| v.iter().map(|s| JsonValue::String(s.into())).collect())
+    prop::collection::vec(any::<String>(), 0..2)
+        .prop_map(|v| v.iter().map(|s| JsonValue::String(s.into())).collect())
+}
+
+#[cfg(any(test, feature = "testing"))]
+fn custom_claims_strategy() -> impl proptest::strategy::Strategy<Value = BTreeMap<String, String>> {
+    prop::collection::btree_map(
+        any::<String>(),
+        arb_json().prop_map(|v| serde_json::to_string(&v).unwrap()),
+        0..2,
+    )
 }
 
 /// This strategy only generates a string (not arbitrary JSON) but
@@ -102,7 +114,7 @@ pub enum ClientMessage {
         new_version: QuerySetVersion,
         #[cfg_attr(
             any(test, feature = "testing"),
-            proptest(strategy = "prop::collection::vec(any::<QuerySetModification>(), 0..8)")
+            proptest(strategy = "prop::collection::vec(any::<QuerySetModification>(), 0..2)")
         )]
         modifications: Vec<QuerySetModification>,
     },
@@ -194,6 +206,11 @@ pub struct UserIdentityAttributes {
     pub address: Option<String>,
     /// Stored as RFC3339 string
     pub updated_at: Option<String>,
+
+    #[cfg_attr(
+        any(test, feature = "testing"),
+        proptest(strategy = "custom_claims_strategy()")
+    )]
     pub custom_claims: BTreeMap<String, String>,
 }