Support private_key_jwt client auth

Which includes having a verifying keystore out of JWKS (and soon out of
a JWKS URI)

crates/jose/src/claims.rs

@@ -0,0 +1,25 @@
+// Copyright 2022 The Matrix.org Foundation C.I.C.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+trait ClaimSet {
+    fn validate(&self) -> anyhow::Result<()>;
+}
+
+struct UnvalidatedClaim<T>(T);
+
+impl<T> ClaimSet for UnvalidatedClaim<T> {
+    fn validate(&self) -> anyhow::Result<()> {
+        Ok(())
+    }
+}

crates/jose/src/iana.rs

@@ -16,9 +16,10 @@
 //!
 //! <https://www.iana.org/assignments/jose/jose.xhtml>
 
+use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, JsonSchema)]
 pub enum JsonWebSignatureAlgorithm {
     /// HMAC using SHA-256
     #[serde(rename = "HS256")]
@@ -157,7 +158,7 @@ pub enum JsonWebSignatureAlgorithm {
     Es256K,
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
 pub enum JsonWebEncryptionAlgorithm {
     /// AES_128_CBC_HMAC_SHA_256 authenticated encryption algorithm
     #[serde(rename = "A128CBC-HS256")]
@@ -184,14 +185,14 @@ pub enum JsonWebEncryptionAlgorithm {
     A256Gcm,
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
 pub enum JsonWebEncryptionCompressionAlgorithm {
     /// DEFLATE
     #[serde(rename = "DEF")]
     Def,
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
 pub enum JsonWebKeyType {
     /// Elliptic Curve
     #[serde(rename = "EC")]
@@ -210,7 +211,7 @@ pub enum JsonWebKeyType {
     Okp,
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, JsonSchema)]
 pub enum JsonWebKeyEcEllipticCurve {
     /// P-256 Curve
     #[serde(rename = "P-256")]
@@ -229,7 +230,7 @@ pub enum JsonWebKeyEcEllipticCurve {
     Secp256K1,
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, JsonSchema)]
 pub enum JsonWebKeyOkpEllipticCurve {
     /// Ed25519 signature algorithm key pairs
     #[serde(rename = "Ed25519")]
@@ -248,7 +249,7 @@ pub enum JsonWebKeyOkpEllipticCurve {
     X448,
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, JsonSchema)]
 pub enum JsonWebKeyUse {
     /// Digital Signature or MAC
     #[serde(rename = "sig")]
@@ -259,7 +260,7 @@ pub enum JsonWebKeyUse {
     Enc,
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, JsonSchema)]
 pub enum JsonWebKeyOperation {
     /// Compute digital signature or MAC
     #[serde(rename = "sign")]

crates/jose/src/jwk.rs

@@ -17,6 +17,7 @@
 use anyhow::bail;
 use p256::NistP256;
 use rsa::{BigUint, PublicKeyParts};
+use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 use serde_with::{
     base64::{Base64, Standard, UrlSafe},
@@ -25,14 +26,17 @@ use serde_with::{
 };
 use url::Url;
 
-use crate::iana::{
-    JsonWebKeyEcEllipticCurve, JsonWebKeyOkpEllipticCurve, JsonWebKeyOperation, JsonWebKeyUse,
-    JsonWebSignatureAlgorithm,
+use crate::{
+    iana::{
+        JsonWebKeyEcEllipticCurve, JsonWebKeyOkpEllipticCurve, JsonWebKeyOperation, JsonWebKeyUse,
+        JsonWebSignatureAlgorithm,
+    },
+    JsonWebKeyType,
 };
 
 #[serde_as]
 #[skip_serializing_none]
-#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
 pub struct JsonWebKey {
     #[serde(flatten)]
     parameters: JsonWebKeyParameters,
@@ -49,17 +53,21 @@ pub struct JsonWebKey {
     #[serde(default)]
     kid: Option<String>,
 
+    #[schemars(with = "Option<String>")]
     #[serde(default)]
     x5u: Option<Url>,
 
+    #[schemars(with = "Vec<String>")]
     #[serde(default)]
     #[serde_as(as = "Option<Vec<Base64<Standard, Padded>>>")]
     x5c: Option<Vec<Vec<u8>>>,
 
+    #[schemars(with = "Option<String>")]
     #[serde(default)]
     #[serde_as(as = "Option<Base64<UrlSafe, Unpadded>>")]
     x5t: Option<Vec<u8>>,
 
+    #[schemars(with = "Option<String>")]
     #[serde(default, rename = "x5t#S256")]
     #[serde_as(as = "Option<Base64<UrlSafe, Unpadded>>")]
     x5t_s256: Option<Vec<u8>>,
@@ -104,13 +112,40 @@ impl JsonWebKey {
         self.kid = Some(kid.into());
         self
     }
+
+    #[must_use]
+    pub fn kty(&self) -> JsonWebKeyType {
+        match self.parameters {
+            JsonWebKeyParameters::Ec { .. } => JsonWebKeyType::Ec,
+            JsonWebKeyParameters::Rsa { .. } => JsonWebKeyType::Rsa,
+            JsonWebKeyParameters::Okp { .. } => JsonWebKeyType::Okp,
+        }
+    }
+
+    #[must_use]
+    pub fn kid(&self) -> Option<&str> {
+        self.kid.as_deref()
+    }
+
+    #[must_use]
+    pub fn params(&self) -> &JsonWebKeyParameters {
+        &self.parameters
+    }
 }
 
-#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
 pub struct JsonWebKeySet {
     keys: Vec<JsonWebKey>,
 }
 
+impl std::ops::Deref for JsonWebKeySet {
+    type Target = Vec<JsonWebKey>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.keys
+    }
+}
+
 impl JsonWebKeySet {
     #[must_use]
     pub fn new(keys: Vec<JsonWebKey>) -> Self {
@@ -119,27 +154,36 @@ impl JsonWebKeySet {
 }
 
 #[serde_as]
-#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
 #[serde(tag = "kty")]
 pub enum JsonWebKeyParameters {
     #[serde(rename = "RSA")]
     Rsa {
+        #[schemars(with = "String")]
         #[serde_as(as = "Base64<UrlSafe, Unpadded>")]
         n: Vec<u8>,
+
+        #[schemars(with = "String")]
         #[serde_as(as = "Base64<UrlSafe, Unpadded>")]
         e: Vec<u8>,
     },
     #[serde(rename = "EC")]
     Ec {
         crv: JsonWebKeyEcEllipticCurve,
+
+        #[schemars(with = "String")]
         #[serde_as(as = "Base64<UrlSafe, Unpadded>")]
         x: Vec<u8>,
+
+        #[schemars(with = "String")]
         #[serde_as(as = "Base64<UrlSafe, Unpadded>")]
         y: Vec<u8>,
     },
     #[serde(rename = "OKP")]
     Okp {
         crv: JsonWebKeyOkpEllipticCurve,
+
+        #[schemars(with = "String")]
         #[serde_as(as = "Base64<UrlSafe, Unpadded>")]
         x: Vec<u8>,
     },

crates/jose/src/keystore/jwks.rs

@@ -0,0 +1,278 @@
+// Copyright 2022 The Matrix.org Foundation C.I.C.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::collections::HashMap;
+
+use anyhow::bail;
+use async_trait::async_trait;
+use chrono::{DateTime, Duration, Utc};
+use digest::Digest;
+use rsa::{PublicKey, RsaPublicKey};
+use sha2::{Sha256, Sha384, Sha512};
+use signature::{Signature, Verifier};
+use tokio::sync::RwLock;
+
+use crate::{
+    ExportJwks, JsonWebKeySet, JsonWebKeyType, JsonWebSignatureAlgorithm, JwtHeader,
+    VerifyingKeystore,
+};
+
+pub struct StaticJwksStore {
+    key_set: JsonWebKeySet,
+    index: HashMap<(JsonWebKeyType, String), usize>,
+}
+
+impl StaticJwksStore {
+    #[must_use]
+    pub fn new(key_set: JsonWebKeySet) -> Self {
+        let index = key_set
+            .iter()
+            .enumerate()
+            .filter_map(|(index, key)| {
+                let kid = key.kid()?.to_string();
+                let kty = key.kty();
+
+                Some(((kty, kid), index))
+            })
+            .collect();
+
+        Self { key_set, index }
+    }
+
+    fn find_rsa_key(&self, kid: String) -> anyhow::Result<RsaPublicKey> {
+        let index = *self
+            .index
+            .get(&(JsonWebKeyType::Rsa, kid))
+            .ok_or_else(|| anyhow::anyhow!("key not found"))?;
+
+        let key = self
+            .key_set
+            .get(index)
+            .ok_or_else(|| anyhow::anyhow!("invalid index"))?;
+
+        let key = key.params().clone().try_into()?;
+
+        Ok(key)
+    }
+
+    fn find_ecdsa_key(&self, kid: String) -> anyhow::Result<ecdsa::VerifyingKey<p256::NistP256>> {
+        let index = *self
+            .index
+            .get(&(JsonWebKeyType::Ec, kid))
+            .ok_or_else(|| anyhow::anyhow!("key not found"))?;
+
+        let key = self
+            .key_set
+            .get(index)
+            .ok_or_else(|| anyhow::anyhow!("invalid index"))?;
+
+        let key = key.params().clone().try_into()?;
+
+        Ok(key)
+    }
+}
+
+#[async_trait]
+impl VerifyingKeystore for &StaticJwksStore {
+    async fn verify(
+        self,
+        header: &JwtHeader,
+        payload: &[u8],
+        signature: &[u8],
+    ) -> anyhow::Result<()> {
+        let kid = header
+            .kid()
+            .ok_or_else(|| anyhow::anyhow!("missing kid"))?
+            .to_string();
+        match header.alg() {
+            JsonWebSignatureAlgorithm::Rs256 => {
+                let key = self.find_rsa_key(kid)?;
+
+                let digest = {
+                    let mut digest = Sha256::new();
+                    digest.update(&payload);
+                    digest.finalize()
+                };
+
+                key.verify(
+                    rsa::PaddingScheme::new_pkcs1v15_sign(Some(rsa::Hash::SHA2_256)),
+                    &digest,
+                    signature,
+                )?;
+            }
+
+            JsonWebSignatureAlgorithm::Rs384 => {
+                let key = self.find_rsa_key(kid)?;
+
+                let digest = {
+                    let mut digest = Sha384::new();
+                    digest.update(&payload);
+                    digest.finalize()
+                };
+
+                key.verify(
+                    rsa::PaddingScheme::new_pkcs1v15_sign(Some(rsa::Hash::SHA2_384)),
+                    &digest,
+                    signature,
+                )?;
+            }
+
+            JsonWebSignatureAlgorithm::Rs512 => {
+                let key = self.find_rsa_key(kid)?;
+
+                let digest = {
+                    let mut digest = Sha512::new();
+                    digest.update(&payload);
+                    digest.finalize()
+                };
+
+                key.verify(
+                    rsa::PaddingScheme::new_pkcs1v15_sign(Some(rsa::Hash::SHA2_512)),
+                    &digest,
+                    signature,
+                )?;
+            }
+
+            JsonWebSignatureAlgorithm::Es256 => {
+                let key = self.find_ecdsa_key(kid)?;
+
+                let signature = ecdsa::Signature::from_bytes(signature)?;
+
+                key.verify(payload, &signature)?;
+            }
+
+            _ => bail!("unsupported algorithm"),
+        };
+
+        Ok(())
+    }
+}
+
+enum RemoteKeySet {
+    Pending,
+    Errored {
+        at: DateTime<Utc>,
+        error: anyhow::Error,
+    },
+    Fulfilled {
+        at: DateTime<Utc>,
+        store: StaticJwksStore,
+    },
+}
+
+impl Default for RemoteKeySet {
+    fn default() -> Self {
+        Self::Pending
+    }
+}
+
+impl RemoteKeySet {
+    fn fullfill(&mut self, key_set: JsonWebKeySet) {
+        *self = Self::Fulfilled {
+            at: Utc::now(),
+            store: StaticJwksStore::new(key_set),
+        }
+    }
+
+    fn error(&mut self, error: anyhow::Error) {
+        *self = Self::Errored {
+            at: Utc::now(),
+            error,
+        }
+    }
+
+    fn should_refresh(&self) -> bool {
+        let now = Utc::now();
+        match self {
+            Self::Pending => true,
+            Self::Errored { at, .. } if *at - now > Duration::minutes(5) => true,
+            Self::Fulfilled { at, .. } if *at - now > Duration::hours(1) => true,
+            _ => false,
+        }
+    }
+
+    fn should_force_refresh(&self) -> bool {
+        let now = Utc::now();
+        match self {
+            Self::Pending => true,
+            Self::Errored { at, .. } | Self::Fulfilled { at, .. }
+                if *at - now > Duration::minutes(5) =>
+            {
+                true
+            }
+            _ => false,
+        }
+    }
+}
+
+pub struct JwksStore<T>
+where
+    T: ExportJwks,
+{
+    exporter: T,
+    cache: RwLock<RemoteKeySet>,
+}
+
+impl<T: ExportJwks> JwksStore<T> {
+    pub fn new(exporter: T) -> Self {
+        Self {
+            exporter,
+            cache: RwLock::default(),
+        }
+    }
+
+    async fn should_refresh(&self) -> bool {
+        let cache = self.cache.read().await;
+        cache.should_refresh()
+    }
+
+    async fn refresh(&self) {
+        let mut cache = self.cache.write().await;
+
+        if cache.should_force_refresh() {
+            let jwks = self.exporter.export_jwks().await;
+
+            match jwks {
+                Ok(jwks) => cache.fullfill(jwks),
+                Err(err) => cache.error(err),
+            }
+        }
+    }
+}
+
+#[async_trait]
+impl<T: ExportJwks + Send + Sync> VerifyingKeystore for &JwksStore<T> {
+    async fn verify(
+        self,
+        header: &JwtHeader,
+        payload: &[u8],
+        signature: &[u8],
+    ) -> anyhow::Result<()> {
+        if self.should_refresh().await {
+            self.refresh().await;
+        }
+
+        let cache = self.cache.read().await;
+        // TODO: we could bubble up the underlying error here
+        let store = match &*cache {
+            RemoteKeySet::Pending => bail!("inconsistent cache state"),
+            RemoteKeySet::Errored { error, .. } => bail!("cache in error state {}", error),
+            RemoteKeySet::Fulfilled { store, .. } => store,
+        };
+
+        store.verify(header, payload, signature).await?;
+
+        Ok(())
+    }
+}

crates/jose/src/keystore/mod.rs

@@ -12,11 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+mod jwks;
 mod shared_secret;
 mod static_keystore;
 mod traits;
 
 pub use self::{
+    jwks::{JwksStore, StaticJwksStore},
     shared_secret::SharedSecret,
     static_keystore::StaticKeystore,
     traits::{ExportJwks, SigningKeystore, VerifyingKeystore},

crates/jose/src/keystore/static_keystore.rs

@@ -27,10 +27,7 @@ use sha2::{Sha256, Sha384, Sha512};
 use signature::{Signature, Signer, Verifier};
 
 use super::{ExportJwks, SigningKeystore, VerifyingKeystore};
-use crate::{
-    iana::{JsonWebKeyOperation, JsonWebSignatureAlgorithm},
-    JsonWebKey, JsonWebKeySet, JwtHeader,
-};
+use crate::{iana::JsonWebSignatureAlgorithm, JsonWebKey, JsonWebKeySet, JwtHeader};
 
 #[derive(Default)]
 pub struct StaticKeystore {
@@ -276,23 +273,13 @@ impl VerifyingKeystore for &StaticKeystore {
 }
 
 #[async_trait]
-impl ExportJwks for &StaticKeystore {
-    async fn export_jwks(self) -> JsonWebKeySet {
-        let rsa = self.rsa_keys.iter().flat_map(|(kid, key)| {
+impl ExportJwks for StaticKeystore {
+    async fn export_jwks(&self) -> anyhow::Result<JsonWebKeySet> {
+        let rsa = self.rsa_keys.iter().map(|(kid, key)| {
             let pubkey = RsaPublicKey::from(key);
-            let basekey = JsonWebKey::new(pubkey.into())
+            JsonWebKey::new(pubkey.into())
                 .with_kid(kid)
                 .with_use(crate::JsonWebKeyUse::Sig)
-                .with_key_ops(vec![JsonWebKeyOperation::Sign]);
-
-            let algs = [
-                JsonWebSignatureAlgorithm::Rs256,
-                JsonWebSignatureAlgorithm::Rs384,
-                JsonWebSignatureAlgorithm::Rs512,
-            ];
-
-            algs.into_iter()
-                .map(move |alg| basekey.clone().with_alg(alg))
         });
 
         let es256 = self.es256_keys.iter().map(|(kid, key)| {
@@ -300,12 +287,11 @@ impl ExportJwks for &StaticKeystore {
             JsonWebKey::new(pubkey.into())
                 .with_kid(kid)
                 .with_use(crate::JsonWebKeyUse::Sig)
-                .with_key_ops(vec![JsonWebKeyOperation::Sign])
                 .with_alg(JsonWebSignatureAlgorithm::Es256)
         });
 
         let keys = rsa.chain(es256).collect();
-        JsonWebKeySet::new(keys)
+        Ok(JsonWebKeySet::new(keys))
     }
 }
 

crates/jose/src/keystore/traits.rs

@@ -14,9 +14,7 @@
 
 use async_trait::async_trait;
 
-use crate::{
-    iana::JsonWebSignatureAlgorithm, JsonWebKeySet, JwtHeader,
-};
+use crate::{iana::JsonWebSignatureAlgorithm, JsonWebKeySet, JwtHeader};
 
 #[async_trait]
 pub trait SigningKeystore {
@@ -32,6 +30,5 @@ pub trait VerifyingKeystore {
 
 #[async_trait]
 pub trait ExportJwks {
-    async fn export_jwks(self) -> JsonWebKeySet;
+    async fn export_jwks(&self) -> anyhow::Result<JsonWebKeySet>;
 }
-

crates/jose/src/lib.rs

@@ -19,6 +19,7 @@
 #![allow(clippy::missing_errors_doc)]
 #![allow(clippy::module_name_repetitions)]
 
+mod claims;
 pub(crate) mod iana;
 pub(crate) mod jwk;
 pub(crate) mod jwt;
@@ -31,5 +32,8 @@ pub use self::{
     },
     jwk::{JsonWebKey, JsonWebKeySet},
     jwt::{DecodedJsonWebToken, JsonWebTokenParts, JwtHeader},
-    keystore::{ExportJwks, SharedSecret, SigningKeystore, StaticKeystore, VerifyingKeystore},
+    keystore::{
+        ExportJwks, JwksStore, SharedSecret, SigningKeystore, StaticJwksStore, StaticKeystore,
+        VerifyingKeystore,
+    },
 };