Pass custom http client through for the KeyStoreManager and jwks disc…

…overy
cduvray · Jun 10, 2024 · 6bd1ef7 · 6bd1ef7
1 parent d87b7ef
commit 6bd1ef7
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 16 deletions.
diff --git a/jwt-authorizer/src/authorizer.rs b/jwt-authorizer/src/authorizer.rs
@@ -56,7 +56,7 @@ where
         refresh: Option<Refresh>,
         validation: crate::validation::Validation,
         jwt_source: JwtSource,
-        http_client: Option<Client>,
+        http_client: Client,
     ) -> Result<Authorizer<C>, InitError> {
         Ok(match key_source_type {
             KeySourceType::RSA(path) => {
@@ -201,7 +201,7 @@ where
             }
             KeySourceType::Jwks(url) => {
                 let jwks_url = Url::parse(url.as_str()).map_err(|e| InitError::JwksUrlError(e.to_string()))?;
-                let key_store_manager = KeyStoreManager::new(jwks_url, refresh.unwrap_or_default());
+                let key_store_manager = KeyStoreManager::new(http_client, jwks_url, refresh.unwrap_or_default());
                 Authorizer {
                     key_source: KeySource::KeyStoreSource(key_store_manager),
                     claims_checker,
@@ -210,10 +210,10 @@ where
                 }
             }
             KeySourceType::Discovery(issuer_url) => {
-                let jwks_url = Url::parse(&oidc::discover_jwks(issuer_url.as_str(), http_client).await?)
+                let jwks_url = Url::parse(&oidc::discover_jwks(issuer_url.as_str(), &http_client).await?)
                     .map_err(|e| InitError::JwksUrlError(e.to_string()))?;
 
-                let key_store_manager = KeyStoreManager::new(jwks_url, refresh.unwrap_or_default());
+                let key_store_manager = KeyStoreManager::new(http_client, jwks_url, refresh.unwrap_or_default());
                 Authorizer {
                     key_source: KeySource::KeyStoreSource(key_store_manager),
                     claims_checker,

diff --git a/jwt-authorizer/src/builder.rs b/jwt-authorizer/src/builder.rs
@@ -233,7 +233,7 @@ where
                 self.refresh,
                 val,
                 self.jwt_source,
-                None,
+                self.http_client.unwrap_or_default(),
             )
             .await?,
         );
@@ -249,7 +249,7 @@ where
             self.refresh,
             val,
             self.jwt_source,
-            self.http_client,
+            self.http_client.unwrap_or_default(),
         )
         .await
     }

diff --git a/jwt-authorizer/src/jwks/key_store_manager.rs b/jwt-authorizer/src/jwks/key_store_manager.rs
@@ -1,5 +1,5 @@
 use jsonwebtoken::{jwk::JwkSet, Algorithm};
-use reqwest::Url;
+use reqwest::{Client, Url};
 use std::{
     sync::Arc,
     time::{Duration, Instant},
@@ -51,6 +51,7 @@ impl Default for Refresh {
 
 #[derive(Clone)]
 pub struct KeyStoreManager {
+    http_client: Client,
     key_url: Url,
     /// in case of fail loading (error or key not found), minimal interval
     refresh: Refresh,
@@ -67,8 +68,9 @@ pub struct KeyStore {
 }
 
 impl KeyStoreManager {
-    pub(crate) fn new(key_url: Url, refresh: Refresh) -> KeyStoreManager {
+    pub(crate) fn new(http_client: Client, key_url: Url, refresh: Refresh) -> KeyStoreManager {
         KeyStoreManager {
+            http_client,
             key_url,
             refresh,
             keystore: Arc::new(Mutex::new(KeyStore {
@@ -85,7 +87,7 @@ impl KeyStoreManager {
         let key = match self.refresh.strategy {
             RefreshStrategy::Interval => {
                 if ks_gard.can_refresh(self.refresh.refresh_interval, self.refresh.retry_interval) {
-                    ks_gard.refresh(&self.key_url, &[]).await?;
+                    ks_gard.refresh(&self.http_client, &self.key_url, &[]).await?;
                 }
                 ks_gard.get_key(header)?
             }
@@ -95,7 +97,7 @@ impl KeyStoreManager {
                     if let Some(jwk) = jwk_opt {
                         jwk
                     } else if ks_gard.can_refresh(self.refresh.refresh_interval, self.refresh.retry_interval) {
-                        ks_gard.refresh(&self.key_url, &[("kid", kid)]).await?;
+                        ks_gard.refresh(&self.http_client, &self.key_url, &[("kid", kid)]).await?;
                         ks_gard.find_kid(kid).ok_or_else(|| AuthError::InvalidKid(kid.to_owned()))?
                     } else {
                         return Err(AuthError::InvalidKid(kid.to_owned()));
@@ -107,6 +109,7 @@ impl KeyStoreManager {
                     } else if ks_gard.can_refresh(self.refresh.refresh_interval, self.refresh.retry_interval) {
                         ks_gard
                             .refresh(
+                                &self.http_client,
                                 &self.key_url,
                                 &[(
                                     "alg",
@@ -127,7 +130,7 @@ impl KeyStoreManager {
                     // if jwks endpoint is down for the loading, respect retry_interval
                     && ks_gard.can_refresh(self.refresh.refresh_interval, self.refresh.retry_interval)
                 {
-                    ks_gard.refresh(&self.key_url, &[]).await?;
+                    ks_gard.refresh(&self.http_client, &self.key_url, &[]).await?;
                 }
                 ks_gard.get_key(header)?
             }
@@ -151,8 +154,8 @@ impl KeyStore {
         }
     }
 
-    async fn refresh(&mut self, key_url: &Url, qparam: &[(&str, &str)]) -> Result<(), AuthError> {
-        reqwest::Client::new()
+    async fn refresh(&mut self, http_client: &Client, key_url: &Url, qparam: &[(&str, &str)]) -> Result<(), AuthError> {
+        http_client
             .get(key_url.as_ref())
             .query(qparam)
             .send()

diff --git a/jwt-authorizer/src/oidc.rs b/jwt-authorizer/src/oidc.rs
@@ -20,9 +20,7 @@ fn discovery_url(issuer: &str) -> Result<Url, InitError> {
     Ok(url)
 }
 
-pub async fn discover_jwks(issuer: &str, client: Option<Client>) -> Result<String, InitError> {
-    let client = client.unwrap_or_default();
-
+pub async fn discover_jwks(issuer: &str, client: &Client) -> Result<String, InitError> {
     client
         .get(discovery_url(issuer)?)
         .send()