Non-recursive DFS token tree revoke

vault/token_store.go

@@ -1214,29 +1214,39 @@ func (ts *TokenStore) RevokeTree(id string) error {
 	return nil
 }
 
-// revokeTreeSalted is used to invalide a given token and all
+// revokeTreeSalted is used to invalidate a given token and all
 // child tokens using a saltedID.
+// Updated to be non-recursive and revoke child tokens
+// before parent tokens(DFS).
 func (ts *TokenStore) revokeTreeSalted(saltedId string) error {
-	// Scan for child tokens
-	path := parentPrefix + saltedId + "/"
-	children, err := ts.view.List(path)
-	if err != nil {
-		return fmt.Errorf("failed to scan for children: %v", err)
-	}
+	var dfs []string
+	dfs = append(dfs, saltedId)
 
-	// Recursively nuke the children. The subtle nuance here is that
-	// we don't have the acutal ID of the child, but we have the salted
-	// value. Turns out, this is good enough!
-	for _, child := range children {
-		if err := ts.revokeTreeSalted(child); err != nil {
-			return err
+	for l := len(dfs); l > 0; l = len(dfs) {
+		id := dfs[0]
+		path := parentPrefix + id + "/"
+		children, err := ts.view.List(path)
+		if err != nil {
+			return fmt.Errorf("failed to scan for children: %v", err)
+		}
+		// If the length of the children array is zero,
+		// then we are at a leaf node.
+		if len(children) == 0 {
+			if err := ts.revokeSalted(id); err != nil {
+				return fmt.Errorf("failed to revoke entry: %v", err)
+			}
+			// If the length of l is equal to 1, then the last token has been deleted
+			if l == 1 {
+				return nil
+			}
+			dfs = dfs[1:]
+		} else {
+			// If we make it here, there are children and they must
+			// be prepended.
+			dfs = append(children, dfs...)
 		}
 	}
 
-	// Revoke this entry
-	if err := ts.revokeSalted(saltedId); err != nil {
-		return fmt.Errorf("failed to revoke entry: %v", err)
-	}
 	return nil
 }
 

vault/token_store_test.go

@@ -770,41 +770,36 @@ func TestTokenStore_Revoke_Orphan(t *testing.T) {
 	}
 }
 
+// This was the original function name, and now it just calls
+// the non recursive version for a variety of depths.
 func TestTokenStore_RevokeTree(t *testing.T) {
-	_, ts, _, _ := TestCoreWithTokenStore(t)
-
-	ent1 := &TokenEntry{}
-	if err := ts.create(ent1); err != nil {
-		t.Fatalf("err: %v", err)
-	}
+	testTokenStore_RevokeTree_NonRecursive(t, 1)
+	testTokenStore_RevokeTree_NonRecursive(t, 2)
+	testTokenStore_RevokeTree_NonRecursive(t, 10)
+}
 
-	ent2 := &TokenEntry{Parent: ent1.ID}
-	if err := ts.create(ent2); err != nil {
-		t.Fatalf("err: %v", err)
-	}
+// Revokes a given Token Store tree non recursively.
+// The second parameter refers to the depth of the tree.
+func testTokenStore_RevokeTree_NonRecursive(t testing.TB, depth uint64) {
+	_, ts, _, _ := TestCoreWithTokenStore(t)
+	root, children := buildTokenTree(t, ts, depth)
+	err := ts.RevokeTree("")
 
-	ent3 := &TokenEntry{Parent: ent2.ID}
-	if err := ts.create(ent3); err != nil {
+	if err.Error() != "cannot tree-revoke blank token" {
 		t.Fatalf("err: %v", err)
 	}
 
-	ent4 := &TokenEntry{Parent: ent2.ID}
-	if err := ts.create(ent4); err != nil {
-		t.Fatalf("err: %v", err)
-	}
+	// Nuke tree non recursively.
+	err = ts.RevokeTree(root.ID)
 
-	err := ts.RevokeTree("")
-	if err.Error() != "cannot tree-revoke blank token" {
-		t.Fatalf("err: %v", err)
-	}
-	err = ts.RevokeTree(ent1.ID)
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
-
-	lookup := []string{ent1.ID, ent2.ID, ent3.ID, ent4.ID}
-	for _, id := range lookup {
-		out, err := ts.Lookup(id)
+	// Append the root to ensure it was successfully
+	// deleted.
+	children = append(children, root)
+	for _, entry := range children {
+		out, err := ts.Lookup(entry.ID)
 		if err != nil {
 			t.Fatalf("err: %v", err)
 		}
@@ -814,6 +809,52 @@ func TestTokenStore_RevokeTree(t *testing.T) {
 	}
 }
 
+// A benchmark function that tests testTokenStore_RevokeTree_NonRecursive
+// for a variety of different depths.
+func BenchmarkTokenStore_RevokeTree(b *testing.B) {
+	benchmarks := []uint64{0, 1, 2, 4, 8, 16, 20}
+	for _, depth := range benchmarks {
+		b.Run(fmt.Sprintf("Tree of Depth %d", depth), func(b *testing.B) {
+			for i := 0; i < b.N; i++ {
+				testTokenStore_RevokeTree_NonRecursive(b, depth)
+			}
+		})
+	}
+}
+
+// Builds a TokenTree of a specified depth, so that
+// we may run revoke tests on it.
+func buildTokenTree(t testing.TB, ts *TokenStore, depth uint64) (root *TokenEntry, children []*TokenEntry) {
+	root = &TokenEntry{}
+	if err := ts.create(root); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	frontier := []*TokenEntry{root}
+	current := uint64(0)
+	for current < depth {
+		next := make([]*TokenEntry, 0, 2*len(frontier))
+		for _, node := range frontier {
+			left := &TokenEntry{Parent: node.ID}
+			if err := ts.create(left); err != nil {
+				t.Fatalf("err: %v", err)
+			}
+
+			right := &TokenEntry{Parent: node.ID}
+			if err := ts.create(right); err != nil {
+				t.Fatalf("err: %v", err)
+			}
+
+			children = append(children, left, right)
+			next = append(next, left, right)
+		}
+		frontier = next
+		current++
+	}
+
+	return root, children
+}
+
 func TestTokenStore_RevokeSelf(t *testing.T) {
 	_, ts, _, _ := TestCoreWithTokenStore(t)