From 71819961045386b23edc18455f1b54764292816c Mon Sep 17 00:00:00 2001
From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com>
Date: Mon, 8 Jul 2024 16:00:36 -0600
Subject: [PATCH] [8.14](backport #40135) [aws] Fix missing metrics bug when
 include_linked_accounts set to false (#40137)

* [aws] Fix missing metrics bug when include_linked_accounts set to false (#40135)

(cherry picked from commit 8da898fb475f142dbdcc0987f79b1015e106b0e7)

* Update CHANGELOG.next.asciidoc

---------

Co-authored-by: kaiyan-sheng <kaiyan.sheng@elastic.co>
---
 CHANGELOG.next.asciidoc                    |  1 +
 x-pack/metricbeat/module/aws/utils.go      | 13 ++---
 x-pack/metricbeat/module/aws/utils_test.go | 56 ++++++++++++++++++++++
 3 files changed, 62 insertions(+), 8 deletions(-)

diff --git a/CHANGELOG.next.asciidoc b/CHANGELOG.next.asciidoc
index f4decc5f937..43dd9f6a677 100644
--- a/CHANGELOG.next.asciidoc
+++ b/CHANGELOG.next.asciidoc
@@ -103,6 +103,7 @@ https://github.com/elastic/beats/compare/v8.8.1\...main[Check the HEAD diff]
 - Normalize AWS RDS CPU Utilization values before making the metadata API call. {pull}39664[39664]
 - Fix behavior of pagetypeinfo metrics {pull}39985[39985]
 - Fix query logic for temp and non-temp tablespaces in Oracle module. {issue}38051[38051] {pull}39787[39787]
+- Fix missing metrics from CloudWatch when include_linked_accounts set to false. {issue}40071[40071] {pull}40135[40135]
 
 *Osquerybeat*
 
diff --git a/x-pack/metricbeat/module/aws/utils.go b/x-pack/metricbeat/module/aws/utils.go
index caf695f1cb9..0ed9d983ba5 100644
--- a/x-pack/metricbeat/module/aws/utils.go
+++ b/x-pack/metricbeat/module/aws/utils.go
@@ -81,15 +81,12 @@ func GetListMetricsOutput(namespace string, regionName string, period time.Durat
 		}
 
 		// when IncludeLinkedAccounts is set to false, ListMetrics API does not return any OwningAccounts
-		if page.OwningAccounts == nil {
-			for _, metric := range page.Metrics {
-				metricWithAccountID = append(metricWithAccountID, MetricWithID{metric, monitoringAccountID})
-			}
-			return metricWithAccountID, nil
-		}
-
 		for i, metric := range page.Metrics {
-			metricWithAccountID = append(metricWithAccountID, MetricWithID{metric, page.OwningAccounts[i]})
+			owningAccount := monitoringAccountID
+			if page.OwningAccounts != nil {
+				owningAccount = page.OwningAccounts[i]
+			}
+			metricWithAccountID = append(metricWithAccountID, MetricWithID{metric, owningAccount})
 		}
 	}
 	return metricWithAccountID, nil
diff --git a/x-pack/metricbeat/module/aws/utils_test.go b/x-pack/metricbeat/module/aws/utils_test.go
index 5fe62b5c80c..e3242c61388 100644
--- a/x-pack/metricbeat/module/aws/utils_test.go
+++ b/x-pack/metricbeat/module/aws/utils_test.go
@@ -6,6 +6,7 @@ package aws
 
 import (
 	"context"
+	"strconv"
 	"testing"
 	"time"
 
@@ -49,6 +50,9 @@ type MockCloudWatchClient struct{}
 // MockCloudwatchClientCrossAccounts struct is used for unit tests.
 type MockCloudwatchClientCrossAccounts struct{}
 
+// MockCloudwatchClientMultiplePages struct is used for unit tests.
+type MockCloudwatchClientMultiplePages struct{}
+
 // GetMetricData implements cloudwatch.GetMetricDataAPIClient interface
 func (m *MockCloudWatchClient) GetMetricData(context.Context, *cloudwatch.GetMetricDataInput, ...func(*cloudwatch.Options)) (*cloudwatch.GetMetricDataOutput, error) {
 	emptyString := ""
@@ -158,6 +162,51 @@ func (m *MockCloudwatchClientCrossAccounts) ListMetrics(context.Context, *cloudw
 	}, nil
 }
 
+func (m *MockCloudwatchClientMultiplePages) ListMetrics(ctx context.Context, input *cloudwatch.ListMetricsInput, opts ...func(*cloudwatch.Options)) (*cloudwatch.ListMetricsOutput, error) {
+	var allMetrics = []cloudwatchtypes.Metric{
+		{
+			MetricName: &metricName,
+			Namespace:  &namespace,
+			Dimensions: []cloudwatchtypes.Dimension{
+				{Name: &dimName, Value: &instanceID1},
+			},
+		},
+		{
+			MetricName: awssdk.String("NetworkIn"),
+			Namespace:  &namespace,
+			Dimensions: []cloudwatchtypes.Dimension{
+				{Name: &dimName, Value: &instanceID1},
+			},
+		},
+	}
+
+	pageSize := 1 // Change this to control the number of metrics per page
+	startIndex := 0
+
+	if input.NextToken != nil {
+		index, err := strconv.Atoi(*input.NextToken)
+		if err != nil {
+			return nil, err
+		}
+		startIndex = index
+	}
+
+	endIndex := startIndex + pageSize
+	if endIndex > len(allMetrics) {
+		endIndex = len(allMetrics)
+	}
+
+	nextToken := ""
+	if endIndex < len(allMetrics) {
+		nextToken = strconv.Itoa(endIndex)
+	}
+
+	return &cloudwatch.ListMetricsOutput{
+		Metrics:   allMetrics[startIndex:endIndex],
+		NextToken: awssdk.String(nextToken),
+	}, nil
+}
+
 // MockResourceGroupsTaggingClient is used for unit tests.
 type MockResourceGroupsTaggingClient struct{}
 
@@ -221,6 +270,13 @@ func TestGetListMetricsCrossAccountsOutput(t *testing.T) {
 	assert.Equal(t, instanceID2, *listMetricsOutput[1].Metric.Dimensions[0].Value)
 }
 
+func TestGetListMetricsOutputWithMultiplePages(t *testing.T) {
+	svcCloudwatch := &MockCloudwatchClientMultiplePages{}
+	listMetricsOutput, err := GetListMetricsOutput("AWS/EC2", "us-west-1", time.Minute*5, false, "123", svcCloudwatch)
+	assert.NoError(t, err)
+	assert.Equal(t, 2, len(listMetricsOutput))
+}
+
 func TestGetListMetricsOutputWithWildcard(t *testing.T) {
 	svcCloudwatch := &MockCloudWatchClient{}
 	listMetricsOutput, err := GetListMetricsOutput("*", "us-west-1", time.Minute*5, false, "123", svcCloudwatch)