From defbbbdfc3c541a19eed88bddec32b96149b2663 Mon Sep 17 00:00:00 2001
From: mozillazg <mozillazg101@gmail.com>
Date: Mon, 18 Nov 2024 13:32:57 +0800
Subject: [PATCH] support new env ACK_RAM_TOOL_TOKEN_EXTRA_KEY_PREFIX

---
 pkg/ctl/credentialplugin/gettoken.go      | 31 +++++++++++++++++-
 pkg/ctl/credentialplugin/gettoken_test.go | 33 +++++++++++++++++++
 pkg/ramauthenticator/token.go             | 40 +++++++++++++++++++++--
 3 files changed, 101 insertions(+), 3 deletions(-)
 create mode 100644 pkg/ctl/credentialplugin/gettoken_test.go

diff --git a/pkg/ctl/credentialplugin/gettoken.go b/pkg/ctl/credentialplugin/gettoken.go
index 65dbd9cc..0dd5f690 100644
--- a/pkg/ctl/credentialplugin/gettoken.go
+++ b/pkg/ctl/credentialplugin/gettoken.go
@@ -8,8 +8,12 @@ import (
 	"github.com/AliyunContainerService/ack-ram-tool/pkg/ramauthenticator"
 	"github.com/AliyunContainerService/ack-ram-tool/pkg/types"
 	"github.com/spf13/cobra"
+	"os"
+	"strings"
 )
 
+const envTokenExtraQueryKeyPrefix = "ACK_RAM_TOOL_TOKEN_EXTRA_KEY_PREFIX"
+
 type GetTokenOpts struct {
 	//clusterId        string
 	privateIpAddress bool
@@ -30,7 +34,9 @@ var getTokenCmd = &cobra.Command{
 		clusterId := ctl.GlobalOption.ClusterId
 		getCredentialOpts.clusterId = clusterId
 
-		token, err := ramauthenticator.GenerateToken(clusterId, client.Credential())
+		generator := ramauthenticator.NewTokenGenerator(clusterId, client.Credential())
+		generator.SetExtraQuery(getExtraTokenQuery())
+		token, err := generator.NewToken()
 		common.ExitIfError(err)
 
 		cred, err := newTokenExecCredential(token)
@@ -42,6 +48,29 @@ var getTokenCmd = &cobra.Command{
 	},
 }
 
+func getExtraTokenQuery() map[string]string {
+	query := make(map[string]string)
+	prefix := os.Getenv(envTokenExtraQueryKeyPrefix)
+	if prefix == "" {
+		return query
+	}
+	for _, item := range os.Environ() {
+		before, after, found := strings.Cut(item, "=")
+		if !found {
+			continue
+		}
+		if after == "" {
+			continue
+		}
+		if !strings.HasPrefix(before, prefix) {
+			continue
+		}
+		k := strings.ToLower(strings.TrimPrefix(before, prefix))
+		query[k] = after
+	}
+	return query
+}
+
 func newTokenExecCredential(token *ramauthenticator.Token) (*types.ExecCredential, error) {
 	version := getApiVersion(getCredentialOpts.apiVersion)
 	var exp *types.KubeTime
diff --git a/pkg/ctl/credentialplugin/gettoken_test.go b/pkg/ctl/credentialplugin/gettoken_test.go
new file mode 100644
index 00000000..e50682cc
--- /dev/null
+++ b/pkg/ctl/credentialplugin/gettoken_test.go
@@ -0,0 +1,33 @@
+package credentialplugin
+
+import (
+	"os"
+	"reflect"
+	"testing"
+)
+
+func Test_getExtraTokenQuery(t *testing.T) {
+	os.Setenv(envTokenExtraQueryKeyPrefix, "FOO_")
+	defer os.Unsetenv(envTokenExtraQueryKeyPrefix)
+	os.Setenv("FOO_BAR", "test1")
+	os.Setenv("FOO_FUZZ", "test2")
+	tests := []struct {
+		name string
+		want map[string]string
+	}{
+		{
+			name: "test",
+			want: map[string]string{
+				"bar":  "test1",
+				"fuzz": "test2",
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := getExtraTokenQuery(); !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("getExtraTokenQuery() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
diff --git a/pkg/ramauthenticator/token.go b/pkg/ramauthenticator/token.go
index 0312a14a..8b2ae2c5 100644
--- a/pkg/ramauthenticator/token.go
+++ b/pkg/ramauthenticator/token.go
@@ -43,12 +43,46 @@ type Token struct {
 	Expiration time.Time `json:"-"`
 }
 
-func GenerateToken(clusterId string, cred credentials.Credential) (*Token, error) {
+type extendOption func(q *openapi.OpenApiRequest)
+
+type TokenGenerator struct {
+	clusterId  string
+	cred       credentials.Credential
+	extraQuery map[string]string
+}
+
+func NewTokenGenerator(clusterId string, cred credentials.Credential) *TokenGenerator {
+	return &TokenGenerator{
+		clusterId:  clusterId,
+		cred:       cred,
+		extraQuery: make(map[string]string),
+	}
+}
+
+func (g *TokenGenerator) NewToken() (*Token, error) {
+	return GenerateToken(g.clusterId, g.cred, g.extendRequest)
+}
+
+func (g *TokenGenerator) SetExtraQuery(extraQuery map[string]string) {
+	g.extraQuery = extraQuery
+}
+
+func (g *TokenGenerator) extendRequest(req *openapi.OpenApiRequest) {
+	for k, v := range g.extraQuery {
+		k = strings.ToLower(k)
+		req.Query[k] = tea.String(v)
+	}
+}
+
+func GenerateToken(clusterId string, cred credentials.Credential, options ...extendOption) (*Token, error) {
 	q := &openapi.OpenApiRequest{
 		Query: map[string]*string{
 			"ACKClusterId": tea.String(clusterId),
 		},
 	}
+	for _, f := range options {
+		f(q)
+	}
 	params := &openapi.Params{
 		Action:      tea.String("GetCallerIdentity"),
 		Version:     tea.String("2015-04-01"),
@@ -82,7 +116,9 @@ func GenerateToken(clusterId string, cred credentials.Credential) (*Token, error
 	}
 	for k, v := range req.Query {
 		if !signParamsWhitelist[strings.ToLower(k)] {
-			continue
+			if q.Query[k] == nil {
+				continue
+			}
 		}
 		t.Query[k] = tea.StringValue(v)
 	}