Cubbyhole cleanup (#6006)

* fix cubbyhole deletion * Fix error handling * Move the cubbyhole tidy logic to token store and track the revocation count * Move fetching of cubby keys before the tidy loop * Fix context getting cancelled * Test the cubbyhole cleanup logic * Add progress counter for cubbyhole cleanup * Minor polish * Use map instead of slice for faster computation * Add test for cubbyhole deletion * Add a log statement for deletion * Add SHA1 hashed tokens into the mix
hashicorp · Jan 9, 2019 · 9a38d5a · 9a38d5a
1 parent 96f1035
commit 9a38d5a
Show file tree

Hide file tree

Showing 2 changed files with 196 additions and 16 deletions.
diff --git a/vault/token_store.go b/vault/token_store.go
@@ -85,16 +85,8 @@ var (
 			return errors.New("nil token entry")
 		}
 
-		tokenNS, err := NamespaceByID(ctx, te.NamespaceID, ts.core)
-		if err != nil {
-			return err
-		}
-		if tokenNS == nil {
-			return namespace.ErrNoNamespace
-		}
-
-		switch tokenNS.ID {
-		case namespace.RootNamespaceID:
+		switch {
+		case te.NamespaceID == namespace.RootNamespaceID && !strings.HasPrefix(te.ID, "s."):
 			saltedID, err := ts.SaltID(ctx, te.ID)
 			if err != nil {
 				return err
@@ -1722,7 +1714,7 @@ func (ts *TokenStore) handleTidy(ctx context.Context, req *logical.Request, data
 
 	ns, err := namespace.FromContext(ctx)
 	if err != nil {
-		return nil, errwrap.Wrapf("failed get namespace from context: {{err}}", err)
+		return nil, errwrap.Wrapf("failed to get namespace from context: {{err}}", err)
 	}
 
 	go func() {
@@ -1751,6 +1743,12 @@ func (ts *TokenStore) handleTidy(ctx context.Context, req *logical.Request, data
 				return errwrap.Wrapf("failed to fetch secondary index entries: {{err}}", err)
 			}
 
+			// List all the cubbyhole storage keys
+			cubbyholeKeys, err := ts.cubbyholeBackend.storageView.List(quitCtx, "")
+			if err != nil {
+				return errwrap.Wrapf("failed to fetch cubbyhole storage keys: {{err}}", err)
+			}
+
 			var countParentEntries, deletedCountParentEntries, countParentList, deletedCountParentList int64
 
 			// Scan through the secondary index entries; if there is an entry
@@ -1824,9 +1822,13 @@ func (ts *TokenStore) handleTidy(ctx context.Context, req *logical.Request, data
 			}
 
 			var countAccessorList,
+				countCubbyholeKeys,
 				deletedCountAccessorEmptyToken,
 				deletedCountAccessorInvalidToken,
-				deletedCountInvalidTokenInAccessor int64
+				deletedCountInvalidTokenInAccessor,
+				deletedCountInvalidCubbyholeKey int64
+
+			validCubbyholeKeys := make(map[string]bool)
 
 			// For each of the accessor, see if the token ID associated with it is
 			// a valid one. If not, delete the leases associated with that token
@@ -1871,10 +1873,12 @@ func (ts *TokenStore) handleTidy(ctx context.Context, req *logical.Request, data
 
 				lock.RUnlock()
 
-				// If token entry is not found assume that the token is not valid any
-				// more and conclude that accessor, leases, and secondary index entries
-				// for this token should not exist as well.
-				if te == nil {
+				switch {
+				case te == nil:
+					// If token entry is not found assume that the token is not valid any
+					// more and conclude that accessor, leases, and secondary index entries
+					// for this token should not exist as well.
+
 					ts.logger.Info("deleting token with nil entry referenced by accessor", "salted_accessor", saltedAccessor)
 
 					// RevokeByToken expects a '*logical.TokenEntry'. For the
@@ -1904,6 +1908,41 @@ func (ts *TokenStore) handleTidy(ctx context.Context, req *logical.Request, data
 						continue
 					}
 					deletedCountAccessorInvalidToken++
+				default:
+					// Cache the cubbyhole storage key when the token is valid
+					switch {
+					case te.NamespaceID == namespace.RootNamespaceID && !strings.HasPrefix(te.ID, "s."):
+						saltedID, err := ts.SaltID(quitCtx, te.ID)
+						if err != nil {
+							tidyErrors = multierror.Append(tidyErrors, errwrap.Wrapf("failed to create salted token id: {{err}}", err))
+							continue
+						}
+						validCubbyholeKeys[salt.SaltID(ts.cubbyholeBackend.saltUUID, saltedID, salt.SHA1Hash)] = true
+					default:
+						if te.CubbyholeID == "" {
+							tidyErrors = multierror.Append(tidyErrors, fmt.Errorf("missing cubbyhole ID for a valid token"))
+							continue
+						}
+						validCubbyholeKeys[te.CubbyholeID] = true
+					}
+				}
+			}
+
+			// Revoke invalid cubbyhole storage keys
+			for _, key := range cubbyholeKeys {
+				countCubbyholeKeys++
+				if countCubbyholeKeys%500 == 0 {
+					ts.logger.Info("checking if there are invalid cubbyholes", "progress", countCubbyholeKeys)
+				}
+
+				key := strings.TrimSuffix(key, "/")
+				if !validCubbyholeKeys[key] {
+					ts.logger.Info("deleting invalid cubbyhole", "key", key)
+					err = ts.cubbyholeBackend.revoke(quitCtx, key)
+					if err != nil {
+						tidyErrors = multierror.Append(tidyErrors, errwrap.Wrapf(fmt.Sprintf("failed to revoke cubbyhole key %q: {{err}}", key), err))
+					}
+					deletedCountInvalidCubbyholeKey++
 				}
 			}
 
@@ -1915,6 +1954,7 @@ func (ts *TokenStore) handleTidy(ctx context.Context, req *logical.Request, data
 			ts.logger.Info("number of deleted accessors which had empty tokens", "count", deletedCountAccessorEmptyToken)
 			ts.logger.Info("number of revoked tokens which were invalid but present in accessors", "count", deletedCountInvalidTokenInAccessor)
 			ts.logger.Info("number of deleted accessors which had invalid tokens", "count", deletedCountAccessorInvalidToken)
+			ts.logger.Info("number of deleted cubbyhole keys that were invalid", "count", deletedCountInvalidCubbyholeKey)
 
 			return tidyErrors.ErrorOrNil()
 		}

diff --git a/vault/token_store_test.go b/vault/token_store_test.go
@@ -22,6 +22,146 @@ import (
 	"github.com/hashicorp/vault/logical"
 )
 
+func TestTokenStore_CubbyholeDeletion(t *testing.T) {
+	c, _, root := TestCoreUnsealed(t)
+	ts := c.tokenStore
+
+	for i := 0; i < 10; i++ {
+		// Create a token
+		tokenReq := &logical.Request{
+			Operation:   logical.UpdateOperation,
+			Path:        "create",
+			ClientToken: root,
+		}
+		// Supplying token ID forces SHA1 hashing to be used
+		if i%2 == 0 {
+			tokenReq.Data = map[string]interface{}{
+				"id": "testroot",
+			}
+		}
+		resp := testMakeTokenViaRequest(t, ts, tokenReq)
+		token := resp.Auth.ClientToken
+
+		// Write data in the token's cubbyhole
+		resp, err := c.HandleRequest(namespace.RootContext(nil), &logical.Request{
+			ClientToken: token,
+			Operation:   logical.UpdateOperation,
+			Path:        "cubbyhole/sample/data",
+			Data: map[string]interface{}{
+				"foo": "bar",
+			},
+		})
+		if err != nil || (resp != nil && resp.IsError()) {
+			t.Fatalf("bad: resp: %#v\nerr: %v", resp, err)
+		}
+
+		// Revoke the token
+		resp, err = ts.HandleRequest(namespace.RootContext(nil), &logical.Request{
+			ClientToken: token,
+			Path:        "revoke-self",
+			Operation:   logical.UpdateOperation,
+		})
+		if err != nil || (resp != nil && resp.IsError()) {
+			t.Fatalf("bad: resp: %#v\nerr: %v", resp, err)
+		}
+	}
+
+	// List the cubbyhole keys
+	cubbyholeKeys, err := ts.cubbyholeBackend.storageView.List(namespace.RootContext(nil), "")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// There should be no entries
+	if len(cubbyholeKeys) != 0 {
+		t.Fatalf("bad: len(cubbyholeKeys); expected: 0, actual: %d", len(cubbyholeKeys))
+	}
+}
+
+func TestTokenStore_CubbyholeTidy(t *testing.T) {
+	c, _, root := TestCoreUnsealed(t)
+	ts := c.tokenStore
+
+	for i := 1; i <= 20; i++ {
+		// Create 20 tokens
+		tokenReq := &logical.Request{
+			Operation:   logical.UpdateOperation,
+			Path:        "create",
+			ClientToken: root,
+		}
+
+		resp := testMakeTokenViaRequest(t, ts, tokenReq)
+		token := resp.Auth.ClientToken
+
+		// Supplying token ID forces SHA1 hashing to be used
+		if i%3 == 0 {
+			tokenReq.Data = map[string]interface{}{
+				"id": "testroot",
+			}
+		}
+
+		// Create 4 junk cubbyhole entries
+		if i%5 == 0 {
+			invalidToken, err := uuid.GenerateUUID()
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			resp, err := ts.cubbyholeBackend.HandleRequest(namespace.RootContext(nil), &logical.Request{
+				ClientToken: invalidToken,
+				Operation:   logical.UpdateOperation,
+				Path:        "cubbyhole/sample/data",
+				Data: map[string]interface{}{
+					"foo": "bar",
+				},
+				Storage: ts.cubbyholeBackend.storageView,
+			})
+			if err != nil || (resp != nil && resp.IsError()) {
+				t.Fatalf("bad: resp: %#v\nerr: %v", resp, err)
+			}
+		}
+
+		// Write into cubbyholes of 10 tokens
+		if i%2 == 0 {
+			continue
+		}
+		resp, err := c.HandleRequest(namespace.RootContext(nil), &logical.Request{
+			ClientToken: token,
+			Operation:   logical.UpdateOperation,
+			Path:        "cubbyhole/sample/data",
+			Data: map[string]interface{}{
+				"foo": "bar",
+			},
+		})
+		if err != nil || (resp != nil && resp.IsError()) {
+			t.Fatalf("bad: resp: %#v\nerr: %v", resp, err)
+		}
+	}
+
+	// Tidy cubbyhole storage
+	resp, err := ts.HandleRequest(namespace.RootContext(nil), &logical.Request{
+		Path:      "tidy",
+		Operation: logical.UpdateOperation,
+	})
+	if err != nil || (resp != nil && resp.IsError()) {
+		t.Fatalf("bad: resp: %#v\nerr: %v", resp, err)
+	}
+
+	// Wait for tidy operation to complete
+	time.Sleep(2 * time.Second)
+
+	// List all the cubbyhole storage keys
+	cubbyholeKeys, err := ts.cubbyholeBackend.storageView.List(namespace.RootContext(nil), "")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// The junk entries must have been cleaned up
+	if len(cubbyholeKeys) != 10 {
+		t.Fatalf("bad: len(cubbyholeKeys); expected: 10, actual: %d", len(cubbyholeKeys))
+	}
+}
+
 func TestTokenStore_Salting(t *testing.T) {
 	c, _, _ := TestCoreUnsealed(t)
 	ts := c.tokenStore