feat: use jwks for external jwt auth (#4089)

* feat: use jwks for external jwts

* fix: build

* update ee ref

* fix: nit error msg

backend/ee-repo-ref.txt

@@ -1 +1 @@
-75e2aa1b2a8cff3cdcd64bab33eadab036e0c3fe
+ba9196530fafbfe017d2592af2830a60b6fa62a0

backend/windmill-api/src/ee.rs

@@ -1,4 +1,8 @@
 use anyhow::anyhow;
+#[cfg(feature = "enterprise")]
+use std::sync::Arc;
+#[cfg(feature = "enterprise")]
+use tokio::sync::RwLock;
 
 pub async fn validate_license_key(_license_key: String) -> anyhow::Result<String> {
     // Implementation is not open source
@@ -9,8 +13,20 @@ pub async fn validate_license_key(_license_key: String) -> anyhow::Result<String
 pub async fn jwt_ext_auth(
     _w_id: Option<&String>,
     _token: &str,
-) -> Option<(crate::db::ApiAuthed, usize)> {
+    _external_jwks: Option<Arc<RwLock<ExternalJwks>>>,
+) -> anyhow::Result<(crate::db::ApiAuthed, usize)> {
     // Implementation is not open source
 
-    None
+    Err(anyhow!("External JWT auth is not open source"))
+}
+
+#[cfg(feature = "enterprise")]
+pub struct ExternalJwks;
+
+#[cfg(feature = "enterprise")]
+impl ExternalJwks {
+    pub async fn load() -> Option<Arc<RwLock<Self>>> {
+        // Implementation is not open source
+        None
+    }
 }

backend/windmill-api/src/lib.rs

@@ -7,6 +7,8 @@
  */
 
 use crate::db::ApiAuthed;
+#[cfg(feature = "enterprise")]
+use crate::ee::ExternalJwks;
 #[cfg(feature = "embedding")]
 use crate::embeddings::load_embeddings_db;
 use crate::oauth2_ee::AllClients;
@@ -164,9 +166,13 @@ pub async fn run_server(
     }
     let user_db = UserDB::new(db.clone());
 
+    #[cfg(feature = "enterprise")]
+    let ext_jwks = ExternalJwks::load().await;
     let auth_cache = Arc::new(users::AuthCache::new(
         db.clone(),
         std::env::var("SUPERADMIN_SECRET").ok(),
+        #[cfg(feature = "enterprise")]
+        ext_jwks,
     ));
     let argon2 = Arc::new(Argon2::default());
 

backend/windmill-api/src/users.rs

@@ -12,6 +12,8 @@ use std::sync::Arc;
 
 use crate::db::ApiAuthed;
 
+#[cfg(feature = "enterprise")]
+use crate::ee::ExternalJwks;
 use crate::oauth2_ee::{check_nb_of_user, InstanceEvent};
 use crate::utils::{
     generate_instance_wide_unique_username, get_and_delete_pending_username_or_generate,
@@ -39,6 +41,8 @@ use regex::Regex;
 use serde::{Deserialize, Serialize};
 use sqlx::FromRow;
 use time::OffsetDateTime;
+#[cfg(feature = "enterprise")]
+use tokio::sync::RwLock;
 use tower_cookies::{Cookie, Cookies};
 use tracing::{Instrument, Span};
 use windmill_audit::audit_ee::{audit_log, AuditAuthor};
@@ -145,11 +149,23 @@ pub struct AuthCache {
     cache: Cache<(String, String), ExpiringAuthCache>,
     db: DB,
     superadmin_secret: Option<String>,
+    #[cfg(feature = "enterprise")]
+    ext_jwks: Option<Arc<RwLock<ExternalJwks>>>,
 }
 
 impl AuthCache {
-    pub fn new(db: DB, superadmin_secret: Option<String>) -> Self {
-        AuthCache { cache: Cache::new(300), db, superadmin_secret }
+    pub fn new(
+        db: DB,
+        superadmin_secret: Option<String>,
+        #[cfg(feature = "enterprise")] ext_jwks: Option<Arc<RwLock<ExternalJwks>>>,
+    ) -> Self {
+        AuthCache {
+            cache: Cache::new(300),
+            db,
+            superadmin_secret,
+            #[cfg(feature = "enterprise")]
+            ext_jwks,
+        }
     }
 
     pub async fn invalidate(&self, w_id: &str, token: String) {
@@ -168,9 +184,19 @@ impl AuthCache {
             }
             #[cfg(feature = "enterprise")]
             _ if token.starts_with("jwt_ext_") => {
-                let authed_and_exp =
-                    crate::ee::jwt_ext_auth(w_id.as_ref(), token.trim_start_matches("jwt_ext_"))
-                        .await;
+                let authed_and_exp = match crate::ee::jwt_ext_auth(
+                    w_id.as_ref(),
+                    token.trim_start_matches("jwt_ext_"),
+                    self.ext_jwks.clone(),
+                )
+                .await
+                {
+                    Ok(r) => Some(r),
+                    Err(e) => {
+                        tracing::error!("JWT_EXT auth error: {:?}", e);
+                        None
+                    }
+                };
 
                 if let Some((authed, exp)) = authed_and_exp.clone() {
                     self.cache.insert(