API: Add context to each raw request call (#4987)

api/auth_token.go

@@ -1,5 +1,7 @@
 package api
 
+import "context"
+
 // TokenAuth is used to perform token backend operations on Vault
 type TokenAuth struct {
 	c *Client
@@ -16,7 +18,9 @@ func (c *TokenAuth) Create(opts *TokenCreateRequest) (*Secret, error) {
 		return nil, err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -31,7 +35,9 @@ func (c *TokenAuth) CreateOrphan(opts *TokenCreateRequest) (*Secret, error) {
 		return nil, err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -46,7 +52,9 @@ func (c *TokenAuth) CreateWithRole(opts *TokenCreateRequest, roleName string) (*
 		return nil, err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -63,7 +71,9 @@ func (c *TokenAuth) Lookup(token string) (*Secret, error) {
 		return nil, err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -79,7 +89,10 @@ func (c *TokenAuth) LookupAccessor(accessor string) (*Secret, error) {
 	}); err != nil {
 		return nil, err
 	}
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -91,7 +104,9 @@ func (c *TokenAuth) LookupAccessor(accessor string) (*Secret, error) {
 func (c *TokenAuth) LookupSelf() (*Secret, error) {
 	r := c.c.NewRequest("GET", "/v1/auth/token/lookup-self")
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -109,7 +124,9 @@ func (c *TokenAuth) Renew(token string, increment int) (*Secret, error) {
 		return nil, err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -126,7 +143,9 @@ func (c *TokenAuth) RenewSelf(increment int) (*Secret, error) {
 		return nil, err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -146,7 +165,9 @@ func (c *TokenAuth) RenewTokenAsSelf(token string, increment int) (*Secret, erro
 		return nil, err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -164,7 +185,10 @@ func (c *TokenAuth) RevokeAccessor(accessor string) error {
 	}); err != nil {
 		return err
 	}
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return err
 	}
@@ -183,7 +207,9 @@ func (c *TokenAuth) RevokeOrphan(token string) error {
 		return err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return err
 	}
@@ -197,7 +223,10 @@ func (c *TokenAuth) RevokeOrphan(token string) error {
 // an effect.
 func (c *TokenAuth) RevokeSelf(token string) error {
 	r := c.c.NewRequest("PUT", "/v1/auth/token/revoke-self")
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return err
 	}
@@ -217,7 +246,9 @@ func (c *TokenAuth) RevokeTree(token string) error {
 		return err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return err
 	}

api/client.go

@@ -608,6 +608,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 +629,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 +650,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, _ = context.WithTimeout(ctx, timeout)
-		ctx, cancelFunc = context.WithTimeout(context.Background(), timeout)
-		req.Request = req.Request.WithContext(ctx)
 	}
+	req.Request = req.Request.WithContext(ctx)
 
 	if backoff == nil {
 		backoff = retryablehttp.LinearJitterBackoff
@@ -667,9 +671,6 @@ START:
 
 	var result *Response
 	resp, err := client.Do(req)
-	if cancelFunc != nil {
-		cancelFunc()
-	}
 	if resp != nil {
 		result = &Response{Response: resp}
 	}

api/help.go

@@ -1,6 +1,7 @@
 package api
 
 import (
+	"context"
 	"fmt"
 )
 
@@ -8,7 +9,10 @@ import (
 func (c *Client) Help(path string) (*Help, error) {
 	r := c.NewRequest("GET", fmt.Sprintf("/v1/%s", path))
 	r.Params.Add("help", "1")
-	resp, err := c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}

api/logical.go

@@ -2,6 +2,7 @@ package api
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"io"
 	"os"
@@ -46,7 +47,10 @@ func (c *Client) Logical() *Logical {
 
 func (c *Logical) Read(path string) (*Secret, error) {
 	r := c.c.NewRequest("GET", "/v1/"+path)
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if resp != nil {
 		defer resp.Body.Close()
 	}
@@ -77,7 +81,10 @@ func (c *Logical) List(path string) (*Secret, error) {
 	// handle the wrapping lookup function
 	r.Method = "GET"
 	r.Params.Set("list", "true")
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if resp != nil {
 		defer resp.Body.Close()
 	}
@@ -108,7 +115,9 @@ func (c *Logical) Write(path string, data map[string]interface{}) (*Secret, erro
 		return nil, err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if resp != nil {
 		defer resp.Body.Close()
 	}
@@ -134,7 +143,10 @@ func (c *Logical) Write(path string, data map[string]interface{}) (*Secret, erro
 
 func (c *Logical) Delete(path string) (*Secret, error) {
 	r := c.c.NewRequest("DELETE", "/v1/"+path)
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if resp != nil {
 		defer resp.Body.Close()
 	}
@@ -175,7 +187,9 @@ func (c *Logical) Unwrap(wrappingToken string) (*Secret, error) {
 		return nil, err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if resp != nil {
 		defer resp.Body.Close()
 	}

api/ssh.go

@@ -1,6 +1,9 @@
 package api
 
-import "fmt"
+import (
+	"context"
+	"fmt"
+)
 
 // SSH is used to return a client to invoke operations on SSH backend.
 type SSH struct {
@@ -28,7 +31,9 @@ func (c *SSH) Credential(role string, data map[string]interface{}) (*Secret, err
 		return nil, err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -45,7 +50,9 @@ func (c *SSH) SignKey(role string, data map[string]interface{}) (*Secret, error)
 		return nil, err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}

api/ssh_agent.go

@@ -1,6 +1,7 @@
 package api
 
 import (
+	"context"
 	"crypto/tls"
 	"crypto/x509"
 	"fmt"
@@ -207,7 +208,9 @@ func (c *SSHHelper) Verify(otp string) (*SSHVerifyResponse, error) {
 		return nil, err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}

api/sys_audit.go

@@ -1,6 +1,7 @@
 package api
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/mitchellh/mapstructure"
@@ -16,7 +17,9 @@ func (c *Sys) AuditHash(path string, input string) (string, error) {
 		return "", err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return "", err
 	}
@@ -37,7 +40,11 @@ func (c *Sys) AuditHash(path string, input string) (string, error) {
 
 func (c *Sys) ListAudit() (map[string]*Audit, error) {
 	r := c.c.NewRequest("GET", "/v1/sys/audit")
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
+
 	if err != nil {
 		return nil, err
 	}
@@ -87,7 +94,10 @@ func (c *Sys) EnableAuditWithOptions(path string, options *EnableAuditOptions) e
 		return err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
+
 	if err != nil {
 		return err
 	}
@@ -98,7 +108,11 @@ func (c *Sys) EnableAuditWithOptions(path string, options *EnableAuditOptions) e
 
 func (c *Sys) DisableAudit(path string) error {
 	r := c.c.NewRequest("DELETE", fmt.Sprintf("/v1/sys/audit/%s", path))
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
+
 	if err == nil {
 		defer resp.Body.Close()
 	}

api/sys_auth.go

@@ -1,6 +1,7 @@
 package api
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/mitchellh/mapstructure"
@@ -8,7 +9,10 @@ import (
 
 func (c *Sys) ListAuth() (map[string]*AuthMount, error) {
 	r := c.c.NewRequest("GET", "/v1/sys/auth")
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -56,7 +60,9 @@ func (c *Sys) EnableAuthWithOptions(path string, options *EnableAuthOptions) err
 		return err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return err
 	}
@@ -67,7 +73,10 @@ func (c *Sys) EnableAuthWithOptions(path string, options *EnableAuthOptions) err
 
 func (c *Sys) DisableAuth(path string) error {
 	r := c.c.NewRequest("DELETE", fmt.Sprintf("/v1/sys/auth/%s", path))
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err == nil {
 		defer resp.Body.Close()
 	}

api/sys_capabilities.go

@@ -1,6 +1,9 @@
 package api
 
-import "fmt"
+import (
+	"context"
+	"fmt"
+)
 
 func (c *Sys) CapabilitiesSelf(path string) ([]string, error) {
 	return c.Capabilities(c.c.Token(), path)
@@ -22,7 +25,9 @@ func (c *Sys) Capabilities(token, path string) ([]string, error) {
 		return nil, err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}

api/sys_config_cors.go

@@ -1,8 +1,13 @@
 package api
 
+import "context"
+
 func (c *Sys) CORSStatus() (*CORSResponse, error) {
 	r := c.c.NewRequest("GET", "/v1/sys/config/cors")
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -19,7 +24,9 @@ func (c *Sys) ConfigureCORS(req *CORSRequest) (*CORSResponse, error) {
 		return nil, err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -33,7 +40,9 @@ func (c *Sys) ConfigureCORS(req *CORSRequest) (*CORSResponse, error) {
 func (c *Sys) DisableCORS() (*CORSResponse, error) {
 	r := c.c.NewRequest("DELETE", "/v1/sys/config/cors")
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}

api/sys_generate_root.go

@@ -1,5 +1,7 @@
 package api
 
+import "context"
+
 func (c *Sys) GenerateRootStatus() (*GenerateRootStatusResponse, error) {
 	return c.generateRootStatusCommon("/v1/sys/generate-root/attempt")
 }
@@ -10,7 +12,10 @@ func (c *Sys) GenerateDROperationTokenStatus() (*GenerateRootStatusResponse, err
 
 func (c *Sys) generateRootStatusCommon(path string) (*GenerateRootStatusResponse, error) {
 	r := c.c.NewRequest("GET", path)
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -40,7 +45,9 @@ func (c *Sys) generateRootInitCommon(path, otp, pgpKey string) (*GenerateRootSta
 		return nil, err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -61,7 +68,10 @@ func (c *Sys) GenerateDROperationTokenCancel() error {
 
 func (c *Sys) generateRootCancelCommon(path string) error {
 	r := c.c.NewRequest("DELETE", path)
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err == nil {
 		defer resp.Body.Close()
 	}
@@ -87,7 +97,9 @@ func (c *Sys) generateRootUpdateCommon(path, shard, nonce string) (*GenerateRoot
 		return nil, err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}

api/sys_health.go

@@ -1,5 +1,7 @@
 package api
 
+import "context"
+
 func (c *Sys) Health() (*HealthResponse, error) {
 	r := c.c.NewRequest("GET", "/v1/sys/health")
 	// If the code is 400 or above it will automatically turn into an error,
@@ -9,7 +11,10 @@ func (c *Sys) Health() (*HealthResponse, error) {
 	r.Params.Add("sealedcode", "299")
 	r.Params.Add("standbycode", "299")
 	r.Params.Add("drsecondarycode", "299")
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}

api/sys_init.go

@@ -1,8 +1,13 @@
 package api
 
+import "context"
+
 func (c *Sys) InitStatus() (bool, error) {
 	r := c.c.NewRequest("GET", "/v1/sys/init")
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return false, err
 	}
@@ -19,7 +24,9 @@ func (c *Sys) Init(opts *InitRequest) (*InitResponse, error) {
 		return nil, err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}

api/sys_leader.go

@@ -1,8 +1,13 @@
 package api
 
+import "context"
+
 func (c *Sys) Leader() (*LeaderResponse, error) {
 	r := c.c.NewRequest("GET", "/v1/sys/leader")
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}

api/sys_leases.go

@@ -1,6 +1,9 @@
 package api
 
-import "errors"
+import (
+	"context"
+	"errors"
+)
 
 func (c *Sys) Renew(id string, increment int) (*Secret, error) {
 	r := c.c.NewRequest("PUT", "/v1/sys/leases/renew")
@@ -13,7 +16,9 @@ func (c *Sys) Renew(id string, increment int) (*Secret, error) {
 		return nil, err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -24,7 +29,10 @@ func (c *Sys) Renew(id string, increment int) (*Secret, error) {
 
 func (c *Sys) Revoke(id string) error {
 	r := c.c.NewRequest("PUT", "/v1/sys/leases/revoke/"+id)
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err == nil {
 		defer resp.Body.Close()
 	}
@@ -33,7 +41,10 @@ func (c *Sys) Revoke(id string) error {
 
 func (c *Sys) RevokePrefix(id string) error {
 	r := c.c.NewRequest("PUT", "/v1/sys/leases/revoke-prefix/"+id)
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err == nil {
 		defer resp.Body.Close()
 	}
@@ -42,7 +53,10 @@ func (c *Sys) RevokePrefix(id string) error {
 
 func (c *Sys) RevokeForce(id string) error {
 	r := c.c.NewRequest("PUT", "/v1/sys/leases/revoke-force/"+id)
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err == nil {
 		defer resp.Body.Close()
 	}
@@ -74,7 +88,9 @@ func (c *Sys) RevokeWithOptions(opts *RevokeOptions) error {
 		}
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err == nil {
 		defer resp.Body.Close()
 	}

api/sys_mounts.go

@@ -1,6 +1,7 @@
 package api
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/mitchellh/mapstructure"
@@ -8,7 +9,10 @@ import (
 
 func (c *Sys) ListMounts() (map[string]*MountOutput, error) {
 	r := c.c.NewRequest("GET", "/v1/sys/mounts")
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -48,7 +52,9 @@ func (c *Sys) Mount(path string, mountInfo *MountInput) error {
 		return err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return err
 	}
@@ -59,7 +65,10 @@ func (c *Sys) Mount(path string, mountInfo *MountInput) error {
 
 func (c *Sys) Unmount(path string) error {
 	r := c.c.NewRequest("DELETE", fmt.Sprintf("/v1/sys/mounts/%s", path))
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err == nil {
 		defer resp.Body.Close()
 	}
@@ -77,7 +86,9 @@ func (c *Sys) Remount(from, to string) error {
 		return err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err == nil {
 		defer resp.Body.Close()
 	}
@@ -90,7 +101,9 @@ func (c *Sys) TuneMount(path string, config MountConfigInput) error {
 		return err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err == nil {
 		defer resp.Body.Close()
 	}
@@ -100,7 +113,9 @@ func (c *Sys) TuneMount(path string, config MountConfigInput) error {
 func (c *Sys) MountConfig(path string) (*MountConfigOutput, error) {
 	r := c.c.NewRequest("GET", fmt.Sprintf("/v1/sys/mounts/%s/tune", path))
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}

api/sys_plugins.go

@@ -1,6 +1,7 @@
 package api
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 )
@@ -19,7 +20,10 @@ type ListPluginsResponse struct {
 func (c *Sys) ListPlugins(i *ListPluginsInput) (*ListPluginsResponse, error) {
 	path := "/v1/sys/plugins/catalog"
 	req := c.c.NewRequest("LIST", path)
-	resp, err := c.c.RawRequest(req)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, req)
 	if err != nil {
 		return nil, err
 	}
@@ -54,7 +58,10 @@ type GetPluginResponse struct {
 func (c *Sys) GetPlugin(i *GetPluginInput) (*GetPluginResponse, error) {
 	path := fmt.Sprintf("/v1/sys/plugins/catalog/%s", i.Name)
 	req := c.c.NewRequest(http.MethodGet, path)
-	resp, err := c.c.RawRequest(req)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, req)
 	if err != nil {
 		return nil, err
 	}
@@ -93,7 +100,9 @@ func (c *Sys) RegisterPlugin(i *RegisterPluginInput) error {
 		return err
 	}
 
-	resp, err := c.c.RawRequest(req)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, req)
 	if err == nil {
 		defer resp.Body.Close()
 	}
@@ -111,7 +120,10 @@ type DeregisterPluginInput struct {
 func (c *Sys) DeregisterPlugin(i *DeregisterPluginInput) error {
 	path := fmt.Sprintf("/v1/sys/plugins/catalog/%s", i.Name)
 	req := c.c.NewRequest(http.MethodDelete, path)
-	resp, err := c.c.RawRequest(req)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, req)
 	if err == nil {
 		defer resp.Body.Close()
 	}

api/sys_policy.go

@@ -1,10 +1,16 @@
 package api
 
-import "fmt"
+import (
+	"context"
+	"fmt"
+)
 
 func (c *Sys) ListPolicies() ([]string, error) {
 	r := c.c.NewRequest("GET", "/v1/sys/policy")
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -33,7 +39,10 @@ func (c *Sys) ListPolicies() ([]string, error) {
 
 func (c *Sys) GetPolicy(name string) (string, error) {
 	r := c.c.NewRequest("GET", fmt.Sprintf("/v1/sys/policy/%s", name))
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if resp != nil {
 		defer resp.Body.Close()
 		if resp.StatusCode == 404 {
@@ -70,7 +79,9 @@ func (c *Sys) PutPolicy(name, rules string) error {
 		return err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return err
 	}
@@ -81,7 +92,10 @@ func (c *Sys) PutPolicy(name, rules string) error {
 
 func (c *Sys) DeletePolicy(name string) error {
 	r := c.c.NewRequest("DELETE", fmt.Sprintf("/v1/sys/policy/%s", name))
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err == nil {
 		defer resp.Body.Close()
 	}

api/sys_rekey.go

@@ -1,8 +1,13 @@
 package api
 
+import "context"
+
 func (c *Sys) RekeyStatus() (*RekeyStatusResponse, error) {
 	r := c.c.NewRequest("GET", "/v1/sys/rekey/init")
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -15,7 +20,10 @@ func (c *Sys) RekeyStatus() (*RekeyStatusResponse, error) {
 
 func (c *Sys) RekeyRecoveryKeyStatus() (*RekeyStatusResponse, error) {
 	r := c.c.NewRequest("GET", "/v1/sys/rekey-recovery-key/init")
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -28,7 +36,10 @@ func (c *Sys) RekeyRecoveryKeyStatus() (*RekeyStatusResponse, error) {
 
 func (c *Sys) RekeyVerificationStatus() (*RekeyVerificationStatusResponse, error) {
 	r := c.c.NewRequest("GET", "/v1/sys/rekey/verify")
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -41,7 +52,10 @@ func (c *Sys) RekeyVerificationStatus() (*RekeyVerificationStatusResponse, error
 
 func (c *Sys) RekeyRecoveryKeyVerificationStatus() (*RekeyVerificationStatusResponse, error) {
 	r := c.c.NewRequest("GET", "/v1/sys/rekey-recovery-key/verify")
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -58,7 +72,9 @@ func (c *Sys) RekeyInit(config *RekeyInitRequest) (*RekeyStatusResponse, error)
 		return nil, err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -75,7 +91,9 @@ func (c *Sys) RekeyRecoveryKeyInit(config *RekeyInitRequest) (*RekeyStatusRespon
 		return nil, err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -88,7 +106,10 @@ func (c *Sys) RekeyRecoveryKeyInit(config *RekeyInitRequest) (*RekeyStatusRespon
 
 func (c *Sys) RekeyCancel() error {
 	r := c.c.NewRequest("DELETE", "/v1/sys/rekey/init")
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err == nil {
 		defer resp.Body.Close()
 	}
@@ -97,7 +118,10 @@ func (c *Sys) RekeyCancel() error {
 
 func (c *Sys) RekeyRecoveryKeyCancel() error {
 	r := c.c.NewRequest("DELETE", "/v1/sys/rekey-recovery-key/init")
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err == nil {
 		defer resp.Body.Close()
 	}
@@ -106,7 +130,10 @@ func (c *Sys) RekeyRecoveryKeyCancel() error {
 
 func (c *Sys) RekeyVerificationCancel() error {
 	r := c.c.NewRequest("DELETE", "/v1/sys/rekey/verify")
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err == nil {
 		defer resp.Body.Close()
 	}
@@ -115,7 +142,10 @@ func (c *Sys) RekeyVerificationCancel() error {
 
 func (c *Sys) RekeyRecoveryKeyVerificationCancel() error {
 	r := c.c.NewRequest("DELETE", "/v1/sys/rekey-recovery-key/verify")
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err == nil {
 		defer resp.Body.Close()
 	}
@@ -133,7 +163,9 @@ func (c *Sys) RekeyUpdate(shard, nonce string) (*RekeyUpdateResponse, error) {
 		return nil, err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -155,7 +187,9 @@ func (c *Sys) RekeyRecoveryKeyUpdate(shard, nonce string) (*RekeyUpdateResponse,
 		return nil, err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -168,7 +202,10 @@ func (c *Sys) RekeyRecoveryKeyUpdate(shard, nonce string) (*RekeyUpdateResponse,
 
 func (c *Sys) RekeyRetrieveBackup() (*RekeyRetrieveResponse, error) {
 	r := c.c.NewRequest("GET", "/v1/sys/rekey/backup")
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -181,7 +218,10 @@ func (c *Sys) RekeyRetrieveBackup() (*RekeyRetrieveResponse, error) {
 
 func (c *Sys) RekeyRetrieveRecoveryBackup() (*RekeyRetrieveResponse, error) {
 	r := c.c.NewRequest("GET", "/v1/sys/rekey/recovery-backup")
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -194,7 +234,10 @@ func (c *Sys) RekeyRetrieveRecoveryBackup() (*RekeyRetrieveResponse, error) {
 
 func (c *Sys) RekeyDeleteBackup() error {
 	r := c.c.NewRequest("DELETE", "/v1/sys/rekey/backup")
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err == nil {
 		defer resp.Body.Close()
 	}
@@ -204,7 +247,10 @@ func (c *Sys) RekeyDeleteBackup() error {
 
 func (c *Sys) RekeyDeleteRecoveryBackup() error {
 	r := c.c.NewRequest("DELETE", "/v1/sys/rekey/recovery-backup")
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err == nil {
 		defer resp.Body.Close()
 	}
@@ -223,7 +269,9 @@ func (c *Sys) RekeyVerificationUpdate(shard, nonce string) (*RekeyVerificationUp
 		return nil, err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}
@@ -245,7 +293,9 @@ func (c *Sys) RekeyRecoveryKeyVerificationUpdate(shard, nonce string) (*RekeyVer
 		return nil, err
 	}
 
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}

api/sys_rotate.go

@@ -1,10 +1,16 @@
 package api
 
-import "time"
+import (
+	"context"
+	"time"
+)
 
 func (c *Sys) Rotate() error {
 	r := c.c.NewRequest("POST", "/v1/sys/rotate")
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err == nil {
 		defer resp.Body.Close()
 	}
@@ -13,7 +19,10 @@ func (c *Sys) Rotate() error {
 
 func (c *Sys) KeyStatus() (*KeyStatus, error) {
 	r := c.c.NewRequest("GET", "/v1/sys/key-status")
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}

api/sys_seal.go

@@ -1,5 +1,7 @@
 package api
 
+import "context"
+
 func (c *Sys) SealStatus() (*SealStatusResponse, error) {
 	r := c.c.NewRequest("GET", "/v1/sys/seal-status")
 	return sealStatusRequest(c, r)
@@ -7,7 +9,10 @@ func (c *Sys) SealStatus() (*SealStatusResponse, error) {
 
 func (c *Sys) Seal() error {
 	r := c.c.NewRequest("PUT", "/v1/sys/seal")
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err == nil {
 		defer resp.Body.Close()
 	}
@@ -37,7 +42,9 @@ func (c *Sys) Unseal(shard string) (*SealStatusResponse, error) {
 }
 
 func sealStatusRequest(c *Sys, r *Request) (*SealStatusResponse, error) {
-	resp, err := c.c.RawRequest(r)
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err != nil {
 		return nil, err
 	}

api/sys_stepdown.go

@@ -1,8 +1,13 @@
 package api
 
+import "context"
+
 func (c *Sys) StepDown() error {
 	r := c.c.NewRequest("PUT", "/v1/sys/step-down")
-	resp, err := c.c.RawRequest(r)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+	resp, err := c.c.RawRequestWithContext(ctx, r)
 	if err == nil {
 		defer resp.Body.Close()
 	}