additional improvements

Author: skotambkar

aws/endpointcreds/provider.go

@@ -87,7 +87,7 @@ func New(cfg aws.Config) *Provider {
 // Retrieve will attempt to request the credentials from the endpoint the Provider
 // was configured for. And error will be returned if the retrieval fails.
 func (p *Provider) retrieveFn(ctx context.Context) (aws.Credentials, error) {
-	resp, err := p.getCredentials()
+	resp, err := p.getCredentials(ctx)
 	if err != nil {
 		return aws.Credentials{},
 			awserr.New("CredentialsEndpointError", "failed to load credentials", err)
@@ -120,7 +120,7 @@ type errorOutput struct {
 	Message string `json:"message"`
 }
 
-func (p *Provider) getCredentials() (*getCredentialsOutput, error) {
+func (p *Provider) getCredentials(ctx context.Context) (*getCredentialsOutput, error) {
 	op := &aws.Operation{
 		Name:       "GetCredentials",
 		HTTPMethod: "GET",
@@ -129,7 +129,7 @@ func (p *Provider) getCredentials() (*getCredentialsOutput, error) {
 	out := &getCredentialsOutput{}
 	req := p.Client.NewRequest(op, nil, out)
 	req.HTTPRequest.Header.Set("Accept", "application/json")
-
+	req.SetContext(ctx)
 	return out, req.Send()
 }
 

aws/signer/v2/v2_test.go

@@ -2,6 +2,7 @@ package v2
 
 import (
 	"bytes"
+	"context"
 	"net/http"
 	"net/url"
 	"os"
@@ -75,7 +76,7 @@ func TestSignRequestWithAndWithoutSession(t *testing.T) {
 
 	signer := builder.BuildSigner()
 
-	err := signer.Sign(nil)
+	err := signer.Sign(context.Background())
 	if err != nil {
 		t.Fatalf("expect no error, got %v", err)
 	}
@@ -122,7 +123,7 @@ func TestSignRequestWithAndWithoutSession(t *testing.T) {
 	builder.SessionToken = "SESSION"
 	signer = builder.BuildSigner()
 
-	err = signer.Sign(nil)
+	err = signer.Sign(context.Background())
 	if err != nil {
 		t.Fatalf("expect no error, got %v", err)
 	}
@@ -165,7 +166,7 @@ func TestMoreComplexSignRequest(t *testing.T) {
 
 	signer := builder.BuildSigner()
 
-	err := signer.Sign(nil)
+	err := signer.Sign(context.Background())
 	if err != nil {
 		t.Fatalf("expect no error, got %v", err)
 	}

aws/static_provider.go

@@ -2,6 +2,7 @@ package aws
 
 import (
 	"context"
+
 	"github.com/aws/aws-sdk-go-v2/aws/awserr"
 )
 
@@ -33,6 +34,12 @@ func NewStaticCredentialsProvider(key, secret, session string) StaticCredentials
 
 // Retrieve returns the credentials or error if the credentials are invalid.
 func (s StaticCredentialsProvider) Retrieve(ctx context.Context) (Credentials, error) {
+	select {
+	case <-ctx.Done():
+		return Credentials{}, awserr.New(ErrCodeRequestCanceled, "context canceled", ctx.Err())
+	default: // do nothing
+	}
+
 	v := s.Value
 	if v.AccessKeyID == "" || v.SecretAccessKey == "" {
 		return Credentials{Source: StaticCredentialsProviderName}, ErrStaticCredentialsEmpty

aws/stscreds/provider.go

@@ -248,9 +248,6 @@ func (p *AssumeRoleProvider) retrieveFn(ctx context.Context) (aws.Credentials, e
 	}
 
 	req := p.Client.AssumeRoleRequest(input)
-	if ctx == nil {
-		ctx = context.Background()
-	}
 	resp, err := req.Send(ctx)
 	if err != nil {
 		return aws.Credentials{Source: ProviderName}, err