Add extra claims to v1/identity/websocket-token (#3705)

# Description of Changes

Due to a limitation around passing headers to a WebSocket connection,
The typescript SDK rely on the endpoint `/v1/identity/websocket-token`
to get a new, short-lived token.
Currently, this endpoint strips all the other claims from the token and
only returns the following claims:

- `hex_identity`
- `sub`
- `iss`
- `aud`
- `iat`
- `exp`

This PR aims to fix this issue by introducing a new member field `extra`
to `SpacetimeIdentityClaims` and `TokenClaims` and letting serde do its
job.

# API and ABI breaking changes

None

# Expected complexity level and risk

2 - The change is trivial (1) but I'm not 100% familiar with all the
places where we would be signing a token (1).

# Testing

1. `curl` the endpoint and checking that the token returned contains all
the expected claims
2. Check that that the endpoint `v1/identity` still correctly issues and
identity and token

---------

Co-authored-by: Jeffrey Dallatezza <jeffreydallatezza@gmail.com>

Author: JulienLavocat

crates/auth/src/identity.rs

@@ -4,6 +4,7 @@ pub use jsonwebtoken::{DecodingKey, EncodingKey};
 use serde::Deserializer;
 use serde::{Deserialize, Serialize};
 use spacetimedb_lib::Identity;
+use std::collections::HashMap;
 use std::time::SystemTime;
 
 #[derive(Debug, Clone)]
@@ -41,6 +42,9 @@ pub struct SpacetimeIdentityClaims {
     pub iat: SystemTime,
     #[serde_as(as = "Option<serde_with::TimestampSeconds>")]
     pub exp: Option<SystemTime>,
+
+    #[serde(flatten)]
+    pub extra: Option<HashMap<String, serde_json::Value>>,
 }
 
 fn deserialize_audience<'de, D>(deserializer: D) -> Result<Vec<String>, D::Error>
@@ -84,6 +88,10 @@ pub struct IncomingClaims {
     pub iat: SystemTime,
     #[serde_as(as = "Option<serde_with::TimestampSeconds>")]
     pub exp: Option<SystemTime>,
+
+    /// All remaining claims from the JWT payload
+    #[serde(flatten)]
+    pub extra: Option<HashMap<String, serde_json::Value>>,
 }
 
 impl TryInto<SpacetimeIdentityClaims> for IncomingClaims {
@@ -122,6 +130,7 @@ impl TryInto<SpacetimeIdentityClaims> for IncomingClaims {
             audience: self.audience,
             iat: self.iat,
             exp: self.exp,
+            extra: self.extra,
         })
     }
 }

crates/client-api/src/auth.rs

@@ -14,6 +14,7 @@ use spacetimedb::auth::token_validation::{
 use spacetimedb::auth::JwtKeys;
 use spacetimedb::energy::EnergyQuanta;
 use spacetimedb::identity::Identity;
+use std::collections::HashMap;
 use std::time::{Duration, SystemTime};
 use uuid::Uuid;
 
@@ -117,15 +118,16 @@ pub struct TokenClaims {
     pub issuer: String,
     pub subject: String,
     pub audience: Vec<String>,
+    pub extra: Option<HashMap<String, serde_json::Value>>,
 }
 
 impl From<SpacetimeAuth> for TokenClaims {
     fn from(auth: SpacetimeAuth) -> Self {
         Self {
             issuer: auth.claims.issuer,
             subject: auth.claims.subject,
-            // This will need to be changed when we care about audiencies.
-            audience: Vec::new(),
+            audience: auth.claims.audience,
+            extra: auth.claims.extra,
         }
     }
 }
@@ -136,6 +138,7 @@ impl TokenClaims {
             issuer,
             subject,
             audience: Vec::new(),
+            extra: None,
         }
     }
 
@@ -159,6 +162,7 @@ impl TokenClaims {
             subject: self.subject.clone(),
             issuer: self.issuer.clone(),
             audience: self.audience.clone(),
+            extra: self.extra.clone(),
             iat,
             exp,
         };
@@ -184,6 +188,7 @@ impl SpacetimeAuth {
             subject: subject.clone(),
             // Placeholder audience.
             audience: vec!["spacetimedb".to_string()],
+            extra: None,
         };
 
         let (claims, token) = claims.encode_and_sign(ctx.jwt_auth_provider()).map_err(log_and_500)?;
@@ -280,7 +285,7 @@ mod tests {
     use anyhow::Ok;
 
     use spacetimedb::auth::{token_validation::TokenValidator, JwtKeys};
-    use std::collections::HashSet;
+    use std::collections::{HashMap, HashSet};
 
     // Make sure that when we encode TokenClaims, we can decode to get the expected identity.
     #[tokio::test]
@@ -291,12 +296,42 @@ mod tests {
             issuer: "localhost".to_string(),
             subject: "test-subject".to_string(),
             audience: vec!["spacetimedb".to_string()],
+            extra: None,
+        };
+        let id = claims.id();
+        let (_, token) = claims.encode_and_sign(&kp.private)?;
+        let decoded = kp.public.validate_token(&token).await?;
+
+        assert_eq!(decoded.identity, id);
+        Ok(())
+    }
+
+    fn to_hashmap(value: serde_json::Value) -> HashMap<String, serde_json::Value> {
+        let mut map = HashMap::new();
+        value.as_object().unwrap().iter().for_each(|(k, v)| {
+            map.insert(k.clone(), v.clone());
+        });
+        map
+    }
+
+    // Make sure that when we encode TokenClaims, we can decode the extra claims.
+    #[tokio::test]
+    async fn decode_encoded_token_with_extra_claims() -> Result<(), anyhow::Error> {
+        let kp = JwtKeys::generate()?;
+
+        let claims = TokenClaims {
+            issuer: "localhost".to_string(),
+            subject: "test-subject".to_string(),
+            audience: vec!["spacetimedb".to_string()],
+            extra: Some(to_hashmap(serde_json::json!({"custom_claim": "value"}))),
         };
         let id = claims.id();
         let (_, token) = claims.encode_and_sign(&kp.private)?;
         let decoded = kp.public.validate_token(&token).await?;
 
         assert_eq!(decoded.identity, id);
+        let custom_claim_value = decoded.extra.as_ref().unwrap().get("custom_claim").unwrap();
+        assert_eq!(custom_claim_value.as_str().unwrap(), "value");
         Ok(())
     }
 
@@ -310,6 +345,7 @@ mod tests {
             issuer: "localhost".to_string(),
             subject: "test-subject".to_string(),
             audience: vec![dummy_audience.clone()],
+            extra: None,
         };
         let (_, token) = claims.encode_and_sign(&kp.private)?;
         let st_creds = SpacetimeCreds::from_signed_token(token);

crates/core/src/auth/token_validation.rs

@@ -349,6 +349,7 @@ mod tests {
             audience: vec![],
             iat: std::time::SystemTime::now(),
             exp: None,
+            extra: None,
         };
         let token = kp.private.sign(&orig_claims)?;
 
@@ -391,6 +392,7 @@ mod tests {
             audience: vec![],
             iat: std::time::SystemTime::now(),
             exp: None,
+            extra: None,
         };
         let token = kp.private.sign(&orig_claims)?;
 
@@ -444,6 +446,7 @@ mod tests {
             audience: vec![],
             iat: std::time::SystemTime::now(),
             exp: None,
+            extra: None,
         };
         let token = kp.private.sign(&orig_claims)?;
 
@@ -609,6 +612,7 @@ mod tests {
             audience: vec![],
             iat: std::time::SystemTime::now(),
             exp: None,
+            extra: None,
         };
         for kp in [kp1, kp2] {
             log::debug!("Testing with key {:?}", kp.kid);

crates/core/src/client/client_connection.rs

@@ -335,6 +335,7 @@ impl ClientConnectionSender {
             audience: vec![],
             iat: SystemTime::now(),
             exp: None,
+            extra: None,
         };
         let sender = Self {
             id,