support new env: ALIBABA_CLOUD_STS_ENDPOINT, ALIBABA_CLOUD_STS_HTTP_S…

…CHEME
AliyunContainerService · Jul 18, 2024 · 4fe3652 · 4fe3652
1 parent efa1ea2
commit 4fe3652
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 5 deletions.
diff --git a/pkg/credentials/provider/endpoint.go b/pkg/credentials/provider/endpoint.go
@@ -1,6 +1,9 @@
 package provider
 
-import "strings"
+import (
+	"os"
+	"strings"
+)
 
 const (
 	regionPlaceholder        = "{region}"
@@ -25,6 +28,9 @@ var stsEndpointsByRegion = map[string][2]string{
 }
 
 func GetSTSEndpoint(region string, vpcNetwork bool) string {
+	if v := os.Getenv(envStsEndpoint); v != "" {
+		return v
+	}
 	endpoints, exist := stsEndpointsByRegion[region]
 	if !exist {
 		endpoints = stsEndpointsByRegion["__default__"]

diff --git a/pkg/credentials/provider/endpoint_test.go b/pkg/credentials/provider/endpoint_test.go
@@ -1,6 +1,7 @@
 package provider
 
 import (
+	"os"
 	"testing"
 )
 
@@ -103,3 +104,13 @@ func TestGetSTSEndpoint(t *testing.T) {
 		})
 	}
 }
+
+func TestGetSTSEndpointFromEnv(t *testing.T) {
+	os.Setenv("ALIBABA_CLOUD_STS_ENDPOINT", "sts.cn-hangzhou.aliyuncs.com")
+	defer os.Unsetenv("ALIBABA_CLOUD_STS_ENDPOINT")
+
+	want := "sts.cn-hangzhou.aliyuncs.com"
+	if got := GetSTSEndpoint("foo", true); got != want {
+		t.Errorf("GetSTSEndpoint() = %v, want %v", got, want)
+	}
+}
diff --git a/pkg/credentials/provider/env.go b/pkg/credentials/provider/env.go
@@ -6,6 +6,9 @@ const (
 	envRoleArn         = "ALIBABA_CLOUD_ROLE_ARN"
 	envOidcProviderArn = "ALIBABA_CLOUD_OIDC_PROVIDER_ARN"
 	envOidcTokenFile   = "ALIBABA_CLOUD_OIDC_TOKEN_FILE"
+
+	envStsEndpoint   = "ALIBABA_CLOUD_STS_ENDPOINT"
+	envStsHttpScheme = "ALIBABA_CLOUD_STS_HTTP_SCHEME"
 )
 
 // https://github.com/aliyun/credentials-go
@@ -142,3 +145,11 @@ func getEnvsValue(keys []string) string {
 func getRoleSessionNameFromEnv() string {
 	return getEnvsValue(roleSessionNameEnvs)
 }
+
+func getStsEndpointFromEnv() string {
+	return getEnvsValue([]string{envStsEndpoint})
+}
+
+func getStsHttpSchemeFromEnv() string {
+	return getEnvsValue([]string{envStsHttpScheme})
+}
diff --git a/pkg/credentials/provider/oidc_provider.go b/pkg/credentials/provider/oidc_provider.go
@@ -13,17 +13,18 @@ import (
 )
 
 const (
-	defaultSTSEndpoint = "sts.aliyuncs.com"
-	defaultSTSScheme   = "HTTPS"
-
 	defaultEnvRoleArn         = "ALIBABA_CLOUD_ROLE_ARN"
 	defaultEnvOIDCProviderArn = "ALIBABA_CLOUD_OIDC_PROVIDER_ARN"
 	defaultEnvOIDCTokenFile   = "ALIBABA_CLOUD_OIDC_TOKEN_FILE"
 
 	defaultExpiryWindowForAssumeRole = time.Minute * 10
 )
 
-var defaultSessionName = "default-session-name"
+var (
+	defaultSessionName = "default-session-name"
+	defaultSTSEndpoint = "sts.aliyuncs.com"
+	defaultSTSScheme   = "HTTPS"
+)
 
 type OIDCProvider struct {
 	u *Updater
@@ -66,6 +67,12 @@ func init() {
 	if sessionName != "" {
 		defaultSessionName = sessionName
 	}
+	if v := getStsEndpointFromEnv(); v != "" {
+		defaultSTSEndpoint = v
+	}
+	if v := getStsHttpSchemeFromEnv(); v != "" {
+		defaultSTSScheme = strings.ToUpper(v)
+	}
 }
 
 func NewOIDCProvider(opts OIDCProviderOptions) *OIDCProvider {