add counters for active service tokens, and identity entities

vault/core_util.go

@@ -91,6 +91,12 @@ func (c *Core) HasFeature(license.Features) bool {
 	return false
 }
 
+func (c *Core) collectNamespaces() []*namespace.Namespace {
+	return []*namespace.Namespace{
+		namespace.RootNamespace,
+	}
+}
+
 func (c *Core) namepaceByPath(string) *namespace.Namespace {
 	return namespace.RootNamespace
 }

vault/counters.go

@@ -173,3 +173,66 @@ func (c *Core) saveCurrentRequestCounters(ctx context.Context, now time.Time) er
 
 	return nil
 }
+
+// ActiveTokens contains the number of active tokens.
+type ActiveTokens struct {
+	// ServiceTokens contains information about the number of active service
+	// tokens.
+	ServiceTokens TokenCounter `json:"service_tokens"`
+}
+
+// TokenCounter counts the number of tokens
+type TokenCounter struct {
+	// Total is the total number of tokens
+	Total int `json:"total"`
+}
+
+// countActiveTokens returns the number of active tokens
+func (c *Core) countActiveTokens(ctx context.Context) (*ActiveTokens, error) {
+
+	// Get all of the namespaces
+	ns := c.collectNamespaces()
+
+	// Count the tokens under each namespace
+	total := 0
+	for i := 0; i < len(ns); i++ {
+		ids, err := c.tokenStore.idView(ns[i]).List(ctx, "")
+		if err != nil {
+			return nil, err
+		}
+		total += len(ids)
+	}
+
+	return &ActiveTokens{
+		ServiceTokens: TokenCounter{
+			Total: total,
+		},
+	}, nil
+}
+
+// ActiveEntities contains the number of active entities.
+type ActiveEntities struct {
+	// Entities contains information about the number of active entities.
+	Entities EntityCounter `json:"entities"`
+}
+
+// EntityCounter counts the number of entities
+type EntityCounter struct {
+	// Total is the total number of entities
+	Total int `json:"total"`
+}
+
+// countActiveEntities returns the number of active entities
+func (c *Core) countActiveEntities(ctx context.Context) (*ActiveEntities, error) {
+
+	count, err := c.identityStore.countEntities()
+	if err != nil {
+		return nil, err
+	}
+
+	return &ActiveEntities{
+		Entities: EntityCounter{
+			Total: count,
+		},
+	}, nil
+}

vault/counters_test.go

@@ -7,6 +7,8 @@ import (
 	"time"
 
 	"github.com/go-test/deep"
+	"github.com/hashicorp/vault/helper/namespace"
+	"github.com/hashicorp/vault/sdk/logical"
 )
 
 //noinspection SpellCheckingInspection
@@ -120,3 +122,121 @@ func TestRequestCounterSaveCurrent(t *testing.T) {
 		t.Errorf("Expected=%v, got=%v, diff=%v", expected2019, all, diff)
 	}
 }
+
+func testCountActiveTokens(t *testing.T, c *Core, root string, expectedServiceTokens int) {
+	rootCtx := namespace.RootContext(nil)
+	resp, err := c.HandleRequest(rootCtx, &logical.Request{
+		ClientToken: root,
+		Operation:   logical.ReadOperation,
+		Path:        "sys/internal/counters/tokens",
+	})
+	if err != nil || (resp != nil && resp.IsError()) {
+		t.Fatalf("bad: resp: %#v\n err: %v", resp, err)
+	}
+
+	if diff := deep.Equal(resp.Data, map[string]interface{}{
+		"counters": &ActiveTokens{
+			ServiceTokens: TokenCounter{
+				Total: expectedServiceTokens,
+			},
+		},
+	}); diff != nil {
+		t.Fatal(diff)
+	}
+}
+
+func TestTokenStore_CountActiveTokens(t *testing.T) {
+	c, _, root := TestCoreUnsealed(t)
+	rootCtx := namespace.RootContext(nil)
+
+	// Count the root token
+	testCountActiveTokens(t, c, root, 1)
+
+	// Create some service tokens
+	req := &logical.Request{
+		ClientToken: root,
+		Operation:   logical.UpdateOperation,
+		Path:        "create",
+	}
+	tokens := make([]string, 10)
+	for i := 0; i < 10; i++ {
+		resp, err := c.tokenStore.HandleRequest(rootCtx, req)
+		if err != nil || (resp != nil && resp.IsError()) {
+			t.Fatalf("bad: resp: %#v\n err: %v", resp, err)
+		}
+		tokens[i] = resp.Auth.ClientToken
+
+		testCountActiveTokens(t, c, root, i+2)
+	}
+
+	// Revoke the service tokens
+	req.Path = "revoke"
+	req.Data = make(map[string]interface{})
+	for i := 0; i < 10; i++ {
+		req.Data["token"] = tokens[i]
+		resp, err := c.tokenStore.HandleRequest(rootCtx, req)
+		if err != nil || (resp != nil && resp.IsError()) {
+			t.Fatalf("bad: resp: %#v\n err: %v", resp, err)
+		}
+
+		testCountActiveTokens(t, c, root, 10-i)
+	}
+}
+
+func testCountActiveEntities(t *testing.T, c *Core, root string, expectedEntities int) {
+	rootCtx := namespace.RootContext(nil)
+	resp, err := c.HandleRequest(rootCtx, &logical.Request{
+		ClientToken: root,
+		Operation:   logical.ReadOperation,
+		Path:        "sys/internal/counters/entities",
+	})
+	if err != nil || (resp != nil && resp.IsError()) {
+		t.Fatalf("bad: resp: %#v\n err: %v", resp, err)
+	}
+
+	if diff := deep.Equal(resp.Data, map[string]interface{}{
+		"counters": &ActiveEntities{
+			Entities: EntityCounter{
+				Total: expectedEntities,
+			},
+		},
+	}); diff != nil {
+		t.Fatal(diff)
+	}
+}
+
+func TestIdentityStore_CountActiveEntities(t *testing.T) {
+	c, _, root := TestCoreUnsealed(t)
+	rootCtx := namespace.RootContext(nil)
+
+	// Count the root token
+	testCountActiveEntities(t, c, root, 0)
+
+	// Create some entities
+	req := &logical.Request{
+		ClientToken: root,
+		Operation:   logical.UpdateOperation,
+		Path:        "entity",
+	}
+	ids := make([]string, 10)
+	for i := 0; i < 10; i++ {
+		resp, err := c.identityStore.HandleRequest(rootCtx, req)
+		if err != nil || (resp != nil && resp.IsError()) {
+			t.Fatalf("bad: resp: %#v\n err: %v", resp, err)
+		}
+		ids[i] = resp.Data["id"].(string)
+
+		testCountActiveEntities(t, c, root, i+1)
+	}
+
+	req.Operation = logical.DeleteOperation
+	for i := 0; i < 10; i++ {
+		req.Path = "entity/id/" + ids[i]
+		resp, err := c.identityStore.HandleRequest(rootCtx, req)
+		if err != nil || (resp != nil && resp.IsError()) {
+			t.Fatalf("bad: resp: %#v\n err: %v", resp, err)
+		}
+
+		testCountActiveEntities(t, c, root, 9-i)
+	}
+}

vault/identity_store_util.go

@@ -2065,3 +2065,21 @@ func (i *IdentityStore) handleAliasListCommon(ctx context.Context, groupAlias bo
 
 	return logical.ListResponseWithInfo(aliasIDs, aliasInfo), nil
 }
+
+func (i *IdentityStore) countEntities() (int, error) {
+	txn := i.db.Txn(false)
+
+	iter, err := txn.Get(entitiesTable, "id")
+	if err != nil {
+		return -1, err
+	}
+
+	count := 0
+	val := iter.Next()
+	for val != nil {
+		count++
+		val = iter.Next()
+	}
+
+	return count, nil
+}

vault/logical_system.go

@@ -3124,6 +3124,36 @@ func (b *SystemBackend) pathInternalCountersRequests(ctx context.Context, req *l
 	return resp, nil
 }
 
+func (b *SystemBackend) pathInternalCountersTokens(ctx context.Context, req *logical.Request, d *framework.FieldData) (*logical.Response, error) {
+	activeTokens, err := b.Core.countActiveTokens(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	resp := &logical.Response{
+		Data: map[string]interface{}{
+			"counters": activeTokens,
+		},
+	}
+
+	return resp, nil
+}
+
+func (b *SystemBackend) pathInternalCountersEntities(ctx context.Context, req *logical.Request, d *framework.FieldData) (*logical.Response, error) {
+	activeEntities, err := b.Core.countActiveEntities(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	resp := &logical.Response{
+		Data: map[string]interface{}{
+			"counters": activeEntities,
+		},
+	}
+
+	return resp, nil
+}
+
 func (b *SystemBackend) pathInternalUIResultantACL(ctx context.Context, req *logical.Request, d *framework.FieldData) (*logical.Response, error) {
 	if req.ClientToken == "" {
 		// 204 -- no ACL
@@ -4052,4 +4082,12 @@ This path responds to the following HTTP methods.
 		"Count of requests seen by this Vault cluster over time.",
 		"Count of requests seen by this Vault cluster over time. Not included in count: health checks, UI asset requests, requests forwarded from another cluster.",
 	},
+	"internal-counters-tokens": {
+		"Count of active tokens in this Vault cluster.",
+		"Count of active tokens in this Vault cluster.",
+	},
+	"internal-counters-entities": {
+		"Count of active entities in this Vault cluster.",
+		"Count of active entities in this Vault cluster.",
+	},
 }

vault/logical_system_paths.go

@@ -871,6 +871,28 @@ func (b *SystemBackend) internalPaths() []*framework.Path {
 			HelpSynopsis:    strings.TrimSpace(sysHelp["internal-counters-requests"][0]),
 			HelpDescription: strings.TrimSpace(sysHelp["internal-counters-requests"][1]),
 		},
+		{
+			Pattern: "internal/counters/tokens",
+			Operations: map[logical.Operation]framework.OperationHandler{
+				logical.ReadOperation: &framework.PathOperation{
+					Callback:    b.pathInternalCountersTokens,
+					Unpublished: true,
+				},
+			},
+			HelpSynopsis:    strings.TrimSpace(sysHelp["internal-counters-tokens"][0]),
+			HelpDescription: strings.TrimSpace(sysHelp["internal-counters-tokens"][1]),
+		},
+		{
+			Pattern: "internal/counters/entities",
+			Operations: map[logical.Operation]framework.OperationHandler{
+				logical.ReadOperation: &framework.PathOperation{
+					Callback:    b.pathInternalCountersEntities,
+					Unpublished: true,
+				},
+			},
+			HelpSynopsis:    strings.TrimSpace(sysHelp["internal-counters-entities"][0]),
+			HelpDescription: strings.TrimSpace(sysHelp["internal-counters-entities"][1]),
+		},
 	}
 }