From 7239288034716865048546a486eba7b8c5723fc3 Mon Sep 17 00:00:00 2001
From: congyi <15605187270@163.com>
Date: Fri, 7 Jun 2024 15:59:48 +0800
Subject: [PATCH] load token from token file

---
 src/azure/storage/config.rs                     | 17 ++++++-----------
 .../storage/workload_identity_credential.rs     | 14 ++++++++------
 2 files changed, 14 insertions(+), 17 deletions(-)

diff --git a/src/azure/storage/config.rs b/src/azure/storage/config.rs
index f962804..2840c29 100644
--- a/src/azure/storage/config.rs
+++ b/src/azure/storage/config.rs
@@ -1,5 +1,5 @@
 use std::env;
-use std::{collections::HashMap, fs};
+use std::{collections::HashMap};
 
 /// Config carries all the configuration for Azure Storage services.
 #[derive(Clone, Default)]
@@ -48,12 +48,12 @@ pub struct Config {
     ///
     /// This is part of use AAD(Azure Active Directory) authenticate on Azure VM
     pub endpoint: Option<String>,
-    /// `federated_token` value will be loaded from:
+    /// `federated_token_file` value will be loaded from:
     ///
     /// - this field if it's `is_some`
-    /// - env value: [`AZURE_FEDERATED_TOKEN`]
-    /// - profile config: `federated_toen_file`
-    pub federated_token: Option<String>,
+    /// - env value: [`AZURE_FEDERATED_TOKEN_FILE`]
+    /// - profile config: `federated_token_file`
+    pub federated_token_file: Option<String>,
     /// `tenant_id` value will be loaded from:
     ///
     /// - this field if it's `is_some`
@@ -68,7 +68,6 @@ pub struct Config {
     pub authority_host: Option<String>,
 }
 
-pub const AZURE_FEDERATED_TOKEN: &str = "AZURE_FEDERATED_TOKEN";
 pub const AZURE_FEDERATED_TOKEN_FILE: &str = "AZURE_FEDERATED_TOKEN_FILE";
 pub const AZURE_TENANT_ID: &str = "AZURE_TENANT_ID";
 pub const AZURE_CLIENT_ID: &str = "AZURE_CLIENT_ID";
@@ -85,11 +84,7 @@ impl Config {
 
         // federated_token can be loaded from both `AZURE_FEDERATED_TOKEN` and `AZURE_FEDERATED_TOKEN_FILE`.
         if let Some(v) = envs.get(AZURE_FEDERATED_TOKEN_FILE) {
-            self.federated_token = Some(fs::read_to_string(v).unwrap_or_default());
-        }
-
-        if let Some(v) = envs.get(AZURE_FEDERATED_TOKEN) {
-            self.federated_token = Some(v.to_string());
+            self.federated_token_file = Some(v.to_string());
         }
 
         if let Some(v) = envs.get(AZURE_TENANT_ID) {
diff --git a/src/azure/storage/workload_identity_credential.rs b/src/azure/storage/workload_identity_credential.rs
index 844690a..a207f2f 100644
--- a/src/azure/storage/workload_identity_credential.rs
+++ b/src/azure/storage/workload_identity_credential.rs
@@ -1,4 +1,4 @@
-use std::str;
+use std::{fs, str};
 
 use http::HeaderValue;
 use http::Method;
@@ -15,17 +15,19 @@ const STORAGE_TOKEN_SCOPE: &str = "https://storage.azure.com/.default";
 ///
 /// See <https://learn.microsoft.com/en-us/azure/app-service/overview-managed-identity?tabs=portal,http#using-the-rest-protocol>
 pub async fn get_workload_identity_token(config: &Config) -> anyhow::Result<Option<LoginResponse>> {
-    let (token, tenant_id, client_id, authority_host) = match (
-        &config.federated_token,
+    let (token_file, tenant_id, client_id, authority_host) = match (
+        &config.federated_token_file,
         &config.tenant_id,
         &config.client_id,
         &config.authority_host,
     ) {
-        (Some(token), Some(tenant_id), Some(client_id), Some(authority_host)) => {
-            (token, tenant_id, client_id, authority_host)
+        (Some(token_file), Some(tenant_id), Some(client_id), Some(authority_host)) => {
+            (token_file, tenant_id, client_id, authority_host)
         }
         _ => return Ok(None),
     };
+
+    let token = fs::read_to_string(token_file)?;
     let url = Url::parse(authority_host)?.join(&format!("/{tenant_id}/oauth2/v2.0/token"))?;
     let scopes: &[&str] = &[STORAGE_TOKEN_SCOPE];
     let encoded_body: String = form_urlencoded::Serializer::new(String::new())
@@ -35,7 +37,7 @@ pub async fn get_workload_identity_token(config: &Config) -> anyhow::Result<Opti
             "client_assertion_type",
             "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",
         )
-        .append_pair("client_assertion", token)
+        .append_pair("client_assertion", &token)
         .append_pair("grant_type", "client_credentials")
         .finish();