github.com/aavshr/aws-sdk-go@v1.41.3/aws/ec2metadata/token_provider.go (about)

     1  package ec2metadata
     2  
     3  import (
     4  	"net/http"
     5  	"sync/atomic"
     6  	"time"
     7  
     8  	"github.com/aavshr/aws-sdk-go/aws/awserr"
     9  	"github.com/aavshr/aws-sdk-go/aws/credentials"
    10  	"github.com/aavshr/aws-sdk-go/aws/request"
    11  )
    12  
    13  // A tokenProvider struct provides access to EC2Metadata client
    14  // and atomic instance of a token, along with configuredTTL for it.
    15  // tokenProvider also provides an atomic flag to disable the
    16  // fetch token operation.
    17  // The disabled member will use 0 as false, and 1 as true.
    18  type tokenProvider struct {
    19  	client        *EC2Metadata
    20  	token         atomic.Value
    21  	configuredTTL time.Duration
    22  	disabled      uint32
    23  }
    24  
    25  // A ec2Token struct helps use of token in EC2 Metadata service ops
    26  type ec2Token struct {
    27  	token string
    28  	credentials.Expiry
    29  }
    30  
    31  // newTokenProvider provides a pointer to a tokenProvider instance
    32  func newTokenProvider(c *EC2Metadata, duration time.Duration) *tokenProvider {
    33  	return &tokenProvider{client: c, configuredTTL: duration}
    34  }
    35  
    36  // fetchTokenHandler fetches token for EC2Metadata service client by default.
    37  func (t *tokenProvider) fetchTokenHandler(r *request.Request) {
    38  
    39  	// short-circuits to insecure data flow if tokenProvider is disabled.
    40  	if v := atomic.LoadUint32(&t.disabled); v == 1 {
    41  		return
    42  	}
    43  
    44  	if ec2Token, ok := t.token.Load().(ec2Token); ok && !ec2Token.IsExpired() {
    45  		r.HTTPRequest.Header.Set(tokenHeader, ec2Token.token)
    46  		return
    47  	}
    48  
    49  	output, err := t.client.getToken(r.Context(), t.configuredTTL)
    50  
    51  	if err != nil {
    52  
    53  		// change the disabled flag on token provider to true,
    54  		// when error is request timeout error.
    55  		if requestFailureError, ok := err.(awserr.RequestFailure); ok {
    56  			switch requestFailureError.StatusCode() {
    57  			case http.StatusForbidden, http.StatusNotFound, http.StatusMethodNotAllowed:
    58  				atomic.StoreUint32(&t.disabled, 1)
    59  			case http.StatusBadRequest:
    60  				r.Error = requestFailureError
    61  			}
    62  
    63  			// Check if request timed out while waiting for response
    64  			if e, ok := requestFailureError.OrigErr().(awserr.Error); ok {
    65  				if e.Code() == request.ErrCodeRequestError {
    66  					atomic.StoreUint32(&t.disabled, 1)
    67  				}
    68  			}
    69  		}
    70  		return
    71  	}
    72  
    73  	newToken := ec2Token{
    74  		token: output.Token,
    75  	}
    76  	newToken.SetExpiration(time.Now().Add(output.TTL), ttlExpirationWindow)
    77  	t.token.Store(newToken)
    78  
    79  	// Inject token header to the request.
    80  	if ec2Token, ok := t.token.Load().(ec2Token); ok {
    81  		r.HTTPRequest.Header.Set(tokenHeader, ec2Token.token)
    82  	}
    83  }
    84  
    85  // enableTokenProviderHandler enables the token provider
    86  func (t *tokenProvider) enableTokenProviderHandler(r *request.Request) {
    87  	// If the error code status is 401, we enable the token provider
    88  	if e, ok := r.Error.(awserr.RequestFailure); ok && e != nil &&
    89  		e.StatusCode() == http.StatusUnauthorized {
    90  		t.token.Store(ec2Token{})
    91  		atomic.StoreUint32(&t.disabled, 0)
    92  	}
    93  }