From 3e6665f65d6f02ff2ec5c83baa15c2dd8b38991b Mon Sep 17 00:00:00 2001
From: Vinny Mannello <94396874+VinnyHC@users.noreply.github.com>
Date: Thu, 14 Apr 2022 09:50:21 -0700
Subject: [PATCH] [Vault-5736] Add (*Client).WithNamespace() for temporary
 namespace handling (#14963)

temporary namespace calls
---
 api/client.go       | 47 ++++++++++++++++++++++++++++++++-
 api/client_test.go  | 64 +++++++++++++++++++++++++++++++++++++++++++++
 changelog/14963.txt |  3 +++
 3 files changed, 113 insertions(+), 1 deletion(-)
 create mode 100644 changelog/14963.txt

diff --git a/api/client.go b/api/client.go
index c9cea7ca6..89160cd9a 100644
--- a/api/client.go
+++ b/api/client.go
@@ -819,10 +819,39 @@ func (c *Client) setNamespace(namespace string) {
 	c.headers.Set(consts.NamespaceHeaderName, namespace)
 }
 
+// ClearNamespace removes the namespace header if set.
 func (c *Client) ClearNamespace() {
 	c.modifyLock.Lock()
 	defer c.modifyLock.Unlock()
-	c.headers.Del(consts.NamespaceHeaderName)
+	if c.headers != nil {
+		c.headers.Del(consts.NamespaceHeaderName)
+	}
+}
+
+// Namespace returns the namespace currently set in this client. It will
+// return an empty string if there is no namespace set.
+func (c *Client) Namespace() string {
+	c.modifyLock.Lock()
+	defer c.modifyLock.Unlock()
+	if c.headers == nil {
+		return ""
+	}
+	return c.headers.Get(consts.NamespaceHeaderName)
+}
+
+// WithNamespace makes a shallow copy of Client, modifies it to use
+// the given namespace, and returns it. Passing an empty string will
+// temporarily unset the namespace.
+func (c *Client) WithNamespace(namespace string) *Client {
+	c2 := *c
+	c2.modifyLock = sync.RWMutex{}
+	c2.headers = c.Headers()
+	if namespace == "" {
+		c2.ClearNamespace()
+	} else {
+		c2.SetNamespace(namespace)
+	}
+	return &c2
 }
 
 // Token returns the access token being used by this client. It will
@@ -1141,12 +1170,22 @@ func (c *Client) rawRequestWithContext(ctx context.Context, r *Request) (*Respon
 	checkRetry := c.config.CheckRetry
 	backoff := c.config.Backoff
 	httpClient := c.config.HttpClient
+	ns := c.headers.Get(consts.NamespaceHeaderName)
 	outputCurlString := c.config.OutputCurlString
 	logger := c.config.Logger
 	c.config.modifyLock.RUnlock()
 
 	c.modifyLock.RUnlock()
 
+	// ensure that the most current namespace setting is used at the time of the call
+	// e.g. calls using (*Client).WithNamespace
+	switch ns {
+	case "":
+		r.Headers.Del(consts.NamespaceHeaderName)
+	default:
+		r.Headers.Set(consts.NamespaceHeaderName, ns)
+	}
+
 	for _, cb := range c.requestCallbacks {
 		cb(r)
 	}
@@ -1278,13 +1317,19 @@ func (c *Client) httpRequestWithContext(ctx context.Context, r *Request) (*Respo
 	limiter := c.config.Limiter
 	httpClient := c.config.HttpClient
 	outputCurlString := c.config.OutputCurlString
+	// add headers
 	if c.headers != nil {
 		for header, vals := range c.headers {
 			for _, val := range vals {
 				req.Header.Add(header, val)
 			}
 		}
+		// explicitly set the namespace header to current client
+		if ns := c.headers.Get(consts.NamespaceHeaderName); ns != "" {
+			r.Headers.Set(consts.NamespaceHeaderName, ns)
+		}
 	}
+
 	c.config.modifyLock.RUnlock()
 	c.modifyLock.RUnlock()
 
diff --git a/api/client_test.go b/api/client_test.go
index 73a306453..58b797d34 100644
--- a/api/client_test.go
+++ b/api/client_test.go
@@ -1136,3 +1136,67 @@ func TestClient_SetCloneToken(t *testing.T) {
 		})
 	}
 }
+
+func TestClientWithNamespace(t *testing.T) {
+	var ns string
+	handler := func(w http.ResponseWriter, req *http.Request) {
+		ns = req.Header.Get(consts.NamespaceHeaderName)
+	}
+	config, ln := testHTTPServer(t, http.HandlerFunc(handler))
+	defer ln.Close()
+
+	// set up a client with a namespace
+	client, err := NewClient(config)
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	ogNS := "test"
+	client.SetNamespace(ogNS)
+	_, err = client.rawRequestWithContext(
+		context.Background(),
+		client.NewRequest(http.MethodGet, "/"))
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if ns != ogNS {
+		t.Fatalf("Expected namespace: \"%s\", got \"%s\"", ogNS, ns)
+	}
+
+	// make a call with a temporary namespace
+	newNS := "new-namespace"
+	_, err = client.WithNamespace(newNS).rawRequestWithContext(
+		context.Background(),
+		client.NewRequest(http.MethodGet, "/"))
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if ns != newNS {
+		t.Fatalf("Expected new namespace: \"%s\", got \"%s\"", newNS, ns)
+	}
+	// ensure client has not been modified
+	_, err = client.rawRequestWithContext(
+		context.Background(),
+		client.NewRequest(http.MethodGet, "/"))
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if ns != ogNS {
+		t.Fatalf("Expected original namespace: \"%s\", got \"%s\"", ogNS, ns)
+	}
+
+	// make call with empty ns
+	_, err = client.WithNamespace("").rawRequestWithContext(
+		context.Background(),
+		client.NewRequest(http.MethodGet, "/"))
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if ns != "" {
+		t.Fatalf("Expected no namespace, got \"%s\"", ns)
+	}
+
+	// ensure client has not been modified
+	if client.Namespace() != ogNS {
+		t.Fatalf("Expected original namespace: \"%s\", got \"%s\"", ogNS, client.Namespace())
+	}
+}
diff --git a/changelog/14963.txt b/changelog/14963.txt
new file mode 100644
index 000000000..9bc3dc86f
--- /dev/null
+++ b/changelog/14963.txt
@@ -0,0 +1,3 @@
+```release-note:improvement
+api: Provide a helper method WithNamespace to create a cloned client with a new NS
+```
\ No newline at end of file