From 25ec3a01109db2620ae8eca5164bf20e09437b91 Mon Sep 17 00:00:00 2001
From: Jeff Mitchell <jeffrey.mitchell@gmail.com>
Date: Fri, 27 Oct 2017 13:55:47 -0400
Subject: [PATCH] Ensure revocation happens before seal/step-down since token
 store isn't available after when using single-use tokens.

Fixes #3497
---
 vault/core.go      | 40 ++++++++++++++++++++--------------------
 vault/core_test.go | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 20 deletions(-)

diff --git a/vault/core.go b/vault/core.go
index 78b2526578cb..0b731cd7d409 100644
--- a/vault/core.go
+++ b/vault/core.go
@@ -1231,16 +1231,6 @@ func (c *Core) sealInitCommon(req *logical.Request) (retErr error) {
 			c.stateLock.RUnlock()
 			return retErr
 		}
-		if te.NumUses == -1 {
-			// Token needs to be revoked
-			defer func(id string) {
-				err = c.tokenStore.Revoke(id)
-				if err != nil {
-					c.logger.Error("core: token needed revocation after seal but failed to revoke", "error", err)
-					retErr = multierror.Append(retErr, ErrInternalError)
-				}
-			}(te.ID)
-		}
 	}
 
 	// Verify that this operation is allowed
@@ -1258,6 +1248,16 @@ func (c *Core) sealInitCommon(req *logical.Request) (retErr error) {
 		return retErr
 	}
 
+	if te != nil && te.NumUses == -1 {
+		// Token needs to be revoked. We do this immediately here because
+		// we won't have a token store after sealing.
+		err = c.tokenStore.Revoke(te.ID)
+		if err != nil {
+			c.logger.Error("core: token needed revocation before seal but failed to revoke", "error", err)
+			retErr = multierror.Append(retErr, ErrInternalError)
+		}
+	}
+
 	// Tell any requests that know about this to stop
 	if c.requestContextCancelFunc != nil {
 		c.requestContextCancelFunc()
@@ -1334,16 +1334,6 @@ func (c *Core) StepDown(req *logical.Request) (retErr error) {
 			retErr = multierror.Append(retErr, logical.ErrPermissionDenied)
 			return retErr
 		}
-		if te.NumUses == -1 {
-			// Token needs to be revoked
-			defer func(id string) {
-				err = c.tokenStore.Revoke(id)
-				if err != nil {
-					c.logger.Error("core: token needed revocation after step-down but failed to revoke", "error", err)
-					retErr = multierror.Append(retErr, ErrInternalError)
-				}
-			}(te.ID)
-		}
 	}
 
 	// Verify that this operation is allowed
@@ -1359,6 +1349,16 @@ func (c *Core) StepDown(req *logical.Request) (retErr error) {
 		return retErr
 	}
 
+	if te != nil && te.NumUses == -1 {
+		// Token needs to be revoked. We do this immediately here because
+		// we won't have a token store after sealing.
+		err = c.tokenStore.Revoke(te.ID)
+		if err != nil {
+			c.logger.Error("core: token needed revocation before step-down but failed to revoke", "error", err)
+			retErr = multierror.Append(retErr, ErrInternalError)
+		}
+	}
+
 	select {
 	case c.manualStepDownCh <- struct{}{}:
 	default:
diff --git a/vault/core_test.go b/vault/core_test.go
index ab3dc746405b..68303b2f0f3a 100644
--- a/vault/core_test.go
+++ b/vault/core_test.go
@@ -272,6 +272,41 @@ func TestCore_Seal_BadToken(t *testing.T) {
 	}
 }
 
+// GH-3497
+func TestCore_Seal_SingleUse(t *testing.T) {
+	c, keys, _ := TestCoreUnsealed(t)
+	c.tokenStore.create(&TokenEntry{
+		ID:       "foo",
+		NumUses:  1,
+		Policies: []string{"root"},
+	})
+	if err := c.Seal("foo"); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if sealed, err := c.Sealed(); err != nil || !sealed {
+		t.Fatalf("err: %v, sealed: %t", err, sealed)
+	}
+	for i, key := range keys {
+		unseal, err := TestCoreUnseal(c, key)
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		if i+1 == len(keys) && !unseal {
+			t.Fatalf("err: should be unsealed")
+		}
+	}
+	if err := c.Seal("foo"); err == nil {
+		t.Fatal("expected error from revoked token")
+	}
+	te, err := c.tokenStore.Lookup("foo")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if te != nil {
+		t.Fatalf("expected nil token entry, got %#v", *te)
+	}
+}
+
 // Ensure we get a LeaseID
 func TestCore_HandleRequest_Lease(t *testing.T) {
 	c, _, root := TestCoreUnsealed(t)