From fa92328423645df3f1a4ad7bd5019312a55a267a Mon Sep 17 00:00:00 2001
From: Conrad Ludgate <conradludgate@gmail.com>
Date: Tue, 30 Jul 2024 14:23:40 +0100
Subject: [PATCH] start stubbing jwt

---
 Cargo.lock                    |   1 +
 proxy/Cargo.toml              |   1 +
 proxy/src/auth/backend.rs     |   1 +
 proxy/src/auth/backend/jwt.rs | 113 ++++++++++++++++++++++++++++++++++
 4 files changed, 116 insertions(+)
 create mode 100644 proxy/src/auth/backend/jwt.rs

diff --git a/Cargo.lock b/Cargo.lock
index 031fae0f37..585cf33183 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4270,6 +4270,7 @@ dependencies = [
  "indexmap 2.0.1",
  "ipnet",
  "itertools 0.10.5",
+ "jsonwebtoken",
  "lasso",
  "md5",
  "measured",
diff --git a/proxy/Cargo.toml b/proxy/Cargo.toml
index b316c53034..0f10a0e8ee 100644
--- a/proxy/Cargo.toml
+++ b/proxy/Cargo.toml
@@ -46,6 +46,7 @@ http-body-util = { version = "0.1" }
 indexmap.workspace = true
 ipnet.workspace = true
 itertools.workspace = true
+jsonwebtoken.workspace = true
 lasso = { workspace = true, features = ["multi-threaded"] }
 md5.workspace = true
 measured = { workspace = true, features = ["lasso"] }
diff --git a/proxy/src/auth/backend.rs b/proxy/src/auth/backend.rs
index 90dea01bf3..3b8e1b53bd 100644
--- a/proxy/src/auth/backend.rs
+++ b/proxy/src/auth/backend.rs
@@ -1,6 +1,7 @@
 mod classic;
 mod hacks;
 mod link;
+mod jwt;
 
 use std::net::IpAddr;
 use std::sync::Arc;
diff --git a/proxy/src/auth/backend/jwt.rs b/proxy/src/auth/backend/jwt.rs
new file mode 100644
index 0000000000..ad66b24a5f
--- /dev/null
+++ b/proxy/src/auth/backend/jwt.rs
@@ -0,0 +1,113 @@
+use std::time::Duration;
+
+use anyhow::{bail, Context};
+use dashmap::{mapref::entry::Entry, DashMap};
+use jsonwebtoken::{jwk::JwkSet, Algorithm, DecodingKey, TokenData, Validation};
+use tokio::time::Instant;
+
+use crate::intern::EndpointIdInt;
+
+pub struct JWKCache {
+    client: reqwest::Client,
+
+    // I don't like this. jsonwebtoken crate is horrendous with the amount of string allocations.. oh well.
+    map: DashMap<EndpointIdInt, JWKCacheEntry>,
+}
+
+pub struct JWKCacheEntry {
+    /// Should refetch at least every hour to verify when old keys have been removed.
+    /// Should refetch when new key IDs are seen only every 5 minutes or so
+    last_retrieved: Instant,
+
+    /// jwks urls
+    urls: Vec<url::Url>,
+
+    /// cplane will return multiple JWKs urls that we need to scrape.
+    key_sets: Vec<JwkSet>,
+}
+
+impl JWKCache {
+    pub async fn check_jwt(
+        &self,
+        endpoint: EndpointIdInt,
+        jwt: String,
+    ) -> Result<(), anyhow::Error> {
+        let header = jsonwebtoken::decode_header(&jwt).context("could not parse JWT")?;
+        let kid = header.kid.context("missing key id")?;
+
+        let entry = match self.map.get(&endpoint) {
+            Some(entry) => entry,
+            None => match self.map.entry(endpoint) {
+                Entry::Occupied(entry) => entry.into_ref().downgrade(),
+                Entry::Vacant(_) => todo!("fetch jwks from cplane"),
+            },
+        };
+
+        let now = Instant::now();
+        if now.duration_since(entry.last_retrieved) > Duration::from_secs(3600) {
+            todo!("refetch jwks from cplane")
+            // drop(entry) - don't want to hold the lock while running long network operations
+            // fetch from cplane
+            // fetch jwks
+            // entry.last_retrieved = now;
+
+            // can we use arcswap compare_and_swap here?
+        }
+
+        let mut jwk = None;
+        for jwks in &entry.key_sets {
+            jwk = jwks.find(&kid);
+            if jwk.is_some_and(|k| k.is_supported()) {
+                break;
+            }
+        }
+
+        let key = match jwk {
+            Some(jwk) if jwk.is_supported() => DecodingKey::from_jwk(jwk).context("invalid key")?,
+            None if now.duration_since(entry.last_retrieved) > Duration::from_secs(300) => {
+                // refetch jwks
+                todo!("fetch jwks again");
+            }
+            _ => {
+                bail!("jwk not found");
+            }
+        };
+
+        let mut validation = Validation::new(Algorithm::ES256);
+        // support all public key algorithms
+        validation.algorithms = vec![
+            Algorithm::ES256,
+            Algorithm::ES384,
+            Algorithm::RS256,
+            Algorithm::RS384,
+            Algorithm::RS512,
+            Algorithm::PS256,
+            Algorithm::PS384,
+            Algorithm::PS512,
+            Algorithm::EdDSA,
+        ];
+
+        // should we require some set of common claims?
+        // eg `exp`?
+        validation.required_spec_claims.clear();
+        validation.leeway = 60;
+
+        validation.validate_exp = true;
+        validation.validate_nbf = true;
+        // we don't have an audience - I want cplane to add this field
+        validation.validate_aud = false;
+
+        // subject can be anything - we support all users
+        validation.sub = None;
+        // don't know the audience - I want cplane to add this field
+        validation.aud = None;
+        // don't know the issuer - I want cplane to add this field
+        validation.iss = None;
+
+        // ew...
+        let token: TokenData<serde_json::Value> =
+            jsonwebtoken::decode(&jwt, &key, &validation).context("invalid token")?;
+
+        Ok(())
+    }
+}