From da87f30d1ff50211599471fd667cffd0c1d2ff48 Mon Sep 17 00:00:00 2001
From: Andrew Burke <31974658+atburke@users.noreply.github.com>
Date: Thu, 2 Jun 2022 16:28:34 -0700
Subject: [PATCH] Filter out invalid EC2 tag keys (#13131)

This change skips over EC2 tag keys that aren't valid Teleport label keys.
---
 lib/labels/ec2/ec2.go      | 19 +++++++++----------
 lib/labels/ec2/ec2_test.go | 24 ++++++++++++++++++++----
 2 files changed, 29 insertions(+), 14 deletions(-)

diff --git a/lib/labels/ec2/ec2.go b/lib/labels/ec2/ec2.go
index 15b3f3acc9e..00cd1819c28 100644
--- a/lib/labels/ec2/ec2.go
+++ b/lib/labels/ec2/ec2.go
@@ -117,13 +117,16 @@ func (l *EC2) Sync(ctx context.Context) error {
 		if err != nil {
 			return trace.Wrap(err)
 		}
-		m[t] = value
+		if types.IsValidLabelKey(t) {
+			m[toAWSLabel(t)] = value
+		} else {
+			l.c.Log.Debugf("Skipping EC2 tag %q, not a valid label key.", t)
+		}
 	}
 
 	l.mu.Lock()
 	defer l.mu.Unlock()
-	l.labels = toAWSLabels(m)
-
+	l.labels = m
 	return nil
 }
 
@@ -148,11 +151,7 @@ func (l *EC2) periodicUpdateLabels(ctx context.Context) {
 	}
 }
 
-// toAWSLabels formats labels coming from EC2.
-func toAWSLabels(labels map[string]string) map[string]string {
-	m := make(map[string]string, len(labels))
-	for k, v := range labels {
-		m[fmt.Sprintf("%s/%s", AWSNamespace, k)] = v
-	}
-	return m
+// toAWSLabel formats labels coming from EC2.
+func toAWSLabel(key string) string {
+	return fmt.Sprintf("%s/%s", AWSNamespace, key)
 }
diff --git a/lib/labels/ec2/ec2_test.go b/lib/labels/ec2/ec2_test.go
index a86a302b11b..89eff2acb4e 100644
--- a/lib/labels/ec2/ec2_test.go
+++ b/lib/labels/ec2/ec2_test.go
@@ -51,6 +51,7 @@ func (m *mockIMDSClient) GetTagValue(ctx context.Context, key string) (string, e
 func TestEC2LabelsSync(t *testing.T) {
 	ctx := context.Background()
 	tags := map[string]string{"a": "1", "b": "2"}
+	expectedTags := map[string]string{"aws/a": "1", "aws/b": "2"}
 	imdsClient := &mockIMDSClient{
 		tags: tags,
 	}
@@ -59,7 +60,7 @@ func TestEC2LabelsSync(t *testing.T) {
 	})
 	require.NoError(t, err)
 	require.NoError(t, ec2Labels.Sync(ctx))
-	require.Equal(t, toAWSLabels(tags), ec2Labels.Get())
+	require.Equal(t, expectedTags, ec2Labels.Get())
 }
 
 func TestEC2LabelsAsync(t *testing.T) {
@@ -91,17 +92,32 @@ func TestEC2LabelsAsync(t *testing.T) {
 	initialTags := map[string]string{"a": "1", "b": "2"}
 	imdsClient.tags = initialTags
 	ec2Labels.Start(ctx)
-	require.Eventually(t, compareLabels(toAWSLabels(initialTags)), time.Second, 100*time.Microsecond)
+	require.Eventually(t, compareLabels(map[string]string{"aws/a": "1", "aws/b": "2"}), time.Second, 100*time.Microsecond)
 
 	// Check that tags are updated over time.
 	updatedTags := map[string]string{"a": "3", "c": "4"}
 	imdsClient.tags = updatedTags
 	clock.Advance(ec2LabelUpdatePeriod)
-	require.Eventually(t, compareLabels(toAWSLabels(updatedTags)), time.Second, 100*time.Millisecond)
+	require.Eventually(t, compareLabels(map[string]string{"aws/a": "3", "aws/c": "4"}), time.Second, 100*time.Millisecond)
 
 	// Check that service stops updating when closed.
 	cancel()
 	imdsClient.tags = map[string]string{"x": "8", "y": "9", "z": "10"}
 	clock.Advance(ec2LabelUpdatePeriod)
-	require.Eventually(t, compareLabels(toAWSLabels(updatedTags)), time.Second, 100*time.Millisecond)
+	require.Eventually(t, compareLabels(map[string]string{"aws/a": "3", "aws/c": "4"}), time.Second, 100*time.Millisecond)
+}
+
+func TestEC2LabelsValidKey(t *testing.T) {
+	ctx := context.Background()
+	tags := map[string]string{"good-label": "1", "bad-l@bel": "2"}
+	expectedTags := map[string]string{"aws/good-label": "1"}
+	imdsClient := &mockIMDSClient{
+		tags: tags,
+	}
+	ec2Labels, err := New(ctx, &Config{
+		Client: imdsClient,
+	})
+	require.NoError(t, err)
+	require.NoError(t, ec2Labels.Sync(ctx))
+	require.Equal(t, expectedTags, ec2Labels.Get())
 }