rust/neon
1
0
Fork 0
mirror of https://github.com/neondatabase/neon.git synced 2026-10-17 09:12:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

extract fetch/update routine

Browse Source
This commit is contained in:
Conrad Ludgate
2024-08-12 09:01:04 +01:00
parent c6b36d8171
commit 20fdf3e19f
1 changed files with 44 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
proxy/src/auth/backend/jwt.rs
+44 -36
View File
@@ -134,6 +134,47 @@ impl JWKCacheEntryLock {
Ok(x)
}
async fn get_or_update_jwk_cache<F: FetchAuthRules>(
self: &Arc<Self>,
client: &reqwest::Client,
fetch: &F,
) -> Result<Arc<JWKCacheEntry>, anyhow::Error> {
let now = Instant::now();
let guard = self.cached.load_full();
// if we have no cached JWKs, try and get some
let Some(cached) = guard else {
let permit = self.acquire_permit().await;
return self.renew_jwks(permit, client, fetch).await;
};
let last_update = now.duration_since(cached.last_retrieved);
// check if the cached JWKs need updating.
if last_update > MAX_RENEW {
let permit = self.acquire_permit().await;
// it's been too long since we checked the keys. wait for them to update.
return self.renew_jwks(permit, client, fetch).await;
}
// every 5 minutes we should spawn a job to eagerly update the token.
if last_update > AUTO_RENEW {
if let Some(permit) = self.try_acquire_permit() {
let entry = self.clone();
let client = client.clone();
let fetch = fetch.clone();
tokio::spawn(async move {
if let Err(e) = entry.renew_jwks(permit, &client, &fetch).await {
tracing::warn!(error=?e, "could not fetch JWKs in background job");
}
});
}
}
Ok(cached)
}
async fn check_jwt<F: FetchAuthRules>(
self: &Arc<Self>,
jwt: String,
@@ -159,41 +200,7 @@ impl JWKCacheEntryLock {
ensure!(header.typ == "JWT");
let kid = header.kid.context("missing key id")?;
// check if the cached JWKs need updating.
let mut guard = {
let now = Instant::now();
let guard = self.cached.load_full();
if let Some(cached) = guard {
let last_update = now.duration_since(cached.last_retrieved);
if last_update > MAX_RENEW {
let permit = self.acquire_permit().await;
// it's been too long since we checked the keys. wait for them to update.
self.renew_jwks(permit, client, fetch).await?
} else if last_update > MIN_RENEW {
// every 5 minutes we should spawn a job to eagerly update the token.
if let Some(permit) = self.try_acquire_permit() {
let entry = self.clone();
let client = client.clone();
let fetch = fetch.clone();
tokio::spawn(async move {
if let Err(e) = entry.renew_jwks(permit, &client, &fetch).await {
tracing::warn!(error=?e, "could not fetch JWKs in background job");
}
});
}
cached
} else {
cached
}
} else {
let permit = self.acquire_permit().await;
self.renew_jwks(permit, client, fetch).await?
}
};
let mut guard = self.get_or_update_jwk_cache(client, fetch).await?;
// get the key from the JWKs if possible. If not, wait for the keys to update.
let jwk = loop {
@@ -238,7 +245,8 @@ impl JWKCacheEntryLock {
}
}
const MIN_RENEW: Duration = Duration::from_secs(300);
const MIN_RENEW: Duration = Duration::from_secs(30);
const AUTO_RENEW: Duration = Duration::from_secs(300);
const MAX_RENEW: Duration = Duration::from_secs(3600);
impl JWKCache {