From 13dd9dff50b050ba8ee3d354461daccc42737abd Mon Sep 17 00:00:00 2001
From: Erik Sundell <erik.sundell@grafana.com>
Date: Tue, 9 Mar 2021 20:50:16 +0100
Subject: [PATCH] CloudWatch: Add support for EC2 IAM role (#31804)

* add support for iam role auth

* add stubbable func

* fix broken test

* goimports

* pr feedback
---
 pkg/tsdb/cloudwatch/cloudwatch.go   | 14 ++++++++++-
 pkg/tsdb/cloudwatch/session.go      |  8 +++++++
 pkg/tsdb/cloudwatch/session_test.go | 36 +++++++++++++++++++++++++++++
 3 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/pkg/tsdb/cloudwatch/cloudwatch.go b/pkg/tsdb/cloudwatch/cloudwatch.go
index 0d329afe342..0ee444a3b73 100644
--- a/pkg/tsdb/cloudwatch/cloudwatch.go
+++ b/pkg/tsdb/cloudwatch/cloudwatch.go
@@ -147,6 +147,13 @@ func (e *cloudWatchExecutor) newSession(region string) (*session.Session, error)
 		})
 	case authTypeDefault:
 		plog.Debug("Authenticating towards AWS with default SDK method", "region", dsInfo.Region)
+	case authTypeEC2IAMRole:
+		plog.Debug("Authenticating towards AWS with IAM Role", "region", dsInfo.Region)
+		sess, err := newSession(cfgs...)
+		if err != nil {
+			return nil, err
+		}
+		cfgs = append(cfgs, &aws.Config{Credentials: newEC2RoleCredentials(sess)})
 	default:
 		panic(fmt.Sprintf("Unrecognized authType: %d", dsInfo.AuthType))
 	}
@@ -396,6 +403,7 @@ const (
 	authTypeDefault authType = iota
 	authTypeSharedCreds
 	authTypeKeys
+	authTypeEC2IAMRole
 )
 
 func (at authType) String() string {
@@ -403,9 +411,11 @@ func (at authType) String() string {
 	case authTypeDefault:
 		return "default"
 	case authTypeSharedCreds:
-		return "sharedCreds"
+		return "credentials"
 	case authTypeKeys:
 		return "keys"
+	case authTypeEC2IAMRole:
+		return "ec2_iam_role"
 	default:
 		panic(fmt.Sprintf("Unrecognized auth type %d", at))
 	}
@@ -432,6 +442,8 @@ func (e *cloudWatchExecutor) getDSInfo(region string) *datasourceInfo {
 		at = authTypeKeys
 	case "default":
 		at = authTypeDefault
+	case "ec2_iam_role":
+		at = authTypeEC2IAMRole
 	case "arn":
 		at = authTypeDefault
 		plog.Warn("Authentication type \"arn\" is deprecated, falling back to default")
diff --git a/pkg/tsdb/cloudwatch/session.go b/pkg/tsdb/cloudwatch/session.go
index e93c25aa05c..0cc3b01c9f1 100644
--- a/pkg/tsdb/cloudwatch/session.go
+++ b/pkg/tsdb/cloudwatch/session.go
@@ -5,6 +5,8 @@ import (
 	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/credentials"
+	"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
 	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
 	"github.com/aws/aws-sdk-go/aws/ec2metadata"
 	"github.com/aws/aws-sdk-go/aws/session"
@@ -34,3 +36,9 @@ var newSTSCredentials = stscreds.NewCredentials
 // Stubbable by tests.
 //nolint:gocritic
 var newEC2Metadata = ec2metadata.New
+
+// EC2 role credentials factory.
+// Stubbable by tests.
+var newEC2RoleCredentials = func(sess *session.Session) *credentials.Credentials {
+	return credentials.NewCredentials(&ec2rolecreds.EC2RoleProvider{Client: ec2metadata.New(sess), ExpiryWindow: stscreds.DefaultDuration})
+}
diff --git a/pkg/tsdb/cloudwatch/session_test.go b/pkg/tsdb/cloudwatch/session_test.go
index 24294c5e11f..ef501128fbd 100644
--- a/pkg/tsdb/cloudwatch/session_test.go
+++ b/pkg/tsdb/cloudwatch/session_test.go
@@ -7,6 +7,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/client"
 	"github.com/aws/aws-sdk-go/aws/credentials"
+	"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
 	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
 	"github.com/aws/aws-sdk-go/aws/ec2metadata"
 	"github.com/aws/aws-sdk-go/aws/session"
@@ -105,3 +106,38 @@ func TestNewSession_AssumeRole(t *testing.T) {
 		assert.Empty(t, diff)
 	})
 }
+
+func TestNewSession_EC2IAMRole(t *testing.T) {
+	newSession = func(cfgs ...*aws.Config) (*session.Session, error) {
+		cfg := aws.Config{}
+		cfg.MergeIn(cfgs...)
+		return &session.Session{
+			Config: &cfg,
+		}, nil
+	}
+	newEC2Metadata = func(p client.ConfigProvider, cfgs ...*aws.Config) *ec2metadata.EC2Metadata {
+		return nil
+	}
+	newEC2RoleCredentials = func(sess *session.Session) *credentials.Credentials {
+		return credentials.NewCredentials(&ec2rolecreds.EC2RoleProvider{Client: newEC2Metadata(nil), ExpiryWindow: stscreds.DefaultDuration})
+	}
+
+	t.Run("Credentials are created", func(t *testing.T) {
+		e := newExecutor(nil)
+		e.DataSource = fakeDataSource()
+		e.DataSource.JsonData.Set("authType", "ec2_iam_role")
+
+		sess, err := e.newSession(defaultRegion)
+		require.NoError(t, err)
+		require.NotNil(t, sess)
+
+		expCreds := credentials.NewCredentials(&ec2rolecreds.EC2RoleProvider{
+			Client: newEC2Metadata(nil), ExpiryWindow: stscreds.DefaultDuration,
+		})
+
+		diff := cmp.Diff(expCreds, sess.Config.Credentials, cmp.Exporter(func(_ reflect.Type) bool {
+			return true
+		}), cmpopts.IgnoreFields(stscreds.AssumeRoleProvider{}, "Expiry"))
+		assert.Empty(t, diff)
+	})
+}