Skip to content

Commit 4d5c77b

Browse files
matheuscscpmatheuscscp
committed
Allow injecting custom HTTP client for AWS, Azure, GCP and HashiCorp
Vault Signed-off-by: Matheus Pimenta <matheuscscp@gmail.com>
1 parent 2356626 commit 4d5c77b

File tree

5 files changed

+76
-13
lines changed

5 files changed

+76
-13
lines changed

‎azkv/keysource.go‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,8 @@ type MasterKey struct {
6060
// using TokenCredential.ApplyToMasterKey.
6161
// If nil, azidentity.NewDefaultAzureCredential is used.
6262
tokenCredential azcore.TokenCredential
63+
// clientOptions contains the azkeys.ClientOptions used by the Azure client.
64+
clientOptions *azkeys.ClientOptions
6365
}
6466

6567
// NewMasterKey creates a new MasterKey from a URL, key name and version,
@@ -118,6 +120,17 @@ func (t TokenCredential) ApplyToMasterKey(key *MasterKey) {
118120
key.tokenCredential = t.token
119121
}
120122

123+
// ClientOptions is a wrapper around azkeys.ClientOptions to allow
124+
// configuration of the Azure Key Vault client.
125+
type ClientOptions struct {
126+
o *azkeys.ClientOptions
127+
}
128+
129+
// ApplyToMasterKey configures the ClientOptions on the provided key.
130+
func (c ClientOptions) ApplyToMasterKey(key *MasterKey) {
131+
key.clientOptions = c.o
132+
}
133+
121134
// Encrypt takes a SOPS data key, encrypts it with Azure Key Vault, and stores
122135
// the result in the EncryptedKey field.
123136
func (key *MasterKey) Encrypt(dataKey []byte) error {
@@ -182,7 +195,7 @@ func (key *MasterKey) Decrypt() ([]byte, error) {
182195
return nil, fmt.Errorf("failed to base64 decode Azure Key Vault encrypted key: %w", err)
183196
}
184197

185-
c, err := azkeys.NewClient(key.VaultURL, token, nil)
198+
c, err := azkeys.NewClient(key.VaultURL, token, key.clientOptions)
186199
if err != nil {
187200
log.WithFields(logrus.Fields{"key": key.Name, "version": key.Version}).Info("Decryption failed")
188201
return nil, fmt.Errorf("failed to construct Azure Key Vault client to decrypt data: %w", err)

‎gcpkms/keysource.go‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,8 @@ type MasterKey struct {
6666
// Mostly useful for testing at present, to wire the client to a mock
6767
// server.
6868
grpcConn *grpc.ClientConn
69+
// grpcDialOpts are the gRPC dial options used to create the gRPC connection.
70+
grpcDialOpts []grpc.DialOption
6971
}
7072

7173
// NewMasterKeyFromResourceID creates a new MasterKey with the provided resource
@@ -116,6 +118,14 @@ func (c CredentialJSON) ApplyToMasterKey(key *MasterKey) {
116118
key.credentialJSON = c
117119
}
118120

121+
// DialOptions are the gRPC dial options used to create the gRPC connection.
122+
type DialOptions []grpc.DialOption
123+
124+
// ApplyToMasterKey configures the DialOptions on the provided key.
125+
func (d DialOptions) ApplyToMasterKey(key *MasterKey) {
126+
key.grpcDialOpts = d
127+
}
128+
119129
// Encrypt takes a SOPS data key, encrypts it with GCP KMS, and stores the
120130
// result in the EncryptedKey field.
121131
func (key *MasterKey) Encrypt(dataKey []byte) error {
@@ -261,8 +271,13 @@ func (key *MasterKey) newKMSClient() (*kms.KeyManagementClient, error) {
261271
}
262272
}
263273

264-
if key.grpcConn != nil {
274+
switch {
275+
case key.grpcConn != nil:
265276
opts = append(opts, option.WithGRPCConn(key.grpcConn))
277+
case len(key.grpcDialOpts) > 0:
278+
for _, opt := range key.grpcDialOpts {
279+
opts = append(opts, option.WithGRPCDialOption(opt))
280+
}
266281
}
267282

268283
ctx := context.Background()

‎hcvault/keysource.go‎

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"errors"
77
"fmt"
88
"io"
9+
"net/http"
910
"net/url"
1011
"os"
1112
"path"
@@ -70,6 +71,8 @@ type MasterKey struct {
7071
// Token.ApplyToMasterKey. If empty, the default client configuration
7172
// is used, before falling back to the token stored in defaultTokenFile.
7273
token string
74+
// httpClient is used to override the default HTTP client used by the Vault client.
75+
httpClient *http.Client
7376
}
7477

7578
// NewMasterKeysFromURIs creates a list of MasterKeys from a list of Vault
@@ -128,12 +131,23 @@ func NewMasterKey(address, enginePath, keyName string) *MasterKey {
128131
return key
129132
}
130133

134+
// HTTPClient is a wrapper around http.Client used for configuring the
135+
// Vault client.
136+
type HTTPClient struct {
137+
hc *http.Client
138+
}
139+
140+
// ApplyToMasterKey configures the HTTP client on the provided key.
141+
func (h HTTPClient) ApplyToMasterKey(key *MasterKey) {
142+
key.httpClient = h.hc
143+
}
144+
131145
// Encrypt takes a SOPS data key, encrypts it with Vault Transit, and stores
132146
// the result in the EncryptedKey field.
133147
func (key *MasterKey) Encrypt(dataKey []byte) error {
134148
fullPath := key.encryptPath()
135149

136-
client, err := vaultClient(key.VaultAddress, key.token)
150+
client, err := vaultClient(key.VaultAddress, key.token, key.httpClient)
137151
if err != nil {
138152
log.WithField("Path", fullPath).Info("Encryption failed")
139153
return err
@@ -178,7 +192,7 @@ func (key *MasterKey) SetEncryptedDataKey(enc []byte) {
178192
func (key *MasterKey) Decrypt() ([]byte, error) {
179193
fullPath := key.decryptPath()
180194

181-
client, err := vaultClient(key.VaultAddress, key.token)
195+
client, err := vaultClient(key.VaultAddress, key.token, key.httpClient)
182196
if err != nil {
183197
log.WithField("Path", fullPath).Info("Decryption failed")
184198
return nil, err
@@ -292,10 +306,14 @@ func dataKeyFromSecret(secret *api.Secret) ([]byte, error) {
292306

293307
// vaultClient returns a new Vault client, configured with the given address
294308
// and token.
295-
func vaultClient(address, token string) (*api.Client, error) {
309+
func vaultClient(address, token string, hc *http.Client) (*api.Client, error) {
296310
cfg := api.DefaultConfig()
297311
cfg.Address = address
298312

313+
if hc != nil {
314+
cfg.HttpClient = hc
315+
}
316+
299317
client, err := api.NewClient(cfg)
300318
if err != nil {
301319
return nil, fmt.Errorf("cannot create Vault client: %w", err)

‎hcvault/keysource_test.go‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -187,7 +187,7 @@ func TestMasterKey_Encrypt(t *testing.T) {
187187
assert.NoError(t, key.Encrypt(dataKey))
188188
assert.NotEmpty(t, key.EncryptedKey)
189189

190-
client, err := vaultClient(key.VaultAddress, key.token)
190+
client, err := vaultClient(key.VaultAddress, key.token, nil)
191191
assert.NoError(t, err)
192192

193193
payload := decryptPayload(key.EncryptedKey)
@@ -230,7 +230,7 @@ func TestMasterKey_Decrypt(t *testing.T) {
230230
(Token(testVaultToken)).ApplyToMasterKey(key)
231231
assert.NoError(t, createVaultKey(key))
232232

233-
client, err := vaultClient(key.VaultAddress, key.token)
233+
client, err := vaultClient(key.VaultAddress, key.token, nil)
234234
assert.NoError(t, err)
235235

236236
dataKey := []byte("the heart of a shrimp is located in its head")
@@ -368,7 +368,7 @@ func Test_vaultClient(t *testing.T) {
368368
t.Setenv("VAULT_TOKEN", "")
369369
t.Setenv("HOME", tmpDir)
370370

371-
got, err := vaultClient(testVaultAddress, "")
371+
got, err := vaultClient(testVaultAddress, "", nil)
372372
assert.NoError(t, err)
373373
assert.NotNil(t, got)
374374
assert.Empty(t, got.Token())
@@ -378,7 +378,7 @@ func Test_vaultClient(t *testing.T) {
378378
token := "test-token"
379379
t.Setenv("VAULT_TOKEN", token)
380380

381-
got, err := vaultClient(testVaultAddress, "")
381+
got, err := vaultClient(testVaultAddress, "", nil)
382382
assert.NoError(t, err)
383383
assert.NotNil(t, got)
384384
assert.Equal(t, token, got.Token())
@@ -388,7 +388,7 @@ func Test_vaultClient(t *testing.T) {
388388
ignored := "test-token"
389389
t.Setenv("VAULT_TOKEN", ignored)
390390

391-
got, err := vaultClient(testVaultAddress, testVaultToken)
391+
got, err := vaultClient(testVaultAddress, testVaultToken, nil)
392392
assert.NoError(t, err)
393393
assert.NotNil(t, got)
394394
assert.Equal(t, testVaultToken, got.Token())
@@ -407,7 +407,7 @@ func Test_vaultClient(t *testing.T) {
407407
t.Setenv("VAULT_TOKEN", "")
408408
t.Setenv("HOME", tmpDir)
409409

410-
got, err := vaultClient(testVaultAddress, "")
410+
got, err := vaultClient(testVaultAddress, "", nil)
411411
assert.NoError(t, err)
412412
assert.NotNil(t, got)
413413
assert.Equal(t, token, got.Token())
@@ -487,7 +487,7 @@ func Test_engineAndKeyFromPath(t *testing.T) {
487487

488488
// enableVaultTransit enables the Vault Transit backend on the given enginePath.
489489
func enableVaultTransit(address, token, enginePath string) error {
490-
client, err := vaultClient(address, token)
490+
client, err := vaultClient(address, token, nil)
491491
if err != nil {
492492
return fmt.Errorf("cannot create Vault client: %w", err)
493493
}
@@ -504,7 +504,7 @@ func enableVaultTransit(address, token, enginePath string) error {
504504
// createVaultKey creates a new RSA-4096 Vault key using the data from the
505505
// provided MasterKey.
506506
func createVaultKey(key *MasterKey) error {
507-
client, err := vaultClient(key.VaultAddress, key.token)
507+
client, err := vaultClient(key.VaultAddress, key.token, nil)
508508
if err != nil {
509509
return fmt.Errorf("cannot create Vault client: %w", err)
510510
}

‎kms/keysource.go‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99
"context"
1010
"encoding/base64"
1111
"fmt"
12+
"net/http"
1213
"os"
1314
"regexp"
1415
"sort"
@@ -79,6 +80,8 @@ type MasterKey struct {
7980
// injected using e.g. an environment variable. The field is not publicly
8081
// exposed, nor configurable.
8182
baseEndpoint string
83+
// httpClient is used to override the default HTTP client used by the AWS client.
84+
httpClient *http.Client
8285
}
8386

8487
// NewMasterKey creates a new MasterKey from an ARN, role and context, setting
@@ -233,6 +236,17 @@ func (c CredentialsProvider) ApplyToMasterKey(key *MasterKey) {
233236
key.credentialsProvider = c.provider
234237
}
235238

239+
// HTTPClient is a wrapper around http.Client used for configuring the
240+
// AWS KMS client.
241+
type HTTPClient struct {
242+
hc *http.Client
243+
}
244+
245+
// ApplyToMasterKey configures the HTTP client on the provided key.
246+
func (h HTTPClient) ApplyToMasterKey(key *MasterKey) {
247+
key.httpClient = h.hc
248+
}
249+
236250
// Encrypt takes a SOPS data key, encrypts it with KMS and stores the result
237251
// in the EncryptedKey field.
238252
func (key *MasterKey) Encrypt(dataKey []byte) error {
@@ -369,6 +383,9 @@ func (key MasterKey) createKMSConfig() (*aws.Config, error) {
369383
lo.SharedConfigProfile = key.AwsProfile
370384
}
371385
lo.Region = region
386+
if key.httpClient != nil {
387+
lo.HTTPClient = key.httpClient
388+
}
372389
return nil
373390
})
374391
if err != nil {

0 commit comments

Comments
 (0)