123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285 |
- package imds
- import (
- "bytes"
- "context"
- "fmt"
- "io/ioutil"
- "net/url"
- "path"
- "time"
- awsmiddleware "github.com/aws/aws-sdk-go-v2/aws/middleware"
- "github.com/aws/aws-sdk-go-v2/aws/retry"
- "github.com/aws/smithy-go/middleware"
- smithyhttp "github.com/aws/smithy-go/transport/http"
- )
- func addAPIRequestMiddleware(stack *middleware.Stack,
- options Options,
- getPath func(interface{}) (string, error),
- getOutput func(*smithyhttp.Response) (interface{}, error),
- ) (err error) {
- err = addRequestMiddleware(stack, options, "GET", getPath, getOutput)
- if err != nil {
- return err
- }
- // Token Serializer build and state management.
- if !options.disableAPIToken {
- err = stack.Finalize.Insert(options.tokenProvider, (*retry.Attempt)(nil).ID(), middleware.After)
- if err != nil {
- return err
- }
- err = stack.Deserialize.Insert(options.tokenProvider, "OperationDeserializer", middleware.Before)
- if err != nil {
- return err
- }
- }
- return nil
- }
- func addRequestMiddleware(stack *middleware.Stack,
- options Options,
- method string,
- getPath func(interface{}) (string, error),
- getOutput func(*smithyhttp.Response) (interface{}, error),
- ) (err error) {
- err = awsmiddleware.AddSDKAgentKey(awsmiddleware.FeatureMetadata, "ec2-imds")(stack)
- if err != nil {
- return err
- }
- // Operation timeout
- err = stack.Initialize.Add(&operationTimeout{
- DefaultTimeout: defaultOperationTimeout,
- }, middleware.Before)
- if err != nil {
- return err
- }
- // Operation Serializer
- err = stack.Serialize.Add(&serializeRequest{
- GetPath: getPath,
- Method: method,
- }, middleware.After)
- if err != nil {
- return err
- }
- // Operation endpoint resolver
- err = stack.Serialize.Insert(&resolveEndpoint{
- Endpoint: options.Endpoint,
- EndpointMode: options.EndpointMode,
- }, "OperationSerializer", middleware.Before)
- if err != nil {
- return err
- }
- // Operation Deserializer
- err = stack.Deserialize.Add(&deserializeResponse{
- GetOutput: getOutput,
- }, middleware.After)
- if err != nil {
- return err
- }
- err = stack.Deserialize.Add(&smithyhttp.RequestResponseLogger{
- LogRequest: options.ClientLogMode.IsRequest(),
- LogRequestWithBody: options.ClientLogMode.IsRequestWithBody(),
- LogResponse: options.ClientLogMode.IsResponse(),
- LogResponseWithBody: options.ClientLogMode.IsResponseWithBody(),
- }, middleware.After)
- if err != nil {
- return err
- }
- err = addSetLoggerMiddleware(stack, options)
- if err != nil {
- return err
- }
- // Retry support
- return retry.AddRetryMiddlewares(stack, retry.AddRetryMiddlewaresOptions{
- Retryer: options.Retryer,
- LogRetryAttempts: options.ClientLogMode.IsRetries(),
- })
- }
- func addSetLoggerMiddleware(stack *middleware.Stack, o Options) error {
- return middleware.AddSetLoggerMiddleware(stack, o.Logger)
- }
- type serializeRequest struct {
- GetPath func(interface{}) (string, error)
- Method string
- }
- func (*serializeRequest) ID() string {
- return "OperationSerializer"
- }
- func (m *serializeRequest) HandleSerialize(
- ctx context.Context, in middleware.SerializeInput, next middleware.SerializeHandler,
- ) (
- out middleware.SerializeOutput, metadata middleware.Metadata, err error,
- ) {
- request, ok := in.Request.(*smithyhttp.Request)
- if !ok {
- return out, metadata, fmt.Errorf("unknown transport type %T", in.Request)
- }
- reqPath, err := m.GetPath(in.Parameters)
- if err != nil {
- return out, metadata, fmt.Errorf("unable to get request URL path, %w", err)
- }
- request.Request.URL.Path = reqPath
- request.Request.Method = m.Method
- return next.HandleSerialize(ctx, in)
- }
- type deserializeResponse struct {
- GetOutput func(*smithyhttp.Response) (interface{}, error)
- }
- func (*deserializeResponse) ID() string {
- return "OperationDeserializer"
- }
- func (m *deserializeResponse) HandleDeserialize(
- ctx context.Context, in middleware.DeserializeInput, next middleware.DeserializeHandler,
- ) (
- out middleware.DeserializeOutput, metadata middleware.Metadata, err error,
- ) {
- out, metadata, err = next.HandleDeserialize(ctx, in)
- if err != nil {
- return out, metadata, err
- }
- resp, ok := out.RawResponse.(*smithyhttp.Response)
- if !ok {
- return out, metadata, fmt.Errorf(
- "unexpected transport response type, %T, want %T", out.RawResponse, resp)
- }
- defer resp.Body.Close()
- // read the full body so that any operation timeouts cleanup will not race
- // the body being read.
- body, err := ioutil.ReadAll(resp.Body)
- if err != nil {
- return out, metadata, fmt.Errorf("read response body failed, %w", err)
- }
- resp.Body = ioutil.NopCloser(bytes.NewReader(body))
- // Anything that's not 200 |< 300 is error
- if resp.StatusCode < 200 || resp.StatusCode >= 300 {
- return out, metadata, &smithyhttp.ResponseError{
- Response: resp,
- Err: fmt.Errorf("request to EC2 IMDS failed"),
- }
- }
- result, err := m.GetOutput(resp)
- if err != nil {
- return out, metadata, fmt.Errorf(
- "unable to get deserialized result for response, %w", err,
- )
- }
- out.Result = result
- return out, metadata, err
- }
- type resolveEndpoint struct {
- Endpoint string
- EndpointMode EndpointModeState
- }
- func (*resolveEndpoint) ID() string {
- return "ResolveEndpoint"
- }
- func (m *resolveEndpoint) HandleSerialize(
- ctx context.Context, in middleware.SerializeInput, next middleware.SerializeHandler,
- ) (
- out middleware.SerializeOutput, metadata middleware.Metadata, err error,
- ) {
- req, ok := in.Request.(*smithyhttp.Request)
- if !ok {
- return out, metadata, fmt.Errorf("unknown transport type %T", in.Request)
- }
- var endpoint string
- if len(m.Endpoint) > 0 {
- endpoint = m.Endpoint
- } else {
- switch m.EndpointMode {
- case EndpointModeStateIPv6:
- endpoint = defaultIPv6Endpoint
- case EndpointModeStateIPv4:
- fallthrough
- case EndpointModeStateUnset:
- endpoint = defaultIPv4Endpoint
- default:
- return out, metadata, fmt.Errorf("unsupported IMDS endpoint mode")
- }
- }
- req.URL, err = url.Parse(endpoint)
- if err != nil {
- return out, metadata, fmt.Errorf("failed to parse endpoint URL: %w", err)
- }
- return next.HandleSerialize(ctx, in)
- }
- const (
- defaultOperationTimeout = 5 * time.Second
- )
- // operationTimeout adds a timeout on the middleware stack if the Context the
- // stack was called with does not have a deadline. The next middleware must
- // complete before the timeout, or the context will be canceled.
- //
- // If DefaultTimeout is zero, no default timeout will be used if the Context
- // does not have a timeout.
- //
- // The next middleware must also ensure that any resources that are also
- // canceled by the stack's context are completely consumed before returning.
- // Otherwise the timeout cleanup will race the resource being consumed
- // upstream.
- type operationTimeout struct {
- DefaultTimeout time.Duration
- }
- func (*operationTimeout) ID() string { return "OperationTimeout" }
- func (m *operationTimeout) HandleInitialize(
- ctx context.Context, input middleware.InitializeInput, next middleware.InitializeHandler,
- ) (
- output middleware.InitializeOutput, metadata middleware.Metadata, err error,
- ) {
- if _, ok := ctx.Deadline(); !ok && m.DefaultTimeout != 0 {
- var cancelFn func()
- ctx, cancelFn = context.WithTimeout(ctx, m.DefaultTimeout)
- defer cancelFn()
- }
- return next.HandleInitialize(ctx, input)
- }
- // appendURIPath joins a URI path component to the existing path with `/`
- // separators between the path components. If the path being added ends with a
- // trailing `/` that slash will be maintained.
- func appendURIPath(base, add string) string {
- reqPath := path.Join(base, add)
- if len(add) != 0 && add[len(add)-1] == '/' {
- reqPath += "/"
- }
- return reqPath
- }
|