nodeadm: retry IMDS 404s and cleanup callers (#1970)

nodeadm/cmd/nodeadm/init/init.go

@@ -5,14 +5,14 @@ import (
 
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/config"
-	"github.com/aws/aws-sdk-go-v2/feature/ec2/imds"
 	"github.com/aws/aws-sdk-go-v2/service/ec2"
 	"github.com/integrii/flaggy"
 	"go.uber.org/zap"
 	"k8s.io/utils/strings/slices"
 
 	"github.com/awslabs/amazon-eks-ami/nodeadm/internal/api"
 	"github.com/awslabs/amazon-eks-ami/nodeadm/internal/aws/ecr"
+	"github.com/awslabs/amazon-eks-ami/nodeadm/internal/aws/imds"
 	"github.com/awslabs/amazon-eks-ami/nodeadm/internal/cli"
 	"github.com/awslabs/amazon-eks-ami/nodeadm/internal/configprovider"
 	"github.com/awslabs/amazon-eks-ami/nodeadm/internal/containerd"
@@ -146,14 +146,18 @@ func (c *initCmd) Run(log *zap.Logger, opts *cli.GlobalOptions) error {
 // perform in-place updates when allowed by the user
 func enrichConfig(log *zap.Logger, cfg *api.NodeConfig) error {
 	log.Info("Fetching instance details..")
-	imdsClient := imds.New(imds.Options{})
-	awsConfig, err := config.LoadDefaultConfig(context.TODO(), config.WithClientLogMode(aws.LogRetries), config.WithEC2IMDSRegion(func(o *config.UseEC2IMDSRegion) {
-		o.Client = imdsClient
-	}))
+	awsConfig, err := config.LoadDefaultConfig(context.TODO(),
+		config.WithClientLogMode(aws.LogRetries),
+		config.WithEC2IMDSRegion(func(o *config.UseEC2IMDSRegion) {
+			// Use our pre-configured IMDS client to avoid hitting common retry
+			// issues with the default config.
+			o.Client = imds.Client
+		}),
+	)
 	if err != nil {
 		return err
 	}
-	instanceDetails, err := api.GetInstanceDetails(context.TODO(), cfg.Spec.FeatureGates, imdsClient, ec2.NewFromConfig(awsConfig))
+	instanceDetails, err := api.GetInstanceDetails(context.TODO(), cfg.Spec.FeatureGates, ec2.NewFromConfig(awsConfig))
 	if err != nil {
 		return err
 	}

nodeadm/internal/api/status.go

@@ -3,34 +3,29 @@ package api
 import (
 	"context"
 	"fmt"
-	"io"
 	"time"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
-	"github.com/aws/aws-sdk-go-v2/feature/ec2/imds"
 	"github.com/aws/aws-sdk-go-v2/service/ec2"
 	ec2extra "github.com/awslabs/amazon-eks-ami/nodeadm/internal/aws/ec2"
+	"github.com/awslabs/amazon-eks-ami/nodeadm/internal/aws/imds"
 )
 
 // Fetch information about the ec2 instance using IMDS data.
 // This information is stored into the internal config to avoid redundant calls
 // to IMDS when looking for instance metadata
-func GetInstanceDetails(ctx context.Context, featureGates map[Feature]bool, imdsClient *imds.Client, ec2Client *ec2.Client) (*InstanceDetails, error) {
-	instanceIdenitityDocument, err := imdsClient.GetInstanceIdentityDocument(ctx, &imds.GetInstanceIdentityDocumentInput{})
+func GetInstanceDetails(ctx context.Context, featureGates map[Feature]bool, ec2Client *ec2.Client) (*InstanceDetails, error) {
+	instanceIdenitityDocument, err := imds.GetInstanceIdentityDocument(ctx)
 	if err != nil {
 		return nil, err
 	}
 
-	macResponse, err := imdsClient.GetMetadata(ctx, &imds.GetMetadataInput{Path: "mac"})
-	if err != nil {
-		return nil, err
-	}
-	mac, err := io.ReadAll(macResponse.Content)
+	mac, err := imds.GetProperty(ctx, "mac")
 	if err != nil {
 		return nil, err
 	}
 
-	privateDNSName := ""
+	var privateDNSName string
 	if !IsFeatureEnabled(InstanceIdNodeName, featureGates) {
 		privateDNSName, err = getPrivateDNSName(ec2Client, instanceIdenitityDocument.InstanceID)
 		if err != nil {

nodeadm/internal/aws/ecr/ecr.go

@@ -39,7 +39,7 @@ func (r *ECRRegistry) GetSandboxImage() string {
 
 func GetEKSRegistry(region string) (ECRRegistry, error) {
 	account, region := getEKSRegistryCoordinates(region)
-	servicesDomain, err := imds.GetProperty(imds.ServicesDomain)
+	servicesDomain, err := imds.GetProperty(context.TODO(), imds.ServicesDomain)
 	if err != nil {
 		return "", err
 	}

nodeadm/internal/aws/imds/imds.go

@@ -9,14 +9,24 @@ import (
 	"github.com/aws/aws-sdk-go-v2/feature/ec2/imds"
 )
 
-var client *imds.Client
+var Client *imds.Client
 
 func init() {
-	client = imds.New(imds.Options{
+	Client = imds.New(imds.Options{
 		DisableDefaultTimeout: true,
 		Retryer: retry.NewStandard(func(so *retry.StandardOptions) {
 			so.MaxAttempts = 15
 			so.MaxBackoff = 1 * time.Second
+			so.Retryables = append(so.Retryables,
+				&retry.RetryableHTTPStatusCode{
+					Codes: map[int]struct{}{
+						// Retry 404s due to the rare occurrence that
+						// credentials take longer to propagate through IMDS and
+						// fail on the first call.
+						404: {},
+					},
+				},
+			)
 		}),
 	})
 }
@@ -27,24 +37,28 @@ const (
 	ServicesDomain IMDSProperty = "services/domain"
 )
 
-func GetUserData() ([]byte, error) {
-	resp, err := client.GetUserData(context.TODO(), &imds.GetUserDataInput{})
+func GetInstanceIdentityDocument(ctx context.Context) (*imds.GetInstanceIdentityDocumentOutput, error) {
+	return Client.GetInstanceIdentityDocument(ctx, &imds.GetInstanceIdentityDocumentInput{})
+}
+
+func GetUserData(ctx context.Context) ([]byte, error) {
+	res, err := Client.GetUserData(ctx, &imds.GetUserDataInput{})
 	if err != nil {
 		return nil, err
 	}
-	return io.ReadAll(resp.Content)
+	return io.ReadAll(res.Content)
 }
 
-func GetProperty(prop IMDSProperty) (string, error) {
-	bytes, err := GetPropertyBytes(prop)
+func GetProperty(ctx context.Context, prop IMDSProperty) (string, error) {
+	bytes, err := GetPropertyBytes(ctx, prop)
 	if err != nil {
 		return "", err
 	}
 	return string(bytes), nil
 }
 
-func GetPropertyBytes(prop IMDSProperty) ([]byte, error) {
-	res, err := client.GetMetadata(context.TODO(), &imds.GetMetadataInput{Path: string(prop)})
+func GetPropertyBytes(ctx context.Context, prop IMDSProperty) ([]byte, error) {
+	res, err := Client.GetMetadata(ctx, &imds.GetMetadataInput{Path: string(prop)})
 	if err != nil {
 		return nil, err
 	}

nodeadm/internal/configprovider/userdata.go

@@ -1,6 +1,7 @@
 package configprovider
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/awslabs/amazon-eks-ami/nodeadm/api"
@@ -22,7 +23,7 @@ type userDataProvider interface {
 type imdsUserDataProvider struct{}
 
 func (p *imdsUserDataProvider) GetUserData() ([]byte, error) {
-	return imds.GetUserData()
+	return imds.GetUserData(context.TODO())
 }
 
 type userDataConfigProvider struct {

nodeadm/internal/kubelet/config.go

@@ -5,7 +5,6 @@ import (
 	_ "embed"
 	"encoding/json"
 	"fmt"
-	"io"
 	"net"
 	"net/url"
 	"os"
@@ -20,10 +19,10 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	k8skubelet "k8s.io/kubelet/config/v1beta1"
 
-	"github.com/aws/aws-sdk-go-v2/feature/ec2/imds"
 	"github.com/aws/smithy-go/ptr"
 
 	"github.com/awslabs/amazon-eks-ami/nodeadm/internal/api"
+	"github.com/awslabs/amazon-eks-ami/nodeadm/internal/aws/imds"
 	"github.com/awslabs/amazon-eks-ami/nodeadm/internal/containerd"
 	"github.com/awslabs/amazon-eks-ami/nodeadm/internal/system"
 	"github.com/awslabs/amazon-eks-ami/nodeadm/internal/util"
@@ -203,7 +202,7 @@ func (ksc *kubeletConfig) withOutpostSetup(cfg *api.NodeConfig) error {
 }
 
 func (ksc *kubeletConfig) withNodeIp(cfg *api.NodeConfig, flags map[string]string) error {
-	nodeIp, err := getNodeIp(context.TODO(), imds.New(imds.Options{}), cfg)
+	nodeIp, err := getNodeIp(context.TODO(), cfg)
 	if err != nil {
 		return err
 	}
@@ -407,36 +406,24 @@ func getProviderId(availabilityZone, instanceId string) string {
 }
 
 // Get the IP of the node depending on the ipFamily configured for the cluster
-func getNodeIp(ctx context.Context, imdsClient *imds.Client, cfg *api.NodeConfig) (string, error) {
+func getNodeIp(ctx context.Context, cfg *api.NodeConfig) (string, error) {
 	ipFamily, err := api.GetCIDRIpFamily(cfg.Spec.Cluster.CIDR)
 	if err != nil {
 		return "", err
 	}
 	switch ipFamily {
 	case api.IPFamilyIPv4:
-		ipv4Response, err := imdsClient.GetMetadata(ctx, &imds.GetMetadataInput{
-			Path: "local-ipv4",
-		})
+		ipv4, err := imds.GetProperty(ctx, "local-ipv4")
 		if err != nil {
 			return "", err
 		}
-		ip, err := io.ReadAll(ipv4Response.Content)
-		if err != nil {
-			return "", err
-		}
-		return string(ip), nil
+		return ipv4, nil
 	case api.IPFamilyIPv6:
-		ipv6Response, err := imdsClient.GetMetadata(ctx, &imds.GetMetadataInput{
-			Path: fmt.Sprintf("network/interfaces/macs/%s/ipv6s", cfg.Status.Instance.MAC),
-		})
-		if err != nil {
-			return "", err
-		}
-		ip, err := io.ReadAll(ipv6Response.Content)
+		ipv6, err := imds.GetProperty(ctx, imds.IMDSProperty(fmt.Sprintf("network/interfaces/macs/%s/ipv6s", cfg.Status.Instance.MAC)))
 		if err != nil {
 			return "", err
 		}
-		return string(ip), nil
+		return ipv6, nil
 	default:
 		return "", fmt.Errorf("invalid ip-family. %s is not one of %v", ipFamily, []api.IPFamily{api.IPFamilyIPv4, api.IPFamilyIPv6})
 	}