Support decrypting API keys encrypted with an encryption context (#77)

DataDog · Jun 7, 2021 · ff5fffb · ff5fffb
1 parent 836741a
commit ff5fffb
Show file tree

Hide file tree

Showing 27 changed files with 141 additions and 34 deletions.
diff --git 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
@@ -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/).
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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,30 +29,58 @@ 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{
 		kmsClient: kms.New(session.New(nil)),
 	}
 }
 
-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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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.