[kms] Support more attributes on CreateKey and deal with InvalidCiphe…

…rtext when decrypting
dzbarsky · Jul 12, 2023 · 1e4dfcb · 1e4dfcb
1 parent 9250c63
commit 1e4dfcb
Show file tree

Hide file tree

Showing 4 changed files with 64 additions and 4 deletions.
diff --git a/services/kms/errors.go b/services/kms/errors.go
@@ -32,6 +32,16 @@ func DisabledException(message string) *awserrors.Error {
 	}
 }
 
+func InvalidCiphertextException(message string) *awserrors.Error {
+	return &awserrors.Error{
+		Code: 400,
+		Body: awserrors.ErrorBody{
+			Type:    "InvalidCiphertextException",
+			Message: message,
+		},
+	}
+}
+
 func KMSInternalException(message string) *awserrors.Error {
 	return &awserrors.Error{
 		Code: 500,

diff --git a/services/kms/kms.go b/services/kms/kms.go
@@ -66,7 +66,12 @@ func (k *KMS) CreateKey(input CreateKeyInput) (*CreateKeyOutput, *awserrors.Erro
 
 	var aesKey *AESKey
 
-	switch input.KeySpec {
+	keySpec := input.KeySpec
+	if keySpec == "" {
+		keySpec = input.CustomerMasterKeySpec
+	}
+
+	switch keySpec {
 	case "", "SYMMETRIC_DEFAULT":
 		aesKey = newAesKey()
 	case "HMAC_224", "HMAC_256", "HMAC_384", "HMAC_512":
@@ -304,7 +309,14 @@ func (k *KMS) Decrypt(input DecryptInput) (*DecryptOutput, *awserrors.Error) {
 
 	// Opposite of Key.Encrypt
 	data := input.CiphertextBlob
+	if len(data) == 0 {
+		return nil, InvalidCiphertextException("")
+	}
+
 	keyArnLen, data := uint8(data[0]), data[1:]
+	if len(data) < 4+int(keyArnLen) {
+		return nil, InvalidCiphertextException("")
+	}
 	keyArn, data := string(data[:keyArnLen]), data[keyArnLen:]
 	version, data := binary.LittleEndian.Uint32(data[:4]), data[4:]
 
@@ -326,7 +338,7 @@ func (k *KMS) Decrypt(input DecryptInput) (*DecryptOutput, *awserrors.Error) {
 
 	plaintext, err := key.Key.Decrypt(data, version, input.EncryptionContext)
 	if err != nil {
-		return nil, KMSInternalException(err.Error())
+		return nil, InvalidCiphertextException(err.Error())
 	}
 
 	return &DecryptOutput{

diff --git a/services/kms/kms_test.go b/services/kms/kms_test.go
@@ -2,6 +2,7 @@ package kms
 
 import (
 	"bytes"
+	"reflect"
 	"strings"
 	"testing"
 
@@ -334,3 +335,38 @@ func TestEncryptDecrypt(t *testing.T) {
 		}
 	}
 }
+
+func TestInvalidCiphertext(t *testing.T) {
+	k, keyId := newKMSWithKey()
+	plaintext := []byte("The quick brown fox jumps over the lazy dog")
+
+	context := map[string]string{"k1": "v1", "k2": "v2"}
+	encryptOutput, err := k.Encrypt(EncryptInput{
+		KeyId:             keyId,
+		EncryptionContext: context,
+		Plaintext:         plaintext,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	ciphertext := encryptOutput.CiphertextBlob
+
+	_, err = k.Decrypt(DecryptInput{
+		KeyId:          keyId,
+		CiphertextBlob: ciphertext,
+		// No context
+	})
+	if !reflect.DeepEqual(err, InvalidCiphertextException("cipher: message authentication failed")) {
+		t.Fatal("bad err", err)
+	}
+
+	_, err = k.Decrypt(DecryptInput{
+		KeyId:             keyId,
+		CiphertextBlob:    []byte("nope"),
+		EncryptionContext: context,
+	})
+	if !reflect.DeepEqual(err, InvalidCiphertextException("")) {
+		t.Fatal("bad err", err)
+	}
+}
diff --git a/services/kms/types.go b/services/kms/types.go
@@ -1,8 +1,10 @@
 package kms
 
 type CreateKeyInput struct {
-	KeySpec string
-	Tags    []APITag
+	Description           string
+	CustomerMasterKeySpec string
+	KeySpec               string
+	Tags                  []APITag
 }
 
 type CreateKeyOutput struct {