add timeout for credential retrieval (#426)

AliyunContainerService · Jul 31, 2024 · 4fa12f6 · 4fa12f6
1 parent 19dbbb4
commit 4fa12f6
Show file tree

Hide file tree

Showing 4 changed files with 51 additions and 16 deletions.
diff --git a/.gitignore b/.gitignore
@@ -26,3 +26,4 @@ kubeconfig
 .terraform/
 .terraform.*
 terraform.tfstate*
+.idea
diff --git a/pkg/credentials/provider/semaphore_provider.go b/pkg/credentials/provider/semaphore_provider.go
@@ -40,3 +40,9 @@ func (o *SemaphoreProviderOptions) applyDefaults() {
 		o.MaxWeight = 1
 	}
 }
+
+func (p *SemaphoreProvider) Stop(ctx context.Context) {
+	if s, ok := p.cp.(Stopper); ok {
+		s.Stop(ctx)
+	}
+}
diff --git a/pkg/credentials/provider/v1sdk.go b/pkg/credentials/provider/v1sdk.go
@@ -3,23 +3,27 @@ package provider
 import (
 	"context"
 	"fmt"
+	"time"
 )
 
 type SignerForV1SDK struct {
-	p      CredentialsProvider
-	Logger Logger
+	p                          CredentialsProvider
+	Logger                     Logger
+	credentialRetrievalTimeout time.Duration
 }
 
 type SignerForV1SDKOptions struct {
-	Logger Logger
+	Logger                     Logger
+	CredentialRetrievalTimeout time.Duration
 }
 
 func NewSignerForV1SDK(p CredentialsProvider, opts SignerForV1SDKOptions) *SignerForV1SDK {
 	opts.applyDefaults()
 
 	return &SignerForV1SDK{
-		p:      p,
-		Logger: opts.Logger,
+		p:                          p,
+		Logger:                     opts.Logger,
+		credentialRetrievalTimeout: opts.CredentialRetrievalTimeout,
 	}
 }
 
@@ -36,15 +40,19 @@ func (s *SignerForV1SDK) GetVersion() string {
 }
 
 func (s *SignerForV1SDK) GetAccessKeyId() (string, error) {
-	cred, err := s.p.Credentials(context.TODO())
+	timeoutCtx, cancel := context.WithTimeout(context.Background(), s.credentialRetrievalTimeout)
+	defer cancel()
+	cred, err := s.p.Credentials(timeoutCtx)
 	if err != nil {
 		return "", err
 	}
 	return cred.AccessKeyId, nil
 }
 
 func (s *SignerForV1SDK) GetExtraParam() map[string]string {
-	cred, err := s.p.Credentials(context.TODO())
+	timeoutCtx, cancel := context.WithTimeout(context.Background(), s.credentialRetrievalTimeout)
+	defer cancel()
+	cred, err := s.p.Credentials(timeoutCtx)
 	if err != nil {
 		s.logger().Error(err, fmt.Sprintf("get credentials failed: %s", err))
 		return nil
@@ -56,7 +64,9 @@ func (s *SignerForV1SDK) GetExtraParam() map[string]string {
 }
 
 func (s *SignerForV1SDK) Sign(stringToSign, secretSuffix string) string {
-	cred, err := s.p.Credentials(context.TODO())
+	timeoutCtx, cancel := context.WithTimeout(context.Background(), s.credentialRetrievalTimeout)
+	defer cancel()
+	cred, err := s.p.Credentials(timeoutCtx)
 	if err != nil {
 		s.logger().Error(err, fmt.Sprintf("get credentials failed: %s", err))
 		return ""
@@ -76,4 +86,7 @@ func (o *SignerForV1SDKOptions) applyDefaults() {
 	if o.Logger == nil {
 		o.Logger = defaultLog
 	}
+	if o.CredentialRetrievalTimeout <= 0 {
+		o.CredentialRetrievalTimeout = defaultTimeout
+	}
 }
diff --git a/pkg/credentials/provider/v2sdk.go b/pkg/credentials/provider/v2sdk.go
@@ -2,44 +2,56 @@ package provider
 
 import (
 	"context"
+	"time"
 )
 
+var defaultTimeout = time.Minute * 10
+
 type CredentialForV2SDK struct {
-	p      CredentialsProvider
-	Logger Logger
+	p                          CredentialsProvider
+	Logger                     Logger
+	credentialRetrievalTimeout time.Duration
 }
 
 type CredentialForV2SDKOptions struct {
-	Logger Logger
+	Logger                     Logger
+	CredentialRetrievalTimeout time.Duration
 }
 
 func NewCredentialForV2SDK(p CredentialsProvider, opts CredentialForV2SDKOptions) *CredentialForV2SDK {
 	opts.applyDefaults()
 
 	return &CredentialForV2SDK{
-		p:      p,
-		Logger: opts.Logger,
+		p:                          p,
+		Logger:                     opts.Logger,
+		credentialRetrievalTimeout: opts.CredentialRetrievalTimeout,
 	}
 }
 
 func (c *CredentialForV2SDK) GetAccessKeyId() (*string, error) {
-	cred, err := c.p.Credentials(context.TODO())
+	timeoutCtx, cancel := context.WithTimeout(context.Background(), c.credentialRetrievalTimeout)
+	defer cancel()
+	cred, err := c.p.Credentials(timeoutCtx)
 	if err != nil {
 		return nil, err
 	}
 	return stringPointer(cred.AccessKeyId), nil
 }
 
 func (c *CredentialForV2SDK) GetAccessKeySecret() (*string, error) {
-	cred, err := c.p.Credentials(context.TODO())
+	timeoutCtx, cancel := context.WithTimeout(context.Background(), c.credentialRetrievalTimeout)
+	defer cancel()
+	cred, err := c.p.Credentials(timeoutCtx)
 	if err != nil {
 		return nil, err
 	}
 	return stringPointer(cred.AccessKeySecret), nil
 }
 
 func (c *CredentialForV2SDK) GetSecurityToken() (*string, error) {
-	cred, err := c.p.Credentials(context.TODO())
+	timeoutCtx, cancel := context.WithTimeout(context.Background(), c.credentialRetrievalTimeout)
+	defer cancel()
+	cred, err := c.p.Credentials(timeoutCtx)
 	if err != nil {
 		return nil, err
 	}
@@ -65,6 +77,9 @@ func (o *CredentialForV2SDKOptions) applyDefaults() {
 	if o.Logger == nil {
 		o.Logger = defaultLog
 	}
+	if o.CredentialRetrievalTimeout <= 0 {
+		o.CredentialRetrievalTimeout = defaultTimeout
+	}
 }
 
 func stringPointer(s string) *string {