diff --git a/vault/token_store.go b/vault/token_store.go index d989393711..b5cbfcf57b 100644 --- a/vault/token_store.go +++ b/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 } diff --git a/vault/token_store_test.go b/vault/token_store_test.go index 1dd1e916c4..1b3d7c346c 100644 --- a/vault/token_store_test.go +++ b/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) { + testTokenStore_RevokeTree_NonRecursive(t, 1) + testTokenStore_RevokeTree_NonRecursive(t, 2) + testTokenStore_RevokeTree_NonRecursive(t, 10) +} + +// 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) - - ent1 := &TokenEntry{} - if err := ts.create(ent1); err != nil { - t.Fatalf("err: %v", err) - } - - ent2 := &TokenEntry{Parent: ent1.ID} - if err := ts.create(ent2); err != nil { - t.Fatalf("err: %v", err) - } - - ent3 := &TokenEntry{Parent: ent2.ID} - if err := ts.create(ent3); err != nil { - t.Fatalf("err: %v", err) - } - - ent4 := &TokenEntry{Parent: ent2.ID} - if err := ts.create(ent4); err != nil { - t.Fatalf("err: %v", err) - } - + root, children := buildTokenTree(t, ts, depth) err := ts.RevokeTree("") + if err.Error() != "cannot tree-revoke blank token" { t.Fatalf("err: %v", err) } - err = ts.RevokeTree(ent1.ID) + + // Nuke tree non recursively. + err = ts.RevokeTree(root.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)