From ff65df43a1bb8ffd6bfaa3c047186ba137345a79 Mon Sep 17 00:00:00 2001
From: Hirokazu Hata <h.hata.ai.t@gmail.com>
Date: Fri, 12 Apr 2024 12:45:38 +0900
Subject: [PATCH] Set expected audience to jsonwebtoken validation option to
 get ID token (#249)

jsonwebtoken validate audience as default validation rule. But audiecen
is not set. So set "target_audience" in the claim as expected audience.

target_audiecen:
https://cloud.google.com/iap/docs/authentication-howto?hl=ja#bash

See for more details:
https://github.com/yoshidan/google-cloud-rust/issues/248
---
 foundation/auth/src/error.rs                  |  3 ++
 foundation/auth/src/token_source/mod.rs       |  7 ++--
 .../service_account_token_source.rs           | 32 ++++++++++++-------
 3 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/foundation/auth/src/error.rs b/foundation/auth/src/error.rs
index 44400854..22d66aa7 100644
--- a/foundation/auth/src/error.rs
+++ b/foundation/auth/src/error.rs
@@ -48,4 +48,7 @@ pub enum Error {
 
     #[error("unexpected impersonation token response : status={0}, detail={1}")]
     UnexpectedImpersonateTokenResponse(u16, String),
+
+    #[error("No target_audience Found in the private claims")]
+    NoTargetAudienceFound,
 }
diff --git a/foundation/auth/src/token_source/mod.rs b/foundation/auth/src/token_source/mod.rs
index fb596576..d0469b81 100644
--- a/foundation/auth/src/token_source/mod.rs
+++ b/foundation/auth/src/token_source/mod.rs
@@ -58,17 +58,18 @@ struct ExpClaim {
 }
 
 impl InternalIdToken {
-    fn to_token(&self) -> Result<Token, Error> {
+    fn to_token(&self, audience: &str) -> Result<Token, Error> {
         Ok(Token {
             access_token: self.id_token.clone(),
             token_type: "Bearer".into(),
-            expiry: time::OffsetDateTime::from_unix_timestamp(self.get_exp()?).ok(),
+            expiry: time::OffsetDateTime::from_unix_timestamp(self.get_exp(audience)?).ok(),
         })
     }
 
-    fn get_exp(&self) -> Result<i64, Error> {
+    fn get_exp(&self, audience: &str) -> Result<i64, Error> {
         let mut validation = jsonwebtoken::Validation::default();
         validation.insecure_disable_signature_validation();
+        validation.set_audience(&[audience]);
         let decoding_key = jsonwebtoken::DecodingKey::from_secret(b"");
         Ok(
             jsonwebtoken::decode::<ExpClaim>(self.id_token.as_str(), &decoding_key, &validation)?
diff --git a/foundation/auth/src/token_source/service_account_token_source.rs b/foundation/auth/src/token_source/service_account_token_source.rs
index 67fd7bec..f747ad40 100644
--- a/foundation/auth/src/token_source/service_account_token_source.rs
+++ b/foundation/auth/src/token_source/service_account_token_source.rs
@@ -176,7 +176,7 @@ impl TokenSource for OAuth2ServiceAccountTokenSource {
         let iat = OffsetDateTime::now_utc();
         let exp = iat + time::Duration::hours(1);
 
-        let request_token = Claims {
+        let claims = Claims {
             iss: self.email.as_ref(),
             sub: self.sub.as_ref().map(|s| s.as_ref()),
             scope: Some(self.scopes.as_ref()),
@@ -184,8 +184,8 @@ impl TokenSource for OAuth2ServiceAccountTokenSource {
             exp: exp.unix_timestamp(),
             iat: iat.unix_timestamp(),
             private_claims: &self.private_claims,
-        }
-        .token(&self.pk, &self.pk_id)?;
+        };
+        let request_token = claims.token(&self.pk, &self.pk_id)?;
 
         let form = [
             ("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer"),
@@ -193,15 +193,23 @@ impl TokenSource for OAuth2ServiceAccountTokenSource {
         ];
 
         match self.use_id_token {
-            true => Ok(self
-                .client
-                .post(self.token_url.as_str())
-                .form(&form)
-                .send()
-                .await?
-                .json::<InternalIdToken>()
-                .await?
-                .to_token()?),
+            true => {
+                let audience = claims
+                    .private_claims
+                    .get("target_audience")
+                    .ok_or(Error::NoTargetAudienceFound)?
+                    .as_str()
+                    .ok_or(Error::NoTargetAudienceFound)?;
+                Ok(self
+                    .client
+                    .post(self.token_url.as_str())
+                    .form(&form)
+                    .send()
+                    .await?
+                    .json::<InternalIdToken>()
+                    .await?
+                    .to_token(audience)?)
+            }
             false => Ok(self
                 .client
                 .post(self.token_url.as_str())