1 files changed, 44 insertions, 9 deletions
diff --git a/vendor/github.com/hashicorp/vault/api/client.go b/vendor/github.com/hashicorp/vault/api/client.go
index 8f0d3f8..c7ced82 100644
--- a/vendor/github.com/hashicorp/vault/api/client.go
+++ b/vendor/github.com/hashicorp/vault/api/client.go
@@ -19,6 +19,7 @@ import (
 	"github.com/hashicorp/go-cleanhttp"
 	retryablehttp "github.com/hashicorp/go-retryablehttp"
 	"github.com/hashicorp/go-rootcerts"
+	"github.com/hashicorp/vault/helper/consts"
 	"github.com/hashicorp/vault/helper/parseutil"
 	"golang.org/x/net/http2"
 	"golang.org/x/time/rate"
@@ -464,6 +465,19 @@ func (c *Client) SetMFACreds(creds []string) {
 	c.mfaCreds = creds
 }
 
+// SetNamespace sets the namespace supplied either via the environment
+// variable or via the command line.
+func (c *Client) SetNamespace(namespace string) {
+	c.modifyLock.Lock()
+	defer c.modifyLock.Unlock()
+
+	if c.headers == nil {
+		c.headers = make(http.Header)
+	}
+
+	c.headers.Set(consts.NamespaceHeaderName, namespace)
+}
+
 // Token returns the access token being used by this client. It will
 // return the empty string if there is no token set.
 func (c *Client) Token() string {
@@ -490,6 +504,26 @@ func (c *Client) ClearToken() {
 	c.token = ""
 }
 
+// Headers gets the current set of headers used for requests. This returns a
+// copy; to modify it make modifications locally and use SetHeaders.
+func (c *Client) Headers() http.Header {
+	c.modifyLock.RLock()
+	defer c.modifyLock.RUnlock()
+
+	if c.headers == nil {
+		return nil
+	}
+
+	ret := make(http.Header)
+	for k, v := range c.headers {
+		for _, val := range v {
+			ret[k] = append(ret[k], val)
+		}
+	}
+
+	return ret
+}
+
 // SetHeaders sets the headers to be used for future requests.
 func (c *Client) SetHeaders(headers http.Header) {
 	c.modifyLock.Lock()
@@ -608,6 +642,13 @@ func (c *Client) NewRequest(method, requestPath string) *Request {
 // a Vault server not configured with this client. This is an advanced operation
 // that generally won't need to be called externally.
 func (c *Client) RawRequest(r *Request) (*Response, error) {
+	return c.RawRequestWithContext(context.Background(), r)
+}
+
+// RawRequestWithContext performs the raw request given. This request may be against
+// a Vault server not configured with this client. This is an advanced operation
+// that generally won't need to be called externally.
+func (c *Client) RawRequestWithContext(ctx context.Context, r *Request) (*Response, error) {
 	c.modifyLock.RLock()
 	token := c.token
 
@@ -622,7 +663,7 @@ func (c *Client) RawRequest(r *Request) (*Response, error) {
 	c.modifyLock.RUnlock()
 
 	if limiter != nil {
-		limiter.Wait(context.Background())
+		limiter.Wait(ctx)
 	}
 
 	// Sanity check the token before potentially erroring from the API
@@ -643,13 +684,10 @@ START:
 		return nil, fmt.Errorf("nil request created")
 	}
 
-	// Set the timeout, if any
-	var cancelFunc context.CancelFunc
 	if timeout != 0 {
-		var ctx context.Context
-		ctx, cancelFunc = context.WithTimeout(context.Background(), timeout)
-		req.Request = req.Request.WithContext(ctx)
+		ctx, _ = context.WithTimeout(ctx, timeout)
 	}
+	req.Request = req.Request.WithContext(ctx)
 
 	if backoff == nil {
 		backoff = retryablehttp.LinearJitterBackoff
@@ -667,9 +705,6 @@ START:
 
 	var result *Response
 	resp, err := client.Do(req)
-	if cancelFunc != nil {
-		cancelFunc()
-	}
 	if resp != nil {
 		result = &Response{Response: resp}
 	}