From e2d713ac057ec69096e8be211f402c5299eaa051 Mon Sep 17 00:00:00 2001
From: Pratyoy Mukhopadhyay <35388175+pmmukh@users.noreply.github.com>
Date: Wed, 19 Jan 2022 09:43:12 -0800
Subject: [PATCH] OSS changes for ent pr (#13696)

---
 api/client.go                   | 17 +++++++++++++++--
 command/agent.go                | 21 +++++++--------------
 command/agent/auth/cert/cert.go |  4 ++++
 command/agent/sink/sink.go      |  4 +++-
 4 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/api/client.go b/api/client.go
index 475899ee4372..0c86d15dcd9f 100644
--- a/api/client.go
+++ b/api/client.go
@@ -928,12 +928,25 @@ func (c *Client) ReadYourWrites() bool {
 // Clone creates a new client with the same configuration. Note that the same
 // underlying http.Client is used; modifying the client from more than one
 // goroutine at once may not be safe, so modify the client as needed and then
-// clone.
+// clone. The headers are cloned based on the CloneHeaders property of the
+// source config
 //
 // Also, only the client's config is currently copied; this means items not in
 // the api.Config struct, such as policy override and wrapping function
 // behavior, must currently then be set as desired on the new client.
 func (c *Client) Clone() (*Client, error) {
+	return c.clone(c.config.CloneHeaders)
+}
+
+// CloneWithHeaders creates a new client similar to Clone, with the difference
+// being that the  headers are always cloned
+func (c *Client) CloneWithHeaders() (*Client, error) {
+	return c.clone(true)
+}
+
+// clone creates a new client, with the headers being cloned based on the
+// passed in cloneheaders boolean
+func (c *Client) clone(cloneHeaders bool) (*Client, error) {
 	c.modifyLock.RLock()
 	defer c.modifyLock.RUnlock()
 
@@ -964,7 +977,7 @@ func (c *Client) Clone() (*Client, error) {
 		return nil, err
 	}
 
-	if config.CloneHeaders {
+	if cloneHeaders {
 		client.SetHeaders(c.Headers().Clone())
 	}
 
diff --git a/command/agent.go b/command/agent.go
index 5cebb729bd7d..27100192ed2c 100644
--- a/command/agent.go
+++ b/command/agent.go
@@ -10,7 +10,6 @@ import (
 	"net"
 	"net/http"
 	"os"
-	"path"
 	"path/filepath"
 	"sort"
 	"strings"
@@ -343,8 +342,12 @@ func (c *AgentCommand) Run(args []string) int {
 
 	var method auth.AuthMethod
 	var sinks []*sink.SinkConfig
-	var namespace string
+	var templateNamespace string
 	if config.AutoAuth != nil {
+		if client.Headers().Get(consts.NamespaceHeaderName) == "" && config.AutoAuth.Method.Namespace != "" {
+			client.SetNamespace(config.AutoAuth.Method.Namespace)
+		}
+		templateNamespace = client.Headers().Get(consts.NamespaceHeaderName)
 		for _, sc := range config.AutoAuth.Sinks {
 			switch sc.Type {
 			case "file":
@@ -371,19 +374,9 @@ func (c *AgentCommand) Run(args []string) int {
 			}
 		}
 
-		// Check if a default namespace has been set
-		mountPath := config.AutoAuth.Method.MountPath
-		if cns := config.AutoAuth.Method.Namespace; cns != "" {
-			namespace = cns
-			// Only set this value if the env var is empty, otherwise we end up with a nested namespace
-			if ens := os.Getenv(api.EnvVaultNamespace); ens == "" {
-				mountPath = path.Join(cns, mountPath)
-			}
-		}
-
 		authConfig := &auth.AuthConfig{
 			Logger:    c.logger.Named(fmt.Sprintf("auth.%s", config.AutoAuth.Method.Type)),
-			MountPath: mountPath,
+			MountPath: config.AutoAuth.Method.MountPath,
 			Config:    config.AutoAuth.Method.Config,
 		}
 		switch config.AutoAuth.Method.Type {
@@ -794,7 +787,7 @@ func (c *AgentCommand) Run(args []string) int {
 			LogLevel:      level,
 			LogWriter:     c.logWriter,
 			AgentConfig:   config,
-			Namespace:     namespace,
+			Namespace:     templateNamespace,
 			ExitAfterAuth: exitAfterAuth,
 		})
 
diff --git a/command/agent/auth/cert/cert.go b/command/agent/auth/cert/cert.go
index 3c8162a5eacd..2703aa8ecd58 100644
--- a/command/agent/auth/cert/cert.go
+++ b/command/agent/auth/cert/cert.go
@@ -9,6 +9,7 @@ import (
 	"github.com/hashicorp/go-hclog"
 	"github.com/hashicorp/vault/api"
 	"github.com/hashicorp/vault/command/agent/auth"
+	"github.com/hashicorp/vault/sdk/helper/consts"
 )
 
 type certMethod struct {
@@ -133,6 +134,9 @@ func (c *certMethod) AuthClient(client *api.Client) (*api.Client, error) {
 		if err != nil {
 			return nil, err
 		}
+		if ns := client.Headers().Get(consts.NamespaceHeaderName); ns != "" {
+			clientToAuth.SetNamespace(ns)
+		}
 
 		// Cache the client for future use
 		c.client = clientToAuth
diff --git a/command/agent/sink/sink.go b/command/agent/sink/sink.go
index 853cc9345f5e..75ea91dc306b 100644
--- a/command/agent/sink/sink.go
+++ b/command/agent/sink/sink.go
@@ -233,14 +233,16 @@ func (s *SinkConfig) encryptToken(token string) (string, error) {
 }
 
 func (s *SinkConfig) wrapToken(client *api.Client, wrapTTL time.Duration, token string) (string, error) {
-	wrapClient, err := client.Clone()
+	wrapClient, err := client.CloneWithHeaders()
 	if err != nil {
 		return "", fmt.Errorf("error deriving client for wrapping, not writing out to sink: %w)", err)
 	}
+
 	wrapClient.SetToken(token)
 	wrapClient.SetWrappingLookupFunc(func(string, string) string {
 		return wrapTTL.String()
 	})
+
 	secret, err := wrapClient.Logical().Write("sys/wrapping/wrap", map[string]interface{}{
 		"token": token,
 	})