api_op_GetToken.go raw

   1  package imds
   2  
   3  import (
   4  	"context"
   5  	"fmt"
   6  	"io"
   7  	"strconv"
   8  	"strings"
   9  	"time"
  10  
  11  	"github.com/aws/smithy-go/middleware"
  12  	smithyhttp "github.com/aws/smithy-go/transport/http"
  13  )
  14  
  15  const getTokenPath = "/latest/api/token"
  16  const tokenTTLHeader = "X-Aws-Ec2-Metadata-Token-Ttl-Seconds"
  17  
  18  // getToken uses the duration to return a token for EC2 IMDS, or an error if
  19  // the request failed.
  20  func (c *Client) getToken(ctx context.Context, params *getTokenInput, optFns ...func(*Options)) (*getTokenOutput, error) {
  21  	if params == nil {
  22  		params = &getTokenInput{}
  23  	}
  24  
  25  	result, metadata, err := c.invokeOperation(ctx, "getToken", params, optFns,
  26  		addGetTokenMiddleware,
  27  	)
  28  	if err != nil {
  29  		return nil, err
  30  	}
  31  
  32  	out := result.(*getTokenOutput)
  33  	out.ResultMetadata = metadata
  34  	return out, nil
  35  }
  36  
  37  type getTokenInput struct {
  38  	TokenTTL time.Duration
  39  }
  40  
  41  type getTokenOutput struct {
  42  	Token    string
  43  	TokenTTL time.Duration
  44  
  45  	ResultMetadata middleware.Metadata
  46  }
  47  
  48  func addGetTokenMiddleware(stack *middleware.Stack, options Options) error {
  49  	err := addRequestMiddleware(stack,
  50  		options,
  51  		"PUT",
  52  		"GetToken",
  53  		buildGetTokenPath,
  54  		buildGetTokenOutput)
  55  	if err != nil {
  56  		return err
  57  	}
  58  
  59  	err = stack.Serialize.Add(&tokenTTLRequestHeader{}, middleware.After)
  60  	if err != nil {
  61  		return err
  62  	}
  63  
  64  	return nil
  65  }
  66  
  67  func buildGetTokenPath(interface{}) (string, error) {
  68  	return getTokenPath, nil
  69  }
  70  
  71  func buildGetTokenOutput(resp *smithyhttp.Response) (v interface{}, err error) {
  72  	defer func() {
  73  		closeErr := resp.Body.Close()
  74  		if err == nil {
  75  			err = closeErr
  76  		} else if closeErr != nil {
  77  			err = fmt.Errorf("response body close error: %v, original error: %w", closeErr, err)
  78  		}
  79  	}()
  80  
  81  	ttlHeader := resp.Header.Get(tokenTTLHeader)
  82  	tokenTTL, err := strconv.ParseInt(ttlHeader, 10, 64)
  83  	if err != nil {
  84  		return nil, fmt.Errorf("unable to parse API token, %w", err)
  85  	}
  86  
  87  	var token strings.Builder
  88  	if _, err = io.Copy(&token, resp.Body); err != nil {
  89  		return nil, fmt.Errorf("unable to read API token, %w", err)
  90  	}
  91  
  92  	return &getTokenOutput{
  93  		Token:    token.String(),
  94  		TokenTTL: time.Duration(tokenTTL) * time.Second,
  95  	}, nil
  96  }
  97  
  98  type tokenTTLRequestHeader struct{}
  99  
 100  func (*tokenTTLRequestHeader) ID() string { return "tokenTTLRequestHeader" }
 101  func (*tokenTTLRequestHeader) HandleSerialize(
 102  	ctx context.Context, in middleware.SerializeInput, next middleware.SerializeHandler,
 103  ) (
 104  	out middleware.SerializeOutput, metadata middleware.Metadata, err error,
 105  ) {
 106  	req, ok := in.Request.(*smithyhttp.Request)
 107  	if !ok {
 108  		return out, metadata, fmt.Errorf("expect HTTP transport, got %T", in.Request)
 109  	}
 110  
 111  	input, ok := in.Parameters.(*getTokenInput)
 112  	if !ok {
 113  		return out, metadata, fmt.Errorf("expect getTokenInput, got %T", in.Parameters)
 114  	}
 115  
 116  	req.Header.Set(tokenTTLHeader, strconv.Itoa(int(input.TokenTTL/time.Second)))
 117  
 118  	return next.HandleSerialize(ctx, in)
 119  }
 120