From 7e313e29fda56427e93d978a77791d081d04bca5 Mon Sep 17 00:00:00 2001
From: akshya96 <87045294+akshya96@users.noreply.github.com>
Date: Wed, 15 Jun 2022 15:41:31 -0700
Subject: [PATCH] Activity Log Filtering Limit Parameter (#16000)

* adding changes from ent branch

* adding fmt changes

* adding changelog
---
 changelog/16000.txt              |  3 +++
 vault/activity_log.go            | 16 +++++++++++++++-
 vault/logical_system_activity.go | 12 +++++++++++-
 3 files changed, 29 insertions(+), 2 deletions(-)
 create mode 100644 changelog/16000.txt

diff --git a/changelog/16000.txt b/changelog/16000.txt
new file mode 100644
index 000000000..fde39b959
--- /dev/null
+++ b/changelog/16000.txt
@@ -0,0 +1,3 @@
+```release-note:improvement
+core: Limit activity log client count usage by namespaces
+```
\ No newline at end of file
diff --git a/vault/activity_log.go b/vault/activity_log.go
index 8ec8e8163..e11c7c483 100644
--- a/vault/activity_log.go
+++ b/vault/activity_log.go
@@ -1505,7 +1505,7 @@ func (a *ActivityLog) DefaultStartTime(endTime time.Time) time.Time {
 	return monthStart.AddDate(0, -a.defaultReportMonths+1, 0)
 }
 
-func (a *ActivityLog) handleQuery(ctx context.Context, startTime, endTime time.Time) (map[string]interface{}, error) {
+func (a *ActivityLog) handleQuery(ctx context.Context, startTime, endTime time.Time, limitNamespaces int) (map[string]interface{}, error) {
 	queryNS, err := namespace.FromContext(ctx)
 	if err != nil {
 		return nil, err
@@ -1557,6 +1557,7 @@ func (a *ActivityLog) handleQuery(ctx context.Context, startTime, endTime time.T
 			} else {
 				displayPath = ns.Path
 			}
+
 			byNamespace = append(byNamespace, &ResponseNamespace{
 				NamespaceID:   nsRecord.NamespaceID,
 				NamespacePath: displayPath,
@@ -1577,6 +1578,19 @@ func (a *ActivityLog) handleQuery(ctx context.Context, startTime, endTime time.T
 	sort.Slice(byNamespace, func(i, j int) bool {
 		return byNamespace[i].Counts.Clients > byNamespace[j].Counts.Clients
 	})
+	if limitNamespaces > 0 {
+		if limitNamespaces > len(byNamespace) {
+			limitNamespaces = len(byNamespace)
+		}
+		byNamespace = byNamespace[:limitNamespaces]
+		// recalculate total entities and tokens
+		totalEntities = 0
+		totalTokens = 0
+		for _, namespaceData := range byNamespace {
+			totalEntities += namespaceData.Counts.DistinctEntities
+			totalTokens += namespaceData.Counts.NonEntityTokens
+		}
+	}
 
 	responseData["by_namespace"] = byNamespace
 	responseData["total"] = &ResponseCounts{
diff --git a/vault/logical_system_activity.go b/vault/logical_system_activity.go
index 77123f767..7b8d67513 100644
--- a/vault/logical_system_activity.go
+++ b/vault/logical_system_activity.go
@@ -27,6 +27,11 @@ func (b *SystemBackend) activityQueryPath() *framework.Path {
 				Type:        framework.TypeTime,
 				Description: "End of query interval",
 			},
+			"limit_namespaces": {
+				Type:        framework.TypeInt,
+				Default:     0,
+				Description: "Limit query output by namespaces",
+			},
 		},
 		HelpSynopsis:    strings.TrimSpace(sysHelp["activity-query"][0]),
 		HelpDescription: strings.TrimSpace(sysHelp["activity-query"][1]),
@@ -198,7 +203,12 @@ func (b *SystemBackend) handleClientMetricQuery(ctx context.Context, req *logica
 		return logical.ErrorResponse(err.Error()), nil
 	}
 
-	results, err := a.handleQuery(ctx, startTime, endTime)
+	var limitNamespaces int
+	if limitNamespacesRaw, ok := d.GetOk("limit_namespaces"); ok {
+		limitNamespaces = limitNamespacesRaw.(int)
+	}
+
+	results, err := a.handleQuery(ctx, startTime, endTime, limitNamespaces)
 	if err != nil {
 		return nil, err
 	}