Skip to content

Commit ecdc889

Browse files
hiddecohiddeco
authored
Merge pull request #1848 from matheuscscp/context
2 parents 60ea3ac + 036ef9f commit ecdc889

File tree

8 files changed

+135
-45
lines changed

8 files changed

+135
-45
lines changed

‎azkv/keysource.go‎

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -120,7 +120,15 @@ func (t TokenCredential) ApplyToMasterKey(key *MasterKey) {
120120

121121
// Encrypt takes a SOPS data key, encrypts it with Azure Key Vault, and stores
122122
// the result in the EncryptedKey field.
123+
//
124+
// Consider using EncryptContext instead.
123125
func (key *MasterKey) Encrypt(dataKey []byte) error {
126+
return key.EncryptContext(context.Background(), dataKey)
127+
}
128+
129+
// EncryptContext takes a SOPS data key, encrypts it with Azure Key Vault, and stores
130+
// the result in the EncryptedKey field.
131+
func (key *MasterKey) EncryptContext(ctx context.Context, dataKey []byte) error {
124132
token, err := key.getTokenCredential()
125133
if err != nil {
126134
log.WithFields(logrus.Fields{"key": key.Name, "version": key.Version}).Info("Encryption failed")
@@ -133,7 +141,7 @@ func (key *MasterKey) Encrypt(dataKey []byte) error {
133141
return fmt.Errorf("failed to construct Azure Key Vault client to encrypt data: %w", err)
134142
}
135143

136-
resp, err := c.Encrypt(context.Background(), key.Name, key.Version, azkeys.KeyOperationParameters{
144+
resp, err := c.Encrypt(ctx, key.Name, key.Version, azkeys.KeyOperationParameters{
137145
Algorithm: to.Ptr(azkeys.EncryptionAlgorithmRSAOAEP256),
138146
Value: dataKey,
139147
}, nil)
@@ -169,7 +177,15 @@ func (key *MasterKey) EncryptIfNeeded(dataKey []byte) error {
169177

170178
// Decrypt decrypts the EncryptedKey field with Azure Key Vault and returns
171179
// the result.
180+
//
181+
// Consider using DecryptContext instead.
172182
func (key *MasterKey) Decrypt() ([]byte, error) {
183+
return key.DecryptContext(context.Background())
184+
}
185+
186+
// DecryptContext decrypts the EncryptedKey field with Azure Key Vault and returns
187+
// the result.
188+
func (key *MasterKey) DecryptContext(ctx context.Context) ([]byte, error) {
173189
token, err := key.getTokenCredential()
174190
if err != nil {
175191
log.WithFields(logrus.Fields{"key": key.Name, "version": key.Version}).Info("Decryption failed")
@@ -188,7 +204,7 @@ func (key *MasterKey) Decrypt() ([]byte, error) {
188204
return nil, fmt.Errorf("failed to construct Azure Key Vault client to decrypt data: %w", err)
189205
}
190206

191-
resp, err := c.Decrypt(context.Background(), key.Name, key.Version, azkeys.KeyOperationParameters{
207+
resp, err := c.Decrypt(ctx, key.Name, key.Version, azkeys.KeyOperationParameters{
192208
Algorithm: to.Ptr(azkeys.EncryptionAlgorithmRSAOAEP256),
193209
Value: rawEncryptedKey,
194210
}, nil)

‎gcpkms/keysource.go‎

Lines changed: 19 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -118,8 +118,16 @@ func (c CredentialJSON) ApplyToMasterKey(key *MasterKey) {
118118

119119
// Encrypt takes a SOPS data key, encrypts it with GCP KMS, and stores the
120120
// result in the EncryptedKey field.
121+
//
122+
// Consider using EncryptContext instead.
121123
func (key *MasterKey) Encrypt(dataKey []byte) error {
122-
service, err := key.newKMSClient()
124+
return key.EncryptContext(context.Background(), dataKey)
125+
}
126+
127+
// EncryptContext takes a SOPS data key, encrypts it with GCP KMS, and stores the
128+
// result in the EncryptedKey field.
129+
func (key *MasterKey) EncryptContext(ctx context.Context, dataKey []byte) error {
130+
service, err := key.newKMSClient(ctx)
123131
if err != nil {
124132
log.WithField("resourceID", key.ResourceID).Info("Encryption failed")
125133
return fmt.Errorf("cannot create GCP KMS service: %w", err)
@@ -134,7 +142,6 @@ func (key *MasterKey) Encrypt(dataKey []byte) error {
134142
Name: key.ResourceID,
135143
Plaintext: dataKey,
136144
}
137-
ctx := context.Background()
138145
resp, err := service.Encrypt(ctx, req)
139146
if err != nil {
140147
log.WithField("resourceID", key.ResourceID).Info("Encryption failed")
@@ -169,8 +176,16 @@ func (key *MasterKey) EncryptIfNeeded(dataKey []byte) error {
169176

170177
// Decrypt decrypts the EncryptedKey field with GCP KMS and returns
171178
// the result.
179+
//
180+
// Consider using DecryptContext instead.
172181
func (key *MasterKey) Decrypt() ([]byte, error) {
173-
service, err := key.newKMSClient()
182+
return key.DecryptContext(context.Background())
183+
}
184+
185+
// DecryptContext decrypts the EncryptedKey field with GCP KMS and returns
186+
// the result.
187+
func (key *MasterKey) DecryptContext(ctx context.Context) ([]byte, error) {
188+
service, err := key.newKMSClient(ctx)
174189
if err != nil {
175190
log.WithField("resourceID", key.ResourceID).Info("Decryption failed")
176191
return nil, fmt.Errorf("cannot create GCP KMS service: %w", err)
@@ -193,7 +208,6 @@ func (key *MasterKey) Decrypt() ([]byte, error) {
193208
Name: key.ResourceID,
194209
Ciphertext: decodedCipher,
195210
}
196-
ctx := context.Background()
197211
resp, err := service.Decrypt(ctx, req)
198212
if err != nil {
199213
log.WithField("resourceID", key.ResourceID).Info("Decryption failed")
@@ -232,7 +246,7 @@ func (key *MasterKey) TypeToIdentifier() string {
232246
// or credentialJSON, and/or grpcConn, falling back to environmental defaults.
233247
// It returns an error if the ResourceID is invalid, or if the setup of the
234248
// client fails.
235-
func (key *MasterKey) newKMSClient() (*kms.KeyManagementClient, error) {
249+
func (key *MasterKey) newKMSClient(ctx context.Context) (*kms.KeyManagementClient, error) {
236250
re := regexp.MustCompile(`^projects/[^/]+/locations/[^/]+/keyRings/[^/]+/cryptoKeys/[^/]+$`)
237251
matches := re.FindStringSubmatch(key.ResourceID)
238252
if matches == nil {
@@ -265,7 +279,6 @@ func (key *MasterKey) newKMSClient() (*kms.KeyManagementClient, error) {
265279
opts = append(opts, option.WithGRPCConn(key.grpcConn))
266280
}
267281

268-
ctx := context.Background()
269282
client, err := kms.NewKeyManagementClient(ctx, opts...)
270283
if err != nil {
271284
return nil, err

‎gcpkms/keysource_test.go‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package gcpkms
22

33
import (
4+
"context"
45
"encoding/base64"
56
"fmt"
67
"net"
@@ -155,7 +156,7 @@ func TestMasterKey_createCloudKMSService_withCredentialsFile(t *testing.T) {
155156
}
156157

157158
for _, tt := range tests {
158-
_, err := tt.key.newKMSClient()
159+
_, err := tt.key.newKMSClient(context.Background())
159160
if tt.errString != "" {
160161
assert.Error(t, err)
161162
assert.ErrorContains(t, err, tt.errString)
@@ -172,7 +173,7 @@ func TestMasterKey_createCloudKMSService_withOauthToken(t *testing.T) {
172173
ResourceID: testResourceID,
173174
}
174175

175-
_, err := masterKey.newKMSClient()
176+
_, err := masterKey.newKMSClient(context.Background())
176177

177178
assert.NoError(t, err)
178179
}
@@ -182,7 +183,7 @@ func TestMasterKey_createCloudKMSService_withoutCredentials(t *testing.T) {
182183
ResourceID: testResourceID,
183184
}
184185

185-
_, err := masterKey.newKMSClient()
186+
_, err := masterKey.newKMSClient(context.Background())
186187

187188
assert.Error(t, err)
188189
assert.ErrorContains(t, err, "credentials: could not find default credentials")

‎hcvault/keysource.go‎

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package hcvault
22

33
import (
44
"bytes"
5+
"context"
56
"encoding/base64"
67
"errors"
78
"fmt"
@@ -130,7 +131,15 @@ func NewMasterKey(address, enginePath, keyName string) *MasterKey {
130131

131132
// Encrypt takes a SOPS data key, encrypts it with Vault Transit, and stores
132133
// the result in the EncryptedKey field.
134+
//
135+
// Consider using EncryptContext instead.
133136
func (key *MasterKey) Encrypt(dataKey []byte) error {
137+
return key.EncryptContext(context.Background(), dataKey)
138+
}
139+
140+
// EncryptContext takes a SOPS data key, encrypts it with Vault Transit, and stores
141+
// the result in the EncryptedKey field.
142+
func (key *MasterKey) EncryptContext(ctx context.Context, dataKey []byte) error {
134143
fullPath := key.encryptPath()
135144

136145
client, err := vaultClient(key.VaultAddress, key.token)
@@ -139,7 +148,7 @@ func (key *MasterKey) Encrypt(dataKey []byte) error {
139148
return err
140149
}
141150

142-
secret, err := client.Logical().Write(fullPath, encryptPayload(dataKey))
151+
secret, err := client.Logical().WriteWithContext(ctx, fullPath, encryptPayload(dataKey))
143152
if err != nil {
144153
log.WithField("Path", fullPath).Info("Encryption failed")
145154
return fmt.Errorf("failed to encrypt sops data key to Vault transit backend '%s': %w", fullPath, err)
@@ -175,7 +184,14 @@ func (key *MasterKey) SetEncryptedDataKey(enc []byte) {
175184
}
176185

177186
// Decrypt decrypts the EncryptedKey field with Vault Transit and returns the result.
187+
//
188+
// Consider using DecryptContext instead.
178189
func (key *MasterKey) Decrypt() ([]byte, error) {
190+
return key.DecryptContext(context.Background())
191+
}
192+
193+
// DecryptContext decrypts the EncryptedKey field with Vault Transit and returns the result.
194+
func (key *MasterKey) DecryptContext(ctx context.Context) ([]byte, error) {
179195
fullPath := key.decryptPath()
180196

181197
client, err := vaultClient(key.VaultAddress, key.token)
@@ -184,7 +200,7 @@ func (key *MasterKey) Decrypt() ([]byte, error) {
184200
return nil, err
185201
}
186202

187-
secret, err := client.Logical().Write(fullPath, decryptPayload(key.EncryptedKey))
203+
secret, err := client.Logical().WriteWithContext(ctx, fullPath, decryptPayload(key.EncryptedKey))
188204
if err != nil {
189205
log.WithField("Path", fullPath).Info("Decryption failed")
190206
return nil, fmt.Errorf("failed to decrypt sops data key from Vault transit backend '%s': %w", fullPath, err)

‎kms/keysource.go‎

Lines changed: 25 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -235,8 +235,16 @@ func (c CredentialsProvider) ApplyToMasterKey(key *MasterKey) {
235235

236236
// Encrypt takes a SOPS data key, encrypts it with KMS and stores the result
237237
// in the EncryptedKey field.
238+
//
239+
// Consider using EncryptContext instead.
238240
func (key *MasterKey) Encrypt(dataKey []byte) error {
239-
cfg, err := key.createKMSConfig()
241+
return key.EncryptContext(context.Background(), dataKey)
242+
}
243+
244+
// EncryptContext takes a SOPS data key, encrypts it with KMS and stores the result
245+
// in the EncryptedKey field.
246+
func (key *MasterKey) EncryptContext(ctx context.Context, dataKey []byte) error {
247+
cfg, err := key.createKMSConfig(ctx)
240248
if err != nil {
241249
log.WithField("arn", key.Arn).Info("Encryption failed")
242250
return err
@@ -247,7 +255,7 @@ func (key *MasterKey) Encrypt(dataKey []byte) error {
247255
Plaintext: dataKey,
248256
EncryptionContext: stringPointerToStringMap(key.EncryptionContext),
249257
}
250-
out, err := client.Encrypt(context.TODO(), input)
258+
out, err := client.Encrypt(ctx, input)
251259
if err != nil {
252260
log.WithField("arn", key.Arn).Info("Encryption failed")
253261
return fmt.Errorf("failed to encrypt sops data key with AWS KMS: %w", err)
@@ -278,13 +286,21 @@ func (key *MasterKey) SetEncryptedDataKey(enc []byte) {
278286

279287
// Decrypt decrypts the EncryptedKey with a newly created AWS KMS config, and
280288
// returns the result.
289+
//
290+
// Consider using DecryptContext instead.
281291
func (key *MasterKey) Decrypt() ([]byte, error) {
292+
return key.DecryptContext(context.Background())
293+
}
294+
295+
// DecryptContext decrypts the EncryptedKey with a newly created AWS KMS config, and
296+
// returns the result.
297+
func (key *MasterKey) DecryptContext(ctx context.Context) ([]byte, error) {
282298
k, err := base64.StdEncoding.DecodeString(key.EncryptedKey)
283299
if err != nil {
284300
log.WithField("arn", key.Arn).Info("Decryption failed")
285301
return nil, fmt.Errorf("error base64-decoding encrypted data key: %s", err)
286302
}
287-
cfg, err := key.createKMSConfig()
303+
cfg, err := key.createKMSConfig(ctx)
288304
if err != nil {
289305
log.WithField("arn", key.Arn).Info("Decryption failed")
290306
return nil, err
@@ -295,7 +311,7 @@ func (key *MasterKey) Decrypt() ([]byte, error) {
295311
CiphertextBlob: k,
296312
EncryptionContext: stringPointerToStringMap(key.EncryptionContext),
297313
}
298-
decrypted, err := client.Decrypt(context.TODO(), input)
314+
decrypted, err := client.Decrypt(ctx, input)
299315
if err != nil {
300316
log.WithField("arn", key.Arn).Info("Decryption failed")
301317
return nil, fmt.Errorf("failed to decrypt sops data key with AWS KMS: %w", err)
@@ -351,15 +367,15 @@ func (key *MasterKey) TypeToIdentifier() string {
351367

352368
// createKMSConfig returns an AWS config with the credentialsProvider of the
353369
// MasterKey, or the default configuration sources.
354-
func (key MasterKey) createKMSConfig() (*aws.Config, error) {
370+
func (key MasterKey) createKMSConfig(ctx context.Context) (*aws.Config, error) {
355371
re := regexp.MustCompile(arnRegex)
356372
matches := re.FindStringSubmatch(key.Arn)
357373
if matches == nil {
358374
return nil, fmt.Errorf("no valid ARN found in '%s'", key.Arn)
359375
}
360376
region := matches[1]
361377

362-
cfg, err := config.LoadDefaultConfig(context.TODO(), func(lo *config.LoadOptions) error {
378+
cfg, err := config.LoadDefaultConfig(ctx, func(lo *config.LoadOptions) error {
363379
// Use the credentialsProvider if present, otherwise default to reading credentials
364380
// from the environment.
365381
if key.credentialsProvider != nil {
@@ -376,7 +392,7 @@ func (key MasterKey) createKMSConfig() (*aws.Config, error) {
376392
}
377393

378394
if key.Role != "" {
379-
return key.createSTSConfig(&cfg)
395+
return key.createSTSConfig(ctx, &cfg)
380396
}
381397
return &cfg, nil
382398
}
@@ -393,7 +409,7 @@ func (key MasterKey) createClient(config *aws.Config) *kms.Client {
393409
// createSTSConfig uses AWS STS to assume a role and returns a config
394410
// configured with that role's credentials. It returns an error if
395411
// it fails to construct a session name, or assume the role.
396-
func (key MasterKey) createSTSConfig(config *aws.Config) (*aws.Config, error) {
412+
func (key MasterKey) createSTSConfig(ctx context.Context, config *aws.Config) (*aws.Config, error) {
397413
name, err := stsSessionName()
398414
if err != nil {
399415
return nil, err
@@ -404,7 +420,7 @@ func (key MasterKey) createSTSConfig(config *aws.Config) (*aws.Config, error) {
404420
}
405421

406422
client := sts.NewFromConfig(*config)
407-
out, err := client.AssumeRole(context.TODO(), input)
423+
out, err := client.AssumeRole(ctx, input)
408424
if err != nil {
409425
return nil, fmt.Errorf("failed to assume role '%s': %w", key.Role, err)
410426
}

‎kms/keysource_test.go‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -535,7 +535,7 @@ aws_secret_access_key = test-secret`), 0600))
535535
if tt.envFunc != nil {
536536
tt.envFunc(t)
537537
}
538-
cfg, err := tt.key.createKMSConfig()
538+
cfg, err := tt.key.createKMSConfig(context.Background())
539539
tt.assertFunc(t, cfg, err)
540540
})
541541
}
@@ -549,7 +549,7 @@ func TestMasterKey_createSTSConfig(t *testing.T) {
549549
return
550550
}
551551
key := NewMasterKeyFromArn(dummyARN, nil, "")
552-
cfg, err := key.createSTSConfig(nil)
552+
cfg, err := key.createSTSConfig(context.Background(), nil)
553553
assert.Error(t, err)
554554
assert.ErrorContains(t, err, "failed to construct STS session name")
555555
assert.Nil(t, cfg)
@@ -558,7 +558,7 @@ func TestMasterKey_createSTSConfig(t *testing.T) {
558558
t.Run("role assumption error", func(t *testing.T) {
559559
key := NewMasterKeyFromArn(dummyARN, nil, "")
560560
key.Role = "role"
561-
got, err := key.createSTSConfig(&aws.Config{})
561+
got, err := key.createSTSConfig(context.Background(), &aws.Config{})
562562
assert.Error(t, err)
563563
assert.ErrorContains(t, err, "failed to assume role 'role'")
564564
assert.Nil(t, got)
@@ -629,7 +629,7 @@ func createTestMasterKey(arn string) MasterKey {
629629
// createTestKMSClient creates a new client with the
630630
// aws.EndpointResolverWithOptions set to epResolver.
631631
func createTestKMSClient(key MasterKey) (*kms.Client, error) {
632-
cfg, err := key.createKMSConfig()
632+
cfg, err := key.createKMSConfig(context.Background())
633633
if err != nil {
634634
return nil, err
635635
}

0 commit comments

Comments
 (0)