Store client credentials in a new system table (#2983)

# Description of Changes

This adds a new system table to store the jwt payloads of connected
clients. I'm planning to use this system table to expose client claims
to modules in subsequent PRs.

The new table is called `st_connection_credentials`. It is a **private**
system table which stores a mapping from `connection_id` to
`jwt_payload`. Note that a jwt payload is a json representation of the
clients claims, not a fully signed token.

The times when we need to insert and delete these rows closely mirrors
that of the existing `st_client` table, with 1.5 exceptions:
1. We weren't previously inserting to `st_client` until after the
`OnConnect` reducer ran (even though it was in the same transaction). We
want `st_connection_credentials` to be populated before calling the
reducer, so that the reducer can use it get the credentials, so I made a
change to insert to `st_client` and `st_connection_credentials` before
calling the reducer.
2. This difference has not actualized, but when clients start sending
refresh tokens, we will probably need to update the credentials stored
in this table.

This also enforces uniqueness of connection ids. A duplicate connection
id will now make the on-connect reducer fail (since it will violate
uniqueness when trying to insert to `st_connection_credentials`).

# Expected complexity level and risk

2.5

Adding a system table is a bit risky. This is almost rollback safe, with
one annoying case that is worth calling out:

If a database is created with this system table, opening it with an
older version of spacetimedb will only work if there is a snapshot of
the database. If we try to load a table without a snapshot, replaying
will fail on the first row for that table. This is because we don't
write the table schema information to the commit log when creating a
database. In practice, this is unlikely to be an issue, because new
databases asynchronously trigger a snapshot immediately after creation.

Migrating existing databases will be fine. On startup this will detect
that there is a missing system table, and add it in a way that writes it
to the commit log. Since it is in the commit log, we can open the
database with an older version and still understand the data for that
table.

# Testing

There are unit tests that cover opening a database created with an older
version (which doesn't have this table).

I manually tested opening a migrated database with an older version of
spacetimedb.

Author: jsdt

.gitignore

@@ -216,3 +216,6 @@ new.json
 # Keys
 *.pem
 .ok.sql
+
+# Test data
+!crates/core/testdata/

Cargo.lock

@@ -11,6 +11,7 @@ spacetimedb-lib = { workspace = true, features = ["serde"] }
 
 anyhow.workspace = true
 serde.workspace = true
+serde_json.workspace = true
 serde_with.workspace = true
 jsonwebtoken.workspace = true
 

crates/auth/Cargo.toml

@@ -6,9 +6,27 @@ use serde::{Deserialize, Serialize};
 use spacetimedb_lib::Identity;
 use std::time::SystemTime;
 
+#[derive(Debug, Clone)]
+pub struct ConnectionAuthCtx {
+    pub claims: SpacetimeIdentityClaims,
+    pub jwt_payload: String,
+}
+
+impl TryFrom<SpacetimeIdentityClaims> for ConnectionAuthCtx {
+    type Error = anyhow::Error;
+    fn try_from(claims: SpacetimeIdentityClaims) -> Result<Self, Self::Error> {
+        let payload =
+            serde_json::to_string(&claims).map_err(|e| anyhow::anyhow!("Failed to serialize claims: {}", e))?;
+        Ok(ConnectionAuthCtx {
+            claims,
+            jwt_payload: payload,
+        })
+    }
+}
+
 // These are the claims that can be attached to a request/connection.
 #[serde_with::serde_as]
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct SpacetimeIdentityClaims {
     #[serde(rename = "hex_identity")]
     pub identity: Identity,

crates/auth/src/identity.rs

@@ -15,6 +15,7 @@ spacetimedb-lib = { workspace = true, features = ["serde"] }
 spacetimedb-paths.workspace = true
 spacetimedb-schema.workspace = true
 
+base64.workspace = true
 tokio = { version = "1.2", features = ["full"] }
 lazy_static = "1.4.0"
 log = "0.4.4"

crates/client-api/Cargo.toml

@@ -1,23 +1,24 @@
-use std::time::{Duration, SystemTime};
-
+use anyhow::anyhow;
 use axum::extract::{Query, Request, State};
 use axum::middleware::Next;
 use axum::response::IntoResponse;
 use axum_extra::typed_header::TypedHeader;
 use headers::{authorization, HeaderMapExt};
 use http::{request, HeaderValue, StatusCode};
 use serde::{Deserialize, Serialize};
-use spacetimedb::auth::identity::SpacetimeIdentityClaims;
+use spacetimedb::auth::identity::{ConnectionAuthCtx, SpacetimeIdentityClaims};
 use spacetimedb::auth::identity::{JwtError, JwtErrorKind};
 use spacetimedb::auth::token_validation::{
     new_validator, DefaultValidator, TokenSigner, TokenValidationError, TokenValidator,
 };
 use spacetimedb::auth::JwtKeys;
 use spacetimedb::energy::EnergyQuanta;
 use spacetimedb::identity::Identity;
+use std::time::{Duration, SystemTime};
 use uuid::Uuid;
 
 use crate::{log_and_500, ControlStateDelegate, NodeDelegate};
+use base64::{engine::general_purpose, Engine};
 
 /// Credentials for login for a spacetime identity, represented as a JWT.
 ///
@@ -41,6 +42,19 @@ impl SpacetimeCreds {
         Self { token }
     }
 
+    fn extract_jwt_payload_string(&self) -> Option<String> {
+        let parts: Vec<&str> = self.token.split('.').collect();
+        if parts.len() != 3 {
+            return None;
+        }
+
+        let payload_encoded = parts[1];
+        let decoded_bytes = general_purpose::URL_SAFE_NO_PAD.decode(payload_encoded).ok()?;
+        let json_str = String::from_utf8(decoded_bytes).ok()?;
+
+        Some(json_str)
+    }
+
     pub fn to_header_value(&self) -> HeaderValue {
         let mut val = HeaderValue::try_from(["Bearer ", self.token()].concat()).unwrap();
         val.set_sensitive(true);
@@ -70,9 +84,31 @@ impl SpacetimeCreds {
 #[derive(Clone)]
 pub struct SpacetimeAuth {
     pub creds: SpacetimeCreds,
-    pub identity: Identity,
-    pub subject: String,
-    pub issuer: String,
+    pub claims: SpacetimeIdentityClaims,
+    /// The JWT payload as a json string (after base64 decoding).
+    pub jwt_payload: String,
+}
+
+impl SpacetimeAuth {
+    pub fn new(creds: SpacetimeCreds, claims: SpacetimeIdentityClaims) -> Result<Self, anyhow::Error> {
+        let payload = creds
+            .extract_jwt_payload_string()
+            .ok_or_else(|| anyhow!("Failed to extract JWT payload"))?;
+        Ok(Self {
+            creds,
+            claims,
+            jwt_payload: payload,
+        })
+    }
+}
+
+impl From<SpacetimeAuth> for ConnectionAuthCtx {
+    fn from(auth: SpacetimeAuth) -> Self {
+        ConnectionAuthCtx {
+            claims: auth.claims,
+            jwt_payload: auth.jwt_payload.clone(),
+        }
+    }
 }
 
 use jsonwebtoken;
@@ -84,10 +120,10 @@ pub struct TokenClaims {
 }
 
 impl From<SpacetimeAuth> for TokenClaims {
-    fn from(claims: SpacetimeAuth) -> Self {
+    fn from(auth: SpacetimeAuth) -> Self {
         Self {
-            issuer: claims.issuer,
-            subject: claims.subject,
+            issuer: auth.claims.issuer,
+            subject: auth.claims.subject,
             // This will need to be changed when we care about audiencies.
             audience: Vec::new(),
         }
@@ -108,11 +144,14 @@ impl TokenClaims {
         Identity::from_claims(&self.issuer, &self.subject)
     }
 
+    /// Encode the claims into a JWT token and sign it with the provided signer.
+    /// This also adds claims for expiry and issued at time.
+    /// Returns an object representing the claims and the signed token.
     pub fn encode_and_sign_with_expiry(
         &self,
         signer: &impl TokenSigner,
         expiry: Option<Duration>,
-    ) -> Result<String, JwtError> {
+    ) -> Result<(SpacetimeIdentityClaims, String), JwtError> {
         let iat = SystemTime::now();
         let exp = expiry.map(|dur| iat + dur);
         let claims = SpacetimeIdentityClaims {
@@ -123,10 +162,14 @@ impl TokenClaims {
             iat,
             exp,
         };
-        signer.sign(&claims)
+        let token = signer.sign(&claims)?;
+        Ok((claims, token))
     }
 
-    pub fn encode_and_sign(&self, signer: &impl TokenSigner) -> Result<String, JwtError> {
+    /// Encode the claims into a JWT token and sign it with the provided signer.
+    /// This also adds a claim for issued at time.
+    /// Returns an object representing the claims and the signed token.
+    pub fn encode_and_sign(&self, signer: &impl TokenSigner) -> Result<(SpacetimeIdentityClaims, String), JwtError> {
         self.encode_and_sign_with_expiry(signer, None)
     }
 }
@@ -143,32 +186,28 @@ impl SpacetimeAuth {
             audience: vec!["spacetimedb".to_string()],
         };
 
-        let identity = claims.id();
-        let creds = {
-            let token = claims.encode_and_sign(ctx.jwt_auth_provider()).map_err(log_and_500)?;
-            SpacetimeCreds::from_signed_token(token)
-        };
-
-        Ok(Self {
-            creds,
-            identity,
-            subject,
-            issuer: ctx.jwt_auth_provider().local_issuer().to_string(),
-        })
+        let (claims, token) = claims.encode_and_sign(ctx.jwt_auth_provider()).map_err(log_and_500)?;
+        let creds = SpacetimeCreds::from_signed_token(token);
+        // Pulling out the payload should never fail, since we just made it.
+        Self::new(creds, claims).map_err(log_and_500)
     }
 
     /// Get the auth credentials as headers to be returned from an endpoint.
     pub fn into_headers(self) -> (TypedHeader<SpacetimeIdentity>, TypedHeader<SpacetimeIdentityToken>) {
         (
-            TypedHeader(SpacetimeIdentity(self.identity)),
+            TypedHeader(SpacetimeIdentity(self.claims.identity)),
             TypedHeader(SpacetimeIdentityToken(self.creds)),
         )
     }
 
     // Sign a new token with the same claims and a new expiry.
     // Note that this will not change the issuer, so the private_key might not match.
     // We do this to create short-lived tokens that we will be able to verify.
-    pub fn re_sign_with_expiry(&self, signer: &impl TokenSigner, expiry: Duration) -> Result<String, JwtError> {
+    pub fn re_sign_with_expiry(
+        &self,
+        signer: &impl TokenSigner,
+        expiry: Duration,
+    ) -> Result<(SpacetimeIdentityClaims, String), JwtError> {
         TokenClaims::from(self.clone()).encode_and_sign_with_expiry(signer, Some(expiry))
     }
 }
@@ -237,9 +276,11 @@ impl<TV: TokenValidator + Send + Sync> JwtAuthProvider for JwtKeyAuthProvider<TV
 
 #[cfg(test)]
 mod tests {
-    use crate::auth::TokenClaims;
+    use crate::auth::{SpacetimeCreds, TokenClaims};
     use anyhow::Ok;
+
     use spacetimedb::auth::{token_validation::TokenValidator, JwtKeys};
+    use std::collections::HashSet;
 
     // Make sure that when we encode TokenClaims, we can decode to get the expected identity.
     #[tokio::test]
@@ -252,12 +293,48 @@ mod tests {
             audience: vec!["spacetimedb".to_string()],
         };
         let id = claims.id();
-        let token = claims.encode_and_sign(&kp.private)?;
+        let (_, token) = claims.encode_and_sign(&kp.private)?;
         let decoded = kp.public.validate_token(&token).await?;
 
         assert_eq!(decoded.identity, id);
         Ok(())
     }
+
+    // Test that extracting a JWT payload from a valid token gets the json representation.
+    #[tokio::test]
+    async fn extract_payload() -> Result<(), anyhow::Error> {
+        let kp = JwtKeys::generate()?;
+
+        let dummy_audience = "spacetimedb".to_string();
+        let claims = TokenClaims {
+            issuer: "localhost".to_string(),
+            subject: "test-subject".to_string(),
+            audience: vec![dummy_audience.clone()],
+        };
+        let (_, token) = claims.encode_and_sign(&kp.private)?;
+        let st_creds = SpacetimeCreds::from_signed_token(token);
+        let payload = st_creds
+            .extract_jwt_payload_string()
+            .ok_or_else(|| anyhow::anyhow!("Failed to extract JWT payload"))?;
+        // Make sure it is valid json.
+        let parsed: serde_json::Value = serde_json::from_str(&payload)?;
+        assert_eq!(parsed.get("iss").unwrap().as_str().unwrap(), claims.issuer);
+        assert_eq!(parsed.get("sub").unwrap().as_str().unwrap(), claims.subject);
+        assert_eq!(
+            parsed.get("aud").unwrap().as_array().unwrap()[0].as_str().unwrap(),
+            dummy_audience
+        );
+        let as_object = parsed
+            .as_object()
+            .ok_or_else(|| anyhow::anyhow!("Failed to parse JWT payload as object"))?;
+        let keys: HashSet<String> = as_object.keys().map(|s| s.to_string()).collect();
+        let expected_keys = vec!["iss", "sub", "aud", "iat", "exp", "hex_identity"]
+            .into_iter()
+            .map(|s| s.to_string())
+            .collect::<HashSet<String>>();
+        assert_eq!(keys, expected_keys);
+        Ok(())
+    }
 }
 
 pub async fn validate_token<S: NodeDelegate>(
@@ -283,11 +360,13 @@ impl<S: NodeDelegate + Send + Sync> axum::extract::FromRequestParts<S> for Space
             .await
             .map_err(AuthorizationRejection::Custom)?;
 
+        let payload = creds.extract_jwt_payload_string().ok_or_else(|| {
+            AuthorizationRejection::Custom(TokenValidationError::Other(anyhow!("Internal error parsing token")))
+        })?;
         let auth = SpacetimeAuth {
             creds,
-            identity: claims.identity,
-            subject: claims.subject,
-            issuer: claims.issuer,
+            claims,
+            jwt_payload: payload,
         };
         Ok(Self { auth: Some(auth) })
     }