(OSS) Path Suffix Support for Rate Limit Quotas (#15989)

* Support for rate limit path suffix quotas

* Support for rate limit path suffix quotas

* Precedence test for support for rate limit path suffix quotas

* Update clone method

* Fix mount determination

* Add changelog

* use constant for mounts

* Fix read endpoint, and remount/disable mount

* update godocs for queryquota

changelog/15989.txt

@@ -0,0 +1,3 @@
+```release-note:improvement
+core/quotas: Added ability to add path suffixes for rate-limit resource quotas
+```

vault/logical_system_quotas.go

@@ -186,15 +186,26 @@ func (b *SystemBackend) handleRateLimitQuotasUpdate() framework.OperationFunc {
 			mountPath = strings.TrimPrefix(mountPath, ns.Path)
 		}
 
+		var pathSuffix string
 		if mountPath != "" {
-			match := b.Core.router.MatchingMount(namespace.ContextWithNamespace(ctx, ns), mountPath)
-			if match == "" {
+			me := b.Core.router.MatchingMountEntry(namespace.ContextWithNamespace(ctx, ns), mountPath)
+			if me == nil {
 				return logical.ErrorResponse("invalid mount path %q", mountPath), nil
 			}
+
+			var newMountPath string
+			if me.Table == mountTableType {
+				newMountPath = me.Path
+			} else {
+				newMountPath = me.Table + "/" + me.Path
+			}
+
+			pathSuffix = strings.TrimSuffix(strings.TrimPrefix(mountPath, newMountPath), "/")
+			mountPath = newMountPath
 		}
 		// Disallow creation of new quota that has properties similar to an
 		// existing quota.
-		quotaByFactors, err := b.Core.quotaManager.QuotaByFactors(ctx, qType, ns.Path, mountPath)
+		quotaByFactors, err := b.Core.quotaManager.QuotaByFactors(ctx, qType, ns.Path, mountPath, pathSuffix)
 		if err != nil {
 			return nil, err
 		}
@@ -210,7 +221,7 @@ func (b *SystemBackend) handleRateLimitQuotasUpdate() framework.OperationFunc {
 
 		switch {
 		case quota == nil:
-			quota = quotas.NewRateLimitQuota(name, ns.Path, mountPath, rate, interval, blockInterval)
+			quota = quotas.NewRateLimitQuota(name, ns.Path, mountPath, pathSuffix, rate, interval, blockInterval)
 		default:
 			// Re-inserting the already indexed object in memdb might cause problems.
 			// So, clone the object. See https://github.com/hashicorp/go-memdb/issues/76.
@@ -218,6 +229,7 @@ func (b *SystemBackend) handleRateLimitQuotasUpdate() framework.OperationFunc {
 			rlq := clonedQuota.(*quotas.RateLimitQuota)
 			rlq.NamespacePath = ns.Path
 			rlq.MountPath = mountPath
+			rlq.PathSuffix = pathSuffix
 			rlq.Rate = rate
 			rlq.Interval = interval
 			rlq.BlockInterval = blockInterval
@@ -264,7 +276,7 @@ func (b *SystemBackend) handleRateLimitQuotasRead() framework.OperationFunc {
 		data := map[string]interface{}{
 			"type":           qType,
 			"name":           rlq.Name,
-			"path":           nsPath + rlq.MountPath,
+			"path":           nsPath + rlq.MountPath + rlq.PathSuffix,
 			"rate":           rlq.Rate,
 			"interval":       int(rlq.Interval.Seconds()),
 			"block_interval": int(rlq.BlockInterval.Seconds()),

vault/quotas/quotas.go

@@ -81,10 +81,11 @@ func (q Type) String() string {
 }
 
 const (
-	indexID             = "id"
-	indexName           = "name"
-	indexNamespace      = "ns"
-	indexNamespaceMount = "ns_mount"
+	indexID                 = "id"
+	indexName               = "name"
+	indexNamespace          = "ns"
+	indexNamespaceMount     = "ns_mount"
+	indexNamespaceMountPath = "ns_mount_path"
 )
 
 const (
@@ -391,7 +392,7 @@ func (m *Manager) QuotaByName(qType string, name string) (Quota, error) {
 }
 
 // QuotaByFactors returns the quota rule that matches the provided factors
-func (m *Manager) QuotaByFactors(ctx context.Context, qType, nsPath, mountPath string) (Quota, error) {
+func (m *Manager) QuotaByFactors(ctx context.Context, qType, nsPath, mountPath, pathSuffix string) (Quota, error) {
 	m.lock.RLock()
 	defer m.lock.RUnlock()
 
@@ -402,10 +403,15 @@ func (m *Manager) QuotaByFactors(ctx context.Context, qType, nsPath, mountPath s
 	}
 
 	idx := indexNamespace
-	args := []interface{}{nsPath, false}
+	args := []interface{}{nsPath, false, false}
 	if mountPath != "" {
-		idx = indexNamespaceMount
-		args = []interface{}{nsPath, mountPath}
+		if pathSuffix != "" {
+			idx = indexNamespaceMountPath
+			args = []interface{}{nsPath, mountPath, pathSuffix}
+		} else {
+			idx = indexNamespaceMount
+			args = []interface{}{nsPath, mountPath, false}
+		}
 	}
 
 	txn := m.db.Txn(false)
@@ -443,6 +449,7 @@ func (m *Manager) QueryQuota(req *Request) (Quota, error) {
 // Priority rules are as follows:
 // - namespace specific quota takes precedence over global quota
 // - mount specific quota takes precedence over namespace specific quota
+// - path suffix specific quota takes precedence over mount specific quota
 func (m *Manager) queryQuota(txn *memdb.Txn, req *Request) (Quota, error) {
 	if txn == nil {
 		txn = m.db.Txn(false)
@@ -478,8 +485,18 @@ func (m *Manager) queryQuota(txn *memdb.Txn, req *Request) (Quota, error) {
 		return quotas[0], nil
 	}
 
+	// Fetch path suffix quota
+	pathSuffix := strings.TrimSuffix(strings.TrimPrefix(strings.TrimPrefix(req.Path, req.NamespacePath), req.MountPath), "/")
+	quota, err := quotaFetchFunc(indexNamespaceMountPath, req.NamespacePath, req.MountPath, pathSuffix)
+	if err != nil {
+		return nil, err
+	}
+	if quota != nil {
+		return quota, nil
+	}
+
 	// Fetch mount quota
-	quota, err := quotaFetchFunc(indexNamespaceMount, req.NamespacePath, req.MountPath)
+	quota, err = quotaFetchFunc(indexNamespaceMount, req.NamespacePath, req.MountPath, false)
 	if err != nil {
 		return nil, err
 	}
@@ -488,7 +505,7 @@ func (m *Manager) queryQuota(txn *memdb.Txn, req *Request) (Quota, error) {
 	}
 
 	// Fetch ns quota. If NamespacePath is root, this will return the global quota.
-	quota, err = quotaFetchFunc(indexNamespace, req.NamespacePath, false)
+	quota, err = quotaFetchFunc(indexNamespace, req.NamespacePath, false, false)
 	if err != nil {
 		return nil, err
 	}
@@ -505,7 +522,7 @@ func (m *Manager) queryQuota(txn *memdb.Txn, req *Request) (Quota, error) {
 	}
 
 	// Fetch global quota
-	quota, err = quotaFetchFunc(indexNamespace, "root", false)
+	quota, err = quotaFetchFunc(indexNamespace, "root", false, false)
 	if err != nil {
 		return nil, err
 	}
@@ -731,6 +748,11 @@ func dbSchema() *memdb.DBSchema {
 							&memdb.FieldSetIndex{
 								Field: "MountPath",
 							},
+							// By sending false as the query parameter, we can
+							// query just the namespace specific quota.
+							&memdb.FieldSetIndex{
+								Field: "PathSuffix",
+							},
 						},
 					},
 				},
@@ -745,6 +767,28 @@ func dbSchema() *memdb.DBSchema {
 							&memdb.StringFieldIndex{
 								Field: "MountPath",
 							},
+							// By sending false as the query parameter, we can
+							// query just the namespace specific quota.
+							&memdb.FieldSetIndex{
+								Field: "PathSuffix",
+							},
+						},
+					},
+				},
+				indexNamespaceMountPath: {
+					Name:         indexNamespaceMountPath,
+					AllowMissing: true,
+					Indexer: &memdb.CompoundMultiIndex{
+						Indexes: []memdb.Indexer{
+							&memdb.StringFieldIndex{
+								Field: "NamespacePath",
+							},
+							&memdb.StringFieldIndex{
+								Field: "MountPath",
+							},
+							&memdb.StringFieldIndex{
+								Field: "PathSuffix",
+							},
 						},
 					},
 				},
@@ -973,35 +1017,49 @@ func (m *Manager) HandleRemount(ctx context.Context, from, to namespace.MountPat
 		toNs = namespace.RootNamespaceID
 	}
 
-	idx := indexNamespaceMount
 	leaseQuotaUpdated := false
-	args := []interface{}{fromNs, from.MountPath}
-	for _, quotaType := range quotaTypes() {
-		iter, err := txn.Get(quotaType, idx, args...)
-		if err != nil {
-			return err
-		}
-		for raw := iter.Next(); raw != nil; raw = iter.Next() {
-			quota := raw.(Quota)
 
-			// Clone the object and update it
-			clonedQuota := quota.Clone()
-			clonedQuota.handleRemount(to.MountPath, toNs)
-			// Update both underlying storage and memdb with the quota change
-			entry, err := logical.StorageEntryJSON(QuotaStoragePath(quotaType, quota.QuotaName()), quota)
+	updateMounts := func(idx string, args ...interface{}) error {
+		for _, quotaType := range quotaTypes() {
+			iter, err := txn.Get(quotaType, idx, args...)
 			if err != nil {
 				return err
 			}
-			if err := m.storage.Put(ctx, entry); err != nil {
-				return err
-			}
-			if err := m.setQuotaLockedWithTxn(ctx, quotaType, clonedQuota, false, txn); err != nil {
-				return err
-			}
-			if quotaType == TypeLeaseCount.String() {
-				leaseQuotaUpdated = true
+			for raw := iter.Next(); raw != nil; raw = iter.Next() {
+				quota := raw.(Quota)
+
+				// Clone the object and update it
+				clonedQuota := quota.Clone()
+				clonedQuota.handleRemount(to.MountPath, toNs)
+				// Update both underlying storage and memdb with the quota change
+				entry, err := logical.StorageEntryJSON(QuotaStoragePath(quotaType, quota.QuotaName()), quota)
+				if err != nil {
+					return err
+				}
+				if err := m.storage.Put(ctx, entry); err != nil {
+					return err
+				}
+				if err := m.setQuotaLockedWithTxn(ctx, quotaType, clonedQuota, false, txn); err != nil {
+					return err
+				}
+				if quotaType == TypeLeaseCount.String() {
+					leaseQuotaUpdated = true
+				}
 			}
 		}
+		return nil
+	}
+
+	// Update mounts for everything without a path prefix
+	err := updateMounts(indexNamespaceMount, fromNs, from.MountPath, false)
+	if err != nil {
+		return err
+	}
+
+	// Update mounts for everything with a path prefix
+	err = updateMounts(indexNamespaceMount, fromNs, from.MountPath, true)
+	if err != nil {
+		return err
 	}
 
 	if leaseQuotaUpdated {
@@ -1031,26 +1089,40 @@ func (m *Manager) HandleBackendDisabling(ctx context.Context, nsPath, mountPath
 		nsPath = "root"
 	}
 
-	idx := indexNamespaceMount
 	leaseQuotaDeleted := false
-	args := []interface{}{nsPath, mountPath}
-	for _, quotaType := range quotaTypes() {
-		iter, err := txn.Get(quotaType, idx, args...)
-		if err != nil {
-			return err
-		}
-		for raw := iter.Next(); raw != nil; raw = iter.Next() {
-			if err := txn.Delete(quotaType, raw); err != nil {
-				return fmt.Errorf("failed to delete quota from db after mount disabling; namespace %q, err %v", nsPath, err)
+
+	updateMounts := func(idx string, args ...interface{}) error {
+		for _, quotaType := range quotaTypes() {
+			iter, err := txn.Get(quotaType, idx, args...)
+			if err != nil {
+				return err
 			}
-			quota := raw.(Quota)
-			if err := m.storage.Delete(ctx, QuotaStoragePath(quotaType, quota.QuotaName())); err != nil {
-				return fmt.Errorf("failed to delete quota from storage after mount disabling; namespace %q, err %v", nsPath, err)
-			}
-			if quotaType == TypeLeaseCount.String() {
-				leaseQuotaDeleted = true
+			for raw := iter.Next(); raw != nil; raw = iter.Next() {
+				if err := txn.Delete(quotaType, raw); err != nil {
+					return fmt.Errorf("failed to delete quota from db after mount disabling; namespace %q, err %v", nsPath, err)
+				}
+				quota := raw.(Quota)
+				if err := m.storage.Delete(ctx, QuotaStoragePath(quotaType, quota.QuotaName())); err != nil {
+					return fmt.Errorf("failed to delete quota from storage after mount disabling; namespace %q, err %v", nsPath, err)
+				}
+				if quotaType == TypeLeaseCount.String() {
+					leaseQuotaDeleted = true
+				}
 			}
 		}
+		return nil
+	}
+
+	// Update mounts for everything without a path prefix
+	err := updateMounts(indexNamespaceMount, nsPath, mountPath, false)
+	if err != nil {
+		return err
+	}
+
+	// Update mounts for everything with a path prefix
+	err = updateMounts(indexNamespaceMount, nsPath, mountPath, true)
+	if err != nil {
+		return err
 	}
 
 	if leaseQuotaDeleted {

vault/quotas/quotas_rate_limit.go

@@ -54,6 +54,9 @@ type RateLimitQuota struct {
 	// MountPath is the path of the mount to which this quota is applicable
 	MountPath string `json:"mount_path"`
 
+	// PathSuffix is the path suffix to which this quota is applicable
+	PathSuffix string `json:"path_suffix"`
+
 	// Rate defines the number of requests allowed per Interval.
 	Rate float64 `json:"rate"`
 
@@ -81,7 +84,7 @@ type RateLimitQuota struct {
 // provided, which will default to 1s when initialized. An optional block
 // duration may be provided, where if set, when a client reaches the rate limit,
 // subsequent requests will fail until the block duration has passed.
-func NewRateLimitQuota(name, nsPath, mountPath string, rate float64, interval, block time.Duration) *RateLimitQuota {
+func NewRateLimitQuota(name, nsPath, mountPath, pathSuffix string, rate float64, interval, block time.Duration) *RateLimitQuota {
 	id, err := uuid.GenerateUUID()
 	if err != nil {
 		// Fall back to generating with a hash of the name, later in initialize
@@ -93,6 +96,7 @@ func NewRateLimitQuota(name, nsPath, mountPath string, rate float64, interval, b
 		Type:          TypeRateLimit,
 		NamespacePath: nsPath,
 		MountPath:     mountPath,
+		PathSuffix:    pathSuffix,
 		Rate:          rate,
 		Interval:      interval,
 		BlockInterval: block,
@@ -108,6 +112,7 @@ func (q *RateLimitQuota) Clone() Quota {
 		MountPath:     q.MountPath,
 		Type:          q.Type,
 		NamespacePath: q.NamespacePath,
+		PathSuffix:    q.PathSuffix,
 		BlockInterval: q.BlockInterval,
 		Rate:          q.Rate,
 		Interval:      q.Interval,

vault/quotas/quotas_rate_limit_test.go

@@ -27,7 +27,7 @@ func TestNewRateLimitQuota(t *testing.T) {
 		rlq       *RateLimitQuota
 		expectErr bool
 	}{
-		{"valid rate", NewRateLimitQuota("test-rate-limiter", "qa", "/foo/bar", 16.7, time.Second, 0), false},
+		{"valid rate", NewRateLimitQuota("test-rate-limiter", "qa", "/foo/bar", "", 16.7, time.Second, 0), false},
 	}
 
 	for _, tc := range testCases {
@@ -44,7 +44,7 @@ func TestNewRateLimitQuota(t *testing.T) {
 }
 
 func TestRateLimitQuota_Close(t *testing.T) {
-	rlq := NewRateLimitQuota("test-rate-limiter", "qa", "/foo/bar", 16.7, time.Second, time.Minute)
+	rlq := NewRateLimitQuota("test-rate-limiter", "qa", "/foo/bar", "", 16.7, time.Second, time.Minute)
 	require.NoError(t, rlq.initialize(logging.NewVaultLogger(log.Trace), metricsutil.BlackholeSink()))
 	require.NoError(t, rlq.close(context.Background()))
 
@@ -218,7 +218,7 @@ func TestRateLimitQuota_Update(t *testing.T) {
 	qm, err := NewManager(logging.NewVaultLogger(log.Trace), nil, metricsutil.BlackholeSink())
 	require.NoError(t, err)
 
-	quota := NewRateLimitQuota("quota1", "", "", 10, time.Second, 0)
+	quota := NewRateLimitQuota("quota1", "", "", "", 10, time.Second, 0)
 	require.NoError(t, qm.SetQuota(context.Background(), TypeRateLimit.String(), quota, true))
 	require.NoError(t, qm.SetQuota(context.Background(), TypeRateLimit.String(), quota, true))
 

vault/quotas/quotas_test.go

@@ -16,7 +16,7 @@ func TestQuotas_MountPathOverwrite(t *testing.T) {
 	qm, err := NewManager(logging.NewVaultLogger(log.Trace), nil, metricsutil.BlackholeSink())
 	require.NoError(t, err)
 
-	quota := NewRateLimitQuota("tq", "", "kv1/", 10, time.Second, 0)
+	quota := NewRateLimitQuota("tq", "", "kv1/", "", 10, time.Second, 0)
 	require.NoError(t, qm.SetQuota(context.Background(), TypeRateLimit.String(), quota, false))
 	quota = quota.Clone().(*RateLimitQuota)
 	quota.MountPath = "kv2/"
@@ -43,19 +43,20 @@ func TestQuotas_Precedence(t *testing.T) {
 	qm, err := NewManager(logging.NewVaultLogger(log.Trace), nil, metricsutil.BlackholeSink())
 	require.NoError(t, err)
 
-	setQuotaFunc := func(t *testing.T, name, nsPath, mountPath string) Quota {
+	setQuotaFunc := func(t *testing.T, name, nsPath, mountPath, pathSuffix string) Quota {
 		t.Helper()
-		quota := NewRateLimitQuota(name, nsPath, mountPath, 10, time.Second, 0)
+		quota := NewRateLimitQuota(name, nsPath, mountPath, pathSuffix, 10, time.Second, 0)
 		require.NoError(t, qm.SetQuota(context.Background(), TypeRateLimit.String(), quota, true))
 		return quota
 	}
 
-	checkQuotaFunc := func(t *testing.T, nsPath, mountPath string, expected Quota) {
+	checkQuotaFunc := func(t *testing.T, nsPath, mountPath, pathSuffix string, expected Quota) {
 		t.Helper()
 		quota, err := qm.QueryQuota(&Request{
 			Type:          TypeRateLimit,
 			NamespacePath: nsPath,
 			MountPath:     mountPath,
+			Path:          nsPath + mountPath + pathSuffix,
 		})
 		require.NoError(t, err)
 
@@ -65,26 +66,34 @@ func TestQuotas_Precedence(t *testing.T) {
 	}
 
 	// No quota present. Expect nil.
-	checkQuotaFunc(t, "", "", nil)
+	checkQuotaFunc(t, "", "", "", nil)
 
 	// Define global quota and expect that to be returned.
-	rateLimitGlobalQuota := setQuotaFunc(t, "rateLimitGlobalQuota", "", "")
-	checkQuotaFunc(t, "", "", rateLimitGlobalQuota)
+	rateLimitGlobalQuota := setQuotaFunc(t, "rateLimitGlobalQuota", "", "", "")
+	checkQuotaFunc(t, "", "", "", rateLimitGlobalQuota)
 
 	// Define a global mount specific quota and expect that to be returned.
-	rateLimitGlobalMountQuota := setQuotaFunc(t, "rateLimitGlobalMountQuota", "", "testmount")
-	checkQuotaFunc(t, "", "testmount", rateLimitGlobalMountQuota)
+	rateLimitGlobalMountQuota := setQuotaFunc(t, "rateLimitGlobalMountQuota", "", "testmount/", "")
+	checkQuotaFunc(t, "", "testmount/", "", rateLimitGlobalMountQuota)
+
+	// Define a global mount + path specific quota and expect that to be returned.
+	rateLimitGlobalMountPathQuota := setQuotaFunc(t, "rateLimitGlobalMountPathQuota", "", "testmount/", "testpath")
+	checkQuotaFunc(t, "", "testmount/", "testpath", rateLimitGlobalMountPathQuota)
 
 	// Define a namespace quota and expect that to be returned.
-	rateLimitNSQuota := setQuotaFunc(t, "rateLimitNSQuota", "testns", "")
-	checkQuotaFunc(t, "testns", "", rateLimitNSQuota)
+	rateLimitNSQuota := setQuotaFunc(t, "rateLimitNSQuota", "testns/", "", "")
+	checkQuotaFunc(t, "testns/", "", "", rateLimitNSQuota)
 
 	// Define a namespace mount specific quota and expect that to be returned.
-	rateLimitNSMountQuota := setQuotaFunc(t, "rateLimitNSMountQuota", "testns", "testmount")
-	checkQuotaFunc(t, "testns", "testmount", rateLimitNSMountQuota)
+	rateLimitNSMountQuota := setQuotaFunc(t, "rateLimitNSMountQuota", "testns/", "testmount/", "")
+	checkQuotaFunc(t, "testns/", "testmount/", "", rateLimitNSMountQuota)
+
+	// Define a namespace mount + path specific quota and expect that to be returned.
+	rateLimitNSMountPathQuota := setQuotaFunc(t, "rateLimitNSMountPathQuota", "testns/", "testmount/", "testpath")
+	checkQuotaFunc(t, "testns/", "testmount/", "testpath", rateLimitNSMountPathQuota)
 
 	// Now that many quota types are defined, verify that the most specific
 	// matches are returned per namespace.
-	checkQuotaFunc(t, "", "", rateLimitGlobalQuota)
-	checkQuotaFunc(t, "testns", "", rateLimitNSQuota)
+	checkQuotaFunc(t, "", "", "", rateLimitGlobalQuota)
+	checkQuotaFunc(t, "testns/", "", "", rateLimitNSQuota)
 }