fix: token validation

src/auth.rs

@@ -10,10 +10,17 @@ pub struct Claims {
     pub exp: usize,
     pub iat: usize,
     pub iss: String,
-    pub aud: String,
+    pub aud: Audience,
 }
 
-#[derive(Debug, Deserialize)]
+#[derive(Debug, Serialize, Deserialize)]
+#[serde(untagged)]
+pub enum Audience {
+    Single(String),
+    Multiple(Vec<String>),
+}
+
+#[derive(Debug, Deserialize, Clone)]
 struct Jwk {
     #[serde(rename = "kty")]
     _kty: String,
@@ -31,13 +38,17 @@ struct Jwks {
 
 pub struct JwksVerifier {
     issuer: String,
+    audience: String,
     jwks_uri: String,
     keys: Arc<RwLock<Vec<Jwk>>>,
     client: Client,
 }
 
 impl JwksVerifier {
-    pub async fn new(issuer: String) -> Result<Self, Box<dyn std::error::Error>> {
+    pub async fn new(
+        issuer: String,
+        audience: String,
+    ) -> Result<Self, Box<dyn std::error::Error>> {
         let client = Client::new();
         // Authentik OIDC discovery
         let discovery_url = format!(
@@ -53,6 +64,7 @@ impl JwksVerifier {
 
         let verifier = Self {
             issuer,
+            audience,
             jwks_uri,
             keys: Arc::new(RwLock::new(Vec::new())),
             client,
@@ -73,18 +85,29 @@ impl JwksVerifier {
         let header = decode_header(token)?;
         let kid = header.kid.ok_or("Missing kid in token header")?;
 
-        let keys = self.keys.read().await;
-        let jwk = keys
-            .iter()
-            .find(|k| k.kid == kid)
-            .ok_or("Key not found in JWKS")?;
+        let jwk = {
+            let keys = self.keys.read().await;
+            keys.iter().find(|k| k.kid == kid).cloned()
+        };
+
+        let jwk = match jwk {
+            Some(key) => key,
+            None => {
+                self.refresh_keys().await?;
+                let keys = self.keys.read().await;
+                keys.iter()
+                    .find(|k| k.kid == kid)
+                    .cloned()
+                    .ok_or("Key not found in JWKS")?
+            }
+        };
 
         let decoding_key = DecodingKey::from_rsa_components(&jwk.n, &jwk.e)?;
 
         let mut validation = Validation::new(Algorithm::RS256);
         validation.set_issuer(&[self.issuer.clone()]);
-        // Aud validation might need careful config, usually it's the client_id
-        validation.validate_aud = false;
+        validation.set_audience(&[self.audience.clone()]);
+        validation.validate_aud = true;
 
         let token_data = decode::<Claims>(token, &decoding_key, &validation)?;
         Ok(token_data.claims)