From 4c7b9bb771d4c6255559ebb871e26d9c227e400c Mon Sep 17 00:00:00 2001
From: Mikko Ylinen <mikko.ylinen@intel.com>
Date: Thu, 8 Aug 2024 15:18:42 +0300
Subject: [PATCH] kbs: token: add verifier with JWKS from OpenID configuration

Signed-off-by: Mikko Ylinen <mikko.ylinen@intel.com>
---
 kbs/Cargo.toml        |  2 +-
 kbs/src/token/mod.rs  | 10 +++++-
 kbs/src/token/oidc.rs | 75 +++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 85 insertions(+), 2 deletions(-)
 create mode 100644 kbs/src/token/oidc.rs

diff --git a/kbs/Cargo.toml b/kbs/Cargo.toml
index a983769c56..e61283c314 100644
--- a/kbs/Cargo.toml
+++ b/kbs/Cargo.toml
@@ -10,7 +10,7 @@ edition.workspace = true
 default = ["coco-as-builtin", "resource", "opa", "rustls"]
 
 # Feature that allows to access resources from KBS
-resource = ["rsa", "dep:openssl", "reqwest", "aes-gcm"]
+resource = ["rsa", "dep:openssl", "reqwest", "aes-gcm", "jsonwebtoken"]
 
 # Support a backend attestation service for KBS
 as = []
diff --git a/kbs/src/token/mod.rs b/kbs/src/token/mod.rs
index 99674c2a46..3963a7988a 100644
--- a/kbs/src/token/mod.rs
+++ b/kbs/src/token/mod.rs
@@ -10,6 +10,7 @@ use strum::EnumString;
 use tokio::sync::RwLock;
 
 mod coco;
+mod oidc;
 
 #[async_trait]
 pub trait AttestationTokenVerifier {
@@ -21,13 +22,16 @@ pub trait AttestationTokenVerifier {
 #[derive(Deserialize, Debug, Clone, EnumString)]
 pub enum AttestationTokenVerifierType {
     CoCo,
+    Oidc,
 }
 
 #[derive(Deserialize, Debug, Clone)]
 pub struct AttestationTokenVerifierConfig {
     pub attestation_token_type: AttestationTokenVerifierType,
 
-    // Trusted Certificates file (PEM format) path to verify Attestation Token Signature.
+    // Trusted Certificates file (PEM format) path or OpenID configuration URLs
+    // that provide "jwks_uri" pointing to JWKS signing certificates to verify
+    // Attestation Token Signature.
     pub trusted_certs_paths: Option<Vec<String>>,
 }
 
@@ -48,5 +52,9 @@ pub async fn create_token_verifier(
             coco::CoCoAttestationTokenVerifier::new(&config)?,
         ))
             as Arc<RwLock<dyn AttestationTokenVerifier + Send + Sync>>),
+        AttestationTokenVerifierType::Oidc => Ok(Arc::new(RwLock::new(
+            oidc::OidcAttestationTokenVerifier::new(&config).await?,
+        ))
+            as Arc<RwLock<dyn AttestationTokenVerifier + Send + Sync>>),
     }
 }
diff --git a/kbs/src/token/oidc.rs b/kbs/src/token/oidc.rs
new file mode 100644
index 0000000000..3361c77f9f
--- /dev/null
+++ b/kbs/src/token/oidc.rs
@@ -0,0 +1,75 @@
+// Copyright (c) 2024 by Intel Corporation
+// Licensed under the Apache License, Version 2.0, see LICENSE for details.
+// SPDX-License-Identifier: Apache-2.0
+
+use crate::token::{AttestationTokenVerifier, AttestationTokenVerifierConfig};
+use anyhow::*;
+use async_trait::async_trait;
+use jsonwebtoken::{decode, decode_header, jwk, Algorithm, DecodingKey, Validation};
+use reqwest::{get, Url};
+use serde_json::Value;
+use std::str::FromStr;
+
+const OPENID_CONFIG_URL_SUFFIX: &str = ".well-known/openid-configuration";
+
+pub struct OidcAttestationTokenVerifier {
+    trusted_certs: Option<jwk::JwkSet>,
+}
+
+impl OidcAttestationTokenVerifier {
+    pub async fn new(config: &AttestationTokenVerifierConfig) -> Result<Self> {
+        let trusted_certs = match &config.trusted_certs_paths {
+            Some(paths) => {
+                let mut keyset = jwk::JwkSet { keys: Vec::new() };
+
+                for url in paths.iter() {
+                    // TODO: accept both "local" jwks and OIDC ones
+                    let oidc_url = Url::parse(url)?.join(OPENID_CONFIG_URL_SUFFIX)?;
+                    let oidc_values = get(oidc_url).await?.json::<Value>().await?;
+
+                    let jwks_uri = oidc_values["jwks_uri"].as_str().ok_or(anyhow!(
+                        "Failed to parse jwks uri from OpenID Configuration"
+                    ))?;
+
+                    let jwkset = get(jwks_uri).await?.json::<jwk::JwkSet>().await?;
+
+                    for jwk in jwkset.keys.iter() {
+                        keyset.keys.push(jwk.clone());
+                    }
+                }
+                Some(keyset)
+            }
+            None => None,
+        };
+
+        Ok(Self { trusted_certs })
+    }
+}
+
+#[async_trait]
+impl AttestationTokenVerifier for OidcAttestationTokenVerifier {
+    async fn verify(&self, token: String) -> Result<String> {
+        let header = decode_header(&token).context("Failed to decode attestation token header")?;
+
+        let Some(keyset) = &self.trusted_certs else {
+            bail!("missing config");
+        };
+
+        let kid = header
+            .kid
+            .ok_or(anyhow!("Failed to decode kid in the token header"))?;
+        let key = keyset
+            .find(&kid)
+            .ok_or(anyhow!("Failed to find kid in trusted certificates"))?;
+
+        let alg = Algorithm::from_str(key.common.key_algorithm.unwrap().to_string().as_str())?;
+
+        let dkey = DecodingKey::from_jwk(key)?;
+        let token_data = decode::<Value>(&token, &dkey, &Validation::new(alg))
+            .context("Failed to decode attestation token")?;
+
+        Ok(serde_json::to_string(&token_data.claims)?)
+    }
+}
+
+// TODO: tests