Make the JWK generic over the parameters

crates/axum-utils/src/client_authorization.rs

@@ -31,8 +31,8 @@ use mas_data_model::{Client, JwksOrJwksUri, StorageBackend};
 use mas_http::HttpServiceExt;
 use mas_iana::oauth::OAuthClientAuthenticationMethod;
 use mas_jose::{
-    DecodedJsonWebToken, DynamicJwksStore, Either, JsonWebKeySet, JsonWebSignatureHeader,
-    JsonWebTokenParts, SharedSecret, StaticJwksStore, VerifyingKeystore,
+    jwk::PublicJsonWebKeySet, DecodedJsonWebToken, DynamicJwksStore, Either,
+    JsonWebSignatureHeader, JsonWebTokenParts, SharedSecret, StaticJwksStore, VerifyingKeystore,
 };
 use mas_storage::{
     oauth2::client::{lookup_client_by_client_id, ClientFetchError},
@@ -185,7 +185,7 @@ fn jwks_key_store(jwks: &JwksOrJwksUri) -> Either<StaticJwksStore, DynamicJwksSt
             // TODO: get the client from somewhere else?
             let exporter = mas_http::client("fetch-jwks")
                 .response_body_to_bytes()
-                .json_response::<JsonWebKeySet>()
+                .json_response::<PublicJsonWebKeySet>()
                 .map_request(move |_: ()| {
                     http::Request::builder()
                         .method("GET")

crates/config/src/sections/clients.rs

@@ -16,7 +16,7 @@ use std::ops::{Deref, DerefMut};
 
 use async_trait::async_trait;
 use mas_iana::oauth::OAuthClientAuthenticationMethod;
-use mas_jose::JsonWebKeySet;
+use mas_jose::jwk::PublicJsonWebKeySet;
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 use serde_with::skip_serializing_none;
@@ -28,12 +28,12 @@ use super::ConfigurationSection;
 #[derive(JsonSchema, Serialize, Deserialize, Clone, Debug)]
 #[serde(rename_all = "snake_case")]
 pub enum JwksOrJwksUri {
-    Jwks(JsonWebKeySet),
+    Jwks(PublicJsonWebKeySet),
     JwksUri(Url),
 }
 
-impl From<JsonWebKeySet> for JwksOrJwksUri {
-    fn from(jwks: JsonWebKeySet) -> Self {
+impl From<PublicJsonWebKeySet> for JwksOrJwksUri {
+    fn from(jwks: PublicJsonWebKeySet) -> Self {
         Self::Jwks(jwks)
     }
 }
@@ -125,7 +125,7 @@ impl ClientConfig {
 
     #[doc(hidden)]
     #[must_use]
-    pub fn jwks(&self) -> Option<&JsonWebKeySet> {
+    pub fn jwks(&self) -> Option<&PublicJsonWebKeySet> {
         match &self.client_auth_method {
             ClientAuthMethodConfig::PrivateKeyJwt(JwksOrJwksUri::Jwks(jwks)) => Some(jwks),
             _ => None,

crates/data-model/src/oauth2/client.rs

@@ -16,7 +16,7 @@ use mas_iana::{
     jose::JsonWebSignatureAlg,
     oauth::{OAuthAuthorizationEndpointResponseType, OAuthClientAuthenticationMethod},
 };
-use mas_jose::JsonWebKeySet;
+use mas_jose::jwk::PublicJsonWebKeySet;
 use oauth2_types::requests::GrantType;
 use serde::Serialize;
 use thiserror::Error;
@@ -28,7 +28,7 @@ use crate::traits::{StorageBackend, StorageBackendMarker};
 #[serde(rename_all = "snake_case")]
 pub enum JwksOrJwksUri {
     /// Client's JSON Web Key Set document, passed by value.
-    Jwks(JsonWebKeySet),
+    Jwks(PublicJsonWebKeySet),
 
     /// URL for the Client's JSON Web Key Set document.
     JwksUri(Url),

crates/jose/src/jwk/mod.rs

@@ -34,6 +34,10 @@ pub(crate) mod public_parameters;
 
 pub use self::public_parameters::JsonWebKeyPublicParameters as JsonWebKeyParameters;
 
+pub trait JwkKty {
+    fn kty(&self) -> JsonWebKeyType;
+}
+
 trait JwkEcCurve {
     const CRV: JsonWebKeyEcEllipticCurve;
 }
@@ -53,9 +57,9 @@ impl JwkEcCurve for k256::Secp256k1 {
 #[serde_as]
 #[skip_serializing_none]
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
-pub struct JsonWebKey {
+pub struct JsonWebKey<P> {
     #[serde(flatten)]
-    parameters: JsonWebKeyParameters,
+    parameters: P,
 
     #[serde(default)]
     r#use: Option<JsonWebKeyUse>,
@@ -89,9 +93,12 @@ pub struct JsonWebKey {
     x5t_s256: Option<Vec<u8>>,
 }
 
-impl JsonWebKey {
+pub type PublicJsonWebKey = JsonWebKey<self::public_parameters::JsonWebKeyPublicParameters>;
+pub type PrivateJsonWebKey = JsonWebKey<self::private_parameters::JsonWebKeyPrivateParameters>;
+
+impl<P> JsonWebKey<P> {
     #[must_use]
-    pub const fn new(parameters: JsonWebKeyParameters) -> Self {
+    pub const fn new(parameters: P) -> Self {
         Self {
             parameters,
             r#use: None,
@@ -135,30 +142,29 @@ impl JsonWebKey {
     }
 
     #[must_use]
-    pub const fn params(&self) -> &JsonWebKeyParameters {
+    pub const fn params(&self) -> &P {
         &self.parameters
     }
 }
 
-impl Constrainable for JsonWebKey {
+impl<P> Constrainable for JsonWebKey<P>
+where
+    P: JwkKty,
+{
     fn kid(&self) -> Option<&str> {
         self.kid.as_deref()
     }
 
     fn kty(&self) -> JsonWebKeyType {
-        match self.parameters {
-            JsonWebKeyParameters::Ec { .. } => JsonWebKeyType::Ec,
-            JsonWebKeyParameters::Rsa { .. } => JsonWebKeyType::Rsa,
-            JsonWebKeyParameters::Okp { .. } => JsonWebKeyType::Okp,
-        }
+        self.parameters.kty()
     }
 
     fn algs(&self) -> Option<Vec<JsonWebSignatureAlg>> {
         if let Some(alg) = self.alg {
             Some(vec![alg])
         } else {
-            match &self.parameters {
-                JsonWebKeyParameters::Rsa { .. } => Some(vec![
+            match self.parameters.kty() {
+                JsonWebKeyType::Rsa => Some(vec![
                     JsonWebSignatureAlg::Rs256,
                     JsonWebSignatureAlg::Rs384,
                     JsonWebSignatureAlg::Rs512,
@@ -166,13 +172,22 @@ impl Constrainable for JsonWebKey {
                     JsonWebSignatureAlg::Ps384,
                     JsonWebSignatureAlg::Ps512,
                 ]),
-                JsonWebKeyParameters::Ec(params) => match params.crv() {
-                    JsonWebKeyEcEllipticCurve::P256 => Some(vec![JsonWebSignatureAlg::Es256]),
-                    JsonWebKeyEcEllipticCurve::P384 => Some(vec![JsonWebSignatureAlg::Es384]),
-                    JsonWebKeyEcEllipticCurve::P521 => Some(vec![JsonWebSignatureAlg::Es512]),
-                    JsonWebKeyEcEllipticCurve::Secp256K1 => Some(vec![JsonWebSignatureAlg::Es256K]),
-                },
-                JsonWebKeyParameters::Okp { .. } => Some(vec![JsonWebSignatureAlg::EdDsa]),
+                JsonWebKeyType::Ec => {
+                    todo!()
+                    /*
+                        match params.crv() {
+                        JsonWebKeyEcEllipticCurve::P256 => Some(vec![JsonWebSignatureAlg::Es256]),
+                        JsonWebKeyEcEllipticCurve::P384 => Some(vec![JsonWebSignatureAlg::Es384]),
+                        JsonWebKeyEcEllipticCurve::P521 => Some(vec![JsonWebSignatureAlg::Es512]),
+                        JsonWebKeyEcEllipticCurve::Secp256K1 => Some(vec![JsonWebSignatureAlg::Es256K]),
+                    }, */
+                }
+                JsonWebKeyType::Okp => Some(vec![JsonWebSignatureAlg::EdDsa]),
+                JsonWebKeyType::Oct => Some(vec![
+                    JsonWebSignatureAlg::Hs256,
+                    JsonWebSignatureAlg::Hs384,
+                    JsonWebSignatureAlg::Hs512,
+                ]),
             }
         }
     }
@@ -183,21 +198,25 @@ impl Constrainable for JsonWebKey {
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
-pub struct JsonWebKeySet {
-    keys: Vec<JsonWebKey>,
+pub struct JsonWebKeySet<P> {
+    keys: Vec<JsonWebKey<P>>,
 }
 
-impl std::ops::Deref for JsonWebKeySet {
-    type Target = Vec<JsonWebKey>;
+pub type PublicJsonWebKeySet = JsonWebKeySet<self::public_parameters::JsonWebKeyPublicParameters>;
+pub type PrivateJsonWebKeySet =
+    JsonWebKeySet<self::private_parameters::JsonWebKeyPrivateParameters>;
+
+impl<P> std::ops::Deref for JsonWebKeySet<P> {
+    type Target = Vec<JsonWebKey<P>>;
 
     fn deref(&self) -> &Self::Target {
         &self.keys
     }
 }
 
-impl JsonWebKeySet {
+impl<P> JsonWebKeySet<P> {
     #[must_use]
-    pub fn new(keys: Vec<JsonWebKey>) -> Self {
+    pub fn new(keys: Vec<JsonWebKey<P>>) -> Self {
         Self { keys }
     }
 }

crates/jose/src/jwk/private_parameters.rs

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use mas_iana::jose::{JsonWebKeyEcEllipticCurve, JsonWebKeyOkpEllipticCurve};
+use mas_iana::jose::{JsonWebKeyEcEllipticCurve, JsonWebKeyOkpEllipticCurve, JsonWebKeyType};
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 use serde_with::{
@@ -21,6 +21,8 @@ use serde_with::{
     serde_as,
 };
 
+use super::JwkKty;
+
 #[serde_as]
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
 #[serde(tag = "kty")]
@@ -38,6 +40,17 @@ pub enum JsonWebKeyPrivateParameters {
     Okp(OkpPrivateParameters),
 }
 
+impl JwkKty for JsonWebKeyPrivateParameters {
+    fn kty(&self) -> JsonWebKeyType {
+        match self {
+            Self::Oct(_) => JsonWebKeyType::Oct,
+            Self::Rsa(_) => JsonWebKeyType::Rsa,
+            Self::Ec(_) => JsonWebKeyType::Ec,
+            Self::Okp(_) => JsonWebKeyType::Okp,
+        }
+    }
+}
+
 #[serde_as]
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
 pub struct OctPrivateParameters {

crates/jose/src/jwk/public_parameters.rs

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use mas_iana::jose::{JsonWebKeyEcEllipticCurve, JsonWebKeyOkpEllipticCurve};
+use mas_iana::jose::{JsonWebKeyEcEllipticCurve, JsonWebKeyOkpEllipticCurve, JsonWebKeyType};
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 use serde_with::{
@@ -21,6 +21,8 @@ use serde_with::{
     serde_as,
 };
 
+use super::JwkKty;
+
 #[serde_as]
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
 #[serde(tag = "kty")]
@@ -35,6 +37,16 @@ pub enum JsonWebKeyPublicParameters {
     Okp(OkpPublicParameters),
 }
 
+impl JwkKty for JsonWebKeyPublicParameters {
+    fn kty(&self) -> JsonWebKeyType {
+        match self {
+            Self::Rsa(_) => JsonWebKeyType::Rsa,
+            Self::Ec(_) => JsonWebKeyType::Ec,
+            Self::Okp(_) => JsonWebKeyType::Okp,
+        }
+    }
+}
+
 #[serde_as]
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
 pub struct RsaPublicParameters {

crates/jose/src/jwt/header.rs

@@ -21,7 +21,7 @@ use serde_with::{
 };
 use url::Url;
 
-use crate::jwk::JsonWebKey;
+use crate::jwk::PublicJsonWebKey;
 
 #[serde_as]
 #[skip_serializing_none]
@@ -33,7 +33,7 @@ pub struct JsonWebSignatureHeader {
     jku: Option<Url>,
 
     #[serde(default)]
-    jwk: Option<JsonWebKey>,
+    jwk: Option<PublicJsonWebKey>,
 
     #[serde(default)]
     kid: Option<String>,
@@ -98,12 +98,12 @@ impl JsonWebSignatureHeader {
     }
 
     #[must_use]
-    pub const fn jwk(&self) -> Option<&JsonWebKey> {
+    pub const fn jwk(&self) -> Option<&PublicJsonWebKey> {
         self.jwk.as_ref()
     }
 
     #[must_use]
-    pub fn with_jwk(mut self, jwk: JsonWebKey) -> Self {
+    pub fn with_jwk(mut self, jwk: PublicJsonWebKey) -> Self {
         self.jwk = Some(jwk);
         self
     }

crates/jose/src/keystore/jwks/dynamic_store.rs

@@ -24,7 +24,7 @@ use tower::{
 };
 
 use super::StaticJwksStore;
-use crate::{JsonWebKeySet, JsonWebSignatureHeader, VerifyingKeystore};
+use crate::{jwk::PublicJsonWebKeySet, JsonWebSignatureHeader, VerifyingKeystore};
 
 #[derive(Debug, Error)]
 pub enum Error {
@@ -60,7 +60,7 @@ impl<E> Default for State<E> {
 }
 
 impl<E> State<E> {
-    fn fullfill(&mut self, key_set: JsonWebKeySet) {
+    fn fullfill(&mut self, key_set: PublicJsonWebKeySet) {
         *self = Self::Fulfilled {
             at: Utc::now(),
             store: StaticJwksStore::new(key_set),
@@ -100,14 +100,14 @@ impl<E> State<E> {
 
 #[derive(Clone)]
 pub struct DynamicJwksStore {
-    exporter: BoxCloneService<(), JsonWebKeySet, BoxError>,
+    exporter: BoxCloneService<(), PublicJsonWebKeySet, BoxError>,
     cache: Arc<RwLock<State<Arc<BoxError>>>>,
 }
 
 impl DynamicJwksStore {
     pub fn new<T>(exporter: T) -> Self
     where
-        T: Service<(), Response = JsonWebKeySet, Error = BoxError> + Send + Clone + 'static,
+        T: Service<(), Response = PublicJsonWebKeySet, Error = BoxError> + Send + Clone + 'static,
         T::Future: Send,
     {
         Self {

crates/jose/src/keystore/jwks/static_store.rs

@@ -22,8 +22,9 @@ use signature::{Signature, Verifier};
 use thiserror::Error;
 
 use crate::{
-    constraints::Constrainable, JsonWebKey, JsonWebKeySet, JsonWebSignatureHeader,
-    VerifyingKeystore,
+    constraints::Constrainable,
+    jwk::{PublicJsonWebKey, PublicJsonWebKeySet},
+    JsonWebSignatureHeader, VerifyingKeystore,
 };
 
 #[derive(Debug, Error)]
@@ -60,7 +61,7 @@ struct KeyConstraint<'a> {
 }
 
 impl<'a> KeyConstraint<'a> {
-    fn matches(&self, key: &'a JsonWebKey) -> bool {
+    fn matches(&self, key: &'a PublicJsonWebKey) -> bool {
         // If a specific KID was asked, match the key only if it has a matching kid
         // field
         if let Some(kid) = self.kid {
@@ -84,22 +85,25 @@ impl<'a> KeyConstraint<'a> {
         true
     }
 
-    fn find_keys(&self, key_set: &'a JsonWebKeySet) -> Vec<&'a JsonWebKey> {
+    fn find_keys(&self, key_set: &'a PublicJsonWebKeySet) -> Vec<&'a PublicJsonWebKey> {
         key_set.iter().filter(|k| self.matches(k)).collect()
     }
 }
 
 pub struct StaticJwksStore {
-    key_set: JsonWebKeySet,
+    key_set: PublicJsonWebKeySet,
 }
 
 impl StaticJwksStore {
     #[must_use]
-    pub fn new(key_set: JsonWebKeySet) -> Self {
+    pub fn new(key_set: PublicJsonWebKeySet) -> Self {
         Self { key_set }
     }
 
-    fn find_key<'a>(&'a self, constraint: &KeyConstraint<'a>) -> Result<&'a JsonWebKey, Error> {
+    fn find_key<'a>(
+        &'a self,
+        constraint: &KeyConstraint<'a>,
+    ) -> Result<&'a PublicJsonWebKey, Error> {
         let keys = constraint.find_keys(&self.key_set);
 
         match &keys[..] {

crates/jose/src/keystore/static_keystore.rs

@@ -34,7 +34,7 @@ use signature::{Signature, Signer, Verifier};
 use tower::Service;
 
 use super::{SigningKeystore, VerifyingKeystore};
-use crate::{JsonWebKey, JsonWebKeySet, JsonWebSignatureHeader};
+use crate::{jwk::PublicJsonWebKeySet, JsonWebKey, JsonWebKeySet, JsonWebSignatureHeader};
 
 // Generate with
 //  openssl genrsa 2048
@@ -365,7 +365,7 @@ impl VerifyingKeystore for StaticKeystore {
 
 impl Service<()> for &StaticKeystore {
     type Future = Ready<Result<Self::Response, Self::Error>>;
-    type Response = JsonWebKeySet;
+    type Response = PublicJsonWebKeySet;
     type Error = Infallible;
 
     fn poll_ready(&mut self, _cx: &mut std::task::Context<'_>) -> Poll<Result<(), Self::Error>> {

crates/jose/src/lib.rs

@@ -20,7 +20,7 @@
 pub mod claims;
 pub mod constraints;
 pub mod hmac;
-pub(crate) mod jwk;
+pub mod jwk;
 pub(crate) mod jwt;
 mod keystore;
 pub(crate) mod rsa;

crates/oauth2-types/src/registration/client_metadata_serde.rs

@@ -20,7 +20,7 @@ use mas_iana::{
     jose::{JsonWebEncryptionAlg, JsonWebEncryptionEnc, JsonWebSignatureAlg},
     oauth::{OAuthAuthorizationEndpointResponseType, OAuthClientAuthenticationMethod},
 };
-use mas_jose::JsonWebKeySet;
+use mas_jose::jwk::PublicJsonWebKeySet;
 use serde::{
     de::{DeserializeOwned, Error},
     ser::SerializeMap,
@@ -99,7 +99,7 @@ pub struct ClientMetadataSerdeHelper {
     application_type: Option<ApplicationType>,
     contacts: Option<Vec<String>>,
     jwks_uri: Option<Url>,
-    jwks: Option<JsonWebKeySet>,
+    jwks: Option<PublicJsonWebKeySet>,
     sector_identifier_uri: Option<Url>,
     subject_type: Option<SubjectType>,
     token_endpoint_auth_method: Option<OAuthClientAuthenticationMethod>,

crates/oauth2-types/src/registration/mod.rs

@@ -20,7 +20,7 @@ use mas_iana::{
     jose::{JsonWebEncryptionAlg, JsonWebEncryptionEnc, JsonWebSignatureAlg},
     oauth::{OAuthAuthorizationEndpointResponseType, OAuthClientAuthenticationMethod},
 };
-use mas_jose::JsonWebKeySet;
+use mas_jose::jwk::PublicJsonWebKeySet;
 use serde::{Deserialize, Serialize};
 use serde_with::{serde_as, skip_serializing_none, TimestampSeconds};
 use thiserror::Error;
@@ -197,7 +197,7 @@ pub struct ClientMetadata {
     /// This field is mutually exclusive with `jwks_uri`.
     ///
     /// [JWK]: https://www.rfc-editor.org/rfc/rfc7517.html
-    pub jwks: Option<JsonWebKeySet>,
+    pub jwks: Option<PublicJsonWebKeySet>,
 
     /// URL to be used in calculating pseudonymous identifiers by the OpenID
     /// Connect provider when [pairwise subject identifiers] are used.
@@ -861,7 +861,7 @@ mod tests {
         jose::{JsonWebEncryptionAlg, JsonWebEncryptionEnc, JsonWebSignatureAlg},
         oauth::{OAuthAuthorizationEndpointResponseType, OAuthClientAuthenticationMethod},
     };
-    use mas_jose::JsonWebKeySet;
+    use mas_jose::jwk::PublicJsonWebKeySet;
     use url::Url;
 
     use super::{ClientMetadata, ClientMetadataVerificationError};
@@ -874,7 +874,7 @@ mod tests {
         }
     }
 
-    fn jwks() -> JsonWebKeySet {
+    fn jwks() -> PublicJsonWebKeySet {
         serde_json::from_value(serde_json::json!({
             "keys": [
                 {

crates/storage/src/oauth2/client.rs

@@ -19,7 +19,7 @@ use mas_iana::{
     jose::JsonWebSignatureAlg,
     oauth::{OAuthAuthorizationEndpointResponseType, OAuthClientAuthenticationMethod},
 };
-use mas_jose::JsonWebKeySet;
+use mas_jose::jwk::PublicJsonWebKeySet;
 use oauth2_types::requests::GrantType;
 use sqlx::{PgConnection, PgExecutor};
 use thiserror::Error;
@@ -331,7 +331,7 @@ pub async fn insert_client(
     policy_uri: Option<&Url>,
     tos_uri: Option<&Url>,
     jwks_uri: Option<&Url>,
-    jwks: Option<&JsonWebKeySet>,
+    jwks: Option<&PublicJsonWebKeySet>,
     id_token_signed_response_alg: Option<JsonWebSignatureAlg>,
     userinfo_signed_response_alg: Option<JsonWebSignatureAlg>,
     token_endpoint_auth_method: Option<OAuthClientAuthenticationMethod>,
@@ -421,7 +421,7 @@ pub async fn insert_client_from_config(
     client_id: &str,
     client_auth_method: OAuthClientAuthenticationMethod,
     encrypted_client_secret: Option<&str>,
-    jwks: Option<&JsonWebKeySet>,
+    jwks: Option<&PublicJsonWebKeySet>,
     jwks_uri: Option<&Url>,
     redirect_uris: &[Url],
 ) -> anyhow::Result<()> {