feat: support OIDCRoleArnProvider beforeRefresh (#252)

* feat: support OIDCRoleArnProvider beforeRefresh
TencentCloud · Feb 28, 2024 · 0f4febc · 0f4febc
1 parent 5d1b2c1
commit 0f4febc
Showing 1 changed file with 47 additions and 25 deletions.
diff --git a/tencentcloud/common/oidc_role_arn_provider.go b/tencentcloud/common/oidc_role_arn_provider.go
@@ -21,6 +21,7 @@ type OIDCRoleArnProvider struct {
 	roleSessionName  string
 	durationSeconds  int64
 	Endpoint         string
+	beforeRefresh    func(provider *OIDCRoleArnProvider) error
 }
 
 type oidcStsRsp struct {
@@ -54,33 +55,46 @@ func NewOIDCRoleArnProvider(region, providerId, webIdentityToken, roleArn, roleS
 //  4. roleSessionName will be "tencentcloud-go-sdk-" + timestamp
 //  5. durationSeconds will be 7200s
 func DefaultTkeOIDCRoleArnProvider() (*OIDCRoleArnProvider, error) {
-	reg := os.Getenv("TKE_REGION")
-	if reg == "" {
-		return nil, errors.New("env TKE_REGION not exist")
+	beforeRefresh := func(provider *OIDCRoleArnProvider) error {
+		reg := os.Getenv("TKE_REGION")
+		if reg == "" {
+			return errors.New("env TKE_REGION not exist")
+		}
+
+		providerId := os.Getenv("TKE_PROVIDER_ID")
+		if providerId == "" {
+			return errors.New("env TKE_PROVIDER_ID not exist")
+		}
+
+		tokenFile := os.Getenv("TKE_WEB_IDENTITY_TOKEN_FILE")
+		if tokenFile == "" {
+			return errors.New("env TKE_WEB_IDENTITY_TOKEN_FILE not exist")
+		}
+		tokenBytes, err := ioutil.ReadFile(tokenFile)
+		if err != nil {
+			return err
+		}
+
+		roleArn := os.Getenv("TKE_ROLE_ARN")
+		if roleArn == "" {
+			return errors.New("env TKE_ROLE_ARN not exist")
+		}
+
+		sessionName := defaultSessionName + strconv.FormatInt(time.Now().UnixNano()/1000, 10)
+
+		provider.region = region
+		provider.providerId = providerId
+		provider.webIdentityToken = string(tokenBytes)
+		provider.roleArn = roleArn
+		provider.roleSessionName = sessionName
+		return nil
 	}
 
-	providerId := os.Getenv("TKE_PROVIDER_ID")
-	if providerId == "" {
-		return nil, errors.New("env TKE_PROVIDER_ID not exist")
+	provider := &OIDCRoleArnProvider{
+		beforeRefresh:   beforeRefresh,
+		durationSeconds: defaultDurationSeconds,
 	}
-
-	tokenFile := os.Getenv("TKE_WEB_IDENTITY_TOKEN_FILE")
-	if tokenFile == "" {
-		return nil, errors.New("env TKE_WEB_IDENTITY_TOKEN_FILE not exist")
-	}
-	tokenBytes, err := ioutil.ReadFile(tokenFile)
-	if err != nil {
-		return nil, err
-	}
-
-	roleArn := os.Getenv("TKE_ROLE_ARN")
-	if roleArn == "" {
-		return nil, errors.New("env TKE_ROLE_ARN not exist")
-	}
-
-	sessionName := defaultSessionName + strconv.FormatInt(time.Now().UnixNano()/1000, 10)
-
-	return NewOIDCRoleArnProvider(reg, providerId, string(tokenBytes), roleArn, sessionName, defaultDurationSeconds), nil
+	return provider, provider.beforeRefresh(provider)
 }
 
 func (r *OIDCRoleArnProvider) GetCredential() (CredentialIface, error) {
@@ -89,6 +103,14 @@ func (r *OIDCRoleArnProvider) GetCredential() (CredentialIface, error) {
 		version = "2018-08-13"
 		action  = "AssumeRoleWithWebIdentity"
 	)
+
+	if r.beforeRefresh != nil {
+		err := r.beforeRefresh(r)
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	if r.durationSeconds > 43200 || r.durationSeconds <= 0 {
 		return nil, tcerr.NewTencentCloudSDKError(creErr, "AssumeRoleWithWebIdentity durationSeconds should be in the range of 0~43200s", "")
 	}
@@ -131,7 +153,7 @@ func (r *OIDCRoleArnProvider) GetCredential() (CredentialIface, error) {
 		roleArn:         r.roleArn,
 		roleSessionName: r.roleSessionName,
 		durationSeconds: r.durationSeconds,
-		expiredTime:     int64(rspSt.Response.ExpiredTime),
+		expiredTime:     int64(rspSt.Response.ExpiredTime) - r.durationSeconds/10,
 		token:           rspSt.Response.Credentials.Token,
 		tmpSecretId:     rspSt.Response.Credentials.TmpSecretId,
 		tmpSecretKey:    rspSt.Response.Credentials.TmpSecretKey,