From bf6793893a4c86d41f8efc5f0f94e91e59323839 Mon Sep 17 00:00:00 2001
From: Nelson Kopliku <nelson.kopliku@suse.com>
Date: Sat, 29 Mar 2025 14:46:41 +0100
Subject: [PATCH 01/11] Support aws imds v2

---
 internal/core/cloud/aws.go      |  74 +++++++-
 internal/core/cloud/aws_test.go | 295 ++++++++++++++++++++++++++++----
 2 files changed, 330 insertions(+), 39 deletions(-)

diff --git a/internal/core/cloud/aws.go b/internal/core/cloud/aws.go
index 6ae6048..db82c7e 100644
--- a/internal/core/cloud/aws.go
+++ b/internal/core/cloud/aws.go
@@ -12,12 +12,22 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"strconv"
 	"strings"
+
+	"github.com/pkg/errors"
+	log "github.com/sirupsen/logrus"
 )
 
 const (
-	awsMetadataURL      = "http://169.254.169.254/latest/"
-	awsMetadataResource = "meta-data"
+	awsMetadataURL                = "http://169.254.169.254/latest/"
+	awsMetadataResource           = "meta-data"
+	metadataTokenTTLHeader string = "X-aws-ec2-metadata-token-ttl-seconds" //nolint
+	metadataTokenHeader    string = "X-aws-ec2-metadata-token"             //nolint
+
+	// https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-instance-metadata-service.html
+	// (TTL) for the token, in seconds, up to a maximum of six hours (21,600 seconds).
+	awsEC2MetadataTokenTTL int = 60 * 2 // 2 minutes should be more than enough for the curren discovery loop
 )
 
 type AWSMetadata struct {
@@ -86,8 +96,13 @@ func NewAWSMetadata(ctx context.Context, client HTTPClient) (*AWSMetadata, error
 		},
 	}
 
+	token, err := requestMetadataToken(ctx, client)
+	if err != nil {
+		return nil, err
+	}
+
 	firstElementsList := []string{fmt.Sprintf("%s/", awsMetadataResource)}
-	metadata, err := buildAWSMetadata(ctx, client, awsMetadataURL, firstElementsList)
+	metadata, err := buildAWSMetadata(ctx, client, awsMetadataURL, firstElementsList, token)
 	if err != nil {
 		return nil, err
 	}
@@ -110,13 +125,14 @@ func buildAWSMetadata(
 	client HTTPClient,
 	url string,
 	elements []string,
-) (map[string]interface{}, error) {
-	metadata := make(map[string]interface{})
+	token string,
+) (map[string]any, error) {
+	metadata := make(map[string]any)
 
 	for _, element := range elements {
 		newURL := url + element
 
-		response, err := requestMetadata(ctx, client, newURL)
+		response, err := requestMetadata(ctx, client, newURL, token)
 		if err != nil {
 			return metadata, err
 		}
@@ -125,7 +141,7 @@ func buildAWSMetadata(
 			currentElement := strings.Trim(element, "/")
 			newElements := strings.Split(fmt.Sprintf("%v", response), "\n")
 
-			metadata[currentElement], err = buildAWSMetadata(ctx, client, newURL, newElements)
+			metadata[currentElement], err = buildAWSMetadata(ctx, client, newURL, newElements, token)
 			if err != nil {
 				return nil, err
 			}
@@ -137,8 +153,14 @@ func buildAWSMetadata(
 	return metadata, nil
 }
 
-func requestMetadata(ctx context.Context, client HTTPClient, url string) (interface{}, error) {
+func requestMetadata(
+	ctx context.Context,
+	client HTTPClient,
+	url string,
+	token string,
+) (any, error) {
 	req, _ := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+	req.Header.Add(metadataTokenHeader, token)
 
 	resp, err := client.Do(req)
 	if err != nil {
@@ -151,11 +173,45 @@ func requestMetadata(ctx context.Context, client HTTPClient, url string) (interf
 		return nil, err
 	}
 
+	if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden {
+		return nil, errors.Errorf("failed to fetch AWS metadata: %s", resp.Status)
+	}
+
 	// The metadata endpoint may return json elements
 	if json.Valid(body) {
-		var jsonData interface{}
+		var jsonData any
 		err := json.Unmarshal(body, &jsonData)
 		return jsonData, err
 	}
 	return string(body), nil
 }
+
+func requestMetadataToken(ctx context.Context, client HTTPClient) (string, error) {
+	log.Debug("Fetching IMDS token...")
+
+	url := fmt.Sprintf("%sapi/token", awsMetadataURL)
+
+	req, _ := http.NewRequestWithContext(ctx, http.MethodPut, url, nil)
+
+	req.Header.Add(metadataTokenTTLHeader, strconv.Itoa(awsEC2MetadataTokenTTL))
+
+	resp, err := client.Do(req)
+	if err != nil {
+		log.Debugf("An error occurred while fetching IMDS token: %s", err)
+		return "", err
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		return "", errors.Errorf("failed to fetch metadata token: %s", resp.Status)
+	}
+
+	defer resp.Body.Close()
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return "", err
+	}
+
+	log.Debug("Metadata token fetched successfully")
+
+	return string(body), nil
+}
diff --git a/internal/core/cloud/aws_test.go b/internal/core/cloud/aws_test.go
index 8504e43..622a366 100644
--- a/internal/core/cloud/aws_test.go
+++ b/internal/core/cloud/aws_test.go
@@ -3,6 +3,7 @@ package cloud_test
 import (
 	"bytes"
 	"context"
+	"fmt"
 	"io"
 	"net/http"
 	"os"
@@ -18,42 +19,134 @@ import (
 
 type AWSMetadataTestSuite struct {
 	suite.Suite
+	mockHTTPClient *mocks.HTTPClient
 }
 
-func TestAWSMetadataTestSuite(t *testing.T) {
-	suite.Run(t, new(AWSMetadataTestSuite))
+type metadataFixtureSpec struct {
+	fixture     string
+	requestPath string
 }
 
-func (suite *AWSMetadataTestSuite) TestNewAWSMetadata() {
-	ctx := context.TODO()
-	clientMock := new(mocks.HTTPClient)
-
-	fixtures := []string{
-		"meta-data",
-		"ami-id",
-		"block-device-mapping",
-		"ebs1",
-		"ebs2",
-		"root",
-		"identity-credentials",
-		"ec2",
-		"info",
-		"instance-id",
-		"instance-type",
-		"network",
-		"interfaces",
-		"macs",
-		"some-mac",
-		"vpc-id",
-		"placement",
-		"availability-zone",
-		"region",
+const mockedIMDSToken string = "awsMetadataToken"
+
+func fetchTokenRequest() any {
+	return mock.MatchedBy(func(req *http.Request) bool {
+		return req.Method == http.MethodPut &&
+			req.URL.Path == "/latest/api/token" &&
+			req.Header.Get("X-aws-ec2-metadata-token-ttl-seconds") == "120"
+	})
+}
+
+func matchesRequestPathAndToken(matchingPath string, matchingToken string) any {
+	return mock.MatchedBy(func(req *http.Request) bool {
+		return req.Method == http.MethodGet &&
+			req.URL.Path == matchingPath &&
+			req.Header.Get("X-aws-ec2-metadata-token") == matchingToken
+	})
+}
+
+func mockSuccessfulTokenResponse(mockHTTPClient *mocks.HTTPClient) *mock.Call {
+	token := []byte(mockedIMDSToken)
+
+	body := io.NopCloser(bytes.NewReader(token))
+	successfulTokenResponse := &http.Response{
+		StatusCode: 200,
+		Body:       body,
+	}
+
+	return mockHTTPClient.On("Do", fetchTokenRequest()).Return(
+		successfulTokenResponse, nil,
+	)
+}
+
+func mockSuccessfulMetadataDiscoveryUntil(mockHTTPClient *mocks.HTTPClient, untilFixture string, failureResponse ...*http.Response) {
+	fixtures := []metadataFixtureSpec{
+		{
+			fixture:     "meta-data",
+			requestPath: "/latest/meta-data/",
+		},
+		{
+			fixture:     "ami-id",
+			requestPath: "/latest/meta-data/ami-id",
+		},
+		{
+			fixture:     "block-device-mapping",
+			requestPath: "/latest/meta-data/block-device-mapping/",
+		},
+		{
+			fixture:     "ebs1",
+			requestPath: "/latest/meta-data/block-device-mapping/ebs1",
+		},
+		{
+			fixture:     "ebs2",
+			requestPath: "/latest/meta-data/block-device-mapping/ebs2",
+		},
+		{
+			fixture:     "root",
+			requestPath: "/latest/meta-data/block-device-mapping/root",
+		},
+		{
+			fixture:     "identity-credentials",
+			requestPath: "/latest/meta-data/identity-credentials/",
+		},
+		{
+			fixture:     "ec2",
+			requestPath: "/latest/meta-data/identity-credentials/ec2/",
+		},
+		{
+			fixture:     "info",
+			requestPath: "/latest/meta-data/identity-credentials/ec2/info",
+		},
+		{
+			fixture:     "instance-id",
+			requestPath: "/latest/meta-data/instance-id",
+		},
+		{
+			fixture:     "instance-type",
+			requestPath: "/latest/meta-data/instance-type",
+		},
+		{
+			fixture:     "network",
+			requestPath: "/latest/meta-data/network/",
+		},
+		{
+			fixture:     "interfaces",
+			requestPath: "/latest/meta-data/network/interfaces/",
+		},
+		{
+			fixture:     "macs",
+			requestPath: "/latest/meta-data/network/interfaces/macs/",
+		},
+		{
+			fixture:     "some-mac",
+			requestPath: "/latest/meta-data/network/interfaces/macs/some-mac/",
+		},
+		{
+			fixture:     "vpc-id",
+			requestPath: "/latest/meta-data/network/interfaces/macs/some-mac/vpc-id",
+		},
+		{
+			fixture:     "placement",
+			requestPath: "/latest/meta-data/placement/",
+		},
+		{
+			fixture:     "availability-zone",
+			requestPath: "/latest/meta-data/placement/availability-zone",
+		},
+		{
+			fixture:     "region",
+			requestPath: "/latest/meta-data/placement/region",
+		},
 	}
 
 	fixturesFolder := helpers.GetFixturePath("discovery/aws")
 
-	for _, fixture := range fixtures {
-		aFile, _ := os.Open(path.Join(fixturesFolder, fixture))
+	matcher := func(fixtureSpec metadataFixtureSpec) any {
+		return matchesRequestPathAndToken(fixtureSpec.requestPath, mockedIMDSToken)
+	}
+
+	for _, fixtureSpec := range fixtures {
+		aFile, _ := os.Open(path.Join(fixturesFolder, fixtureSpec.fixture))
 		bodyText, _ := io.ReadAll(aFile)
 		body := io.NopCloser(bytes.NewReader(bodyText))
 
@@ -62,12 +155,154 @@ func (suite *AWSMetadataTestSuite) TestNewAWSMetadata() {
 			Body:       body,
 		}
 
-		clientMock.On("Do", mock.AnythingOfType("*http.Request")).Return(
+		if fixtureSpec.fixture == untilFixture {
+			if len(failureResponse) > 0 {
+				response = failureResponse[0]
+			} else {
+				response = &http.Response{
+					StatusCode: 403,
+					Status:     "403 Forbidden",
+					Body:       io.NopCloser(bytes.NewReader([]byte(""))),
+				}
+			}
+		}
+
+		mockHTTPClient.On("Do", matcher(fixtureSpec)).Return(
 			response, nil,
 		).Once()
 	}
+}
+
+func mockSuccessfulMetadataDiscovery(mockHTTPClient *mocks.HTTPClient) {
+	mockSuccessfulMetadataDiscoveryUntil(mockHTTPClient, "the-end")
+}
+
+func withEmptyBody(responses ...*http.Response) []*http.Response {
+	emptyBodies := [][]byte{
+		[]byte(""),
+		[]byte(`
+		`),
+		[]byte(`
+	
+		`),
+	}
+
+	resultingResponses := []*http.Response{}
+	for _, errorResponse := range responses {
+		for _, emptyBody := range emptyBodies {
+			errorResponseWithEmptyBody := errorResponse
+			errorResponseWithEmptyBody.Body = io.NopCloser(bytes.NewReader(emptyBody))
+			resultingResponses = append(resultingResponses, errorResponseWithEmptyBody)
+		}
+	}
+	return resultingResponses
+}
+
+func TestAWSMetadataTestSuite(t *testing.T) {
+	suite.Run(t, new(AWSMetadataTestSuite))
+}
+
+func (suite *AWSMetadataTestSuite) SetupTest() {
+	suite.mockHTTPClient = new(mocks.HTTPClient)
+}
 
-	m, err := cloud.NewAWSMetadata(ctx, clientMock)
+func (suite *AWSMetadataTestSuite) TestUnableToFetchMetadataToken() {
+	responses := []*http.Response{
+		{
+			StatusCode: 400,
+			Status:     "400 Bad Request",
+		},
+		{
+			StatusCode: 401,
+			Status:     "401 Unauthorized",
+		},
+		{
+			StatusCode: 403,
+			Status:     "403 Forbidden",
+		},
+		{
+			StatusCode: 404,
+			Status:     "404 Not Found",
+		},
+		{
+			StatusCode: 500,
+			Status:     "500 Internal Server Error",
+		},
+	}
+
+	for _, response := range withEmptyBody(responses...) {
+		ctx := context.TODO()
+
+		suite.mockHTTPClient.On("Do", fetchTokenRequest()).Return(
+			response, nil,
+		).Once()
+
+		extractedMetadata, err := cloud.NewAWSMetadata(ctx, suite.mockHTTPClient)
+
+		suite.ErrorContains(err, fmt.Sprintf("failed to fetch metadata token: %s", response.Status))
+		suite.Nil(extractedMetadata)
+	}
+
+}
+
+func (suite *AWSMetadataTestSuite) TestUnableToFetchRootMetadata() {
+	responses := []*http.Response{
+		{
+			StatusCode: 401,
+			Status:     "401 Unauthorized",
+		},
+		{
+			StatusCode: 403,
+			Status:     "403 Forbidden",
+		},
+	}
+
+	for _, response := range withEmptyBody(responses...) {
+		mockSuccessfulTokenResponse(suite.mockHTTPClient).Once()
+		mockSuccessfulMetadataDiscoveryUntil(suite.mockHTTPClient, "meta-data", response)
+
+		ctx := context.TODO()
+
+		extractedMetadata, err := cloud.NewAWSMetadata(ctx, suite.mockHTTPClient)
+
+		suite.ErrorContains(err, fmt.Sprintf("failed to fetch AWS metadata: %s", response.Status))
+
+		suite.Nil(extractedMetadata)
+	}
+}
+
+func (suite *AWSMetadataTestSuite) TestUnableToFetchMetadata() {
+	responses := []*http.Response{
+		{
+			StatusCode: 401,
+			Status:     "401 Unauthorized",
+		},
+		{
+			StatusCode: 403,
+			Status:     "403 Forbidden",
+		},
+	}
+
+	for _, response := range withEmptyBody(responses...) {
+		mockSuccessfulTokenResponse(suite.mockHTTPClient).Once()
+		mockSuccessfulMetadataDiscoveryUntil(suite.mockHTTPClient, "info", response)
+
+		ctx := context.TODO()
+
+		extractedMetadata, err := cloud.NewAWSMetadata(ctx, suite.mockHTTPClient)
+
+		suite.ErrorContains(err, fmt.Sprintf("failed to fetch AWS metadata: %s", response.Status))
+
+		suite.Nil(extractedMetadata)
+	}
+}
+
+func (suite *AWSMetadataTestSuite) TestNewAWSMetadata() {
+	mockSuccessfulTokenResponse(suite.mockHTTPClient).Once()
+	mockSuccessfulMetadataDiscovery(suite.mockHTTPClient)
+
+	ctx := context.TODO()
+	m, err := cloud.NewAWSMetadata(ctx, suite.mockHTTPClient)
 
 	suite.NoError(err)
 
-- 
2.35.3

