1package s3crypto
2
3import (
4	"fmt"
5
6	"github.com/aws/aws-sdk-go/aws"
7	"github.com/aws/aws-sdk-go/service/kms"
8	"github.com/aws/aws-sdk-go/service/kms/kmsiface"
9)
10
11const (
12	// KMSContextWrap is a constant used during decryption to build a kms+context key handler
13	KMSContextWrap      = "kms+context"
14	kmsAWSCEKContextKey = "aws:" + cekAlgorithmHeader
15
16	kmsReservedKeyConflictErrMsg = "conflict in reserved KMS Encryption Context key %s. This value is reserved for the S3 Encryption Client and cannot be set by the user"
17	kmsMismatchCEKAlg            = "the content encryption algorithm used at encryption time does not match the algorithm stored for decryption time. The object may be altered or corrupted"
18)
19
20// NewKMSContextKeyGenerator builds a new kms+context key provider using the customer key ID and material
21// description.
22//
23// Example:
24//	sess := session.Must(session.NewSession())
25//	cmkID := "KMS Key ARN"
26//	var matdesc s3crypto.MaterialDescription
27//	handler := s3crypto.NewKMSContextKeyGenerator(kms.New(sess), cmkID, matdesc)
28func NewKMSContextKeyGenerator(client kmsiface.KMSAPI, cmkID string, matdesc MaterialDescription) CipherDataGeneratorWithCEKAlg {
29	return newKMSContextKeyHandler(client, cmkID, matdesc)
30}
31
32// RegisterKMSContextWrapWithCMK registers the kms+context wrapping algorithm to the given WrapRegistry. The wrapper
33// will be configured to only call KMS Decrypt using the provided CMK.
34//
35// Example:
36//	cr := s3crypto.NewCryptoRegistry()
37//	if err := RegisterKMSContextWrapWithCMK(); err != nil {
38//		panic(err) // handle error
39//	}
40func RegisterKMSContextWrapWithCMK(registry *CryptoRegistry, client kmsiface.KMSAPI, cmkID string) error {
41	if registry == nil {
42		return errNilCryptoRegistry
43	}
44	return registry.AddWrap(KMSContextWrap, newKMSContextWrapEntryWithCMK(client, cmkID))
45}
46
47// RegisterKMSContextWrapWithAnyCMK registers the kms+context wrapping algorithm to the given WrapRegistry. The wrapper
48// will be configured to call KMS decrypt without providing a CMK.
49//
50// Example:
51//	sess := session.Must(session.NewSession())
52//	cr := s3crypto.NewCryptoRegistry()
53//	if err := s3crypto.RegisterKMSContextWrapWithAnyCMK(cr, kms.New(sess)); err != nil {
54//		panic(err) // handle error
55//	}
56func RegisterKMSContextWrapWithAnyCMK(registry *CryptoRegistry, client kmsiface.KMSAPI) error {
57	if registry == nil {
58		return errNilCryptoRegistry
59	}
60	return registry.AddWrap(KMSContextWrap, newKMSContextWrapEntryWithAnyCMK(client))
61}
62
63// newKMSContextWrapEntryWithCMK builds returns a new kms+context key provider and its decrypt handler.
64// The returned handler will be configured to calls KMS Decrypt API without specifying a specific KMS CMK.
65func newKMSContextWrapEntryWithCMK(kmsClient kmsiface.KMSAPI, cmkID string) WrapEntry {
66	// These values are read only making them thread safe
67	kp := &kmsContextKeyHandler{
68		kms:   kmsClient,
69		cmkID: &cmkID,
70	}
71
72	return kp.decryptHandler
73}
74
75// newKMSContextWrapEntryWithAnyCMK builds returns a new kms+context key provider and its decrypt handler.
76// The returned handler will be configured to calls KMS Decrypt API without specifying a specific KMS CMK.
77func newKMSContextWrapEntryWithAnyCMK(kmsClient kmsiface.KMSAPI) WrapEntry {
78	// These values are read only making them thread safe
79	kp := &kmsContextKeyHandler{
80		kms: kmsClient,
81	}
82
83	return kp.decryptHandler
84}
85
86// kmsContextKeyHandler wraps the kmsKeyHandler to explicitly make this type incompatible with the v1 client
87// by not exposing the old interface implementations.
88type kmsContextKeyHandler struct {
89	kms   kmsiface.KMSAPI
90	cmkID *string
91
92	CipherData
93}
94
95func (kp *kmsContextKeyHandler) isAWSFixture() bool {
96	return true
97}
98
99func newKMSContextKeyHandler(client kmsiface.KMSAPI, cmkID string, matdesc MaterialDescription) *kmsContextKeyHandler {
100	kp := &kmsContextKeyHandler{
101		kms:   client,
102		cmkID: &cmkID,
103	}
104
105	if matdesc == nil {
106		matdesc = MaterialDescription{}
107	}
108
109	kp.CipherData.WrapAlgorithm = KMSContextWrap
110	kp.CipherData.MaterialDescription = matdesc
111
112	return kp
113}
114
115func (kp *kmsContextKeyHandler) GenerateCipherDataWithCEKAlg(ctx aws.Context, keySize int, ivSize int, cekAlgorithm string) (CipherData, error) {
116	cd := kp.CipherData.Clone()
117
118	if len(cekAlgorithm) == 0 {
119		return CipherData{}, fmt.Errorf("cek algorithm identifier must not be empty")
120	}
121
122	if _, ok := cd.MaterialDescription[kmsAWSCEKContextKey]; ok {
123		return CipherData{}, fmt.Errorf(kmsReservedKeyConflictErrMsg, kmsAWSCEKContextKey)
124	}
125	cd.MaterialDescription[kmsAWSCEKContextKey] = &cekAlgorithm
126
127	out, err := kp.kms.GenerateDataKeyWithContext(ctx,
128		&kms.GenerateDataKeyInput{
129			EncryptionContext: cd.MaterialDescription,
130			KeyId:             kp.cmkID,
131			KeySpec:           aws.String("AES_256"),
132		})
133	if err != nil {
134		return CipherData{}, err
135	}
136
137	iv, err := generateBytes(ivSize)
138	if err != nil {
139		return CipherData{}, err
140	}
141
142	cd.Key = out.Plaintext
143	cd.IV = iv
144	cd.EncryptedKey = out.CiphertextBlob
145
146	return cd, nil
147}
148
149// decryptHandler initializes a KMS keyprovider with a material description. This
150// is used with Decrypting kms content, due to the cmkID being in the material description.
151func (kp kmsContextKeyHandler) decryptHandler(env Envelope) (CipherDataDecrypter, error) {
152	if env.WrapAlg != KMSContextWrap {
153		return nil, fmt.Errorf("%s value `%s` did not match the expected algorithm `%s` for this handler", cekAlgorithmHeader, env.WrapAlg, KMSContextWrap)
154	}
155
156	m := MaterialDescription{}
157	err := m.decodeDescription([]byte(env.MatDesc))
158	if err != nil {
159		return nil, err
160	}
161
162	if v, ok := m[kmsAWSCEKContextKey]; !ok {
163		return nil, fmt.Errorf("required key %v is missing from encryption context", kmsAWSCEKContextKey)
164	} else if v == nil || *v != env.CEKAlg {
165		return nil, fmt.Errorf(kmsMismatchCEKAlg)
166	}
167
168	kp.MaterialDescription = m
169	kp.WrapAlgorithm = KMSContextWrap
170
171	return &kp, nil
172}
173
174// DecryptKey makes a call to KMS to decrypt the key.
175func (kp *kmsContextKeyHandler) DecryptKey(key []byte) ([]byte, error) {
176	return kp.DecryptKeyWithContext(aws.BackgroundContext(), key)
177}
178
179// DecryptKeyWithContext makes a call to KMS to decrypt the key with request context.
180func (kp *kmsContextKeyHandler) DecryptKeyWithContext(ctx aws.Context, key []byte) ([]byte, error) {
181	out, err := kp.kms.DecryptWithContext(ctx,
182		&kms.DecryptInput{
183			KeyId:             kp.cmkID, // will be nil and not serialized if created with the AnyCMK constructor
184			EncryptionContext: kp.MaterialDescription,
185			CiphertextBlob:    key,
186			GrantTokens:       []*string{},
187		})
188	if err != nil {
189		return nil, err
190	}
191	return out.Plaintext, nil
192}
193
194var (
195	_ CipherDataGeneratorWithCEKAlg  = (*kmsContextKeyHandler)(nil)
196	_ CipherDataDecrypter            = (*kmsContextKeyHandler)(nil)
197	_ CipherDataDecrypterWithContext = (*kmsContextKeyHandler)(nil)
198	_ awsFixture                     = (*kmsContextKeyHandler)(nil)
199)
200