From 9ed568e37c29394498f374cdd02cfa520362da4a Mon Sep 17 00:00:00 2001
From: congyi wang <58715567+wcy-fdu@users.noreply.github.com>
Date: Fri, 28 Jun 2024 16:19:11 +0800
Subject: [PATCH] feat(oss): support define sts endpoint for oss (#448)

as title

---------

Co-authored-by: Xuanwo <github@xuanwo.io>
---
 src/aliyun/config.rs     | 9 +++++++++
 src/aliyun/constants.rs  | 1 +
 src/aliyun/credential.rs | 9 ++++++++-
 3 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/src/aliyun/config.rs b/src/aliyun/config.rs
index 6432e75e..db596b4d 100644
--- a/src/aliyun/config.rs
+++ b/src/aliyun/config.rs
@@ -40,6 +40,11 @@ pub struct Config {
     /// - this field if it's `is_some`
     /// - env value: [`ALIBABA_CLOUD_OIDC_TOKEN_FILE`]
     pub oidc_token_file: Option<String>,
+    /// `sts_endpoint` will be loaded from
+    ///
+    /// - this field if it's `is_some`
+    /// - env value: [`ALIBABA_CLOUD_STS_ENDPOINT`]
+    pub sts_endpoint: Option<String>,
 }
 
 impl Default for Config {
@@ -52,6 +57,7 @@ impl Default for Config {
             role_session_name: "resign".to_string(),
             oidc_provider_arn: None,
             oidc_token_file: None,
+            sts_endpoint: None,
         }
     }
 }
@@ -79,6 +85,9 @@ impl Config {
         if let Some(v) = envs.get(ALIBABA_CLOUD_OIDC_TOKEN_FILE) {
             self.oidc_token_file.get_or_insert(v.clone());
         }
+        if let Some(v) = envs.get(ALIBABA_CLOUD_STS_ENDPOINT) {
+            self.sts_endpoint.get_or_insert(v.clone());
+        }
 
         self
     }
diff --git a/src/aliyun/constants.rs b/src/aliyun/constants.rs
index 31c59aa7..90de2f75 100644
--- a/src/aliyun/constants.rs
+++ b/src/aliyun/constants.rs
@@ -5,3 +5,4 @@ pub const ALIBABA_CLOUD_SECURITY_TOKEN: &str = "ALIBABA_CLOUD_SECURITY_TOKEN";
 pub const ALIBABA_CLOUD_ROLE_ARN: &str = "ALIBABA_CLOUD_ROLE_ARN";
 pub const ALIBABA_CLOUD_OIDC_PROVIDER_ARN: &str = "ALIBABA_CLOUD_OIDC_PROVIDER_ARN";
 pub const ALIBABA_CLOUD_OIDC_TOKEN_FILE: &str = "ALIBABA_CLOUD_OIDC_TOKEN_FILE";
+pub const ALIBABA_CLOUD_STS_ENDPOINT: &str = "ALIBABA_CLOUD_STS_ENDPOINT";
diff --git a/src/aliyun/credential.rs b/src/aliyun/credential.rs
index 3efebbac..2edf02b7 100644
--- a/src/aliyun/credential.rs
+++ b/src/aliyun/credential.rs
@@ -138,7 +138,7 @@ impl Loader {
         let role_session_name = &self.config.role_session_name;
 
         // Construct request to Aliyun STS Service.
-        let url = format!("https://sts.aliyuncs.com/?Action=AssumeRoleWithOIDC&OIDCProviderArn={}&RoleArn={}&RoleSessionName={}&Format=JSON&Version=2015-04-01&Timestamp={}&OIDCToken={}", provider_arn, role_arn, role_session_name, format_rfc3339(now()), token);
+        let url = format!("{}/?Action=AssumeRoleWithOIDC&OIDCProviderArn={}&RoleArn={}&RoleSessionName={}&Format=JSON&Version=2015-04-01&Timestamp={}&OIDCToken={}", self.get_sts_endpoint(), provider_arn, role_arn, role_session_name, format_rfc3339(now()), token);
 
         let req = self.client.get(&url).header(
             http::header::CONTENT_TYPE.as_str(),
@@ -163,6 +163,13 @@ impl Loader {
 
         Ok(Some(cred))
     }
+
+    fn get_sts_endpoint(&self) -> String {
+        match &self.config.sts_endpoint {
+            Some(defined_sts_endpoint) => format!("https://{}", defined_sts_endpoint),
+            None => "https://sts.aliyuncs.com".to_string(),
+        }
+    }
 }
 
 #[derive(Default, Debug, Deserialize)]