1package imds
2
3import (
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
15const getTokenPath = "/latest/api/token"
16const 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.
20func (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
37type getTokenInput struct {
38 TokenTTL time.Duration
39}
40
41type getTokenOutput struct {
42 Token string
43 TokenTTL time.Duration
44
45 ResultMetadata middleware.Metadata
46}
47
48func 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
67func buildGetTokenPath(interface{}) (string, error) {
68 return getTokenPath, nil
69}
70
71func 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
98type tokenTTLRequestHeader struct{}
99
100func (*tokenTTLRequestHeader) ID() string { return "tokenTTLRequestHeader" }
101func (*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}