Change UseToken mechanics.

Add locking around UseToken and Lookup. Have UseToken flag an entry that
needs to be revoked so that it can be done at the appropriate time, but
so that Lookup in the interm doesn't return a value.

The locking is a map of 4096 locks keyed off of the first three
characters of the token ID which should provide good distribution.

meta/meta_test.go

@@ -18,7 +18,7 @@ func TestFlagSet(t *testing.T) {
 		},
 		{
 			FlagSetServer,
-			[]string{"address", "ca-cert", "ca-path", "client-cert", "client-key", "insecure", "tls-skip-verify"},
+			[]string{"address", "ca-cert", "ca-path", "client-cert", "client-key", "insecure", "tls-skip-verify", "wrap-ttl"},
 		},
 	}
 

vault/core.go

@@ -411,7 +411,7 @@ func (c *Core) checkToken(req *logical.Request) (*logical.Auth, *TokenEntry, err
 
 	acl, te, err := c.fetchACLandTokenEntry(req)
 	if err != nil {
-		return nil, nil, err
+		return nil, te, err
 	}
 
 	// Check if this is a root protected path
@@ -450,13 +450,14 @@ func (c *Core) checkToken(req *logical.Request) (*logical.Auth, *TokenEntry, err
 		}
 	}
 
-	// Check the standard non-root ACLs
+	// Check the standard non-root ACLs. Return the token entry if it's not
+	// allowed so we can decrement the use count.
 	allowed, rootPrivs := acl.AllowOperation(req.Operation, req.Path)
 	if !allowed {
-		return nil, nil, logical.ErrPermissionDenied
+		return nil, te, logical.ErrPermissionDenied
 	}
 	if rootPath && !rootPrivs {
-		return nil, nil, logical.ErrPermissionDenied
+		return nil, te, logical.ErrPermissionDenied
 	}
 
 	// Create the auth response
@@ -687,12 +688,20 @@ func (c *Core) Seal(token string) (retErr error) {
 		return err
 	}
 	// Attempt to use the token (decrement num_uses)
-	// If we can't, we still continue attempting the seal, so long as the token
+	// On error bail out; if the token has been revoked, bail out too
-	// has appropriate permissions
 	if te != nil {
-		if err := c.tokenStore.UseToken(te); err != nil {
+		te, err = c.tokenStore.UseToken(te)
+		if err != nil {
 			c.logger.Printf("[ERR] core: failed to use token: %v", err)
-			retErr = ErrInternalError
+			return ErrInternalError
+		}
+		if te == nil {
+			// Token is no longer valid
+			return logical.ErrPermissionDenied
+		}
+		if te.NumUses == -1 {
+			// Token needs to be revoked
+			return c.tokenStore.Revoke(te.ID)
 		}
 	}
 
@@ -744,10 +753,19 @@ func (c *Core) StepDown(token string) error {
 	}
 	// Attempt to use the token (decrement num_uses)
 	if te != nil {
-		if err := c.tokenStore.UseToken(te); err != nil {
+		te, err = c.tokenStore.UseToken(te)
+		if err != nil {
 			c.logger.Printf("[ERR] core: failed to use token: %v", err)
 			return err
 		}
+		if te == nil {
+			// Token has been revoked
+			return logical.ErrPermissionDenied
+		}
+		if te.NumUses == -1 {
+			// Token needs to be revoked
+			return c.tokenStore.Revoke(te.ID)
+		}
 	}
 
 	// Verify that this operation is allowed

vault/request_handling.go

@@ -146,10 +146,7 @@ func (c *Core) HandleRequest(req *logical.Request) (resp *logical.Response, err
 	// wrapped information, since the original response has been audit logged
 	if wrapping {
 		wrappingResp := &logical.Response{
-			WrapInfo: logical.WrapInfo{
+			WrapInfo: resp.WrapInfo,
-				Token: resp.WrapInfo.Token,
-				TTL:   resp.WrapInfo.TTL,
-			},
 		}
 		wrappingResp.CloneWarnings(resp)
 		resp = wrappingResp
@@ -162,45 +159,58 @@ func (c *Core) handleRequest(req *logical.Request) (retResp *logical.Response, r
 	defer metrics.MeasureSince([]string{"core", "handle_request"}, time.Now())
 
 	// Validate the token
-	auth, te, err := c.checkToken(req)
+	auth, te, ctErr := c.checkToken(req)
+	// We run this logic first because we want to decrement the use count even in the case of an error
 	if te != nil {
-		defer func() {
+		// Attempt to use the token (decrement NumUses)
-			// Attempt to use the token (decrement num_uses)
+		var err error
-			// If a secret was generated and num_uses is currently 1, it will be
+		te, err = c.tokenStore.UseToken(te)
-			// immediately revoked; in that case, don't return the leased
+		if err != nil {
-			// credentials as they are now invalid.
+			c.logger.Printf("[ERR] core: failed to use token: %v", err)
-			if retResp != nil &&
+			return nil, nil, ErrInternalError
-				te != nil && te.NumUses == 1 &&
+		}
-				retResp.Secret != nil &&
+		if te == nil {
-				// Some backends return a TTL even without a Lease ID
+			// Token has been revoked by this point
-				retResp.Secret.LeaseID != "" {
+			return nil, nil, logical.ErrPermissionDenied
-				retResp = logical.ErrorResponse("Secret cannot be returned; token had one use left, so leased credentials were immediately revoked.")
+		}
-			}
+		if te.NumUses == -1 {
-			if err := c.tokenStore.UseToken(te); err != nil {
+			// We defer a revocation until after logic has run, since this is a
-				c.logger.Printf("[ERR] core: failed to use token: %v", err)
+			// valid request (this is the token's final use). We pass the ID in
-				retResp = nil
+			// directly just to be safe in case something else modifies te later.
-				retAuth = nil
+			defer func(id string) {
-				retErr = ErrInternalError
+				err = c.tokenStore.Revoke(id)
-			}
+				if err != nil {
-		}()
+					c.logger.Printf("[ERR] core: failed to revoke token: %v", err)
+					retResp = nil
+					retAuth = nil
+					retErr = ErrInternalError
+				}
+				if retResp != nil && retResp.Secret != nil &&
+					// Some backends return a TTL even without a Lease ID
+					retResp.Secret.LeaseID != "" {
+					retResp = logical.ErrorResponse("Secret cannot be returned; token had one use left, so leased credentials were immediately revoked.")
+					return
+				}
+			}(te.ID)
+		}
 	}
-	if err != nil {
+	if ctErr != nil {
 		// If it is an internal error we return that, otherwise we
 		// return invalid request so that the status codes can be correct
 		var errType error
-		switch err {
+		switch ctErr {
 		case ErrInternalError, logical.ErrPermissionDenied:
-			errType = err
+			errType = ctErr
 		default:
 			errType = logical.ErrInvalidRequest
 		}
 
-		if err := c.auditBroker.LogRequest(auth, req, err); err != nil {
+		if err := c.auditBroker.LogRequest(auth, req, ctErr); err != nil {
 			c.logger.Printf("[ERR] core: failed to audit request with path (%s): %v",
 				req.Path, err)
 		}
 
-		return logical.ErrorResponse(err.Error()), nil, errType
+		return logical.ErrorResponse(ctErr.Error()), nil, errType
 	}
 
 	// Attach the display name

vault/token_store.go

@@ -5,7 +5,9 @@ import (
 	"fmt"
 	"regexp"
 	"sort"
+	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/armon/go-metrics"
@@ -61,6 +63,8 @@ type TokenStore struct {
 	cubbyholeBackend *CubbyholeBackend
 
 	policyLookupFunc func(string) (*Policy, error)
+
+	tokenLocks map[string]*sync.RWMutex
 }
 
 // NewTokenStore is used to construct a token store that is
@@ -87,6 +91,19 @@ func NewTokenStore(c *Core, config *logical.BackendConfig) (*TokenStore, error)
 	}
 	t.salt = salt
 
+	t.tokenLocks = map[string]*sync.RWMutex{}
+	for i := int64(0); i < 16; i++ {
+		for j := int64(0); j < 16; j++ {
+			for k := int64(0); k < 16; k++ {
+				t.tokenLocks[fmt.Sprintf("%s%s%s",
+					strconv.FormatInt(i, 16),
+					strconv.FormatInt(j, 16),
+					strconv.FormatInt(k, 16))] = &sync.RWMutex{}
+			}
+		}
+	}
+	t.tokenLocks["global"] = &sync.RWMutex{}
+
 	// Setup the framework endpoints
 	t.Backend = &framework.Backend{
 		AuthRenew: t.authRenew,
@@ -547,33 +564,74 @@ func (ts *TokenStore) storeCommon(entry *TokenEntry, writeSecondary bool) error
 	return nil
 }
 
-// UseToken is used to manage restricted use tokens and decrement
+func (ts *TokenStore) getTokenLock(id string) *sync.RWMutex {
-// their available uses. Note: this is potentially racy, but the simple
+	// Find our multilevel lock, or fall back to global
-// solution of a global lock would be severely detrimental to performance. Also
+	var lock *sync.RWMutex
-// note the specific revoke case below.
+	var ok bool
-func (ts *TokenStore) UseToken(te *TokenEntry) error {
+	if len(id) >= 3 {
+		lock, ok = ts.tokenLocks[id[0:3]]
+	}
+	if !ok || lock == nil {
+		lock = ts.tokenLocks["global"]
+	}
+
+	return lock
+}
+
+// UseToken is used to manage restricted use tokens and decrement their
+// available uses. Returns two values: a potentially updated entry or, if the
+// token has been revoked, nil; and whether an error was encountered. The
+// locking here isn't perfect, as other parts of the code may update an entry,
+// but usually none after the entry is already created...so this is pretty
+// good.
+func (ts *TokenStore) UseToken(te *TokenEntry) (*TokenEntry, error) {
+	if te == nil {
+		return nil, fmt.Errorf("invalid token entry provided for use count decrementing")
+	}
+
+	lock := ts.getTokenLock(te.ID)
+
+	lock.RLock()
 	// If the token is not restricted, there is nothing to do
 	if te.NumUses == 0 {
-		return nil
+		lock.RUnlock()
+		return te, nil
+	}
+
+	lock.RUnlock()
+	lock.Lock()
+	defer lock.Unlock()
+
+	// Call lookupSalted to avoid lock contention
+	te, err := ts.lookupSalted(ts.SaltID(te.ID))
+	if err != nil {
+		return nil, fmt.Errorf("failed to refresh entry: %v", err)
+	}
+	if te == nil {
+		return nil, fmt.Errorf("token entry nil after refreshing to decrement use count; token has likely been used already")
+	}
+
+	// Check if it's already been revoked at this point, if so, don't return an
+	// error, but do indicate with a nil response that the token is no longer
+	// valid
+	if te.NumUses == -1 {
+		return nil, nil
 	}
 
 	// Decrement the count
 	te.NumUses -= 1
 
-	// Revoke the token if there are no remaining uses.
+	// We need to indicate that this is no longer valid, but revocation is
-	// XXX: There is a race condition here with parallel
+	// deferred to the end of the call, so this will make sure that any Lookup
-	// requests using the same token. This would require
+	// that happens doesn't return an entry
-	// some global coordination to avoid, as we must ensure
-	// no requests using the same restricted token are handled
-	// in parallel.
 	if te.NumUses == 0 {
-		return ts.Revoke(te.ID)
+		te.NumUses = -1
 	}
 
 	// Marshal the entry
 	enc, err := json.Marshal(te)
 	if err != nil {
-		return fmt.Errorf("failed to encode entry: %v", err)
+		return nil, fmt.Errorf("failed to encode entry: %v", err)
 	}
 
 	// Write under the primary ID
@@ -581,17 +639,23 @@ func (ts *TokenStore) UseToken(te *TokenEntry) error {
 	path := lookupPrefix + saltedId
 	le := &logical.StorageEntry{Key: path, Value: enc}
 	if err := ts.view.Put(le); err != nil {
-		return fmt.Errorf("failed to persist entry: %v", err)
+		return nil, fmt.Errorf("failed to persist entry: %v", err)
 	}
-	return nil
+
+	return te, nil
 }
 
-// Lookup is used to find a token given its ID
+// Lookup is used to find a token given its ID. It acquires a read lock, then calls lookupSalted.
 func (ts *TokenStore) Lookup(id string) (*TokenEntry, error) {
 	defer metrics.MeasureSince([]string{"token", "lookup"}, time.Now())
 	if id == "" {
 		return nil, fmt.Errorf("cannot lookup blank token")
 	}
+
+	lock := ts.getTokenLock(id)
+	lock.RLock()
+	defer lock.RUnlock()
+
 	return ts.lookupSalted(ts.SaltID(id))
 }
 
@@ -614,6 +678,12 @@ func (ts *TokenStore) lookupSalted(saltedId string) (*TokenEntry, error) {
 	if err := json.Unmarshal(raw.Value, entry); err != nil {
 		return nil, fmt.Errorf("failed to decode entry: %v", err)
 	}
+
+	// This is a token that is awaiting deferred revocation
+	if entry.NumUses == -1 {
+		return nil, nil
+	}
+
 	return entry, nil
 }
 

vault/token_store_test.go

@@ -244,10 +244,13 @@ func TestTokenStore_UseToken(t *testing.T) {
 	}
 
 	// Root is an unlimited use token, should be a no-op
-	err = ts.UseToken(ent)
+	te, err := ts.UseToken(ent)
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
+	if te == nil {
+		t.Fatalf("token entry after use was nil")
+	}
 
 	// Lookup the root token again
 	ent2, err := ts.Lookup(root)
@@ -266,10 +269,13 @@ func TestTokenStore_UseToken(t *testing.T) {
 	}
 
 	// Use the token
-	err = ts.UseToken(ent)
+	te, err = ts.UseToken(ent)
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
+	if te == nil {
+		t.Fatalf("token entry for use #1 was nil")
+	}
 
 	// Lookup the token
 	ent2, err = ts.Lookup(ent.ID)
@@ -283,10 +289,17 @@ func TestTokenStore_UseToken(t *testing.T) {
 	}
 
 	// Use the token
-	err = ts.UseToken(ent)
+	te, err = ts.UseToken(ent)
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
+	if te == nil {
+		t.Fatalf("token entry for use #2 was nil")
+	}
+	if te.NumUses != -1 {
+		t.Fatalf("token entry after use #2 did not have revoke flag")
+	}
+	ts.Revoke(te.ID)
 
 	// Lookup the token
 	ent2, err = ts.Lookup(ent.ID)