From b5af87ebf33bbf98a5cce086ef0aa0a5ee714953 Mon Sep 17 00:00:00 2001
From: Tim Gross <tgross@hashicorp.com>
Date: Tue, 14 Nov 2023 09:54:36 -0500
Subject: [PATCH] set Vault namespace from task in `vault_hook` JWT login
 (#19080)

The JWT login codepath for the `vault_hook` was missing the Vault namespace, so
the login request for non-default namespaces would fail.
---
 client/allocrunner/taskrunner/vault_hook.go |  5 ++--
 client/vaultclient/vaultclient.go           | 26 +++++++++++++++------
 client/vaultclient/vaultclient_test.go      |  8 ++++---
 3 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/client/allocrunner/taskrunner/vault_hook.go b/client/allocrunner/taskrunner/vault_hook.go
index a594474e3..b62ca4a4e 100644
--- a/client/allocrunner/taskrunner/vault_hook.go
+++ b/client/allocrunner/taskrunner/vault_hook.go
@@ -431,8 +431,9 @@ func (h *vaultHook) deriveVaultTokenJWT() (string, error) {
 
 	// Derive Vault token with signed identity.
 	token, err := h.client.DeriveTokenWithJWT(h.ctx, vaultclient.JWTLoginRequest{
-		JWT:  signed.JWT,
-		Role: role,
+		JWT:       signed.JWT,
+		Role:      role,
+		Namespace: h.vaultBlock.Namespace,
 	})
 	if err != nil {
 		return "", structs.WrapRecoverable(
diff --git a/client/vaultclient/vaultclient.go b/client/vaultclient/vaultclient.go
index 75d36335e..f5b1b4f97 100644
--- a/client/vaultclient/vaultclient.go
+++ b/client/vaultclient/vaultclient.go
@@ -41,6 +41,10 @@ type JWTLoginRequest struct {
 	// Nomad client's create_from_role value is used, or the Vault cluster
 	// default role.
 	Role string
+
+	// Namespace is the Vault namespace to use for the login request. If empty,
+	// the Nomad client's Vault configuration namespace will be used.
+	Namespace string
 }
 
 // VaultClient is the interface which nomad client uses to interact with vault and
@@ -250,10 +254,14 @@ func (c *vaultClient) Stop() {
 	close(c.stopCh)
 }
 
-// unlockAndUnset is used to unset the vault token on the client and release the
-// lock. Helper method for deferring a call that does both.
-func (c *vaultClient) unlockAndUnset() {
+// unlockAndUnset is used to unset the vault token on the client, restore the
+// client's namespace, and release the lock. Helper method for deferring a call
+// that does both.
+func (c *vaultClient) unlockAndUnset(previousNs string) {
 	c.client.SetToken("")
+	if previousNs != "" {
+		c.client.SetNamespace(previousNs)
+	}
 	c.lock.Unlock()
 }
 
@@ -270,7 +278,7 @@ func (c *vaultClient) DeriveToken(alloc *structs.Allocation, taskNames []string)
 	}
 
 	c.lock.Lock()
-	defer c.unlockAndUnset()
+	defer c.unlockAndUnset(c.client.Namespace())
 
 	// Use the token supplied to interact with vault
 	c.client.SetToken("")
@@ -294,10 +302,14 @@ func (c *vaultClient) DeriveTokenWithJWT(ctx context.Context, req JWTLoginReques
 	}
 
 	c.lock.Lock()
-	defer c.unlockAndUnset()
+	defer c.unlockAndUnset(c.client.Namespace())
 
-	// Make sure the login request is not passing any token.
+	// Make sure the login request is not passing any token and that we're using
+	// the expected namespace to login
 	c.client.SetToken("")
+	if req.Namespace != "" {
+		c.client.SetNamespace(req.Namespace)
+	}
 
 	jwtLoginPath := fmt.Sprintf("auth/%s/login", c.config.JWTAuthBackendPath)
 	s, err := c.client.Logical().WriteWithContext(ctx, jwtLoginPath,
@@ -337,7 +349,7 @@ func (c *vaultClient) GetConsulACL(token, path string) (*vaultapi.Secret, error)
 	}
 
 	c.lock.Lock()
-	defer c.unlockAndUnset()
+	defer c.unlockAndUnset(c.client.Namespace())
 
 	// Use the token supplied to interact with vault
 	c.client.SetToken(token)
diff --git a/client/vaultclient/vaultclient_test.go b/client/vaultclient/vaultclient_test.go
index bc38fb7c5..2085afe08 100644
--- a/client/vaultclient/vaultclient_test.go
+++ b/client/vaultclient/vaultclient_test.go
@@ -218,7 +218,8 @@ func TestVaultClient_DeriveTokenWithJWT(t *testing.T) {
 	// Derive Vault token using signed JWT.
 	jwtStr := signedWIDs[0].JWT
 	token, err := c.DeriveTokenWithJWT(context.Background(), JWTLoginRequest{
-		JWT: jwtStr,
+		JWT:       jwtStr,
+		Namespace: "default",
 	})
 	must.NoError(t, err)
 	must.NotEq(t, "", token)
@@ -257,8 +258,9 @@ func TestVaultClient_DeriveTokenWithJWT(t *testing.T) {
 
 	// Derive Vault token with non-existing role.
 	token, err = c.DeriveTokenWithJWT(context.Background(), JWTLoginRequest{
-		JWT:  jwtStr,
-		Role: "test",
+		JWT:       jwtStr,
+		Role:      "test",
+		Namespace: "default",
 	})
 	must.ErrorContains(t, err, `role "test" could not be found`)
 }