VAULT-6613 Add role support for rate limit quotas (OSS Changes)

changelog/16115.txt

@@ -0,0 +1,3 @@
+```release-note:improvement
+core/quotas: Added ability to add role information for rate-limit resource quotas, to limit login requests on auth mounts made using that role
+```

http/util.go

@@ -1,7 +1,10 @@
 package http
 
 import (
+	"bytes"
+	"errors"
 	"fmt"
+	"io/ioutil"
 	"net"
 	"net/http"
 	"strings"
@@ -47,11 +50,21 @@ func rateLimitQuotaWrapping(handler http.Handler, core *vault.Core) http.Handler
 			respondError(w, status, err)
 			return
 		}
+		mountPath := strings.TrimPrefix(core.MatchingMount(r.Context(), path), ns.Path)
+
+		// Clone body, so we do not close the request body reader
+		bodyBytes, err := ioutil.ReadAll(r.Body)
+		if err != nil {
+			respondError(w, http.StatusInternalServerError, errors.New("failed to read request body"))
+			return
+		}
+		r.Body = ioutil.NopCloser(bytes.NewBuffer(bodyBytes))
 
 		quotaResp, err := core.ApplyRateLimitQuota(r.Context(), &quotas.Request{
 			Type:          quotas.TypeRateLimit,
 			Path:          path,
-			MountPath:     strings.TrimPrefix(core.MatchingMount(r.Context(), path), ns.Path),
+			MountPath:     mountPath,
+			Role:          core.DetermineRoleFromLoginRequest(mountPath, bodyBytes, r.Context()),
 			NamespacePath: ns.Path,
 			ClientAddress: parseRemoteIPAddress(r),
 		})

vault/core.go

@@ -3300,3 +3300,32 @@ func (c *Core) CheckPluginPerms(pluginName string) (err error) {
 	}
 	return err
 }
+
+// DetermineRoleFromLoginRequest will determine the role that should be applied to a quota for a given
+// login request
+func (c *Core) DetermineRoleFromLoginRequest(mountPoint string, payload []byte, ctx context.Context) string {
+	matchingBackend := c.router.MatchingBackend(ctx, mountPoint)
+	if matchingBackend == nil || matchingBackend.Type() != logical.TypeCredential {
+		// Role based quotas do not apply to this request
+		return ""
+	}
+
+	data := make(map[string]interface{})
+	err := jsonutil.DecodeJSON(payload, &data)
+	if err != nil {
+		// Cannot discern a role from a request we cannot parse
+		return ""
+	}
+
+	resp, err := matchingBackend.HandleRequest(ctx, &logical.Request{
+		MountPoint: mountPoint,
+		Path:       "login",
+		Operation:  logical.ResolveRoleOperation,
+		Data:       data,
+		Storage:    c.router.MatchingStorageByAPIPath(ctx, mountPoint+"login"),
+	})
+	if err != nil || resp.Data["role"] == nil {
+		return ""
+	}
+	return resp.Data["role"].(string)
+}

vault/logical_system_quotas.go

@@ -202,15 +202,6 @@ func (b *SystemBackend) handleRateLimitQuotasUpdate() framework.OperationFunc {
 			pathSuffix = strings.TrimSuffix(strings.TrimPrefix(mountPath, mountAPIPath), "/")
 			mountPath = mountAPIPath
 		}
-		// Disallow creation of new quota that has properties similar to an
-		// existing quota.
-		quotaByFactors, err := b.Core.quotaManager.QuotaByFactors(ctx, qType, ns.Path, mountPath, pathSuffix)
-		if err != nil {
-			return nil, err
-		}
-		if quotaByFactors != nil && quotaByFactors.QuotaName() != name {
-			return logical.ErrorResponse("quota rule with similar properties exists under the name %q", quotaByFactors.QuotaName()), nil
-		}
 
 		role := d.Get("role").(string)
 		// If this is a quota with a role, ensure the backend supports role resolution
@@ -232,6 +223,16 @@ func (b *SystemBackend) handleRateLimitQuotasUpdate() framework.OperationFunc {
 			}
 		}
 
+		// Disallow creation of new quota that has properties similar to an
+		// existing quota.
+		quotaByFactors, err := b.Core.quotaManager.QuotaByFactors(ctx, qType, ns.Path, mountPath, pathSuffix, role)
+		if err != nil {
+			return nil, err
+		}
+		if quotaByFactors != nil && quotaByFactors.QuotaName() != name {
+			return logical.ErrorResponse("quota rule with similar properties exists under the name %q", quotaByFactors.QuotaName()), nil
+		}
+
 		// If a quota already exists, fetch and update it.
 		quota, err := b.Core.quotaManager.QuotaByName(qType, name)
 		if err != nil {
@@ -240,7 +241,7 @@ func (b *SystemBackend) handleRateLimitQuotasUpdate() framework.OperationFunc {
 
 		switch {
 		case quota == nil:
-			quota = quotas.NewRateLimitQuota(name, ns.Path, mountPath, pathSuffix, rate, interval, blockInterval)
+			quota = quotas.NewRateLimitQuota(name, ns.Path, mountPath, pathSuffix, role, 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.
@@ -296,6 +297,7 @@ func (b *SystemBackend) handleRateLimitQuotasRead() framework.OperationFunc {
 			"type":           qType,
 			"name":           rlq.Name,
 			"path":           nsPath + rlq.MountPath + rlq.PathSuffix,
+			"role":           rlq.Role,
 			"rate":           rlq.Rate,
 			"interval":       int(rlq.Interval.Seconds()),
 			"block_interval": int(rlq.BlockInterval.Seconds()),

vault/quotas/quotas.go

@@ -86,6 +86,7 @@ const (
 	indexNamespace          = "ns"
 	indexNamespaceMount     = "ns_mount"
 	indexNamespaceMountPath = "ns_mount_path"
+	indexNamespaceMountRole = "ns_mount_role"
 )
 
 const (
@@ -233,6 +234,9 @@ type Request struct {
 	// Path is the request path to which quota rules are being queried for
 	Path string
 
+	// Role is the role given as part of the request to a login endpoint
+	Role string
+
 	// NamespacePath is the namespace path to which the request belongs
 	NamespacePath string
 
@@ -392,7 +396,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, pathSuffix string) (Quota, error) {
+func (m *Manager) QuotaByFactors(ctx context.Context, qType, nsPath, mountPath, pathSuffix, role string) (Quota, error) {
 	m.lock.RLock()
 	defer m.lock.RUnlock()
 
@@ -403,14 +407,17 @@ func (m *Manager) QuotaByFactors(ctx context.Context, qType, nsPath, mountPath,
 	}
 
 	idx := indexNamespace
-	args := []interface{}{nsPath, false, false}
+	args := []interface{}{nsPath, false, false, false}
 	if mountPath != "" {
 		if pathSuffix != "" {
 			idx = indexNamespaceMountPath
-			args = []interface{}{nsPath, mountPath, pathSuffix}
+			args = []interface{}{nsPath, mountPath, pathSuffix, false}
+		} else if role != "" {
+			idx = indexNamespaceMountRole
+			args = []interface{}{nsPath, mountPath, false, role}
 		} else {
 			idx = indexNamespaceMount
-			args = []interface{}{nsPath, mountPath, false}
+			args = []interface{}{nsPath, mountPath, false, false}
 		}
 	}
 
@@ -450,6 +457,7 @@ func (m *Manager) QueryQuota(req *Request) (Quota, error) {
 // - 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
+// - role based quota takes precedence over path suffix/mount specific quota
 func (m *Manager) queryQuota(txn *memdb.Txn, req *Request) (Quota, error) {
 	if txn == nil {
 		txn = m.db.Txn(false)
@@ -485,9 +493,18 @@ func (m *Manager) queryQuota(txn *memdb.Txn, req *Request) (Quota, error) {
 		return quotas[0], nil
 	}
 
+	// Fetch role suffix quota
+	quota, err := quotaFetchFunc(indexNamespaceMountRole, req.NamespacePath, req.MountPath, false, req.Role)
+	if err != nil {
+		return nil, err
+	}
+	if quota != nil {
+		return quota, 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)
+	quota, err = quotaFetchFunc(indexNamespaceMountPath, req.NamespacePath, req.MountPath, pathSuffix, false)
 	if err != nil {
 		return nil, err
 	}
@@ -496,7 +513,7 @@ func (m *Manager) queryQuota(txn *memdb.Txn, req *Request) (Quota, error) {
 	}
 
 	// Fetch mount quota
-	quota, err = quotaFetchFunc(indexNamespaceMount, req.NamespacePath, req.MountPath, false)
+	quota, err = quotaFetchFunc(indexNamespaceMount, req.NamespacePath, req.MountPath, false, false)
 	if err != nil {
 		return nil, err
 	}
@@ -505,7 +522,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, false)
+	quota, err = quotaFetchFunc(indexNamespace, req.NamespacePath, false, false, false)
 	if err != nil {
 		return nil, err
 	}
@@ -522,7 +539,7 @@ func (m *Manager) queryQuota(txn *memdb.Txn, req *Request) (Quota, error) {
 	}
 
 	// Fetch global quota
-	quota, err = quotaFetchFunc(indexNamespace, "root", false, false)
+	quota, err = quotaFetchFunc(indexNamespace, "root", false, false, false)
 	if err != nil {
 		return nil, err
 	}
@@ -753,6 +770,11 @@ func dbSchema() *memdb.DBSchema {
 							&memdb.FieldSetIndex{
 								Field: "PathSuffix",
 							},
+							// By sending false as the query parameter, we can
+							// query just the namespace specific quota.
+							&memdb.FieldSetIndex{
+								Field: "Role",
+							},
 						},
 					},
 				},
@@ -772,6 +794,33 @@ func dbSchema() *memdb.DBSchema {
 							&memdb.FieldSetIndex{
 								Field: "PathSuffix",
 							},
+							// By sending false as the query parameter, we can
+							// query just the namespace specific quota.
+							&memdb.FieldSetIndex{
+								Field: "Role",
+							},
+						},
+					},
+				},
+				indexNamespaceMountRole: {
+					Name:         indexNamespaceMountRole,
+					AllowMissing: true,
+					Indexer: &memdb.CompoundMultiIndex{
+						Indexes: []memdb.Indexer{
+							&memdb.StringFieldIndex{
+								Field: "NamespacePath",
+							},
+							&memdb.StringFieldIndex{
+								Field: "MountPath",
+							},
+							// By sending false as the query parameter, we can
+							// query just the role specific quota.
+							&memdb.FieldSetIndex{
+								Field: "PathSuffix",
+							},
+							&memdb.StringFieldIndex{
+								Field: "Role",
+							},
 						},
 					},
 				},
@@ -789,6 +838,11 @@ func dbSchema() *memdb.DBSchema {
 							&memdb.StringFieldIndex{
 								Field: "PathSuffix",
 							},
+							// By sending false as the query parameter, we can
+							// query just the namespace specific quota.
+							&memdb.FieldSetIndex{
+								Field: "Role",
+							},
 						},
 					},
 				},
@@ -1050,14 +1104,20 @@ func (m *Manager) HandleRemount(ctx context.Context, from, to namespace.MountPat
 		return nil
 	}
 
-	// Update mounts for everything without a path prefix
-	err := updateMounts(indexNamespaceMount, fromNs, from.MountPath, false)
+	// Update mounts for everything without a path prefix or role
+	err := updateMounts(indexNamespaceMount, fromNs, from.MountPath, false, false)
 	if err != nil {
 		return err
 	}
 
 	// Update mounts for everything with a path prefix
-	err = updateMounts(indexNamespaceMount, fromNs, from.MountPath, true)
+	err = updateMounts(indexNamespaceMount, fromNs, from.MountPath, true, false)
+	if err != nil {
+		return err
+	}
+
+	// Update mounts for everything with a role
+	err = updateMounts(indexNamespaceMount, fromNs, from.MountPath, false, true)
 	if err != nil {
 		return err
 	}
@@ -1113,14 +1173,20 @@ func (m *Manager) HandleBackendDisabling(ctx context.Context, nsPath, mountPath
 		return nil
 	}
 
-	// Update mounts for everything without a path prefix
-	err := updateMounts(indexNamespaceMount, nsPath, mountPath, false)
+	// Update mounts for everything without a path prefix or role
+	err := updateMounts(indexNamespaceMount, nsPath, mountPath, false, false)
 	if err != nil {
 		return err
 	}
 
 	// Update mounts for everything with a path prefix
-	err = updateMounts(indexNamespaceMount, nsPath, mountPath, true)
+	err = updateMounts(indexNamespaceMount, nsPath, mountPath, true, false)
+	if err != nil {
+		return err
+	}
+
+	// Update mounts for everything with a role
+	err = updateMounts(indexNamespaceMount, nsPath, mountPath, false, true)
 	if err != nil {
 		return err
 	}

vault/quotas/quotas_rate_limit.go

@@ -54,6 +54,10 @@ type RateLimitQuota struct {
 	// MountPath is the path of the mount to which this quota is applicable
 	MountPath string `json:"mount_path"`
 
+	// Role is the role on an auth mount to apply the quota to upon /login requests
+	// Not applicable for use with path suffixes
+	Role string `json:"role"`
+
 	// PathSuffix is the path suffix to which this quota is applicable
 	PathSuffix string `json:"path_suffix"`
 
@@ -84,7 +88,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, pathSuffix string, rate float64, interval, block time.Duration) *RateLimitQuota {
+func NewRateLimitQuota(name, nsPath, mountPath, pathSuffix, role 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
@@ -96,6 +100,7 @@ func NewRateLimitQuota(name, nsPath, mountPath, pathSuffix string, rate float64,
 		Type:          TypeRateLimit,
 		NamespacePath: nsPath,
 		MountPath:     mountPath,
+		Role:          role,
 		PathSuffix:    pathSuffix,
 		Rate:          rate,
 		Interval:      interval,
@@ -110,6 +115,7 @@ func (q *RateLimitQuota) Clone() Quota {
 		ID:            q.ID,
 		Name:          q.Name,
 		MountPath:     q.MountPath,
+		Role:          q.Role,
 		Type:          q.Type,
 		NamespacePath: q.NamespacePath,
 		PathSuffix:    q.PathSuffix,

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))