service/s3/s3crypto: Added X-Ray support to s3crypto decryption client (

#2912)

The s3crypto package was not passing a context object down to the KMS client calls so panics were being emitted by the X-Ray SDK. This PR adds support for passing a context to the s3crypto client's PutObjectWithContext and GetObjectWithContext API calls so that the underlying KMS client's API calls are made with the passed in context.

CHANGELOG_PENDING.md

@@ -1,5 +1,7 @@
 ### SDK Features
 
 ### SDK Enhancements
+* `service/s3/s3crypto`: Added X-Ray support to encrypt/decrypt clients ([#2912](https://github.com/aws/aws-sdk-go/pull/2912))
+  * Adds support for passing Context down to the crypto client's KMS client enabling tracing for tools like X-Ray, and metrics.
 
 ### SDK Bugs

service/s3/s3crypto/aes_gcm_content_cipher.go

@@ -2,6 +2,8 @@ package s3crypto
 
 import (
 	"io"
+
+	"github.com/aws/aws-sdk-go/aws"
 )
 
 const (
@@ -20,7 +22,18 @@ func AESGCMContentCipherBuilder(generator CipherDataGenerator) ContentCipherBuil
 }
 
 func (builder gcmContentCipherBuilder) ContentCipher() (ContentCipher, error) {
-	cd, err := builder.generator.GenerateCipherData(gcmKeySize, gcmNonceSize)
+	return builder.ContentCipherWithContext(aws.BackgroundContext())
+}
+
+func (builder gcmContentCipherBuilder) ContentCipherWithContext(ctx aws.Context) (ContentCipher, error) {
+	var cd CipherData
+	var err error
+
+	if v, ok := builder.generator.(CipherDataGeneratorWithContext); ok {
+		cd, err = v.GenerateCipherDataWithContext(ctx, gcmKeySize, gcmNonceSize)
+	} else {
+		cd, err = builder.generator.GenerateCipherData(gcmKeySize, gcmNonceSize)
+	}
 	if err != nil {
 		return nil, err
 	}

service/s3/s3crypto/cipher_builder.go

@@ -1,13 +1,23 @@
 package s3crypto
 
-import "io"
+import (
+	"io"
+
+	"github.com/aws/aws-sdk-go/aws"
+)
 
 // ContentCipherBuilder is a builder interface that builds
 // ciphers for each request.
 type ContentCipherBuilder interface {
 	ContentCipher() (ContentCipher, error)
 }
 
+// ContentCipherBuilderWithContext is a builder interface that builds
+// ciphers for each request.
+type ContentCipherBuilderWithContext interface {
+	ContentCipherWithContext(aws.Context) (ContentCipher, error)
+}
+
 // ContentCipher deals with encrypting and decrypting content
 type ContentCipher interface {
 	EncryptContents(io.Reader) (io.Reader, error)

service/s3/s3crypto/cipher_util.go

@@ -5,16 +5,17 @@ import (
 	"strconv"
 	"strings"
 
+	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/awserr"
 )
 
-func (client *DecryptionClient) contentCipherFromEnvelope(env Envelope) (ContentCipher, error) {
+func (client *DecryptionClient) contentCipherFromEnvelope(ctx aws.Context, env Envelope) (ContentCipher, error) {
 	wrap, err := client.wrapFromEnvelope(env)
 	if err != nil {
 		return nil, err
 	}
 
-	return client.cekFromEnvelope(env, wrap)
+	return client.cekFromEnvelope(ctx, env, wrap)
 }
 
 func (client *DecryptionClient) wrapFromEnvelope(env Envelope) (CipherDataDecrypter, error) {
@@ -36,7 +37,7 @@ const AESGCMNoPadding = "AES/GCM/NoPadding"
 // AESCBC is the string constant that signifies the AES CBC algorithm cipher.
 const AESCBC = "AES/CBC"
 
-func (client *DecryptionClient) cekFromEnvelope(env Envelope, decrypter CipherDataDecrypter) (ContentCipher, error) {
+func (client *DecryptionClient) cekFromEnvelope(ctx aws.Context, env Envelope, decrypter CipherDataDecrypter) (ContentCipher, error) {
 	f, ok := client.CEKRegistry[env.CEKAlg]
 	if !ok || f == nil {
 		return nil, awserr.New(
@@ -55,7 +56,13 @@ func (client *DecryptionClient) cekFromEnvelope(env Envelope, decrypter CipherDa
 	if err != nil {
 		return nil, err
 	}
-	key, err = decrypter.DecryptKey(key)
+
+	if d, ok := decrypter.(CipherDataDecrypterWithContext); ok {
+		key, err = d.DecryptKeyWithContext(ctx, key)
+	} else {
+		key, err = decrypter.DecryptKey(key)
+	}
+
 	if err != nil {
 		return nil, err
 	}

service/s3/s3crypto/cipher_util_test.go

@@ -140,7 +140,7 @@ func TestCEKFactory(t *testing.T) {
 		MatDesc:   `{"kms_cmk_id":""}`,
 	}
 	wrap, err := c.wrapFromEnvelope(env)
-	cek, err := c.cekFromEnvelope(env, wrap)
+	cek, err := c.cekFromEnvelope(aws.BackgroundContext(), env, wrap)
 
 	if err != nil {
 		t.Errorf("expected no error, but received %v", err)
@@ -199,7 +199,7 @@ func TestCEKFactoryNoCEK(t *testing.T) {
 		MatDesc:   `{"kms_cmk_id":""}`,
 	}
 	wrap, err := c.wrapFromEnvelope(env)
-	cek, err := c.cekFromEnvelope(env, wrap)
+	cek, err := c.cekFromEnvelope(aws.BackgroundContext(), env, wrap)
 
 	if err == nil {
 		t.Error("expected error, but received none")
@@ -256,7 +256,7 @@ func TestCEKFactoryCustomEntry(t *testing.T) {
 		MatDesc:   `{"kms_cmk_id":""}`,
 	}
 	wrap, err := c.wrapFromEnvelope(env)
-	cek, err := c.cekFromEnvelope(env, wrap)
+	cek, err := c.cekFromEnvelope(aws.BackgroundContext(), env, wrap)
 
 	if err != nil {
 		t.Errorf("expected no error, but received %v", err)

service/s3/s3crypto/decryption_client.go

@@ -96,7 +96,7 @@ func (c *DecryptionClient) GetObjectRequest(input *s3.GetObjectInput) (*request.
 
 		// If KMS should return the correct CEK algorithm with the proper
 		// KMS key provider
-		cipher, err := c.contentCipherFromEnvelope(env)
+		cipher, err := c.contentCipherFromEnvelope(r.Context(), env)
 		if err != nil {
 			r.Error = err
 			out.Body.Close()

service/s3/s3crypto/encryption_client.go

@@ -84,12 +84,20 @@ func (c *EncryptionClient) PutObjectRequest(input *s3.PutObjectInput) (*request.
 		return req, out
 	}
 
-	encryptor, err := c.ContentCipherBuilder.ContentCipher()
 	req.Handlers.Build.PushFront(func(r *request.Request) {
 		if err != nil {
 			r.Error = err
 			return
 		}
+		var encryptor ContentCipher
+		if v, ok := c.ContentCipherBuilder.(ContentCipherBuilderWithContext); ok {
+			encryptor, err = v.ContentCipherWithContext(r.Context())
+		} else {
+			encryptor, err = c.ContentCipherBuilder.ContentCipher()
+		}
+		if err != nil {
+			r.Error = err
+		}
 
 		md5 := newMD5Reader(input.Body)
 		sha := newSHA256Writer(dst)

service/s3/s3crypto/key_handler.go

@@ -1,6 +1,10 @@
 package s3crypto
 
-import "crypto/rand"
+import (
+	"crypto/rand"
+
+	"github.com/aws/aws-sdk-go/aws"
+)
 
 // CipherDataGenerator handles generating proper key and IVs of proper size for the
 // content cipher. CipherDataGenerator will also encrypt the key and store it in
@@ -9,11 +13,23 @@ type CipherDataGenerator interface {
 	GenerateCipherData(int, int) (CipherData, error)
 }
 
+// CipherDataGeneratorWithContext handles generating proper key and IVs of
+// proper size for the content cipher. CipherDataGenerator will also encrypt
+// the key and store it in the CipherData.
+type CipherDataGeneratorWithContext interface {
+	GenerateCipherDataWithContext(aws.Context, int, int) (CipherData, error)
+}
+
 // CipherDataDecrypter is a handler to decrypt keys from the envelope.
 type CipherDataDecrypter interface {
 	DecryptKey([]byte) ([]byte, error)
 }
 
+// CipherDataDecrypterWithContext is a handler to decrypt keys from the envelope with request context.
+type CipherDataDecrypterWithContext interface {
+	DecryptKeyWithContext(aws.Context, []byte) ([]byte, error)
+}
+
 func generateBytes(n int) []byte {
 	b := make([]byte, n)
 	rand.Read(b)

service/s3/s3crypto/kms_key_handler.go

@@ -98,11 +98,17 @@ func (kp kmsKeyHandler) decryptHandler(env Envelope) (CipherDataDecrypter, error
 
 // DecryptKey makes a call to KMS to decrypt the key.
 func (kp *kmsKeyHandler) DecryptKey(key []byte) ([]byte, error) {
-	out, err := kp.kms.Decrypt(&kms.DecryptInput{
-		EncryptionContext: map[string]*string(kp.CipherData.MaterialDescription),
-		CiphertextBlob:    key,
-		GrantTokens:       []*string{},
-	})
+	return kp.DecryptKeyWithContext(aws.BackgroundContext(), key)
+}
+
+// DecryptKeyWithContext makes a call to KMS to decrypt the key with request context.
+func (kp *kmsKeyHandler) DecryptKeyWithContext(ctx aws.Context, key []byte) ([]byte, error) {
+	out, err := kp.kms.DecryptWithContext(ctx,
+		&kms.DecryptInput{
+			EncryptionContext: kp.CipherData.MaterialDescription,
+			CiphertextBlob:    key,
+			GrantTokens:       []*string{},
+		})
 	if err != nil {
 		return nil, err
 	}
@@ -112,11 +118,18 @@ func (kp *kmsKeyHandler) DecryptKey(key []byte) ([]byte, error) {
 // GenerateCipherData makes a call to KMS to generate a data key, Upon making
 // the call, it also sets the encrypted key.
 func (kp *kmsKeyHandler) GenerateCipherData(keySize, ivSize int) (CipherData, error) {
-	out, err := kp.kms.GenerateDataKey(&kms.GenerateDataKeyInput{
-		EncryptionContext: kp.CipherData.MaterialDescription,
-		KeyId:             kp.cmkID,
-		KeySpec:           aws.String("AES_256"),
-	})
+	return kp.GenerateCipherDataWithContext(aws.BackgroundContext(), keySize, ivSize)
+}
+
+// GenerateCipherDataWithContext makes a call to KMS to generate a data key,
+// Upon making the call, it also sets the encrypted key.
+func (kp *kmsKeyHandler) GenerateCipherDataWithContext(ctx aws.Context, keySize, ivSize int) (CipherData, error) {
+	out, err := kp.kms.GenerateDataKeyWithContext(ctx,
+		&kms.GenerateDataKeyInput{
+			EncryptionContext: kp.CipherData.MaterialDescription,
+			KeyId:             kp.cmkID,
+			KeySpec:           aws.String("AES_256"),
+		})
 	if err != nil {
 		return CipherData{}, err
 	}