From ff5fffbace038f9a2017e70ad7ce5c09ff601b9e Mon Sep 17 00:00:00 2001
From: Nicolas Hinsch <nick.hinsch@datadoghq.com>
Date: Mon, 7 Jun 2021 18:49:51 -0400
Subject: [PATCH] Support decrypting API keys encrypted with an encryption
 context (#77)

---
 LICENSE                                |  2 +-
 NOTICE                                 |  2 +-
 README.md                              |  2 +-
 ddlambda.go                            |  2 +-
 ddlambda_test.go                       |  2 +-
 internal/metrics/api.go                |  2 +-
 internal/metrics/api_test.go           |  2 +-
 internal/metrics/batcher.go            |  2 +-
 internal/metrics/batcher_test.go       |  2 +-
 internal/metrics/constants.go          |  2 +-
 internal/metrics/context.go            |  2 +-
 internal/metrics/context_test.go       |  2 +-
 internal/metrics/kms_decrypter.go      | 47 +++++++++++++---
 internal/metrics/kms_decrypter_test.go | 76 ++++++++++++++++++++++++++
 internal/metrics/listener.go           |  2 +-
 internal/metrics/listener_test.go      |  2 +-
 internal/metrics/model.go              |  2 +-
 internal/metrics/processor.go          |  2 +-
 internal/metrics/processor_test.go     |  2 +-
 internal/metrics/time.go               |  4 +-
 internal/trace/constants.go            |  2 +-
 internal/trace/context.go              |  2 +-
 internal/trace/context_test.go         |  2 +-
 internal/trace/listener.go             |  2 +-
 internal/trace/listener_test.go        |  2 +-
 internal/wrapper/wrap_handler.go       |  2 +-
 internal/wrapper/wrap_handler_test.go  |  2 +-
 27 files changed, 141 insertions(+), 34 deletions(-)
 create mode 100644 internal/metrics/kms_decrypter_test.go

diff --git a/LICENSE b/LICENSE
index 82633259..eac65501 100644
--- a/LICENSE
+++ b/LICENSE
@@ -186,7 +186,7 @@
       same "printed page" as the copyright notice for easier
       identification within third-party archives.
 
-   Copyright 2018 Datadog, Inc.
+   Copyright 2021 Datadog, Inc.
 
    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.
diff --git a/NOTICE b/NOTICE
index e5b59f6f..fe488423 100644
--- a/NOTICE
+++ b/NOTICE
@@ -1,4 +1,4 @@
 Datadog datadog-lambda-go
-Copyright 2019 Datadog, Inc.
+Copyright 2021 Datadog, Inc.
 
 This product includes software developed at Datadog (https://www.datadoghq.com/).
\ No newline at end of file
diff --git a/README.md b/README.md
index e5e33871..49e981fe 100644
--- a/README.md
+++ b/README.md
@@ -110,4 +110,4 @@ For product feedback and questions, join the `#serverless` channel in the [Datad
 
 Unless explicitly stated otherwise all files in this repository are licensed under the Apache License Version 2.0.
 
-This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2020 Datadog, Inc.
+This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2021 Datadog, Inc.
diff --git a/ddlambda.go b/ddlambda.go
index 0fb1c6ef..52b8f541 100644
--- a/ddlambda.go
+++ b/ddlambda.go
@@ -3,7 +3,7 @@
  * under the Apache License Version 2.0.
  *
  * This product includes software developed at Datadog (https://www.datadoghq.com/).
- * Copyright 2019 Datadog, Inc.
+ * Copyright 2021 Datadog, Inc.
  */
 
 package ddlambda
diff --git a/ddlambda_test.go b/ddlambda_test.go
index bc65f236..b9f82ef6 100644
--- a/ddlambda_test.go
+++ b/ddlambda_test.go
@@ -3,7 +3,7 @@
  * under the Apache License Version 2.0.
  *
  * This product includes software developed at Datadog (https://www.datadoghq.com/).
- * Copyright 2019 Datadog, Inc.
+ * Copyright 2021 Datadog, Inc.
  */
 package ddlambda
 
diff --git a/internal/metrics/api.go b/internal/metrics/api.go
index 60e6e19f..f394946d 100644
--- a/internal/metrics/api.go
+++ b/internal/metrics/api.go
@@ -3,7 +3,7 @@
  * under the Apache License Version 2.0.
  *
  * This product includes software developed at Datadog (https://www.datadoghq.com/).
- * Copyright 2019 Datadog, Inc.
+ * Copyright 2021 Datadog, Inc.
  */
 
 package metrics
diff --git a/internal/metrics/api_test.go b/internal/metrics/api_test.go
index 7d1b40b8..bcefa0d5 100644
--- a/internal/metrics/api_test.go
+++ b/internal/metrics/api_test.go
@@ -3,7 +3,7 @@
  * under the Apache License Version 2.0.
  *
  * This product includes software developed at Datadog (https://www.datadoghq.com/).
- * Copyright 2019 Datadog, Inc.
+ * Copyright 2021 Datadog, Inc.
  */
 
 package metrics
diff --git a/internal/metrics/batcher.go b/internal/metrics/batcher.go
index 69097e28..8e141787 100644
--- a/internal/metrics/batcher.go
+++ b/internal/metrics/batcher.go
@@ -3,7 +3,7 @@
  * under the Apache License Version 2.0.
  *
  * This product includes software developed at Datadog (https://www.datadoghq.com/).
- * Copyright 2019 Datadog, Inc.
+ * Copyright 2021 Datadog, Inc.
  */
 
 package metrics
diff --git a/internal/metrics/batcher_test.go b/internal/metrics/batcher_test.go
index 182e40c3..75d26234 100644
--- a/internal/metrics/batcher_test.go
+++ b/internal/metrics/batcher_test.go
@@ -3,7 +3,7 @@
  * under the Apache License Version 2.0.
  *
  * This product includes software developed at Datadog (https://www.datadoghq.com/).
- * Copyright 2019 Datadog, Inc.
+ * Copyright 2021 Datadog, Inc.
  */
 
 package metrics
diff --git a/internal/metrics/constants.go b/internal/metrics/constants.go
index fd7f1738..e77ffe6e 100644
--- a/internal/metrics/constants.go
+++ b/internal/metrics/constants.go
@@ -3,7 +3,7 @@
  * under the Apache License Version 2.0.
  *
  * This product includes software developed at Datadog (https://www.datadoghq.com/).
- * Copyright 2019 Datadog, Inc.
+ * Copyright 2021 Datadog, Inc.
  */
 
 package metrics
diff --git a/internal/metrics/context.go b/internal/metrics/context.go
index 7d7ee19a..744cc8ca 100644
--- a/internal/metrics/context.go
+++ b/internal/metrics/context.go
@@ -3,7 +3,7 @@
  * under the Apache License Version 2.0.
  *
  * This product includes software developed at Datadog (https://www.datadoghq.com/).
- * Copyright 2019 Datadog, Inc.
+ * Copyright 2021 Datadog, Inc.
  */
 
 package metrics
diff --git a/internal/metrics/context_test.go b/internal/metrics/context_test.go
index 3c616348..ef675dc1 100644
--- a/internal/metrics/context_test.go
+++ b/internal/metrics/context_test.go
@@ -3,7 +3,7 @@
  * under the Apache License Version 2.0.
  *
  * This product includes software developed at Datadog (https://www.datadoghq.com/).
- * Copyright 2019 Datadog, Inc.
+ * Copyright 2021 Datadog, Inc.
  */
 
 package metrics
diff --git a/internal/metrics/kms_decrypter.go b/internal/metrics/kms_decrypter.go
index bbcc0a83..bda57b08 100644
--- a/internal/metrics/kms_decrypter.go
+++ b/internal/metrics/kms_decrypter.go
@@ -3,16 +3,19 @@
  * under the Apache License Version 2.0.
  *
  * This product includes software developed at Datadog (https://www.datadoghq.com/).
- * Copyright 2019 Datadog, Inc.
+ * Copyright 2021 Datadog, Inc.
  */
 package metrics
 
 import (
 	"encoding/base64"
 	"fmt"
+	"os"
 
+	"github.com/DataDog/datadog-lambda-go/internal/logger"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/kms"
+	"github.com/aws/aws-sdk-go/service/kms/kmsiface"
 )
 
 type (
@@ -26,6 +29,12 @@ type (
 	}
 )
 
+// functionNameEnvVar is the environment variable that stores the Lambda function name
+const functionNameEnvVar = "AWS_LAMBDA_FUNCTION_NAME"
+
+// encryptionContextKey is the key added to the encryption context by the Lambda console UI
+const encryptionContextKey = "LambdaFunctionName"
+
 // MakeKMSDecrypter creates a new decrypter which uses the AWS KMS service to decrypt variables
 func MakeKMSDecrypter() Decrypter {
 	return &kmsDecrypter{
@@ -33,23 +42,45 @@ func MakeKMSDecrypter() Decrypter {
 	}
 }
 
-func (kd *kmsDecrypter) Decrypt(cipherText string) (string, error) {
+func (kd *kmsDecrypter) Decrypt(ciphertext string) (string, error) {
+	return decryptKMS(kd.kmsClient, ciphertext)
+}
 
-	decodedBytes, err := base64.StdEncoding.DecodeString(cipherText)
+// decryptKMS decodes and deciphers the base64-encoded ciphertext given as a parameter using KMS.
+// For this to work properly, the Lambda function must have the appropriate IAM permissions.
+func decryptKMS(kmsClient kmsiface.KMSAPI, ciphertext string) (string, error) {
+	decodedBytes, err := base64.StdEncoding.DecodeString(ciphertext)
 	if err != nil {
 		return "", fmt.Errorf("Failed to encode cipher text to base64: %v", err)
 	}
 
+	// The Lambda console UI changed the way it encrypts environment variables.
+	// The current behavior as of May 2021 is to encrypt environment variables using the function name as an encryption context.
+	// Previously, the behavior was to encrypt environment variables without an encryption context.
+	// We need to try both, as supplying the incorrect encryption context will cause decryption to fail.
+
+	// Try with encryption context
+	functionName := os.Getenv(functionNameEnvVar)
 	params := &kms.DecryptInput{
 		CiphertextBlob: decodedBytes,
+		EncryptionContext: map[string]*string{
+			encryptionContextKey: &functionName,
+		},
 	}
+	response, err := kmsClient.Decrypt(params)
 
-	response, err := kd.kmsClient.Decrypt(params)
 	if err != nil {
-		return "", fmt.Errorf("Failed to decrypt ciphertext with kms: %v", err)
+		logger.Debug("Failed to decrypt ciphertext with encryption context, retrying without encryption context")
+		// Try without encryption context
+		params = &kms.DecryptInput{
+			CiphertextBlob: decodedBytes,
+		}
+		response, err = kmsClient.Decrypt(params)
+		if err != nil {
+			return "", fmt.Errorf("Failed to decrypt ciphertext with kms: %v", err)
+		}
 	}
-	// Plaintext is a byte array, so convert to string
-	decrypted := string(response.Plaintext[:])
 
-	return decrypted, nil
+	plaintext := string(response.Plaintext)
+	return plaintext, nil
 }
diff --git a/internal/metrics/kms_decrypter_test.go b/internal/metrics/kms_decrypter_test.go
new file mode 100644
index 00000000..52434d23
--- /dev/null
+++ b/internal/metrics/kms_decrypter_test.go
@@ -0,0 +1,76 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2016-present Datadog, Inc.
+
+package metrics
+
+import (
+	"bytes"
+	"errors"
+	"os"
+	"testing"
+
+	"github.com/aws/aws-sdk-go/service/kms"
+	"github.com/aws/aws-sdk-go/service/kms/kmsiface"
+	"github.com/stretchr/testify/assert"
+)
+
+// mockEncryptedAPIKeyBase64 represents an API key encrypted with KMS and encoded as a base64 string
+const mockEncryptedAPIKeyBase64 = "MjIyMjIyMjIyMjIyMjIyMg=="
+
+// mockDecodedEncryptedAPIKey represents the encrypted API key after it has been decoded from base64
+const mockDecodedEncryptedAPIKey = "2222222222222222"
+
+// expectedDecryptedAPIKey represents the true value of the API key after decryption by KMS
+const expectedDecryptedAPIKey = "1111111111111111"
+
+// mockFunctionName represents the name of the current function
+var mockFunctionName = "my-Function"
+
+type mockKMSClientWithEncryptionContext struct {
+	kmsiface.KMSAPI
+}
+
+func (mockKMSClientWithEncryptionContext) Decrypt(params *kms.DecryptInput) (*kms.DecryptOutput, error) {
+	if *params.EncryptionContext[encryptionContextKey] != mockFunctionName {
+		return nil, errors.New("InvalidCiphertextExeption")
+	}
+	if bytes.Equal(params.CiphertextBlob, []byte(mockDecodedEncryptedAPIKey)) {
+		return &kms.DecryptOutput{
+			Plaintext: []byte(expectedDecryptedAPIKey),
+		}, nil
+	}
+	return nil, errors.New("KMS error")
+}
+
+type mockKMSClientNoEncryptionContext struct {
+	kmsiface.KMSAPI
+}
+
+func (mockKMSClientNoEncryptionContext) Decrypt(params *kms.DecryptInput) (*kms.DecryptOutput, error) {
+	if params.EncryptionContext[encryptionContextKey] != nil {
+		return nil, errors.New("InvalidCiphertextExeption")
+	}
+	if bytes.Equal(params.CiphertextBlob, []byte(mockDecodedEncryptedAPIKey)) {
+		return &kms.DecryptOutput{
+			Plaintext: []byte(expectedDecryptedAPIKey),
+		}, nil
+	}
+	return nil, errors.New("KMS error")
+}
+
+func TestDecryptKMSWithEncryptionContext(t *testing.T) {
+	os.Setenv(functionNameEnvVar, mockFunctionName)
+	defer os.Setenv(functionNameEnvVar, "")
+
+	client := mockKMSClientWithEncryptionContext{}
+	result, _ := decryptKMS(client, mockEncryptedAPIKeyBase64)
+	assert.Equal(t, expectedDecryptedAPIKey, result)
+}
+
+func TestDecryptKMSNoEncryptionContext(t *testing.T) {
+	client := mockKMSClientNoEncryptionContext{}
+	result, _ := decryptKMS(client, mockEncryptedAPIKeyBase64)
+	assert.Equal(t, expectedDecryptedAPIKey, result)
+}
diff --git a/internal/metrics/listener.go b/internal/metrics/listener.go
index 3ae9116a..cca4d40a 100644
--- a/internal/metrics/listener.go
+++ b/internal/metrics/listener.go
@@ -3,7 +3,7 @@
  * under the Apache License Version 2.0.
  *
  * This product includes software developed at Datadog (https://www.datadoghq.com/).
- * Copyright 2019 Datadog, Inc.
+ * Copyright 2021 Datadog, Inc.
  */
 
 package metrics
diff --git a/internal/metrics/listener_test.go b/internal/metrics/listener_test.go
index 1216e384..261815f0 100644
--- a/internal/metrics/listener_test.go
+++ b/internal/metrics/listener_test.go
@@ -3,7 +3,7 @@
  * under the Apache License Version 2.0.
  *
  * This product includes software developed at Datadog (https://www.datadoghq.com/).
- * Copyright 2019 Datadog, Inc.
+ * Copyright 2021 Datadog, Inc.
  */
 
 package metrics
diff --git a/internal/metrics/model.go b/internal/metrics/model.go
index 2082ffd0..d4e3aea0 100644
--- a/internal/metrics/model.go
+++ b/internal/metrics/model.go
@@ -3,7 +3,7 @@
  * under the Apache License Version 2.0.
  *
  * This product includes software developed at Datadog (https://www.datadoghq.com/).
- * Copyright 2019 Datadog, Inc.
+ * Copyright 2021 Datadog, Inc.
  */
 
 package metrics
diff --git a/internal/metrics/processor.go b/internal/metrics/processor.go
index c2084e38..960d6d4b 100644
--- a/internal/metrics/processor.go
+++ b/internal/metrics/processor.go
@@ -3,7 +3,7 @@
  * under the Apache License Version 2.0.
  *
  * This product includes software developed at Datadog (https://www.datadoghq.com/).
- * Copyright 2019 Datadog, Inc.
+ * Copyright 2021 Datadog, Inc.
  */
 
 package metrics
diff --git a/internal/metrics/processor_test.go b/internal/metrics/processor_test.go
index 462d57e8..3ef25d17 100644
--- a/internal/metrics/processor_test.go
+++ b/internal/metrics/processor_test.go
@@ -3,7 +3,7 @@
  * under the Apache License Version 2.0.
  *
  * This product includes software developed at Datadog (https://www.datadoghq.com/).
- * Copyright 2019 Datadog, Inc.
+ * Copyright 2021 Datadog, Inc.
  */
 
 package metrics
diff --git a/internal/metrics/time.go b/internal/metrics/time.go
index f563e6bf..ed71525e 100644
--- a/internal/metrics/time.go
+++ b/internal/metrics/time.go
@@ -1,9 +1,9 @@
 /*
  * Unless explicitly stated otherwise all files in this repository are licensed
  * under the Apache License Version 2.0.
- * 
+ *
  * This product includes software developed at Datadog (https://www.datadoghq.com/).
- * Copyright 2019 Datadog, Inc.
+ * Copyright 2021 Datadog, Inc.
  */
 
 package metrics
diff --git a/internal/trace/constants.go b/internal/trace/constants.go
index 73a2381d..09b08605 100644
--- a/internal/trace/constants.go
+++ b/internal/trace/constants.go
@@ -3,7 +3,7 @@
  * under the Apache License Version 2.0.
  *
  * This product includes software developed at Datadog (https://www.datadoghq.com/).
- * Copyright 2019 Datadog, Inc.
+ * Copyright 2021 Datadog, Inc.
  */
 
 package trace
diff --git a/internal/trace/context.go b/internal/trace/context.go
index c7a9769d..fb3d6a70 100644
--- a/internal/trace/context.go
+++ b/internal/trace/context.go
@@ -3,7 +3,7 @@
  * under the Apache License Version 2.0.
  *
  * This product includes software developed at Datadog (https://www.datadoghq.com/).
- * Copyright 2019 Datadog, Inc.
+ * Copyright 2021 Datadog, Inc.
  */
 
 package trace
diff --git a/internal/trace/context_test.go b/internal/trace/context_test.go
index 3bfa2ef1..559580e3 100644
--- a/internal/trace/context_test.go
+++ b/internal/trace/context_test.go
@@ -3,7 +3,7 @@
  * under the Apache License Version 2.0.
  *
  * This product includes software developed at Datadog (https://www.datadoghq.com/).
- * Copyright 2019 Datadog, Inc.
+ * Copyright 2021 Datadog, Inc.
  */
 
 package trace
diff --git a/internal/trace/listener.go b/internal/trace/listener.go
index 3d64e809..d632b65b 100644
--- a/internal/trace/listener.go
+++ b/internal/trace/listener.go
@@ -3,7 +3,7 @@
  * under the Apache License Version 2.0.
  *
  * This product includes software developed at Datadog (https://www.datadoghq.com/).
- * Copyright 2019 Datadog, Inc.
+ * Copyright 2021 Datadog, Inc.
  */
 
 package trace
diff --git a/internal/trace/listener_test.go b/internal/trace/listener_test.go
index 4266b865..93630851 100644
--- a/internal/trace/listener_test.go
+++ b/internal/trace/listener_test.go
@@ -3,7 +3,7 @@
  * under the Apache License Version 2.0.
  *
  * This product includes software developed at Datadog (https://www.datadoghq.com/).
- * Copyright 2019 Datadog, Inc.
+ * Copyright 2021 Datadog, Inc.
  */
 
 package trace
diff --git a/internal/wrapper/wrap_handler.go b/internal/wrapper/wrap_handler.go
index 7a5322f2..dfd95a61 100644
--- a/internal/wrapper/wrap_handler.go
+++ b/internal/wrapper/wrap_handler.go
@@ -3,7 +3,7 @@
  * under the Apache License Version 2.0.
  *
  * This product includes software developed at Datadog (https://www.datadoghq.com/).
- * Copyright 2019 Datadog, Inc.
+ * Copyright 2021 Datadog, Inc.
  */
 
 package wrapper
diff --git a/internal/wrapper/wrap_handler_test.go b/internal/wrapper/wrap_handler_test.go
index 3d5b679d..2c31e998 100644
--- a/internal/wrapper/wrap_handler_test.go
+++ b/internal/wrapper/wrap_handler_test.go
@@ -3,7 +3,7 @@
  * under the Apache License Version 2.0.
  *
  * This product includes software developed at Datadog (https://www.datadoghq.com/).
- * Copyright 2019 Datadog, Inc.
+ * Copyright 2021 Datadog, Inc.
  */
 
 package wrapper