Add ability to use KMS encrypted endpoint

go.mod

@@ -4,6 +4,10 @@ go 1.15
 
 require (
 	github.com/aws/aws-sdk-go v1.35.23
+	github.com/aws/aws-sdk-go-v2 v1.17.1
+	github.com/aws/aws-sdk-go-v2/config v1.17.11
+	github.com/aws/aws-sdk-go-v2/service/kms v1.18.16
 	github.com/google/uuid v1.1.2
 	github.com/sirupsen/logrus v1.7.0
+	golang.org/x/sys v0.2.0 // indirect
 )

go.sum

@@ -1,8 +1,36 @@
 github.com/aws/aws-sdk-go v1.35.23 h1:SCP0d0XvyJTDmfnHEQPvBaYi3kea1VNUo7uQmkVgFts=
 github.com/aws/aws-sdk-go v1.35.23/go.mod h1:tlPOdRjfxPBpNIwqDj61rmsnA85v9jc0Ps9+muhnW+k=
+github.com/aws/aws-sdk-go-v2 v1.17.1 h1:02c72fDJr87N8RAC2s3Qu0YuvMRZKNZJ9F+lAehCazk=
+github.com/aws/aws-sdk-go-v2 v1.17.1/go.mod h1:JLnGeGONAyi2lWXI1p0PCIOIy333JMVK1U7Hf0aRFLw=
+github.com/aws/aws-sdk-go-v2/config v1.17.11 h1:9JQUKwRN8oUqeOFIrNaH6RSPmmcNk1+bQrDka/f/bPc=
+github.com/aws/aws-sdk-go-v2/config v1.17.11/go.mod h1:cw6HIEr0FaZQfcoyRWYZpMfv4qAH19hZFZ5mglwWo3g=
+github.com/aws/aws-sdk-go-v2/credentials v1.12.24 h1:yz4fhoMfgwymG0rU6q5eCydFhYNQxk9yrNjMA7L7xmg=
+github.com/aws/aws-sdk-go-v2/credentials v1.12.24/go.mod h1:prZpUfBu1KZLBLVX482Sq4DpDXGugAre08TPEc21GUg=
+github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.12.19 h1:E3PXZSI3F2bzyj6XxUXdTIfvp425HHhwKsFvmzBwHgs=
+github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.12.19/go.mod h1:VihW95zQpeKQWVPGkwT+2+WJNQV8UXFfMTWdU6VErL8=
+github.com/aws/aws-sdk-go-v2/internal/configsources v1.1.25 h1:nBO/RFxeq/IS5G9Of+ZrgucRciie2qpLy++3UGZ+q2E=
+github.com/aws/aws-sdk-go-v2/internal/configsources v1.1.25/go.mod h1:Zb29PYkf42vVYQY6pvSyJCJcFHlPIiY+YKdPtwnvMkY=
+github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.4.19 h1:oRHDrwCTVT8ZXi4sr9Ld+EXk7N/KGssOr2ygNeojEhw=
+github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.4.19/go.mod h1:6Q0546uHDp421okhmmGfbxzq2hBqbXFNpi4k+Q1JnQA=
+github.com/aws/aws-sdk-go-v2/internal/ini v1.3.26 h1:Mza+vlnZr+fPKFKRq/lKGVvM6B/8ZZmNdEopOwSQLms=
+github.com/aws/aws-sdk-go-v2/internal/ini v1.3.26/go.mod h1:Y2OJ+P+MC1u1VKnavT+PshiEuGPyh/7DqxoDNij4/bg=
+github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.9.19 h1:GE25AWCdNUPh9AOJzI9KIJnja7IwUc1WyUqz/JTyJ/I=
+github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.9.19/go.mod h1:02CP6iuYP+IVnBX5HULVdSAku/85eHB2Y9EsFhrkEwU=
+github.com/aws/aws-sdk-go-v2/service/kms v1.18.16 h1:KHzeOb0G5ZvaIOewRSs3iyHR5MeAKkIZ75tUJCO9ijg=
+github.com/aws/aws-sdk-go-v2/service/kms v1.18.16/go.mod h1:kZodDPTQjSH/qM6/OvyTfM5mms5JHB/EKYp5dhn/vI4=
+github.com/aws/aws-sdk-go-v2/service/sso v1.11.25 h1:GFZitO48N/7EsFDt8fMa5iYdmWqkUDDB3Eje6z3kbG0=
+github.com/aws/aws-sdk-go-v2/service/sso v1.11.25/go.mod h1:IARHuzTXmj1C0KS35vboR0FeJ89OkEy1M9mWbK2ifCI=
+github.com/aws/aws-sdk-go-v2/service/ssooidc v1.13.8 h1:jcw6kKZrtNfBPJkaHrscDOZoe5gvi9wjudnxvozYFJo=
+github.com/aws/aws-sdk-go-v2/service/ssooidc v1.13.8/go.mod h1:er2JHN+kBY6FcMfcBBKNGCT3CarImmdFzishsqBmSRI=
+github.com/aws/aws-sdk-go-v2/service/sts v1.17.2 h1:tpwEMRdMf2UsplengAOnmSIRdvAxf75oUFR+blBr92I=
+github.com/aws/aws-sdk-go-v2/service/sts v1.17.2/go.mod h1:bXcN3koeVYiJcdDU89n3kCYILob7Y34AeLopUbZgLT4=
+github.com/aws/smithy-go v1.13.4 h1:/RN2z1txIJWeXeOkzX+Hk/4Uuvv7dWtCjbmVJcrskyk=
+github.com/aws/smithy-go v1.13.4/go.mod h1:Tg+OJXh4MB2R/uN61Ko2f6hTZwB/ZYGOtib8J3gBHzA=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/google/go-cmp v0.5.8 h1:e6P7q2lk1O+qJJb4BtCQXlK8vWEO8V1ZeuEdJNOqZyg=
+github.com/google/go-cmp v0.5.8/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/google/uuid v1.1.2 h1:EVhdT+1Kseyi1/pUmXKaFxYsDNy9RQYkMWRH68J/W7Y=
 github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/jmespath/go-jmespath v0.4.0 h1:BEgLn5cpjn8UN1mAw4NjwDrS35OdebyEtFe+9YPoQUg=
@@ -21,8 +49,9 @@ golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACk
 golang.org/x/net v0.0.0-20200202094626-16171245cfb2 h1:CCH4IOTTfewWjGOlSp+zGcjutRKlBEZQ6wTn8ozI/nI=
 golang.org/x/net v0.0.0-20200202094626-16171245cfb2/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
-golang.org/x/sys v0.0.0-20191026070338-33540a1f6037 h1:YyJpGZS1sBuBCzLAR1VEpK193GlqGZbnPFnPV/5Rsb4=
 golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.2.0 h1:ljd4t30dBnAvMZaQCevtY0xLLD0A+bRZXbgLMLU1F/A=
+golang.org/x/sys v0.2.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/text v0.3.0 h1:g61tztE5qeGQ89tm6NTjjM9VPIm088od1l6aSorWRWg=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=

lambda-extensions/config/config.go

@@ -17,6 +17,7 @@ import (
 // LambdaExtensionConfig config for storing all configurable parameters
 type LambdaExtensionConfig struct {
 	SumoHTTPEndpoint       string
+	KMSKeyId               string
 	EnableFailover         bool
 	S3BucketName           string
 	S3BucketRegion         string
@@ -36,6 +37,7 @@ type LambdaExtensionConfig struct {
 	SourceCategoryOverride string
 	EnhanceJsonLogs        bool
 	EnableSpanDrops        bool
+	KmsCacheSeconds        int64
 }
 
 var defaultLogTypes = []string{"platform", "function"}
@@ -46,6 +48,7 @@ func GetConfig() (*LambdaExtensionConfig, error) {
 
 	config := &LambdaExtensionConfig{
 		SumoHTTPEndpoint:       os.Getenv("SUMO_HTTP_ENDPOINT"),
+		KMSKeyId:               os.Getenv("KMS_KEY_ID"),
 		S3BucketName:           os.Getenv("SUMO_S3_BUCKET_NAME"),
 		S3BucketRegion:         os.Getenv("SUMO_S3_BUCKET_REGION"),
 		AWSLambdaRuntimeAPI:    os.Getenv("AWS_LAMBDA_RUNTIME_API"),
@@ -78,40 +81,54 @@ func (cfg *LambdaExtensionConfig) setDefaults() {
 	logTypes := os.Getenv("SUMO_LOG_TYPES")
 	enhanceJsonLogs := os.Getenv("SUMO_ENHANCE_JSON_LOGS")
 	enableSpanDrops := os.Getenv("SUMO_SPAN_DROP")
+	kmsCacheSeconds := os.Getenv("KMS_CACHE_SECONDS")
 
 	if numRetry == "" {
 		cfg.NumRetry = 3
 	}
+
 	if logLevel == "" {
 		cfg.LogLevel = logrus.InfoLevel
 	}
+
 	if maxDataQueueLength == "" {
 		cfg.MaxDataQueueLength = 20
 	}
+
 	if maxConcurrentRequests == "" {
 		cfg.MaxConcurrentRequests = 3
 	}
+
 	if enableFailover == "" {
 		cfg.EnableFailover = false
 	}
+
 	if cfg.AWSLambdaRuntimeAPI == "" {
 		cfg.AWSLambdaRuntimeAPI = "127.0.0.1:9001"
 	}
+
 	if logTypes == "" {
 		cfg.LogTypes = defaultLogTypes
 	} else {
 		cfg.LogTypes = strings.Split(logTypes, ",")
 	}
+
 	if retrySleepTime == "" {
 		cfg.RetrySleepTime = 300 * time.Millisecond
 	}
+
 	if enhanceJsonLogs == "" {
 		cfg.EnhanceJsonLogs = true
 	}
+
 	if enableSpanDrops == "" {
 		// by default, spans will not be dropped if user did not configure the env variable
 		cfg.EnableSpanDrops = false
 	}
+
+	if kmsCacheSeconds == "" {
+		cfg.KmsCacheSeconds = 5
+	}
 }
 
 func (cfg *LambdaExtensionConfig) validateConfig() error {
@@ -123,6 +140,7 @@ func (cfg *LambdaExtensionConfig) validateConfig() error {
 	retrySleepTime := os.Getenv("SUMO_RETRY_SLEEP_TIME_MS")
 	enhanceJsonLogs := os.Getenv("SUMO_ENHANCE_JSON_LOGS")
 	enableSpanDrops := os.Getenv("SUMO_SPAN_DROP")
+	kmsCacheSeconds := os.Getenv("KMS_CACHE_SECONDS")
 
 	var allErrors []string
 	var err error
@@ -132,7 +150,7 @@ func (cfg *LambdaExtensionConfig) validateConfig() error {
 	}
 
 	// Todo test url valid
-	if cfg.SumoHTTPEndpoint != "" {
+	if cfg.SumoHTTPEndpoint != "" && cfg.KMSKeyId == "" {
 		_, err = url.ParseRequestURI(cfg.SumoHTTPEndpoint)
 		if err != nil {
 			allErrors = append(allErrors, "SUMO_HTTP_ENDPOINT is not Valid")
@@ -214,6 +232,13 @@ func (cfg *LambdaExtensionConfig) validateConfig() error {
 		}
 	}
 
+	if kmsCacheSeconds != "" {
+		cfg.KmsCacheSeconds, err = strconv.ParseInt(kmsCacheSeconds, 10, 32)
+		if err != nil {
+			allErrors = append(allErrors, fmt.Sprintf("Unable to parse KMS_CACHE_SECONDS: %v", err))
+		}
+	}
+
 	// test valid log format type
 	for _, logType := range cfg.LogTypes {
 		if !utils.StringInSlice(strings.TrimSpace(logType), validLogTypes) {

lambda-extensions/sumoclient/sumoclient.go

@@ -3,6 +3,7 @@ package sumoclient
 import (
 	"bytes"
 	"context"
+	b64 "encoding/base64"
 	"encoding/binary"
 	"encoding/json"
 	"fmt"
@@ -14,11 +15,18 @@ import (
 
 	"github.com/SumoLogic/sumologic-lambda-extensions/lambda-extensions/config"
 
+	"github.com/aws/aws-sdk-go-v2/aws"
+	awsConfig "github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/service/kms"
+
 	uuid "github.com/google/uuid"
 	"github.com/sirupsen/logrus"
 )
 
-var isColdStart = true
+var isColdStart bool = true
+
+var decryptedSumoHttpEndpoint string
+var kmsEndpointCacheTime = time.Now().Add(-5 * time.Minute)
 
 // LogSender interface which needs to be implemented to send logs
 type LogSender interface {
@@ -37,6 +45,12 @@ type sumoLogicClient struct {
 // It is assumed that logs will be array of json objects and all channel payloads satisfy this format
 type responseBody []map[string]interface{}
 
+type KMSDecryptAPI interface {
+	Decrypt(ctx context.Context,
+		params *kms.DecryptInput,
+		optFns ...func(*kms.Options)) (*kms.DecryptOutput, error)
+}
+
 // NewLogSenderClient returns interface pointing to the concrete version of LogSender client
 func NewLogSenderClient(logger *logrus.Entry, cfg *config.LambdaExtensionConfig) LogSender {
 	// setting the cold start variable here since this function is called
@@ -56,8 +70,12 @@ func (s *sumoLogicClient) getColdStart() bool {
 }
 
 func (s *sumoLogicClient) makeRequest(ctx context.Context, buf *bytes.Buffer) (*http.Response, error) {
+	endpoint, err := s.getHttpEndpoint()
+	if err != nil {
+		err = fmt.Errorf("Failed to get SUMO HTTP Endpoint", err)
+	}
 
-	request, err := http.NewRequestWithContext(ctx, "POST", s.config.SumoHTTPEndpoint, buf)
+	request, err := http.NewRequestWithContext(ctx, "POST", endpoint, buf)
 	if err != nil {
 		err = fmt.Errorf("http.NewRequest() error: %v", err)
 		return nil, err
@@ -74,6 +92,56 @@ func (s *sumoLogicClient) makeRequest(ctx context.Context, buf *bytes.Buffer) (*
 	return response, err
 }
 
+// Use cached KMS decrypted endpoint, refresh the cached endpoint, or return unencrypted endpoint
+func (s *sumoLogicClient) getHttpEndpoint() (string, error) {
+	if s.config.KMSKeyId == "" {
+		return s.config.SumoHTTPEndpoint, nil
+	}
+
+	if s.config.KMSKeyId != "" && time.Until(kmsEndpointCacheTime) > 0 {
+		return decryptedSumoHttpEndpoint, nil
+	}
+
+	if s.config.KMSKeyId != "" && (time.Until(kmsEndpointCacheTime) <= 0 || s.config.KmsCacheSeconds == 0) {
+
+		cfg, err := awsConfig.LoadDefaultConfig(context.TODO())
+		if err != nil {
+			fmt.Errorf("Configuration error in aws client,", err)
+		}
+
+		client := kms.NewFromConfig(cfg)
+
+		blob, err := b64.StdEncoding.DecodeString(s.config.SumoHTTPEndpoint)
+		if err != nil {
+			fmt.Errorf("Error converting string to blob,", err)
+		}
+
+		input := &kms.DecryptInput{
+			CiphertextBlob: blob,
+			KeyId:          aws.String(s.config.KMSKeyId),
+		}
+
+		result, err := DecodeData(context.TODO(), client, input)
+
+		if err != nil {
+			fmt.Errorf("Got error decrypting data: ", err)
+			return "", err
+		}
+
+		// Set the decrypted endpoint var as decrypted string to use as cache
+		decryptedSumoHttpEndpoint := string(result.Plaintext)
+
+		// Set new cache time
+		kmsEndpointCacheTime = time.Now()
+
+		return decryptedSumoHttpEndpoint, nil
+	}
+
+	err := fmt.Errorf("Failed to select a valid Sumo HTTP endpoint")
+
+	return "", err
+}
+
 // getS3KeyName returns the key by combining function name, version, date and uuid(version 1)
 func (s *sumoLogicClient) getS3KeyName() (string, error) {
 	currentTime := time.Now()
@@ -414,3 +482,7 @@ func (s *sumoLogicClient) postToSumo(ctx context.Context, logStringToSend *strin
 
 	return nil
 }
+
+func DecodeData(c context.Context, api KMSDecryptAPI, input *kms.DecryptInput) (*kms.DecryptOutput, error) {
+	return api.Decrypt(c, input)
+}