From 2e6edb843d995a29e6ce63cdfe22dbafb97dee47 Mon Sep 17 00:00:00 2001
From: Louis Ruch <louisruch@gmail.com>
Date: Tue, 9 Apr 2024 10:51:53 -0700
Subject: [PATCH] feat(vault): Add headers to client interface

---
 internal/credential/vault/vault.go      |  6 ++++++
 internal/credential/vault/vault_test.go | 28 ++++++++++++++++++++++++-
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/internal/credential/vault/vault.go b/internal/credential/vault/vault.go
index 136806b62d..9b8165afad 100644
--- a/internal/credential/vault/vault.go
+++ b/internal/credential/vault/vault.go
@@ -30,6 +30,7 @@ type vaultClient interface {
 	get(context.Context, string) (*vault.Secret, error)
 	post(context.Context, string, []byte) (*vault.Secret, error)
 	capabilities(context.Context, []string) (pathCapabilities, error)
+	headers(ctx context.Context) (http.Header, error)
 }
 
 var vaultClientFactoryFn = vaultClientFactory
@@ -270,3 +271,8 @@ func (c *client) capabilities(ctx context.Context, paths []string) (pathCapabili
 
 	return newPathCapabilities(res), nil
 }
+
+// headers returns the underlying Vault Client http headers
+func (c *client) headers(_ context.Context) (http.Header, error) {
+	return c.cl.Headers(), nil
+}
diff --git a/internal/credential/vault/vault_test.go b/internal/credential/vault/vault_test.go
index 231bf7c30e..753fc7a3e1 100644
--- a/internal/credential/vault/vault_test.go
+++ b/internal/credential/vault/vault_test.go
@@ -89,7 +89,9 @@ func Test_newClient(t *testing.T) {
 			require.NoError(err)
 			assert.NotNil(client)
 
-			corIdHeader := client.cl.Headers().Get(globals.CorrelationIdKey)
+			headers, err := client.headers(context.Background())
+			require.NoError(err)
+			corIdHeader := headers.Get(globals.CorrelationIdKey)
 			assert.Equal(tt.wantCorId, corIdHeader)
 		})
 	}
@@ -314,3 +316,27 @@ func TestClient_revokeLease(t *testing.T) {
 	// verify the database credentials no longer work
 	assert.Error(testDatabase.ValidateCredential(t, cred))
 }
+
+func Test_headers(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	v := NewTestVaultServer(t)
+
+	clientConfig := &clientConfig{
+		Addr:  v.Addr,
+		Token: TokenSecret(v.RootToken),
+	}
+
+	client, err := newClient(ctx, clientConfig)
+	require.NoError(t, err)
+	assert.NotNil(t, client)
+
+	// Add header to underlying vault client
+	client.cl.AddHeader("test-header", "test-header-value")
+
+	// Get headers from client
+	headers, err := client.headers(context.Background())
+	require.NoError(t, err)
+	got := headers.Get("test-header")
+	assert.Equal(t, "test-header-value", got)
+}