Skip to content

Commit d0c2ac4

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 ecdc889 commit d0c2ac4

File tree

5 files changed

+93
-14
lines changed

5 files changed

+93
-14
lines changed

‎azkv/keysource.go‎

Lines changed: 21 additions & 2 deletions
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,23 @@ 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+
// NewClientOptions creates a new ClientOptions with the provided
130+
// azkeys.ClientOptions.
131+
func NewClientOptions(o *azkeys.ClientOptions) *ClientOptions {
132+
return &ClientOptions{o: o}
133+
}
134+
135+
// ApplyToMasterKey configures the ClientOptions on the provided key.
136+
func (c ClientOptions) ApplyToMasterKey(key *MasterKey) {
137+
key.clientOptions = c.o
138+
}
139+
121140
// Encrypt takes a SOPS data key, encrypts it with Azure Key Vault, and stores
122141
// the result in the EncryptedKey field.
123142
//
@@ -135,7 +154,7 @@ func (key *MasterKey) EncryptContext(ctx context.Context, dataKey []byte) error
135154
return fmt.Errorf("failed to get Azure token credential to encrypt data: %w", err)
136155
}
137156

138-
c, err := azkeys.NewClient(key.VaultURL, token, nil)
157+
c, err := azkeys.NewClient(key.VaultURL, token, key.clientOptions)
139158
if err != nil {
140159
log.WithFields(logrus.Fields{"key": key.Name, "version": key.Version}).Info("Encryption failed")
141160
return fmt.Errorf("failed to construct Azure Key Vault client to encrypt data: %w", err)
@@ -198,7 +217,7 @@ func (key *MasterKey) DecryptContext(ctx context.Context) ([]byte, error) {
198217
return nil, fmt.Errorf("failed to base64 decode Azure Key Vault encrypted key: %w", err)
199218
}
200219

201-
c, err := azkeys.NewClient(key.VaultURL, token, nil)
220+
c, err := azkeys.NewClient(key.VaultURL, token, key.clientOptions)
202221
if err != nil {
203222
log.WithFields(logrus.Fields{"key": key.Name, "version": key.Version}).Info("Decryption failed")
204223
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
//
@@ -275,8 +285,13 @@ func (key *MasterKey) newKMSClient(ctx context.Context) (*kms.KeyManagementClien
275285
}
276286
}
277287

278-
if key.grpcConn != nil {
288+
switch {
289+
case key.grpcConn != nil:
279290
opts = append(opts, option.WithGRPCConn(key.grpcConn))
291+
case len(key.grpcDialOpts) > 0:
292+
for _, opt := range key.grpcDialOpts {
293+
opts = append(opts, option.WithGRPCDialOption(opt))
294+
}
280295
}
281296

282297
client, err := kms.NewKeyManagementClient(ctx, opts...)

‎hcvault/keysource.go‎

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

7679
// NewMasterKeysFromURIs creates a list of MasterKeys from a list of Vault
@@ -129,6 +132,22 @@ func NewMasterKey(address, enginePath, keyName string) *MasterKey {
129132
return key
130133
}
131134

135+
// HTTPClient is a wrapper around http.Client used for configuring the
136+
// Vault client.
137+
type HTTPClient struct {
138+
hc *http.Client
139+
}
140+
141+
// NewHTTPClient creates a new HTTPClient with the provided http.Client.
142+
func NewHTTPClient(hc *http.Client) *HTTPClient {
143+
return &HTTPClient{hc: hc}
144+
}
145+
146+
// ApplyToMasterKey configures the HTTP client on the provided key.
147+
func (h HTTPClient) ApplyToMasterKey(key *MasterKey) {
148+
key.httpClient = h.hc
149+
}
150+
132151
// Encrypt takes a SOPS data key, encrypts it with Vault Transit, and stores
133152
// the result in the EncryptedKey field.
134153
//
@@ -142,7 +161,7 @@ func (key *MasterKey) Encrypt(dataKey []byte) error {
142161
func (key *MasterKey) EncryptContext(ctx context.Context, dataKey []byte) error {
143162
fullPath := key.encryptPath()
144163

145-
client, err := vaultClient(key.VaultAddress, key.token)
164+
client, err := vaultClient(key.VaultAddress, key.token, key.httpClient)
146165
if err != nil {
147166
log.WithField("Path", fullPath).Info("Encryption failed")
148167
return err
@@ -194,7 +213,7 @@ func (key *MasterKey) Decrypt() ([]byte, error) {
194213
func (key *MasterKey) DecryptContext(ctx context.Context) ([]byte, error) {
195214
fullPath := key.decryptPath()
196215

197-
client, err := vaultClient(key.VaultAddress, key.token)
216+
client, err := vaultClient(key.VaultAddress, key.token, key.httpClient)
198217
if err != nil {
199218
log.WithField("Path", fullPath).Info("Decryption failed")
200219
return nil, err
@@ -308,10 +327,14 @@ func dataKeyFromSecret(secret *api.Secret) ([]byte, error) {
308327

309328
// vaultClient returns a new Vault client, configured with the given address
310329
// and token.
311-
func vaultClient(address, token string) (*api.Client, error) {
330+
func vaultClient(address, token string, hc *http.Client) (*api.Client, error) {
312331
cfg := api.DefaultConfig()
313332
cfg.Address = address
314333

334+
if hc != nil {
335+
cfg.HttpClient = hc
336+
}
337+
315338
client, err := api.NewClient(cfg)
316339
if err != nil {
317340
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: 22 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,22 @@ 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+
// NewHTTPClient creates a new HTTPClient with the provided http.Client.
246+
func NewHTTPClient(hc *http.Client) *HTTPClient {
247+
return &HTTPClient{hc: hc}
248+
}
249+
250+
// ApplyToMasterKey configures the HTTP client on the provided key.
251+
func (h HTTPClient) ApplyToMasterKey(key *MasterKey) {
252+
key.httpClient = h.hc
253+
}
254+
236255
// Encrypt takes a SOPS data key, encrypts it with KMS and stores the result
237256
// in the EncryptedKey field.
238257
//
@@ -385,6 +404,9 @@ func (key MasterKey) createKMSConfig(ctx context.Context) (*aws.Config, error) {
385404
lo.SharedConfigProfile = key.AwsProfile
386405
}
387406
lo.Region = region
407+
if key.httpClient != nil {
408+
lo.HTTPClient = key.httpClient
409+
}
388410
return nil
389411
})
390412
if err != nil {

0 commit comments

Comments
 (0)