From c1fb6b5d087a8799fa319691946713293c61b837 Mon Sep 17 00:00:00 2001
From: dacozai <eden.chen@bath.edu>
Date: Mon, 14 Oct 2024 17:00:00 +0800
Subject: [PATCH 1/3] add with_audience method with the custom service account

---
 src/custom_service_account.rs | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/src/custom_service_account.rs b/src/custom_service_account.rs
index 8304c93..7149930 100644
--- a/src/custom_service_account.rs
+++ b/src/custom_service_account.rs
@@ -31,6 +31,7 @@ pub struct CustomServiceAccount {
     signer: Signer,
     tokens: RwLock<HashMap<Vec<String>, Arc<Token>>>,
     subject: Option<String>,
+    audience: Option<String>,
 }
 
 impl CustomServiceAccount {
@@ -59,6 +60,12 @@ impl CustomServiceAccount {
         self
     }
 
+    /// Set the `Audience` to impersonate a user
+    pub fn with_audience(mut self, audience: String) -> Self {
+        self.audience = Some(audience);
+        self
+    }
+
     fn new(credentials: ServiceAccountKey, client: HttpClient) -> Result<Self, Error> {
         debug!(project = ?credentials.project_id, email = credentials.client_email, "found credentials");
         Ok(Self {
@@ -67,13 +74,14 @@ impl CustomServiceAccount {
             credentials,
             tokens: RwLock::new(HashMap::new()),
             subject: None,
+            audience: None,
         })
     }
 
     #[instrument(level = Level::DEBUG, skip(self))]
     async fn fetch_token(&self, scopes: &[&str]) -> Result<Arc<Token>, Error> {
         let jwt =
-            Claims::new(&self.credentials, scopes, self.subject.as_deref()).to_jwt(&self.signer)?;
+            Claims::new(&self.credentials, scopes, self.subject.as_deref(), self.audience.as_deref()).to_jwt(&self.signer)?;
         let body = Bytes::from(
             form_urlencoded::Serializer::new(String::new())
                 .extend_pairs(&[("grant_type", GRANT_TYPE), ("assertion", jwt.as_str())])
@@ -156,7 +164,7 @@ pub(crate) struct Claims<'a> {
 }
 
 impl<'a> Claims<'a> {
-    pub(crate) fn new(key: &'a ServiceAccountKey, scopes: &[&str], sub: Option<&'a str>) -> Self {
+    pub(crate) fn new(key: &'a ServiceAccountKey, scopes: &[&str], sub: Option<&'a str>, aud: Option<&'a str>) -> Self {
         let mut scope = String::with_capacity(16);
         for (i, s) in scopes.iter().enumerate() {
             if i != 0 {
@@ -167,9 +175,14 @@ impl<'a> Claims<'a> {
         }
 
         let iat = Utc::now().timestamp();
+        let aud = if let Some(aud_str) = aud {
+            aud_str
+        } else {
+            &key.token_uri
+        };
         Claims {
             iss: &key.client_email,
-            aud: &key.token_uri,
+            aud,
             exp: iat + 3600 - 5, // Max validity is 1h
             iat,
             sub,

From c5ee02818410d6b65ee072f74b16d4be42f8f573 Mon Sep 17 00:00:00 2001
From: dacozai <eden.chen@bath.edu>
Date: Mon, 14 Oct 2024 17:25:13 +0800
Subject: [PATCH 2/3] fix linter

---
 src/custom_service_account.rs | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/src/custom_service_account.rs b/src/custom_service_account.rs
index 7149930..00fac31 100644
--- a/src/custom_service_account.rs
+++ b/src/custom_service_account.rs
@@ -80,8 +80,13 @@ impl CustomServiceAccount {
 
     #[instrument(level = Level::DEBUG, skip(self))]
     async fn fetch_token(&self, scopes: &[&str]) -> Result<Arc<Token>, Error> {
-        let jwt =
-            Claims::new(&self.credentials, scopes, self.subject.as_deref(), self.audience.as_deref()).to_jwt(&self.signer)?;
+        let jwt = Claims::new(
+            &self.credentials,
+            scopes,
+            self.subject.as_deref(),
+            self.audience.as_deref(),
+        )
+        .to_jwt(&self.signer)?;
         let body = Bytes::from(
             form_urlencoded::Serializer::new(String::new())
                 .extend_pairs(&[("grant_type", GRANT_TYPE), ("assertion", jwt.as_str())])
@@ -164,7 +169,12 @@ pub(crate) struct Claims<'a> {
 }
 
 impl<'a> Claims<'a> {
-    pub(crate) fn new(key: &'a ServiceAccountKey, scopes: &[&str], sub: Option<&'a str>, aud: Option<&'a str>) -> Self {
+    pub(crate) fn new(
+        key: &'a ServiceAccountKey,
+        scopes: &[&str],
+        sub: Option<&'a str>,
+        aud: Option<&'a str>,
+    ) -> Self {
         let mut scope = String::with_capacity(16);
         for (i, s) in scopes.iter().enumerate() {
             if i != 0 {

From 347cdfd7f2c111b1569a471a60feaa887a09c5c7 Mon Sep 17 00:00:00 2001
From: dacozai <eden.chen@bath.edu>
Date: Mon, 14 Oct 2024 17:26:39 +0800
Subject: [PATCH 3/3] update the audience to inline style

---
 src/custom_service_account.rs | 7 +------
 1 file changed, 1 insertion(+), 6 deletions(-)

diff --git a/src/custom_service_account.rs b/src/custom_service_account.rs
index 00fac31..ec61020 100644
--- a/src/custom_service_account.rs
+++ b/src/custom_service_account.rs
@@ -185,14 +185,9 @@ impl<'a> Claims<'a> {
         }
 
         let iat = Utc::now().timestamp();
-        let aud = if let Some(aud_str) = aud {
-            aud_str
-        } else {
-            &key.token_uri
-        };
         Claims {
             iss: &key.client_email,
-            aud,
+            aud: aud.unwrap_or(&key.token_uri),
             exp: iat + 3600 - 5, // Max validity is 1h
             iat,
             sub,