core: revoke the proper token on partial failures from token-related requests (#7835)

* core: revoke the proper token on partial failures from token-related requests

* move test to vault package, move test trigger to expiration manager

* update logging messages for clarity

* docstring fix

vault/core_test.go

@@ -2386,3 +2386,63 @@ func TestCore_HandleRequest_Headers_denyList(t *testing.T) {
 		t.Fatalf("did not expect %q to be in the headers map", consts.AuthHeaderName)
 	}
 }
+
+func TestCore_HandleRequest_TokenCreate_RegisterAuthFailure(t *testing.T) {
+	core, _, root := TestCoreUnsealed(t)
+
+	// Create a root token and use that for subsequent requests
+	req := logical.TestRequest(t, logical.CreateOperation, "auth/token/create")
+	req.Data = map[string]interface{}{
+		"policies": []string{"root"},
+	}
+	req.ClientToken = root
+	resp, err := core.HandleRequest(namespace.RootContext(nil), req)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if resp == nil || resp.Auth == nil || resp.Auth.ClientToken == "" {
+		t.Fatalf("expected a response from token creation, got: %#v", resp)
+	}
+	tokenWithRootPolicy := resp.Auth.ClientToken
+
+	// Use new token to create yet a new token, this should succeed
+	req = logical.TestRequest(t, logical.CreateOperation, "auth/token/create")
+	req.ClientToken = tokenWithRootPolicy
+	_, err = core.HandleRequest(namespace.RootContext(nil), req)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Try again but force failure on RegisterAuth to simulate a network failure
+	// when registering the lease (e.g. a storage failure). This should trigger
+	// an expiration manager cleanup on the newly created token
+	core.expiration.testRegisterAuthFailure.Store(true)
+	req = logical.TestRequest(t, logical.CreateOperation, "auth/token/create")
+	req.ClientToken = tokenWithRootPolicy
+	resp, err = core.HandleRequest(namespace.RootContext(nil), req)
+	if err == nil {
+		t.Fatalf("expected error, got a response: %#v", resp)
+	}
+	core.expiration.testRegisterAuthFailure.Store(false)
+
+	// Do a lookup against the client token that we used for the failed request.
+	// It should still be present
+	req = logical.TestRequest(t, logical.UpdateOperation, "auth/token/lookup")
+	req.Data = map[string]interface{}{
+		"token": tokenWithRootPolicy,
+	}
+	req.ClientToken = root
+	_, err = core.HandleRequest(namespace.RootContext(nil), req)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Do a token creation request with the token to ensure that it's still
+	// valid, should succeed.
+	req = logical.TestRequest(t, logical.CreateOperation, "auth/token/create")
+	req.ClientToken = tokenWithRootPolicy
+	resp, err = core.HandleRequest(namespace.RootContext(nil), req)
+	if err != nil {
+		t.Fatal(err)
+	}
+}

vault/expiration.go

@@ -23,6 +23,7 @@ import (
 	"github.com/hashicorp/vault/sdk/helper/jsonutil"
 	"github.com/hashicorp/vault/sdk/helper/locksutil"
 	"github.com/hashicorp/vault/sdk/logical"
+	uberAtomic "go.uber.org/atomic"
 )
 
 const (
@@ -48,7 +49,7 @@ const (
 	// defaultLeaseDuration is the default lease duration used when no lease is specified
 	defaultLeaseTTL = maxLeaseTTL
 
-	//maxLeaseThreshold is the maximum lease count before generating log warning
+	// maxLeaseThreshold is the maximum lease count before generating log warning
 	maxLeaseThreshold = 256000
 )
 
@@ -87,6 +88,11 @@ type ExpirationManager struct {
 
 	logLeaseExpirations bool
 	expireFunc          ExpireLeaseStrategy
+
+	// testRegisterAuthFailure, if set to true, triggers an explicit failure on
+	// RegisterAuth to simulate a partial failure during a token creation
+	// request. This value should only be set by tests.
+	testRegisterAuthFailure uberAtomic.Bool
 }
 
 type ExpireLeaseStrategy func(context.Context, *ExpirationManager, *leaseEntry)
@@ -1148,6 +1154,12 @@ func (m *ExpirationManager) Register(ctx context.Context, req *logical.Request,
 func (m *ExpirationManager) RegisterAuth(ctx context.Context, te *logical.TokenEntry, auth *logical.Auth) error {
 	defer metrics.MeasureSince([]string{"expire", "register-auth"}, time.Now())
 
+	// Triggers failure of RegisterAuth. This should only be set and triggered
+	// by tests to simulate partial failure during a token creation request.
+	if m.testRegisterAuthFailure.Load() {
+		return fmt.Errorf("failing explicitly on RegisterAuth")
+	}
+
 	authExpirationTime := auth.ExpirationTime()
 
 	if te.TTL == 0 && authExpirationTime.IsZero() && (len(te.Policies) != 1 || te.Policies[0] != "root") {

vault/request_handling.go

@@ -860,7 +860,11 @@ func (c *Core) handleRequest(ctx context.Context, req *logical.Request) (retResp
 
 		_, identityPolicies, err := c.fetchEntityAndDerivedPolicies(ctx, tokenNS, resp.Auth.EntityID)
 		if err != nil {
-			c.tokenStore.revokeOrphan(ctx, te.ID)
+			// Best-effort clean up on error, so we log the cleanup error as a
+			// warning but still return as internal error.
+			if err := c.tokenStore.revokeOrphan(ctx, resp.Auth.ClientToken); err != nil {
+				c.logger.Warn("failed to clean up token lease from entity and policy lookup failure", "request_path", req.Path, "error", err)
+			}
 			return nil, nil, ErrInternalError
 		}
 
@@ -874,8 +878,12 @@ func (c *Core) handleRequest(ctx context.Context, req *logical.Request) (retResp
 				Path:        resp.Auth.CreationPath,
 				NamespaceID: ns.ID,
 			}, resp.Auth); err != nil {
-				c.tokenStore.revokeOrphan(ctx, te.ID)
-				c.logger.Error("failed to register token lease", "request_path", req.Path, "error", err)
+				// Best-effort clean up on error, so we log the cleanup error as
+				// a warning but still return as internal error.
+				if err := c.tokenStore.revokeOrphan(ctx, resp.Auth.ClientToken); err != nil {
+					c.logger.Warn("failed to clean up token lease during auth/token/ request", "request_path", req.Path, "error", err)
+				}
+				c.logger.Error("failed to register token lease during auth/token/ request", "request_path", req.Path, "error", err)
 				retErr = multierror.Append(retErr, ErrInternalError)
 				return nil, auth, retErr
 			}
@@ -1159,6 +1167,8 @@ func (c *Core) handleLoginRequest(ctx context.Context, req *logical.Request) (re
 	return resp, auth, routeErr
 }
 
+// RegisterAuth uses a logical.Auth object to create a token entry in the token
+// store, and registers a corresponding token lease to the expiration manager.
 func (c *Core) RegisterAuth(ctx context.Context, tokenTTL time.Duration, path string, auth *logical.Auth) error {
 	// We first assign token policies to what was returned from the backend
 	// via auth.Policies. Then, we get the full set of policies into
@@ -1209,8 +1219,10 @@ func (c *Core) RegisterAuth(ctx context.Context, tokenTTL time.Duration, path st
 	case logical.TokenTypeService:
 		// Register with the expiration manager
 		if err := c.expiration.RegisterAuth(ctx, &te, auth); err != nil {
-			c.tokenStore.revokeOrphan(ctx, te.ID)
-			c.logger.Error("failed to register token lease", "request_path", path, "error", err)
+			if err := c.tokenStore.revokeOrphan(ctx, te.ID); err != nil {
+				c.logger.Warn("failed to clean up token lease during login request", "request_path", path, "error", err)
+			}
+			c.logger.Error("failed to register token lease during login request", "request_path", path, "error", err)
 			return ErrInternalError
 		}
 	}