Clean up auth related code with traits, and reuse existing key genera…

…tion code. (#1988)

crates/cli/src/util.rs

@@ -5,7 +5,7 @@ use base64::{
 };
 use reqwest::RequestBuilder;
 use serde::Deserialize;
-use spacetimedb::auth::identity::{IncomingClaims, SpacetimeIdentityClaims2};
+use spacetimedb::auth::identity::{IncomingClaims, SpacetimeIdentityClaims};
 use spacetimedb_client_api_messages::name::{DnsLookupResponse, RegisterTldResult, ReverseDNSResponse};
 use spacetimedb_data_structures::map::HashMap;
 use spacetimedb_lib::{AlgebraicType, Identity};
@@ -277,7 +277,7 @@ pub fn decode_identity(config: &Config) -> anyhow::Result<String> {
     let decoded_string = String::from_utf8(decoded_bytes)?;
 
     let claims_data: IncomingClaims = serde_json::from_str(decoded_string.as_str())?;
-    let claims_data: SpacetimeIdentityClaims2 = claims_data.try_into()?;
+    let claims_data: SpacetimeIdentityClaims = claims_data.try_into()?;
 
     Ok(claims_data.identity.to_string())
 }

crates/client-api/Cargo.toml

@@ -41,10 +41,8 @@ tokio-tungstenite.workspace = true
 itoa.workspace = true
 derive_more = "0.99.17"
 uuid.workspace = true
-blake3.workspace = true
 jsonwebtoken.workspace = true
 scopeguard.workspace = true
 
 [dev-dependencies]
-jsonwebkey = { version = "0.3.5", features = ["generate","jwt-convert"] }
 jsonwebtoken.workspace = true

crates/client-api/src/auth.rs

@@ -6,12 +6,13 @@ use axum::response::IntoResponse;
 use axum_extra::typed_header::TypedHeader;
 use headers::{authorization, HeaderMapExt};
 use http::{request, HeaderValue, StatusCode};
-use serde::Deserialize;
-use spacetimedb::auth::identity::SpacetimeIdentityClaims2;
-use spacetimedb::auth::identity::{
-    decode_token, encode_token, DecodingKey, EncodingKey, JwtError, JwtErrorKind, SpacetimeIdentityClaims,
+use serde::{Deserialize, Serialize};
+use spacetimedb::auth::identity::SpacetimeIdentityClaims;
+use spacetimedb::auth::identity::{JwtError, JwtErrorKind};
+use spacetimedb::auth::token_validation::{
+    new_validator, DefaultValidator, TokenSigner, TokenValidationError, TokenValidator,
 };
-use spacetimedb::auth::token_validation::{validate_token, TokenValidationError};
+use spacetimedb::auth::JwtKeys;
 use spacetimedb::energy::EnergyQuanta;
 use spacetimedb::identity::Identity;
 use uuid::Uuid;
@@ -57,18 +58,10 @@ impl SpacetimeCreds {
     pub fn token(&self) -> &str {
         &self.token
     }
-    /// Decode this token into auth claims.
-    pub fn decode_token(&self, public_key: &DecodingKey) -> Result<SpacetimeIdentityClaims, JwtError> {
-        decode_token(public_key, self.token()).map(|x| x.claims)
-    }
+
     pub fn from_signed_token(token: String) -> Self {
         Self { token }
     }
-    /// Mint a new credentials JWT for an identity.
-    pub fn encode_token(private_key: &EncodingKey, identity: Identity) -> Result<Self, JwtError> {
-        let token = encode_token(private_key, identity)?;
-        Ok(Self { token })
-    }
 
     /// Extract credentials from the headers or else query string of a request.
     fn from_request_parts(parts: &request::Parts) -> Result<Option<Self>, headers::Error> {
@@ -137,25 +130,24 @@ impl TokenClaims {
 
     pub fn encode_and_sign_with_expiry(
         &self,
-        private_key: &EncodingKey,
+        signer: &impl TokenSigner,
         expiry: Option<Duration>,
     ) -> Result<String, JwtError> {
         let iat = SystemTime::now();
         let exp = expiry.map(|dur| iat + dur);
-        let claims = SpacetimeIdentityClaims2 {
+        let claims = SpacetimeIdentityClaims {
             identity: self.id(),
             subject: self.subject.clone(),
             issuer: self.issuer.clone(),
             audience: self.audience.clone(),
             iat,
             exp,
         };
-        let header = jsonwebtoken::Header::new(jsonwebtoken::Algorithm::ES256);
-        jsonwebtoken::encode(&header, &claims, private_key)
+        signer.sign(&claims)
     }
 
-    pub fn encode_and_sign(&self, private_key: &EncodingKey) -> Result<String, JwtError> {
-        self.encode_and_sign_with_expiry(private_key, None)
+    pub fn encode_and_sign(&self, signer: &impl TokenSigner) -> Result<String, JwtError> {
+        self.encode_and_sign_with_expiry(signer, None)
     }
 }
 
@@ -165,23 +157,23 @@ impl SpacetimeAuth {
         // Generate claims with a random subject.
         let subject = Uuid::new_v4().to_string();
         let claims = TokenClaims {
-            issuer: ctx.local_issuer(),
+            issuer: ctx.jwt_auth_provider().local_issuer().to_owned(),
             subject: subject.clone(),
             // Placeholder audience.
             audience: vec!["spacetimedb".to_string()],
         };
 
         let identity = claims.id();
         let creds = {
-            let token = claims.encode_and_sign(ctx.private_key()).map_err(log_and_500)?;
+            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.local_issuer(),
+            issuer: ctx.jwt_auth_provider().local_issuer().to_string(),
         })
     }
 
@@ -196,52 +188,94 @@ impl SpacetimeAuth {
     // 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, private_key: &EncodingKey, expiry: Duration) -> Result<String, JwtError> {
-        TokenClaims::from(self.clone()).encode_and_sign_with_expiry(private_key, Some(expiry))
+    pub fn re_sign_with_expiry(&self, signer: &impl TokenSigner, expiry: Duration) -> Result<String, JwtError> {
+        TokenClaims::from(self.clone()).encode_and_sign_with_expiry(signer, Some(expiry))
     }
 }
 
-#[cfg(test)]
-mod tests {
-    use crate::auth::TokenClaims;
-    use anyhow::Ok;
-    use jsonwebkey as jwk;
-    use jsonwebtoken::{DecodingKey, EncodingKey};
-    use spacetimedb::auth::identity;
+// JwtAuthProvider is used for signing and verifying JWT tokens.
+pub trait JwtAuthProvider: Sync + Send + TokenSigner {
+    type TV: TokenValidator + Send + Sync;
+    /// Used to validate incoming JWTs.
+    fn validator(&self) -> &Self::TV;
+
+    /// The issuer to use when signing JWTs.
+    fn local_issuer(&self) -> &str;
+
+    /// Return the public key used to verify JWTs, as the bytes of a PEM public key file.
+    ///
+    /// The `/identity/public-key` route calls this method to return the public key to callers.
+    fn public_key_bytes(&self) -> &[u8];
+}
+
+pub struct JwtKeyAuthProvider<TV: TokenValidator + Send + Sync> {
+    keys: JwtKeys,
+    local_issuer: String,
+    validator: TV,
+}
+
+pub type DefaultJwtAuthProvider = JwtKeyAuthProvider<DefaultValidator>;
+
+// Create a new AuthEnvironment using the default caching validator.
+pub fn default_auth_environment(keys: JwtKeys, local_issuer: String) -> JwtKeyAuthProvider<DefaultValidator> {
+    let validator = new_validator(keys.public.clone(), local_issuer.clone());
+    JwtKeyAuthProvider::new(keys, local_issuer, validator)
+}
+
+impl<TV: TokenValidator + Send + Sync> JwtKeyAuthProvider<TV> {
+    fn new(keys: JwtKeys, local_issuer: String, validator: TV) -> Self {
+        Self {
+            keys,
+            local_issuer,
+            validator,
+        }
+    }
+}
 
-    // TODO: this keypair stuff is duplicated. We should create a test-only crate with helpers.
-    struct KeyPair {
-        pub public_key: DecodingKey,
-        pub private_key: EncodingKey,
+impl<TV: TokenValidator + Send + Sync> TokenSigner for JwtKeyAuthProvider<TV> {
+    fn sign<T: Serialize>(&self, claims: &T) -> Result<String, JwtError> {
+        let header = jsonwebtoken::Header::new(jsonwebtoken::Algorithm::ES256);
+        jsonwebtoken::encode(&header, &claims, &self.keys.private)
     }
+}
 
-    fn new_keypair() -> anyhow::Result<KeyPair> {
-        let mut my_jwk = jwk::JsonWebKey::new(jwk::Key::generate_p256());
+impl<TV: TokenValidator + Send + Sync> JwtAuthProvider for JwtKeyAuthProvider<TV> {
+    type TV = TV;
 
-        my_jwk.set_algorithm(jwk::Algorithm::ES256).unwrap();
-        let public_key = jsonwebtoken::DecodingKey::from_ec_pem(my_jwk.key.to_public().unwrap().to_pem().as_bytes())?;
-        let private_key = jsonwebtoken::EncodingKey::from_ec_pem(my_jwk.key.try_to_pem()?.as_bytes())?;
-        Ok(KeyPair {
-            public_key,
-            private_key,
-        })
+    fn local_issuer(&self) -> &str {
+        &self.local_issuer
+    }
+
+    fn public_key_bytes(&self) -> &[u8] {
+        &self.keys.public_pem
     }
 
+    fn validator(&self) -> &Self::TV {
+        &self.validator
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::auth::TokenClaims;
+    use anyhow::Ok;
+    use spacetimedb::auth::{token_validation::TokenValidator, JwtKeys};
+
     // Make sure that when we encode TokenClaims, we can decode to get the expected identity.
-    #[test]
-    fn decode_encoded_token() -> Result<(), anyhow::Error> {
-        let kp = new_keypair()?;
+    #[tokio::test]
+    async fn decode_encoded_token() -> Result<(), anyhow::Error> {
+        let kp = JwtKeys::generate()?;
 
         let claims = TokenClaims {
             issuer: "localhost".to_string(),
             subject: "test-subject".to_string(),
             audience: vec!["spacetimedb".to_string()],
         };
         let id = claims.id();
-        let token = claims.encode_and_sign(&kp.private_key)?;
+        let token = claims.encode_and_sign(&kp.private)?;
+        let decoded = kp.public.validate_token(&token).await?;
 
-        let decoded = identity::decode_token(&kp.public_key, &token)?;
-        assert_eq!(decoded.claims.identity, id);
+        assert_eq!(decoded.identity, id);
         Ok(())
     }
 }
@@ -258,7 +292,10 @@ impl<S: NodeDelegate + Send + Sync> axum::extract::FromRequestParts<S> for Space
             return Ok(Self { auth: None });
         };
 
-        let claims = validate_token(state.public_key().clone(), &state.local_issuer(), &creds.token)
+        let claims = state
+            .jwt_auth_provider()
+            .validator()
+            .validate_token(&creds.token)
             .await
             .map_err(AuthorizationRejection::Custom)?;
 

crates/client-api/src/lib.rs

@@ -4,7 +4,6 @@ use async_trait::async_trait;
 use axum::response::ErrorResponse;
 use http::StatusCode;
 
-use spacetimedb::auth::identity::{DecodingKey, EncodingKey};
 use spacetimedb::client::ClientActorIndex;
 use spacetimedb::energy::{EnergyBalance, EnergyQuanta};
 use spacetimedb::host::{HostController, UpdateDatabaseResult};
@@ -26,19 +25,8 @@ pub trait NodeDelegate: Send + Sync {
     fn host_controller(&self) -> &HostController;
     fn client_actor_index(&self) -> &ClientActorIndex;
 
-    /// Return a JWT decoding key for verifying credentials.
-    fn public_key(&self) -> &DecodingKey;
-
-    // The issuer to use when signing JWTs.
-    fn local_issuer(&self) -> String;
-
-    /// Return the public key used to verify JWTs, as the bytes of a PEM public key file.
-    ///
-    /// The `/identity/public-key` route calls this method to return the public key to callers.
-    fn public_key_bytes(&self) -> &[u8];
-
-    /// Return a JWT encoding key for signing credentials.
-    fn private_key(&self) -> &EncodingKey;
+    type JwtAuthProviderT: auth::JwtAuthProvider;
+    fn jwt_auth_provider(&self) -> &Self::JwtAuthProviderT;
 }
 
 /// Parameters for publishing a database.
@@ -222,14 +210,11 @@ impl<T: ControlStateWriteAccess + ?Sized> ControlStateWriteAccess for Arc<T> {
 }
 
 impl<T: NodeDelegate + ?Sized> NodeDelegate for Arc<T> {
+    type JwtAuthProviderT = T::JwtAuthProviderT;
     fn gather_metrics(&self) -> Vec<prometheus::proto::MetricFamily> {
         (**self).gather_metrics()
     }
 
-    fn local_issuer(&self) -> String {
-        (**self).local_issuer()
-    }
-
     fn host_controller(&self) -> &HostController {
         (**self).host_controller()
     }
@@ -238,16 +223,8 @@ impl<T: NodeDelegate + ?Sized> NodeDelegate for Arc<T> {
         (**self).client_actor_index()
     }
 
-    fn public_key(&self) -> &DecodingKey {
-        (**self).public_key()
-    }
-
-    fn public_key_bytes(&self) -> &[u8] {
-        (**self).public_key_bytes()
-    }
-
-    fn private_key(&self) -> &EncodingKey {
-        (**self).private_key()
+    fn jwt_auth_provider(&self) -> &Self::JwtAuthProviderT {
+        (**self).jwt_auth_provider()
     }
 }
 

crates/client-api/src/routes/identity.rs

@@ -9,7 +9,7 @@ use serde::{Deserialize, Serialize};
 use spacetimedb_lib::de::serde::DeserializeWrapper;
 use spacetimedb_lib::Identity;
 
-use crate::auth::{SpacetimeAuth, SpacetimeAuthRequired};
+use crate::auth::{JwtAuthProvider, SpacetimeAuth, SpacetimeAuthRequired};
 use crate::{log_and_500, ControlStateDelegate, NodeDelegate};
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -104,7 +104,7 @@ pub async fn create_websocket_token<S: NodeDelegate>(
 ) -> axum::response::Result<impl IntoResponse> {
     let expiry = Duration::from_secs(60);
     let token = auth
-        .re_sign_with_expiry(ctx.private_key(), expiry)
+        .re_sign_with_expiry(ctx.jwt_auth_provider(), expiry)
         .map_err(log_and_500)?;
     // let token = encode_token_with_expiry(ctx.private_key(), auth.identity, Some(expiry)).map_err(log_and_500)?;
     Ok(axum::Json(WebsocketTokenResponse { token }))
@@ -131,7 +131,7 @@ pub async fn validate_token(
 pub async fn get_public_key<S: NodeDelegate>(State(ctx): State<S>) -> axum::response::Result<impl IntoResponse> {
     Ok((
         [(CONTENT_TYPE, "application/pem-certificate-chain")],
-        ctx.public_key_bytes().to_owned(),
+        ctx.jwt_auth_provider().public_key_bytes().to_owned(),
     ))
 }
 

crates/core/Cargo.toml

@@ -45,7 +45,6 @@ clap.workspace = true
 crossbeam-channel.workspace = true
 derive_more.workspace = true
 dirs.workspace = true
-email_address.workspace = true
 enum-as-inner.workspace = true
 enum-map.workspace = true
 flate2.workspace = true
@@ -127,7 +126,5 @@ proptest-derive.workspace = true
 rand.workspace = true
 env_logger.workspace = true
 pretty_assertions.workspace = true
-jsonwebkey = { version = "0.3.5", features = ["generate", "jwt-convert"] }
 jsonwebtoken.workspace = true
-axum-test = "16.2.0"
 axum.workspace = true