Implement private_key_jwks client authentication

This involves a lot of things, including:
 - better VerifyingKeystore trait
 - better errors in the JOSE crate
 - getting rid of async_trait in some JOSE traits

crates/jose/src/jwt.rs

@@ -30,7 +30,7 @@ use crate::{jwk::JsonWebKey, SigningKeystore, VerifyingKeystore};
 
 #[serde_as]
 #[skip_serializing_none]
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct JwtHeader {
     alg: JsonWebSignatureAlg,
 
@@ -163,7 +163,7 @@ where
         Ok(format!("{}.{}", header, payload))
     }
 
-    pub async fn sign<S: SigningKeystore>(&self, store: S) -> anyhow::Result<JsonWebTokenParts> {
+    pub async fn sign<S: SigningKeystore>(&self, store: &S) -> anyhow::Result<JsonWebTokenParts> {
         let payload = self.serialize()?;
         let signature = store.sign(&self.header, payload.as_bytes()).await?;
         Ok(JsonWebTokenParts { payload, signature })
@@ -205,22 +205,24 @@ impl JsonWebTokenParts {
         Ok(decoded)
     }
 
-    pub async fn verify<T, S: VerifyingKeystore>(
+    pub fn verify<T, S: VerifyingKeystore>(
         &self,
         decoded: &DecodedJsonWebToken<T>,
-        store: S,
-    ) -> anyhow::Result<()> {
-        store
-            .verify(&decoded.header, self.payload.as_bytes(), &self.signature)
-            .await?;
-
-        Ok(())
+        store: &S,
+    ) -> S::Future
+    where
+        S::Error: std::error::Error + Send + Sync + 'static,
+    {
+        store.verify(&decoded.header, self.payload.as_bytes(), &self.signature)
     }
 
     pub async fn decode_and_verify<T: DeserializeOwned, S: VerifyingKeystore>(
         &self,
-        store: S,
-    ) -> anyhow::Result<DecodedJsonWebToken<T>> {
+        store: &S,
+    ) -> anyhow::Result<DecodedJsonWebToken<T>>
+    where
+        S::Error: std::error::Error + Send + Sync + 'static,
+    {
         let decoded = self.decode()?;
         self.verify(&decoded, store).await?;
         Ok(decoded)

crates/jose/src/keystore/jwks.rs

@@ -1,276 +0,0 @@
-// 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 mas_iana::jose::{JsonWebKeyType, JsonWebSignatureAlg};
-use rsa::{PublicKey, RsaPublicKey};
-use sha2::{Sha256, Sha384, Sha512};
-use signature::{Signature, Verifier};
-use tokio::sync::RwLock;
-
-use crate::{ExportJwks, JsonWebKeySet, 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() {
-            JsonWebSignatureAlg::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,
-                )?;
-            }
-
-            JsonWebSignatureAlg::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,
-                )?;
-            }
-
-            JsonWebSignatureAlg::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,
-                )?;
-            }
-
-            JsonWebSignatureAlg::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/jwks/dynamic_store.rs

@@ -0,0 +1,160 @@
+// 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::sync::Arc;
+
+use chrono::{DateTime, Duration, Utc};
+use futures_util::future::BoxFuture;
+use thiserror::Error;
+use tokio::sync::RwLock;
+use tower::{
+    util::{BoxCloneService, ServiceExt},
+    BoxError, Service,
+};
+
+use super::StaticJwksStore;
+use crate::{JsonWebKeySet, JwtHeader, VerifyingKeystore};
+
+#[derive(Debug, Error)]
+pub enum Error {
+    #[error("cache in inconsistent state")]
+    InconsistentCache,
+
+    #[error(transparent)]
+    Cached(Arc<BoxError>),
+
+    #[error("todo")]
+    Todo,
+
+    #[error(transparent)]
+    Verification(#[from] super::static_store::Error),
+}
+
+enum State<E> {
+    Pending,
+    Errored {
+        at: DateTime<Utc>,
+        error: E,
+    },
+    Fulfilled {
+        at: DateTime<Utc>,
+        store: StaticJwksStore,
+    },
+}
+
+impl<E> Default for State<E> {
+    fn default() -> Self {
+        Self::Pending
+    }
+}
+
+impl<E> State<E> {
+    fn fullfill(&mut self, key_set: JsonWebKeySet) {
+        *self = Self::Fulfilled {
+            at: Utc::now(),
+            store: StaticJwksStore::new(key_set),
+        }
+    }
+
+    fn error(&mut self, error: E) {
+        *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,
+        }
+    }
+}
+
+#[derive(Clone)]
+pub struct DynamicJwksStore {
+    exporter: BoxCloneService<(), JsonWebKeySet, 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::Future: Send,
+    {
+        Self {
+            exporter: exporter.boxed_clone(),
+            cache: Arc::default(),
+        }
+    }
+}
+
+impl VerifyingKeystore for DynamicJwksStore {
+    type Error = Error;
+    type Future = BoxFuture<'static, Result<(), Self::Error>>;
+
+    fn verify(&self, header: &JwtHeader, payload: &[u8], signature: &[u8]) -> Self::Future {
+        let cache = self.cache.clone();
+        let exporter = self.exporter.clone();
+        let header = header.clone();
+        let payload = payload.to_owned();
+        let signature = signature.to_owned();
+
+        let fut = async move {
+            if cache.read().await.should_refresh() {
+                let mut cache = cache.write().await;
+
+                if cache.should_force_refresh() {
+                    let jwks = async move { exporter.ready_oneshot().await?.call(()).await }.await;
+
+                    match jwks {
+                        Ok(jwks) => cache.fullfill(jwks),
+                        Err(err) => cache.error(Arc::new(err)),
+                    }
+                }
+            }
+
+            let cache = cache.read().await;
+            // TODO: we could bubble up the underlying error here
+            let store = match &*cache {
+                State::Pending => return Err(Error::InconsistentCache),
+                State::Errored { error, .. } => return Err(Error::Cached(error.clone())),
+                State::Fulfilled { store, .. } => store,
+            };
+
+            store.verify(&header, &payload, &signature).await?;
+
+            Ok(())
+        };
+
+        Box::pin(fut)
+    }
+}

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

@@ -0,0 +1,18 @@
+// 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.
+
+mod dynamic_store;
+mod static_store;
+
+pub use self::{dynamic_store::DynamicJwksStore, static_store::StaticJwksStore};

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

@@ -0,0 +1,196 @@
+// 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, future::Ready};
+
+use digest::Digest;
+use mas_iana::jose::{JsonWebKeyType, JsonWebSignatureAlg};
+use rsa::{PublicKey, RsaPublicKey};
+use sha2::{Sha256, Sha384, Sha512};
+use signature::{Signature, Verifier};
+use thiserror::Error;
+
+use crate::{JsonWebKeySet, JwtHeader, VerifyingKeystore};
+
+#[derive(Debug, Error)]
+pub enum Error {
+    #[error("key not found")]
+    KeyNotFound,
+
+    #[error("invalid index")]
+    InvalidIndex,
+
+    #[error(r#"missing "kid" field in header"#)]
+    MissingKid,
+
+    #[error(transparent)]
+    Rsa(#[from] rsa::errors::Error),
+
+    #[error("unsupported algorithm {alg}")]
+    UnsupportedAlgorithm { alg: JsonWebSignatureAlg },
+
+    #[error(transparent)]
+    Signature(#[from] signature::Error),
+
+    #[error("invalid {kty} key {kid}")]
+    InvalidKey {
+        kty: JsonWebKeyType,
+        kid: String,
+        source: anyhow::Error,
+    },
+}
+
+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) -> Result<RsaPublicKey, Error> {
+        let index = *self
+            .index
+            .get(&(JsonWebKeyType::Rsa, kid.clone()))
+            .ok_or(Error::KeyNotFound)?;
+
+        let key = self.key_set.get(index).ok_or(Error::InvalidIndex)?;
+
+        let key = key
+            .params()
+            .clone()
+            .try_into()
+            .map_err(|source| Error::InvalidKey {
+                kty: JsonWebKeyType::Rsa,
+                kid,
+                source,
+            })?;
+
+        Ok(key)
+    }
+
+    fn find_ecdsa_key(&self, kid: String) -> Result<ecdsa::VerifyingKey<p256::NistP256>, Error> {
+        let index = *self
+            .index
+            .get(&(JsonWebKeyType::Ec, kid.clone()))
+            .ok_or(Error::KeyNotFound)?;
+
+        let key = self.key_set.get(index).ok_or(Error::InvalidIndex)?;
+
+        let key = key
+            .params()
+            .clone()
+            .try_into()
+            .map_err(|source| Error::InvalidKey {
+                kty: JsonWebKeyType::Ec,
+                kid,
+                source,
+            })?;
+
+        Ok(key)
+    }
+
+    fn verify_sync(
+        &self,
+        header: &JwtHeader,
+        payload: &[u8],
+        signature: &[u8],
+    ) -> Result<(), Error> {
+        let kid = header.kid().ok_or(Error::MissingKid)?.to_string();
+        match header.alg() {
+            JsonWebSignatureAlg::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,
+                )?;
+            }
+
+            JsonWebSignatureAlg::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,
+                )?;
+            }
+
+            JsonWebSignatureAlg::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,
+                )?;
+            }
+
+            JsonWebSignatureAlg::Es256 => {
+                let key = self.find_ecdsa_key(kid)?;
+
+                let signature = ecdsa::Signature::from_bytes(signature)?;
+
+                key.verify(payload, &signature)?;
+            }
+
+            alg => return Err(Error::UnsupportedAlgorithm { alg }),
+        };
+
+        Ok(())
+    }
+}
+
+impl VerifyingKeystore for StaticJwksStore {
+    type Error = Error;
+    type Future = Ready<Result<(), Self::Error>>;
+
+    fn verify(&self, header: &JwtHeader, payload: &[u8], signature: &[u8]) -> Self::Future {
+        std::future::ready(self.verify_sync(header, payload, signature))
+    }
+}

crates/jose/src/keystore/mod.rs

@@ -18,8 +18,8 @@ mod static_keystore;
 mod traits;
 
 pub use self::{
-    jwks::{JwksStore, StaticJwksStore},
+    jwks::{DynamicJwksStore, StaticJwksStore},
     shared_secret::SharedSecret,
     static_keystore::StaticKeystore,
-    traits::{ExportJwks, SigningKeystore, VerifyingKeystore},
+    traits::{SigningKeystore, VerifyingKeystore},
 };

crates/jose/src/keystore/shared_secret.rs

@@ -12,17 +12,31 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::collections::HashSet;
+use std::{collections::HashSet, future::Ready};
 
 use anyhow::bail;
 use async_trait::async_trait;
+use digest::{InvalidLength, MacError};
 use hmac::{Hmac, Mac};
 use mas_iana::jose::JsonWebSignatureAlg;
 use sha2::{Sha256, Sha384, Sha512};
+use thiserror::Error;
 
 use super::{SigningKeystore, VerifyingKeystore};
 use crate::JwtHeader;
 
+#[derive(Debug, Error)]
+pub enum Error {
+    #[error("invalid key")]
+    InvalidKey(#[from] InvalidLength),
+
+    #[error("unsupported algorithm {alg}")]
+    UnsupportedAlgorithm { alg: JsonWebSignatureAlg },
+
+    #[error("signature verification failed")]
+    Verification(#[from] MacError),
+}
+
 pub struct SharedSecret<'a> {
     inner: &'a [u8],
 }
@@ -33,11 +47,42 @@ impl<'a> SharedSecret<'a> {
             inner: source.as_ref(),
         }
     }
+
+    fn verify_sync(
+        &self,
+        header: &JwtHeader,
+        payload: &[u8],
+        signature: &[u8],
+    ) -> Result<(), Error> {
+        match header.alg() {
+            JsonWebSignatureAlg::Hs256 => {
+                let mut mac = Hmac::<Sha256>::new_from_slice(self.inner)?;
+                mac.update(payload);
+                mac.verify(signature.into())?;
+            }
+
+            JsonWebSignatureAlg::Hs384 => {
+                let mut mac = Hmac::<Sha384>::new_from_slice(self.inner)?;
+                mac.update(payload);
+                mac.verify(signature.into())?;
+            }
+
+            JsonWebSignatureAlg::Hs512 => {
+                let mut mac = Hmac::<Sha512>::new_from_slice(self.inner)?;
+                mac.update(payload);
+                mac.verify(signature.into())?;
+            }
+
+            alg => return Err(Error::UnsupportedAlgorithm { alg }),
+        };
+
+        Ok(())
+    }
 }
 
 #[async_trait]
-impl<'a> SigningKeystore for &SharedSecret<'a> {
-    fn supported_algorithms(self) -> HashSet<JsonWebSignatureAlg> {
+impl<'a> SigningKeystore for SharedSecret<'a> {
+    fn supported_algorithms(&self) -> HashSet<JsonWebSignatureAlg> {
         let mut algorithms = HashSet::with_capacity(3);
 
         algorithms.insert(JsonWebSignatureAlg::Hs256);
@@ -47,7 +92,7 @@ impl<'a> SigningKeystore for &SharedSecret<'a> {
         algorithms
     }
 
-    async fn prepare_header(self, alg: JsonWebSignatureAlg) -> anyhow::Result<JwtHeader> {
+    async fn prepare_header(&self, alg: JsonWebSignatureAlg) -> anyhow::Result<JwtHeader> {
         if !matches!(
             alg,
             JsonWebSignatureAlg::Hs256 | JsonWebSignatureAlg::Hs384 | JsonWebSignatureAlg::Hs512,
@@ -58,7 +103,7 @@ impl<'a> SigningKeystore for &SharedSecret<'a> {
         Ok(JwtHeader::new(alg))
     }
 
-    async fn sign(self, header: &JwtHeader, msg: &[u8]) -> anyhow::Result<Vec<u8>> {
+    async fn sign(&self, header: &JwtHeader, msg: &[u8]) -> anyhow::Result<Vec<u8>> {
         // TODO: do the signing in a blocking task
         // TODO: should we bail out if the key is too small?
         let signature = match header.alg() {
@@ -87,38 +132,12 @@ impl<'a> SigningKeystore for &SharedSecret<'a> {
     }
 }
 
-#[async_trait]
-impl<'a> VerifyingKeystore for &SharedSecret<'a> {
-    async fn verify(
-        self,
-        header: &JwtHeader,
-        payload: &[u8],
-        signature: &[u8],
-    ) -> anyhow::Result<()> {
-        // TODO: do the verification in a blocking task
-        match header.alg() {
-            JsonWebSignatureAlg::Hs256 => {
-                let mut mac = Hmac::<Sha256>::new_from_slice(self.inner)?;
-                mac.update(payload);
-                mac.verify(signature.try_into()?)?;
-            }
+impl<'a> VerifyingKeystore for SharedSecret<'a> {
+    type Error = Error;
+    type Future = Ready<Result<(), Self::Error>>;
 
-            JsonWebSignatureAlg::Hs384 => {
-                let mut mac = Hmac::<Sha384>::new_from_slice(self.inner)?;
-                mac.update(payload);
-                mac.verify(signature.try_into()?)?;
-            }
-
-            JsonWebSignatureAlg::Hs512 => {
-                let mut mac = Hmac::<Sha512>::new_from_slice(self.inner)?;
-                mac.update(payload);
-                mac.verify(signature.try_into()?)?;
-            }
-
-            _ => bail!("unsupported algorithm"),
-        };
-
-        Ok(())
+    fn verify(&self, header: &JwtHeader, payload: &[u8], signature: &[u8]) -> Self::Future {
+        std::future::ready(self.verify_sync(header, payload, signature))
     }
 }
 

crates/jose/src/keystore/static_keystore.rs

@@ -12,7 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::collections::{HashMap, HashSet};
+use std::{
+    collections::{HashMap, HashSet},
+    convert::Infallible,
+    future::Ready,
+    task::Poll,
+};
 
 use anyhow::bail;
 use async_trait::async_trait;
@@ -26,8 +31,9 @@ use pkcs8::{DecodePrivateKey, EncodePublicKey};
 use rsa::{PublicKey as _, RsaPrivateKey, RsaPublicKey};
 use sha2::{Sha256, Sha384, Sha512};
 use signature::{Signature, Signer, Verifier};
+use tower::Service;
 
-use super::{ExportJwks, SigningKeystore, VerifyingKeystore};
+use super::{SigningKeystore, VerifyingKeystore};
 use crate::{JsonWebKey, JsonWebKeySet, JwtHeader};
 
 // Generate with
@@ -123,135 +129,9 @@ impl StaticKeystore {
         self.es256_keys.insert(kid, key);
         Ok(())
     }
-}
 
-#[async_trait]
-impl SigningKeystore for &StaticKeystore {
-    fn supported_algorithms(self) -> HashSet<JsonWebSignatureAlg> {
-        let has_rsa = !self.rsa_keys.is_empty();
-        let has_es256 = !self.es256_keys.is_empty();
-
-        let capacity = (if has_rsa { 3 } else { 0 }) + (if has_es256 { 1 } else { 0 });
-        let mut algorithms = HashSet::with_capacity(capacity);
-
-        if has_rsa {
-            algorithms.insert(JsonWebSignatureAlg::Rs256);
-            algorithms.insert(JsonWebSignatureAlg::Rs384);
-            algorithms.insert(JsonWebSignatureAlg::Rs512);
-        }
-
-        if has_es256 {
-            algorithms.insert(JsonWebSignatureAlg::Es256);
-        }
-
-        algorithms
-    }
-
-    async fn prepare_header(self, alg: JsonWebSignatureAlg) -> anyhow::Result<JwtHeader> {
-        let header = JwtHeader::new(alg);
-
-        let kid = match alg {
-            JsonWebSignatureAlg::Rs256
-            | JsonWebSignatureAlg::Rs384
-            | JsonWebSignatureAlg::Rs512 => self
-                .rsa_keys
-                .keys()
-                .next()
-                .ok_or_else(|| anyhow::anyhow!("no RSA keys in keystore"))?,
-            JsonWebSignatureAlg::Es256 => self
-                .es256_keys
-                .keys()
-                .next()
-                .ok_or_else(|| anyhow::anyhow!("no ECDSA keys in keystore"))?,
-            _ => bail!("unsupported algorithm"),
-        };
-
-        Ok(header.with_kid(kid))
-    }
-
-    async fn sign(self, header: &JwtHeader, msg: &[u8]) -> anyhow::Result<Vec<u8>> {
-        let kid = header
-            .kid()
-            .ok_or_else(|| anyhow::anyhow!("missing kid from the JWT header"))?;
-
-        // TODO: do the signing in a blocking task
-        let signature = match header.alg() {
-            JsonWebSignatureAlg::Rs256 => {
-                let key = self
-                    .rsa_keys
-                    .get(kid)
-                    .ok_or_else(|| anyhow::anyhow!("RSA key not found in key store"))?;
-
-                let digest = {
-                    let mut digest = Sha256::new();
-                    digest.update(&msg);
-                    digest.finalize()
-                };
-
-                key.sign(
-                    rsa::PaddingScheme::new_pkcs1v15_sign(Some(rsa::Hash::SHA2_256)),
-                    &digest,
-                )?
-            }
-
-            JsonWebSignatureAlg::Rs384 => {
-                let key = self
-                    .rsa_keys
-                    .get(kid)
-                    .ok_or_else(|| anyhow::anyhow!("RSA key not found in key store"))?;
-
-                let digest = {
-                    let mut digest = Sha384::new();
-                    digest.update(&msg);
-                    digest.finalize()
-                };
-
-                key.sign(
-                    rsa::PaddingScheme::new_pkcs1v15_sign(Some(rsa::Hash::SHA2_384)),
-                    &digest,
-                )?
-            }
-
-            JsonWebSignatureAlg::Rs512 => {
-                let key = self
-                    .rsa_keys
-                    .get(kid)
-                    .ok_or_else(|| anyhow::anyhow!("RSA key not found in key store"))?;
-
-                let digest = {
-                    let mut digest = Sha512::new();
-                    digest.update(&msg);
-                    digest.finalize()
-                };
-
-                key.sign(
-                    rsa::PaddingScheme::new_pkcs1v15_sign(Some(rsa::Hash::SHA2_512)),
-                    &digest,
-                )?
-            }
-
-            JsonWebSignatureAlg::Es256 => {
-                let key = self
-                    .es256_keys
-                    .get(kid)
-                    .ok_or_else(|| anyhow::anyhow!("ECDSA key not found in key store"))?;
-
-                let signature = key.try_sign(msg)?;
-                let signature: &[u8] = signature.as_ref();
-                signature.to_vec()
-            }
-
-            _ => bail!("Unsupported algorithm"),
-        };
-
-        Ok(signature)
-    }
-}
-
-#[async_trait]
-impl VerifyingKeystore for &StaticKeystore {
-    async fn verify(
-        self,
+    fn verify_sync(
+        &self,
         header: &JwtHeader,
         payload: &[u8],
         signature: &[u8],
@@ -344,8 +224,147 @@ impl VerifyingKeystore for &StaticKeystore {
 }
 
 #[async_trait]
-impl ExportJwks for StaticKeystore {
-    async fn export_jwks(&self) -> anyhow::Result<JsonWebKeySet> {
+impl SigningKeystore for StaticKeystore {
+    fn supported_algorithms(&self) -> HashSet<JsonWebSignatureAlg> {
+        let has_rsa = !self.rsa_keys.is_empty();
+        let has_es256 = !self.es256_keys.is_empty();
+
+        let capacity = (if has_rsa { 3 } else { 0 }) + (if has_es256 { 1 } else { 0 });
+        let mut algorithms = HashSet::with_capacity(capacity);
+
+        if has_rsa {
+            algorithms.insert(JsonWebSignatureAlg::Rs256);
+            algorithms.insert(JsonWebSignatureAlg::Rs384);
+            algorithms.insert(JsonWebSignatureAlg::Rs512);
+        }
+
+        if has_es256 {
+            algorithms.insert(JsonWebSignatureAlg::Es256);
+        }
+
+        algorithms
+    }
+
+    async fn prepare_header(&self, alg: JsonWebSignatureAlg) -> anyhow::Result<JwtHeader> {
+        let header = JwtHeader::new(alg);
+
+        let kid = match alg {
+            JsonWebSignatureAlg::Rs256
+            | JsonWebSignatureAlg::Rs384
+            | JsonWebSignatureAlg::Rs512 => self
+                .rsa_keys
+                .keys()
+                .next()
+                .ok_or_else(|| anyhow::anyhow!("no RSA keys in keystore"))?,
+            JsonWebSignatureAlg::Es256 => self
+                .es256_keys
+                .keys()
+                .next()
+                .ok_or_else(|| anyhow::anyhow!("no ECDSA keys in keystore"))?,
+            _ => bail!("unsupported algorithm"),
+        };
+
+        Ok(header.with_kid(kid))
+    }
+
+    async fn sign(&self, header: &JwtHeader, msg: &[u8]) -> anyhow::Result<Vec<u8>> {
+        let kid = header
+            .kid()
+            .ok_or_else(|| anyhow::anyhow!("missing kid from the JWT header"))?;
+
+        // TODO: do the signing in a blocking task
+        let signature = match header.alg() {
+            JsonWebSignatureAlg::Rs256 => {
+                let key = self
+                    .rsa_keys
+                    .get(kid)
+                    .ok_or_else(|| anyhow::anyhow!("RSA key not found in key store"))?;
+
+                let digest = {
+                    let mut digest = Sha256::new();
+                    digest.update(&msg);
+                    digest.finalize()
+                };
+
+                key.sign(
+                    rsa::PaddingScheme::new_pkcs1v15_sign(Some(rsa::Hash::SHA2_256)),
+                    &digest,
+                )?
+            }
+
+            JsonWebSignatureAlg::Rs384 => {
+                let key = self
+                    .rsa_keys
+                    .get(kid)
+                    .ok_or_else(|| anyhow::anyhow!("RSA key not found in key store"))?;
+
+                let digest = {
+                    let mut digest = Sha384::new();
+                    digest.update(&msg);
+                    digest.finalize()
+                };
+
+                key.sign(
+                    rsa::PaddingScheme::new_pkcs1v15_sign(Some(rsa::Hash::SHA2_384)),
+                    &digest,
+                )?
+            }
+
+            JsonWebSignatureAlg::Rs512 => {
+                let key = self
+                    .rsa_keys
+                    .get(kid)
+                    .ok_or_else(|| anyhow::anyhow!("RSA key not found in key store"))?;
+
+                let digest = {
+                    let mut digest = Sha512::new();
+                    digest.update(&msg);
+                    digest.finalize()
+                };
+
+                key.sign(
+                    rsa::PaddingScheme::new_pkcs1v15_sign(Some(rsa::Hash::SHA2_512)),
+                    &digest,
+                )?
+            }
+
+            JsonWebSignatureAlg::Es256 => {
+                let key = self
+                    .es256_keys
+                    .get(kid)
+                    .ok_or_else(|| anyhow::anyhow!("ECDSA key not found in key store"))?;
+
+                let signature = key.try_sign(msg)?;
+                let signature: &[u8] = signature.as_ref();
+                signature.to_vec()
+            }
+
+            _ => bail!("Unsupported algorithm"),
+        };
+
+        Ok(signature)
+    }
+}
+
+impl VerifyingKeystore for StaticKeystore {
+    type Error = anyhow::Error;
+    type Future = Ready<Result<(), Self::Error>>;
+
+    fn verify(&self, header: &JwtHeader, msg: &[u8], signature: &[u8]) -> Self::Future {
+        std::future::ready(self.verify_sync(header, msg, signature))
+    }
+}
+
+impl Service<()> for &StaticKeystore {
+    type Future = Ready<Result<Self::Response, Self::Error>>;
+    type Response = JsonWebKeySet;
+    type Error = Infallible;
+
+    fn poll_ready(&mut self, _cx: &mut std::task::Context<'_>) -> Poll<Result<(), Self::Error>> {
+        Poll::Ready(Ok(()))
+    }
+
+    fn call(&mut self, _req: ()) -> Self::Future {
         let rsa = self.rsa_keys.iter().map(|(kid, key)| {
             let pubkey = RsaPublicKey::from(key);
             JsonWebKey::new(pubkey.into())
@@ -362,7 +381,7 @@ impl ExportJwks for StaticKeystore {
         });
 
         let keys = rsa.chain(es256).collect();
-        Ok(JsonWebKeySet::new(keys))
+        std::future::ready(Ok(JsonWebKeySet::new(keys)))
     }
 }
 

crates/jose/src/keystore/traits.rs

@@ -12,28 +12,78 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::collections::HashSet;
+use std::{collections::HashSet, future::Future, sync::Arc};
 
 use async_trait::async_trait;
+use futures_util::{
+    future::{Either, MapErr},
+    TryFutureExt,
+};
 use mas_iana::jose::JsonWebSignatureAlg;
+use thiserror::Error;
 
-use crate::{JsonWebKeySet, JwtHeader};
+use crate::JwtHeader;
 
 #[async_trait]
 pub trait SigningKeystore {
-    fn supported_algorithms(self) -> HashSet<JsonWebSignatureAlg>;
+    fn supported_algorithms(&self) -> HashSet<JsonWebSignatureAlg>;
 
-    async fn prepare_header(self, alg: JsonWebSignatureAlg) -> anyhow::Result<JwtHeader>;
+    async fn prepare_header(&self, alg: JsonWebSignatureAlg) -> anyhow::Result<JwtHeader>;
 
-    async fn sign(self, header: &JwtHeader, msg: &[u8]) -> anyhow::Result<Vec<u8>>;
+    async fn sign(&self, header: &JwtHeader, msg: &[u8]) -> anyhow::Result<Vec<u8>>;
 }
 
-#[async_trait]
 pub trait VerifyingKeystore {
-    async fn verify(self, header: &JwtHeader, msg: &[u8], signature: &[u8]) -> anyhow::Result<()>;
+    type Error;
+    type Future: Future<Output = Result<(), Self::Error>>;
+
+    fn verify(&self, header: &JwtHeader, msg: &[u8], signature: &[u8]) -> Self::Future;
 }
 
-#[async_trait]
-pub trait ExportJwks {
-    async fn export_jwks(&self) -> anyhow::Result<JsonWebKeySet>;
+#[derive(Debug, Error)]
+pub enum EitherError<A, B> {
+    #[error(transparent)]
+    Left(A),
+    #[error(transparent)]
+    Right(B),
+}
+
+impl<L, R> VerifyingKeystore for Either<L, R>
+where
+    L: VerifyingKeystore,
+    R: VerifyingKeystore,
+{
+    type Error = EitherError<L::Error, R::Error>;
+
+    #[allow(clippy::type_complexity)]
+    type Future = Either<
+        MapErr<L::Future, fn(L::Error) -> Self::Error>,
+        MapErr<R::Future, fn(R::Error) -> Self::Error>,
+    >;
+
+    fn verify(&self, header: &JwtHeader, msg: &[u8], signature: &[u8]) -> Self::Future {
+        match self {
+            Either::Left(left) => Either::Left(
+                left.verify(header, msg, signature)
+                    .map_err(EitherError::Left),
+            ),
+            Either::Right(right) => Either::Right(
+                right
+                    .verify(header, msg, signature)
+                    .map_err(EitherError::Right),
+            ),
+        }
+    }
+}
+
+impl<T> VerifyingKeystore for Arc<T>
+where
+    T: VerifyingKeystore,
+{
+    type Error = T::Error;
+    type Future = T::Future;
+
+    fn verify(&self, header: &JwtHeader, msg: &[u8], signature: &[u8]) -> Self::Future {
+        self.as_ref().verify(header, msg, signature)
+    }
 }

crates/jose/src/lib.rs

@@ -26,7 +26,7 @@ pub use self::{
     jwk::{JsonWebKey, JsonWebKeySet},
     jwt::{DecodedJsonWebToken, JsonWebTokenParts, JwtHeader},
     keystore::{
-        ExportJwks, JwksStore, SharedSecret, SigningKeystore, StaticJwksStore, StaticKeystore,
+        DynamicJwksStore, SharedSecret, SigningKeystore, StaticJwksStore, StaticKeystore,
         VerifyingKeystore,
     },
 };