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

http/logical_test.go

@@ -437,3 +437,48 @@ func TestLogical_Audit_invalidWrappingToken(t *testing.T) {
 		}
 	}
 }
+
+func TestLogical_CreateToken_Failure(t *testing.T) {
+	core, _, root := vault.TestCoreUnsealed(t)
+
+	ln, addr := TestServer(t, core)
+	defer ln.Close()
+	TestServerAuth(t, addr, root)
+
+	// Create a root token and use that for subsequent requests
+	resp := testHttpPut(t, root, addr+"/v1/auth/token/create", map[string]interface{}{
+		"policies": []string{"root"},
+	})
+	testResponseStatus(t, resp, 200)
+
+	var actual map[string]interface{}
+	testResponseBody(t, resp, &actual)
+	tokenWithRootPolicy := actual["auth"].(map[string]interface{})["client_token"].(string)
+	if tokenWithRootPolicy == "" {
+		t.Fatalf("expected a client_token string but not found\n%#v", actual)
+	}
+
+	// Use new token to create yet a new token, this should succeed
+	resp = testHttpPut(t, tokenWithRootPolicy, addr+"/v1/auth/token/create", nil)
+	testResponseStatus(t, resp, 200)
+
+	// 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.SetTestRegisterAuthFailure(true)
+	resp = testHttpPut(t, tokenWithRootPolicy, addr+"/v1/auth/token/create", nil)
+	testResponseStatus(t, resp, 500)
+	core.SetTestRegisterAuthFailure(false)
+
+	// Do a lookup against the client token that we used for the failed request.
+	// It should still be present.
+	resp = testHttpPost(t, root, addr+"/v1/auth/token/lookup", map[string]interface{}{
+		"token": tokenWithRootPolicy,
+	})
+	testResponseStatus(t, resp, 200)
+
+	// Do a token creation request with the token to ensure that it's still
+	// valid, should succeed.
+	resp = testHttpPut(t, tokenWithRootPolicy, addr+"/v1/auth/token/create", nil)
+	testResponseStatus(t, resp, 200)
+}

vault/core.go

@@ -44,6 +44,7 @@ import (
 	"github.com/hashicorp/vault/vault/seal"
 	shamirseal "github.com/hashicorp/vault/vault/seal/shamir"
 	cache "github.com/patrickmn/go-cache"
+	uberAtomic "go.uber.org/atomic"
 	"google.golang.org/grpc"
 )
 
@@ -478,6 +479,11 @@ type Core struct {
 	secureRandomReader io.Reader
 
 	recoveryMode bool
+
+	// 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
 }
 
 // CoreConfig is used to parameterize a core
@@ -2175,6 +2181,14 @@ func (c *Core) MetricsHelper() *metricsutil.MetricsHelper {
 	return c.metricsHelper
 }
 
+// SetTestRegisterAuthFailure sets the an internal atomic bool within
+// core to trigger failure on its expiration manager's RegisterAuth
+// method to simulate partial failure during a token creation request.
+// This method should only be used by tests.
+func (c *Core) SetTestRegisterAuthFailure(active bool) {
+	c.testRegisterAuthFailure.Store(active)
+}
+
 // BuiltinRegistry is an interface that allows the "vault" package to use
 // the registry of builtin plugins without getting an import cycle. It
 // also allows for mocking the registry easily.

vault/expiration.go

@@ -48,7 +48,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
 )
 
@@ -1147,6 +1147,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 triggered by tests
+	// to simulate partial failure during a token creation request.
+	if m.core.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", "request_path", req.Path, "error", err)
+			}
 			return nil, nil, ErrInternalError
 		}
 
@@ -874,7 +878,11 @@ 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)
+				// 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", "request_path", req.Path, "error", err)
+				}
 				c.logger.Error("failed to register token lease", "request_path", req.Path, "error", err)
 				retErr = multierror.Append(retErr, ErrInternalError)
 				return nil, auth, retErr
@@ -1209,7 +1217,9 @@ 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)
+			if err := c.tokenStore.revokeOrphan(ctx, te.ID); err != nil {
+				c.logger.Warn("failed to clean up token lease", "request_path", path, "error", err)
+			}
 			c.logger.Error("failed to register token lease", "request_path", path, "error", err)
 			return ErrInternalError
 		}