EC2 IMDS IPv6 Support (#1337)

Author: skmcgrail

config/env_config.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"github.com/aws/aws-sdk-go-v2/feature/ec2/imds"
 	"io"
 	"io/ioutil"
 	"os"
@@ -42,14 +43,20 @@ const (
 
 	awsCustomCABundleEnvVar = "AWS_CA_BUNDLE"
 
-	awsWebIdentityTokenFilePathEnvKey = "AWS_WEB_IDENTITY_TOKEN_FILE"
+	awsWebIdentityTokenFilePathEnvVar = "AWS_WEB_IDENTITY_TOKEN_FILE"
 
-	awsRoleARNEnvKey         = "AWS_ROLE_ARN"
-	awsRoleSessionNameEnvKey = "AWS_ROLE_SESSION_NAME"
+	awsRoleARNEnvVar         = "AWS_ROLE_ARN"
+	awsRoleSessionNameEnvVar = "AWS_ROLE_SESSION_NAME"
 
-	awsEnableEndpointDiscoveryEnvKey = "AWS_ENABLE_ENDPOINT_DISCOVERY"
+	awsEnableEndpointDiscoveryEnvVar = "AWS_ENABLE_ENDPOINT_DISCOVERY"
 
 	awsS3UseARNRegionEnvVar = "AWS_S3_USE_ARN_REGION"
+
+	awsEc2MetadataServiceEndpointModeEnvVar = "AWS_EC2_METADATA_SERVICE_ENDPOINT_MODE"
+
+	awsEc2MetadataServiceEndpointEnvVar = "AWS_EC2_METADATA_SERVICE_ENDPOINT"
+
+	awsEc2MetadataDisabled = "AWS_EC2_METADATA_DISABLED"
 )
 
 var (
@@ -180,6 +187,21 @@ type EnvConfig struct {
 	//
 	// AWS_S3_USE_ARN_REGION=true
 	S3UseARNRegion *bool
+
+	// Specifies if the EC2 IMDS service client is enabled.
+	//
+	// AWS_EC2_METADATA_DISABLED=true
+	EC2IMDSClientEnableState imds.ClientEnableState
+
+	// Specifies the EC2 Instance Metadata Service default endpoint selection mode (IPv4 or IPv6)
+	//
+	// AWS_EC2_METADATA_SERVICE_ENDPOINT_MODE=IPv6
+	EC2IMDSEndpointMode imds.EndpointModeState
+
+	// Specifies the EC2 Instance Metadata Service endpoint to use. If specified it overrides EC2IMDSEndpointMode.
+	//
+	// AWS_EC2_METADATA_SERVICE_ENDPOINT=http://fd00:ec2::254
+	EC2IMDSEndpoint string
 }
 
 // loadEnvConfig reads configuration values from the OS's environment variables.
@@ -215,22 +237,59 @@ func NewEnvConfig() (EnvConfig, error) {
 
 	cfg.CustomCABundle = os.Getenv(awsCustomCABundleEnvVar)
 
-	cfg.WebIdentityTokenFilePath = os.Getenv(awsWebIdentityTokenFilePathEnvKey)
+	cfg.WebIdentityTokenFilePath = os.Getenv(awsWebIdentityTokenFilePathEnvVar)
 
-	cfg.RoleARN = os.Getenv(awsRoleARNEnvKey)
-	cfg.RoleSessionName = os.Getenv(awsRoleSessionNameEnvKey)
+	cfg.RoleARN = os.Getenv(awsRoleARNEnvVar)
+	cfg.RoleSessionName = os.Getenv(awsRoleSessionNameEnvVar)
 
-	if err := setEndpointDiscoveryTypeFromEnvVal(&cfg.EnableEndpointDiscovery, []string{awsEnableEndpointDiscoveryEnvKey}); err != nil {
+	if err := setEndpointDiscoveryTypeFromEnvVal(&cfg.EnableEndpointDiscovery, []string{awsEnableEndpointDiscoveryEnvVar}); err != nil {
 		return cfg, err
 	}
 
 	if err := setBoolPtrFromEnvVal(&cfg.S3UseARNRegion, []string{awsS3UseARNRegionEnvVar}); err != nil {
 		return cfg, err
 	}
 
+	setEC2IMDSClientEnableState(&cfg.EC2IMDSClientEnableState, []string{awsEc2MetadataDisabled})
+	if err := setEC2IMDSEndpointMode(&cfg.EC2IMDSEndpointMode, []string{awsEc2MetadataServiceEndpointModeEnvVar}); err != nil {
+		return cfg, err
+	}
+	cfg.EC2IMDSEndpoint = os.Getenv(awsEc2MetadataServiceEndpointEnvVar)
+
 	return cfg, nil
 }
 
+func setEC2IMDSClientEnableState(state *imds.ClientEnableState, keys []string) {
+	for _, k := range keys {
+		value := os.Getenv(k)
+		if len(value) == 0 {
+			continue
+		}
+		switch {
+		case strings.EqualFold(value, "true"):
+			*state = imds.ClientDisabled
+		case strings.EqualFold(value, "false"):
+			*state = imds.ClientEnabled
+		default:
+			continue
+		}
+		break
+	}
+}
+
+func setEC2IMDSEndpointMode(mode *imds.EndpointModeState, keys []string) error {
+	for _, k := range keys {
+		value := os.Getenv(k)
+		if len(value) == 0 {
+			continue
+		}
+		if err := mode.SetFromString(value); err != nil {
+			return fmt.Errorf("invalid value for environment variable, %s=%s, %v", k, value, err)
+		}
+	}
+	return nil
+}
+
 // GetRegion returns the AWS Region if set in the environment. Returns an empty
 // string if not set.
 func (c EnvConfig) getRegion(ctx context.Context) (string, bool, error) {
@@ -371,3 +430,30 @@ func (c EnvConfig) GetEnableEndpointDiscovery(ctx context.Context) (value aws.En
 
 	return c.EnableEndpointDiscovery, true, nil
 }
+
+// GetEC2IMDSClientEnableState implements a EC2IMDSClientEnableState options resolver interface.
+func (c EnvConfig) GetEC2IMDSClientEnableState() (imds.ClientEnableState, bool, error) {
+	if c.EC2IMDSClientEnableState == imds.ClientDefaultEnableState {
+		return imds.ClientDefaultEnableState, false, nil
+	}
+
+	return c.EC2IMDSClientEnableState, true, nil
+}
+
+// GetEC2IMDSEndpointMode implements a EC2IMDSEndpointMode option resolver interface.
+func (c EnvConfig) GetEC2IMDSEndpointMode() (imds.EndpointModeState, bool, error) {
+	if c.EC2IMDSEndpointMode == imds.EndpointModeStateUnset {
+		return imds.EndpointModeStateUnset, false, nil
+	}
+
+	return c.EC2IMDSEndpointMode, true, nil
+}
+
+// GetEC2IMDSEndpoint implements a EC2IMDSEndpoint option resolver interface.
+func (c EnvConfig) GetEC2IMDSEndpoint() (string, bool, error) {
+	if len(c.EC2IMDSEndpoint) == 0 {
+		return "", false, nil
+	}
+
+	return c.EC2IMDSEndpoint, true, nil
+}

config/env_config_test.go

@@ -1,6 +1,8 @@
 package config
 
 import (
+	"github.com/aws/aws-sdk-go-v2/feature/ec2/imds"
+	"github.com/google/go-cmp/cmp"
 	"os"
 	"reflect"
 	"strconv"
@@ -105,8 +107,9 @@ func TestNewEnvConfig(t *testing.T) {
 	defer awstesting.PopEnv(restoreEnv)
 
 	cases := []struct {
-		Env    map[string]string
-		Config EnvConfig
+		Env     map[string]string
+		Config  EnvConfig
+		WantErr bool
 	}{
 		0: {
 			Env: map[string]string{
@@ -250,6 +253,69 @@ func TestNewEnvConfig(t *testing.T) {
 			Env:    map[string]string{},
 			Config: EnvConfig{},
 		},
+		16: {
+			Env: map[string]string{
+				"AWS_EC2_METADATA_SERVICE_ENDPOINT_MODE": "IPv6",
+			},
+			Config: EnvConfig{
+				EC2IMDSEndpointMode: imds.EndpointModeStateIPv6,
+			},
+		},
+		17: {
+			Env: map[string]string{
+				"AWS_EC2_METADATA_SERVICE_ENDPOINT_MODE": "IPv4",
+			},
+			Config: EnvConfig{
+				EC2IMDSEndpointMode: imds.EndpointModeStateIPv4,
+			},
+		},
+		18: {
+			Env: map[string]string{
+				"AWS_EC2_METADATA_SERVICE_ENDPOINT_MODE": "foobar",
+			},
+			Config:  EnvConfig{},
+			WantErr: true,
+		},
+		19: {
+			Env: map[string]string{
+				"AWS_EC2_METADATA_SERVICE_ENDPOINT": "http://endpoint.localhost",
+			},
+			Config: EnvConfig{
+				EC2IMDSEndpoint: "http://endpoint.localhost",
+			},
+		},
+		20: {
+			Env: map[string]string{
+				"AWS_EC2_METADATA_SERVICE_ENDPOINT_MODE": "IPv6",
+				"AWS_EC2_METADATA_SERVICE_ENDPOINT":      "http://endpoint.localhost",
+			},
+			Config: EnvConfig{
+				EC2IMDSEndpoint:     "http://endpoint.localhost",
+				EC2IMDSEndpointMode: imds.EndpointModeStateIPv6,
+			},
+		},
+		21: {
+			Env: map[string]string{
+				"AWS_EC2_METADATA_DISABLED": "false",
+			},
+			Config: EnvConfig{
+				EC2IMDSClientEnableState: imds.ClientEnabled,
+			},
+		},
+		22: {
+			Env: map[string]string{
+				"AWS_EC2_METADATA_DISABLED": "true",
+			},
+			Config: EnvConfig{
+				EC2IMDSClientEnableState: imds.ClientDisabled,
+			},
+		},
+		23: {
+			Env: map[string]string{
+				"AWS_EC2_METADATA_DISABLED": "foobar",
+			},
+			Config: EnvConfig{},
+		},
 	}
 
 	for i, c := range cases {
@@ -261,13 +327,13 @@ func TestNewEnvConfig(t *testing.T) {
 			}
 
 			cfg, err := NewEnvConfig()
-			if err != nil {
-				t.Fatalf("expect no error, got %v", err)
+			if (err != nil) != c.WantErr {
+				t.Fatalf("WantErr=%v, got err=%v", c.WantErr, err)
 			}
 
-			if !reflect.DeepEqual(c.Config, cfg) {
+			if diff := cmp.Diff(c.Config, cfg); len(diff) > 0 {
 				t.Errorf("expect config to match.\n%s",
-					awstesting.SprintExpectActual(c.Config, cfg))
+					diff)
 			}
 		})
 	}

config/load_options.go

@@ -126,6 +126,17 @@ type LoadOptions struct {
 	// EnableEndpointDiscovery specifies if endpoint discovery is enable for
 	// the client.
 	EnableEndpointDiscovery aws.EndpointDiscoveryEnableState
+
+	// Specifies if the EC2 IMDS service client is enabled.
+	//
+	// AWS_EC2_METADATA_DISABLED=true
+	EC2IMDSClientEnableState imds.ClientEnableState
+
+	// Specifies the EC2 Instance Metadata Service default endpoint selection mode (IPv4 or IPv6)
+	EC2IMDSEndpointMode imds.EndpointModeState
+
+	// Specifies the EC2 Instance Metadata Service endpoint to use. If specified it overrides EC2IMDSEndpointMode.
+	EC2IMDSEndpoint string
 }
 
 // getRegion returns Region from config's LoadOptions
@@ -642,3 +653,54 @@ func WithSSOProviderOptions(v func(*ssocreds.Options)) LoadOptionsFunc {
 		return nil
 	}
 }
+
+// GetEC2IMDSClientEnableState implements a EC2IMDSClientEnableState options resolver interface.
+func (o LoadOptions) GetEC2IMDSClientEnableState() (imds.ClientEnableState, bool, error) {
+	if o.EC2IMDSClientEnableState == imds.ClientDefaultEnableState {
+		return imds.ClientDefaultEnableState, false, nil
+	}
+
+	return o.EC2IMDSClientEnableState, true, nil
+}
+
+// GetEC2IMDSEndpointMode implements a EC2IMDSEndpointMode option resolver interface.
+func (o LoadOptions) GetEC2IMDSEndpointMode() (imds.EndpointModeState, bool, error) {
+	if o.EC2IMDSEndpointMode == imds.EndpointModeStateUnset {
+		return imds.EndpointModeStateUnset, false, nil
+	}
+
+	return o.EC2IMDSEndpointMode, true, nil
+}
+
+// GetEC2IMDSEndpoint implements a EC2IMDSEndpoint option resolver interface.
+func (o LoadOptions) GetEC2IMDSEndpoint() (string, bool, error) {
+	if len(o.EC2IMDSEndpoint) == 0 {
+		return "", false, nil
+	}
+
+	return o.EC2IMDSEndpoint, true, nil
+}
+
+// WithEC2IMDSClientEnableState is a helper function to construct functional options that sets the EC2IMDSClientEnableState.
+func WithEC2IMDSClientEnableState(v imds.ClientEnableState) LoadOptionsFunc {
+	return func(o *LoadOptions) error {
+		o.EC2IMDSClientEnableState = v
+		return nil
+	}
+}
+
+// WithEC2IMDSEndpointMode is a helper function to construct functional options that sets the EC2IMDSEndpointMode.
+func WithEC2IMDSEndpointMode(v imds.EndpointModeState) LoadOptionsFunc {
+	return func(o *LoadOptions) error {
+		o.EC2IMDSEndpointMode = v
+		return nil
+	}
+}
+
+// WithEC2IMDSEndpoint is a helper function to construct functional options that sets the EC2IMDSEndpoint.
+func WithEC2IMDSEndpoint(v string) LoadOptionsFunc {
+	return func(o *LoadOptions) error {
+		o.EC2IMDSEndpoint = v
+		return nil
+	}
+}

config/resolve_credentials.go

@@ -289,14 +289,8 @@ func resolveEC2RoleCredentials(ctx context.Context, cfg *aws.Config, configs con
 
 	optFns = append(optFns, func(o *ec2rolecreds.Options) {
 		// Only define a client from config if not already defined.
-		if o.Client != nil {
-			options := imds.Options{
-				HTTPClient: cfg.HTTPClient,
-			}
-			if cfg.Retryer != nil {
-				options.Retryer = cfg.Retryer()
-			}
-			o.Client = imds.New(options)
+		if o.Client == nil {
+			o.Client = imds.NewFromConfig(*cfg)
 		}
 	})