From 86decbdbaea53344dd16b3144c9b2e4b07ec7ba5 Mon Sep 17 00:00:00 2001
From: Andrea Spacca <andrea.spacca@elastic.co>
Date: Wed, 13 Oct 2021 16:13:05 +0200
Subject: [PATCH] =?UTF-8?q?[Libbeat]=20Security=20-=20fetch=20IMDSv2=20tok?=
 =?UTF-8?q?en=20for=20add=5Fcloud=5Fmetadata=20suppor=E2=80=A6=20(#28285)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* [Libbeat] Security - fetch IMDSv2 token for add_cloud_metadata support on aws
---
 CHANGELOG.next.asciidoc                       |  1 +
 .../add_cloud_metadata/provider_aws_ec2.go    | 86 ++++++++++++++++++-
 2 files changed, 85 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.next.asciidoc b/CHANGELOG.next.asciidoc
index f7a40745bab..90c8e764521 100644
--- a/CHANGELOG.next.asciidoc
+++ b/CHANGELOG.next.asciidoc
@@ -23,6 +23,7 @@ https://github.com/elastic/beats/compare/v7.0.0-alpha2...master[Check the HEAD d
 - Remove deprecated/undocumented IncludeCreatorMetadata setting from kubernetes metadata config options {pull}28006[28006]
 - Remove deprecated fields from kubernetes module {pull}28046[28046]
 - Remove deprecated config option aws_partition. {pull}28120[28120]
+- Enable IMDSv2 support for `add_cloud_metadata` processor on AWS. {issue}22101[22101] {pull}28285[28285]
 - Update kubernetes.namespace from keyword to group field and add name, labels, annotations, uuid as its fields {pull}27917[27917]
 
 *Auditbeat*
diff --git a/libbeat/processors/add_cloud_metadata/provider_aws_ec2.go b/libbeat/processors/add_cloud_metadata/provider_aws_ec2.go
index 350ec63b8c6..8978d51ae32 100644
--- a/libbeat/processors/add_cloud_metadata/provider_aws_ec2.go
+++ b/libbeat/processors/add_cloud_metadata/provider_aws_ec2.go
@@ -18,12 +18,88 @@
 package add_cloud_metadata
 
 import (
+	"fmt"
+	"io"
+	"io/ioutil"
+	"net"
+	"net/http"
+
+	"github.com/elastic/beats/v7/libbeat/logp"
+
 	"github.com/elastic/beats/v7/libbeat/common"
 	s "github.com/elastic/beats/v7/libbeat/common/schema"
 	c "github.com/elastic/beats/v7/libbeat/common/schema/mapstriface"
+	"github.com/elastic/beats/v7/libbeat/common/transport/tlscommon"
 )
 
-const ec2InstanceIdentityURI = "/2014-02-25/dynamic/instance-identity/document"
+const (
+	ec2InstanceIdentityURI            = "/2014-02-25/dynamic/instance-identity/document"
+	ec2InstanceIMDSv2TokenValueHeader = "X-aws-ec2-metadata-token"
+	ec2InstanceIMDSv2TokenTTLHeader   = "X-aws-ec2-metadata-token-ttl-seconds"
+	ec2InstanceIMDSv2TokenTTLValue    = "21600"
+	ec2InstanceIMDSv2TokenURI         = "/latest/api/token"
+)
+
+// fetches IMDSv2 token, returns empty one on errors
+func getIMDSv2Token(c *common.Config) string {
+	logger := logp.NewLogger("add_cloud_metadata")
+
+	config := defaultConfig()
+	if err := c.Unpack(&config); err != nil {
+		logger.Warnf("error when load config for getting IMDSv2 token: %s. No token in the metadata request will be used.", err)
+		return ""
+	}
+
+	tlsConfig, err := tlscommon.LoadTLSConfig(config.TLS)
+	if err != nil {
+		logger.Warnf("error when load TLS config for getting IMDSv2 token: %s. No token in the metadata request will be used.", err)
+		return ""
+	}
+
+	client := http.Client{
+		Timeout: config.Timeout,
+		Transport: &http.Transport{
+			DisableKeepAlives: true,
+			DialContext: (&net.Dialer{
+				Timeout:   config.Timeout,
+				KeepAlive: 0,
+			}).DialContext,
+			TLSClientConfig: tlsConfig.ToConfig(),
+		},
+	}
+
+	tokenReq, err := http.NewRequest("PUT", fmt.Sprintf("http://%s%s", metadataHost, ec2InstanceIMDSv2TokenURI), nil)
+	if err != nil {
+		logger.Warnf("error when make token request for getting IMDSv2 token: %s. No token in the metadata request will be used.", err)
+		return ""
+	}
+
+	tokenReq.Header.Add(ec2InstanceIMDSv2TokenTTLHeader, ec2InstanceIMDSv2TokenTTLValue)
+	rsp, err := client.Do(tokenReq)
+	defer func(body io.ReadCloser) {
+		if body != nil {
+			body.Close()
+		}
+	}(rsp.Body)
+
+	if err != nil {
+		logger.Warnf("error when read token request for getting IMDSv2 token: %s. No token in the metadata request will be used.", err)
+		return ""
+	}
+
+	if rsp.StatusCode != http.StatusOK {
+		logger.Warnf("error when check request status for getting IMDSv2 token: http request status %d. No token in the metadata request will be used.", rsp.StatusCode)
+		return ""
+	}
+
+	all, err := ioutil.ReadAll(rsp.Body)
+	if err != nil {
+		logger.Warnf("error when reading token request for getting IMDSv2 token: %s. No token in the metadata request will be used.", err)
+		return ""
+	}
+
+	return string(all)
+}
 
 // AWS EC2 Metadata Service
 var ec2MetadataFetcher = provider{
@@ -48,7 +124,13 @@ var ec2MetadataFetcher = provider{
 			return common.MapStr{"cloud": out}
 		}
 
-		fetcher, err := newMetadataFetcher(config, "aws", nil, metadataHost, ec2Schema, ec2InstanceIdentityURI)
+		headers := make(map[string]string, 1)
+		token := getIMDSv2Token(config)
+		if len(token) > 0 {
+			headers[ec2InstanceIMDSv2TokenValueHeader] = token
+		}
+
+		fetcher, err := newMetadataFetcher(config, "aws", headers, metadataHost, ec2Schema, ec2InstanceIdentityURI)
 		return fetcher, err
 	},
 }