ec2 client credential init fix for AWS IRSA

pkg/util/ec2/ec2_tags.go

@@ -15,6 +15,7 @@ import (
 	"time"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go-v2/credentials"
 	"github.com/aws/aws-sdk-go-v2/service/ec2"
 	"github.com/aws/aws-sdk-go-v2/service/ec2/types"
@@ -97,7 +98,7 @@ func fetchEc2TagsFromAPI(ctx context.Context) ([]string, error) {
 	// except when a more specific role (e.g. task role in ECS) does not have
 	// EC2:DescribeTags permission, but a more general role (e.g. instance role)
 	// does have it.
-	tags, err := getTagsWithCreds(ctx, instanceIdentity, nil)
+	tags, err := getTagsWithCreds(ctx, instanceIdentity, ec2Connection(ctx, instanceIdentity.Region, nil))
 	if err == nil {
 		return tags, nil
 	}
@@ -111,15 +112,29 @@ func fetchEc2TagsFromAPI(ctx context.Context) ([]string, error) {
 	}
 
 	awsCreds := credentials.NewStaticCredentialsProvider(iamParams.AccessKeyID, iamParams.SecretAccessKey, iamParams.Token)
-	return getTagsWithCreds(ctx, instanceIdentity, awsCreds)
+	return getTagsWithCreds(ctx, instanceIdentity, ec2Connection(ctx, instanceIdentity.Region, awsCreds))
 }
 
-func getTagsWithCreds(ctx context.Context, instanceIdentity *EC2Identity, awsCreds aws.CredentialsProvider) ([]string, error) {
-	connection := ec2.New(ec2.Options{
-		Region:      instanceIdentity.Region,
-		Credentials: awsCreds,
-	})
+type ec2ClientInterface interface {
+	DescribeTags(ctx context.Context, params *ec2.DescribeTagsInput, optFns ...func(*ec2.Options)) (*ec2.DescribeTagsOutput, error)
+}
+
+// ec2Connection creates an ec2 client with the given region and credentials
+func ec2Connection(ctx context.Context, region string, awsCreds aws.CredentialsProvider) ec2ClientInterface {
+	// using aws config to read the build in credentials to set up the ec2 client.
+	cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region), config.WithCredentialsProvider(awsCreds))
+	if err != nil {
+		log.Warnf("unable to get aws configurations: %s", err)
+		return nil
+	}
+	log.Debug("aws config loaded successfully")
+	return ec2.NewFromConfig(cfg)
+}
 
+func getTagsWithCreds(ctx context.Context, instanceIdentity *EC2Identity, connection ec2ClientInterface) ([]string, error) {
+	if connection == nil {
+		return nil, fmt.Errorf("ec2 client is not set, unable to get tags")
+	}
 	// We want to use 'ec2_metadata_timeout' here instead of current context. 'ctx' comes from the agent main and will
 	// only be canceled if the agent is stopped. The default timeout for the AWS SDK is 1 minutes (20s timeout with
 	// 3 retries). Since we call getTagsWithCreds twice in a row, it can be a 2 minutes latency.

pkg/util/ec2/ec2_tags_test.go

@@ -9,19 +9,22 @@ package ec2
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
 	"net/http/httptest"
 	"os"
 	"testing"
 
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-
 	configmock "github.com/DataDog/datadog-agent/pkg/config/mock"
 	pkgconfigsetup "github.com/DataDog/datadog-agent/pkg/config/setup"
 	"github.com/DataDog/datadog-agent/pkg/util/cache"
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/service/ec2"
+	"github.com/aws/aws-sdk-go-v2/service/ec2/types"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestGetIAMRole(t *testing.T) {
@@ -208,3 +211,112 @@ func TestGetTagsFullWorkflow(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, []string{"tag1", "tag2"}, tags)
 }
+
+// Mock implementation of ec2ClientInterface
+type mockEC2Client struct {
+	DescribeTagsFunc func(ctx context.Context, params *ec2.DescribeTagsInput, optFns ...func(*ec2.Options)) (*ec2.DescribeTagsOutput, error)
+}
+
+func (m *mockEC2Client) DescribeTags(ctx context.Context, params *ec2.DescribeTagsInput, optFns ...func(*ec2.Options)) (*ec2.DescribeTagsOutput, error) {
+	return m.DescribeTagsFunc(ctx, params, optFns...)
+}
+
+// Helper function to compare slices of strings
+func equalStringSlices(a, b []string) bool {
+	if len(a) != len(b) {
+		return false
+	}
+	aMap := make(map[string]struct{}, len(a))
+	for _, v := range a {
+		aMap[v] = struct{}{}
+	}
+	for _, v := range b {
+		if _, ok := aMap[v]; !ok {
+			return false
+		}
+	}
+	return true
+}
+
+func TestGetTagsWithCreds(t *testing.T) {
+	tests := []struct {
+		name             string
+		instanceIdentity *EC2Identity
+		mockEC2Client    ec2ClientInterface
+		expectedTags     []string
+		expectedError    assert.ErrorAssertionFunc
+	}{
+		{
+			name: "Successful retrieval of tags",
+			instanceIdentity: &EC2Identity{
+				InstanceID: "i-1234567890abcdef0",
+			},
+			mockEC2Client: setupMockEC2Client(&ec2.DescribeTagsOutput{
+				Tags: []types.TagDescription{
+					{Key: aws.String("Name"), Value: aws.String("TestInstance")},
+					{Key: aws.String("Env"), Value: aws.String("Production")},
+				}}, nil),
+			expectedTags:  []string{"Name:TestInstance", "Env:Production"},
+			expectedError: assert.NoError,
+		},
+		{
+			name: "Excluded tags are filtered out",
+			instanceIdentity: &EC2Identity{
+				InstanceID: "i-1234567890abcdef0",
+			},
+			mockEC2Client: setupMockEC2Client(&ec2.DescribeTagsOutput{
+				Tags: []types.TagDescription{
+					{Key: aws.String("Name"), Value: aws.String("TestInstance")},
+					{Key: aws.String("aws:cloudformation:stack-name"), Value: aws.String("MyStack")},
+				}}, nil),
+			expectedTags:  []string{"Name:TestInstance", "aws:cloudformation:stack-name:MyStack"},
+			expectedError: assert.NoError,
+		},
+		{
+			name: "DescribeTags returns an error",
+			instanceIdentity: &EC2Identity{
+				InstanceID: "i-1234567890abcdef0",
+			},
+			mockEC2Client: setupMockEC2Client(&ec2.DescribeTagsOutput{
+				Tags: nil,
+			}, errors.New("DescribeTags failed")),
+			expectedTags:  nil,
+			expectedError: assert.Error,
+		},
+		{
+			name: "ec2 config failure, connection is nil",
+			instanceIdentity: &EC2Identity{
+				InstanceID: "i-1234567890abcdef0",
+			},
+			mockEC2Client: nil,
+			expectedTags:  nil,
+			expectedError: assert.Error,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Create a background context
+			ctx := context.Background()
+
+			// Call the function under test
+			tags, err := getTagsWithCreds(ctx, tt.instanceIdentity, tt.mockEC2Client)
+
+			// Validate the error
+			tt.expectedError(t, err)
+			// Validate the tags
+			if !equalStringSlices(tags, tt.expectedTags) {
+				t.Fatalf("Expected tags '%v', got '%v'", tt.expectedTags, tags)
+			}
+		})
+	}
+}
+
+func setupMockEC2Client(mockOutput *ec2.DescribeTagsOutput, mockError error) *mockEC2Client {
+	return &mockEC2Client{
+		DescribeTagsFunc: func(_ context.Context, _ *ec2.DescribeTagsInput, _ ...func(*ec2.Options)) (*ec2.DescribeTagsOutput, error) {
+			return mockOutput, mockError
+		},
+	}
+
+}

releasenotes/notes/ec2-cred-fix-9326728a245381e7.yaml

@@ -0,0 +1,12 @@
+# Each section from every release note are combined when the
+# CHANGELOG.rst is rendered. So the text needs to be worded so that
+# it does not depend on any information only available in another
+# section. This may mean repeating some details, but each section
+# must be readable independently of the other.
+#
+# Each section note must be formatted as reStructuredText.
+---
+enhancements:
+  - |
+    Modifies the initialization of the Amazon EC2 client to use the AWS configuration package.
+    This change enables support for built-in credentials, such as AWS IAM Roles for Service Accounts (IRSA).