Add context handling to Consul operations (#4739)

hashicorp · Jun 11, 2018 · c707039 · c707039
1 parent 24b25cd
commit c707039
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 10 deletions.
diff --git a/builtin/logical/consul/path_token.go b/builtin/logical/consul/path_token.go
@@ -59,12 +59,15 @@ func (b *backend) pathTokenRead(ctx context.Context, req *logical.Request, d *fr
 	// Generate a name for the token
 	tokenName := fmt.Sprintf("Vault %s %s %d", role, req.DisplayName, time.Now().UnixNano())
 
+	writeOpts := &api.WriteOpts{}
+	writeOpts = writeOpts.WithContext(ctx)
+
 	// Create it
 	token, _, err := c.ACL().Create(&api.ACLEntry{
 		Name:  tokenName,
 		Type:  result.TokenType,
 		Rules: result.Policy,
-	}, nil)
+	}, writeOpts)
 	if err != nil {
 		return logical.ErrorResponse(err.Error()), nil
 	}

diff --git a/physical/consul/consul.go b/physical/consul/consul.go
@@ -394,7 +394,10 @@ func (c *ConsulBackend) Transaction(ctx context.Context, txns []*physical.TxnEnt
 	c.permitPool.Acquire()
 	defer c.permitPool.Release()
 
-	ok, resp, _, err := c.kv.Txn(ops, nil)
+	queryOpts := &api.QueryOptions{}
+	queryOpts = queryOpts.WithContext(ctx)
+
+	ok, resp, _, err := c.kv.Txn(ops, queryOpts)
 	if err != nil {
 		return err
 	}
@@ -422,7 +425,10 @@ func (c *ConsulBackend) Put(ctx context.Context, entry *physical.Entry) error {
 		Value: entry.Value,
 	}
 
-	_, err := c.kv.Put(pair, nil)
+	writeOpts := &api.WriteOptions{}
+	writeOpts = writeOpts.WithContext(ctx)
+
+	_, err := c.kv.Put(pair, writeOpts)
 	return err
 }
 
@@ -433,14 +439,14 @@ func (c *ConsulBackend) Get(ctx context.Context, key string) (*physical.Entry, e
 	c.permitPool.Acquire()
 	defer c.permitPool.Release()
 
-	var queryOptions *api.QueryOptions
+	queryOpts := &api.QueryOptions{}
+	queryOpts = queryOpts.WithContext(ctx)
+
 	if c.consistencyMode == consistencyModeStrong {
-		queryOptions = &api.QueryOptions{
-			RequireConsistent: true,
-		}
+		queryOpts.RequireConsistent = true
 	}
 
-	pair, _, err := c.kv.Get(c.path+key, queryOptions)
+	pair, _, err := c.kv.Get(c.path+key, queryOpts)
 	if err != nil {
 		return nil, err
 	}
@@ -461,7 +467,10 @@ func (c *ConsulBackend) Delete(ctx context.Context, key string) error {
 	c.permitPool.Acquire()
 	defer c.permitPool.Release()
 
-	_, err := c.kv.Delete(c.path+key, nil)
+	writeOpts := &api.WriteOptions{}
+	writeOpts = writeOpts.WithContext(ctx)
+
+	_, err := c.kv.Delete(c.path+key, writeOpts)
 	return err
 }
 
@@ -481,7 +490,10 @@ func (c *ConsulBackend) List(ctx context.Context, prefix string) ([]string, erro
 	c.permitPool.Acquire()
 	defer c.permitPool.Release()
 
-	out, _, err := c.kv.Keys(scan, "/", nil)
+	queryOpts := &api.QueryOptions{}
+	queryOpts = queryOpts.WithContext(ctx)
+
+	out, _, err := c.kv.Keys(scan, "/", queryOpts)
 	for idx, val := range out {
 		out[idx] = strings.TrimPrefix(val, scan)
 	}