Only resolve roles for role quotas and leases (#22597)

changelog/22597.txt

@@ -0,0 +1,3 @@
+```release-note:bug
+core/quotas: Only perform ResolveRoleOperation for role-based quotas and lease creation.
+```

http/util.go

@@ -60,22 +60,33 @@ func rateLimitQuotaWrapping(handler http.Handler, core *vault.Core) http.Handler
 		}
 		r.Body = ioutil.NopCloser(bytes.NewBuffer(bodyBytes))
 
-		role := core.DetermineRoleFromLoginRequestFromBytes(mountPath, bodyBytes, r.Context())
-
-		// add an entry to the context to prevent recalculating request role unnecessarily
-		r = r.WithContext(context.WithValue(r.Context(), logical.CtxKeyRequestRole{}, role))
-
-		quotaResp, err := core.ApplyRateLimitQuota(r.Context(), &quotas.Request{
+		quotaReq := &quotas.Request{
 			Type:          quotas.TypeRateLimit,
 			Path:          path,
 			MountPath:     mountPath,
-			Role:          role,
 			NamespacePath: ns.Path,
 			ClientAddress: parseRemoteIPAddress(r),
-		})
+		}
+		requiresResolveRole, err := core.ResolveRoleForQuotas(r.Context(), quotaReq)
+		if err != nil {
+			core.Logger().Error("failed to lookup quotas", "path", path, "error", err)
+			respondError(w, http.StatusInternalServerError, err)
+			return
+		}
+
+		// If any role-based quotas are enabled for this namespace/mount, just
+		// do the role resolution once here.
+		if requiresResolveRole {
+			role := core.DetermineRoleFromLoginRequestFromBytes(r.Context(), mountPath, bodyBytes)
+			// add an entry to the context to prevent recalculating request role unnecessarily
+			r = r.WithContext(context.WithValue(r.Context(), logical.CtxKeyRequestRole{}, role))
+			quotaReq.Role = role
+		}
+
+		quotaResp, err := core.ApplyRateLimitQuota(r.Context(), quotaReq)
 		if err != nil {
 			core.Logger().Error("failed to apply quota", "path", path, "error", err)
-			respondError(w, http.StatusUnprocessableEntity, err)
+			respondError(w, http.StatusInternalServerError, err)
 			return
 		}
 

vault/core.go

@@ -3939,20 +3939,20 @@ func (c *Core) LoadNodeID() (string, error) {
 
 // DetermineRoleFromLoginRequestFromBytes will determine the role that should be applied to a quota for a given
 // login request, accepting a byte payload
-func (c *Core) DetermineRoleFromLoginRequestFromBytes(mountPoint string, payload []byte, ctx context.Context) string {
+func (c *Core) DetermineRoleFromLoginRequestFromBytes(ctx context.Context, mountPoint string, payload []byte) string {
 	data := make(map[string]interface{})
 	err := jsonutil.DecodeJSON(payload, &data)
 	if err != nil {
 		// Cannot discern a role from a request we cannot parse
 		return ""
 	}
 
-	return c.DetermineRoleFromLoginRequest(mountPoint, data, ctx)
+	return c.DetermineRoleFromLoginRequest(ctx, mountPoint, data)
 }
 
 // DetermineRoleFromLoginRequest will determine the role that should be applied to a quota for a given
 // login request
-func (c *Core) DetermineRoleFromLoginRequest(mountPoint string, data map[string]interface{}, ctx context.Context) string {
+func (c *Core) DetermineRoleFromLoginRequest(ctx context.Context, mountPoint string, data map[string]interface{}) string {
 	c.authLock.RLock()
 	defer c.authLock.RUnlock()
 	matchingBackend := c.router.MatchingBackend(ctx, mountPoint)
@@ -3971,9 +3971,19 @@ func (c *Core) DetermineRoleFromLoginRequest(mountPoint string, data map[string]
 	if err != nil || resp.Data["role"] == nil {
 		return ""
 	}
+
 	return resp.Data["role"].(string)
 }
 
+// ResolveRoleForQuotas looks for any quotas requiring a role for early
+// computation in the RateLimitQuotaWrapping handler.
+func (c *Core) ResolveRoleForQuotas(ctx context.Context, req *quotas.Request) (bool, error) {
+	if c.quotaManager == nil {
+		return false, nil
+	}
+	return c.quotaManager.QueryResolveRoleQuotas(req)
+}
+
 // aliasNameFromLoginRequest will determine the aliasName from the login Request
 func (c *Core) aliasNameFromLoginRequest(ctx context.Context, req *logical.Request) (string, error) {
 	c.authLock.RLock()

vault/quotas/quotas.go

@@ -617,6 +617,37 @@ func (m *Manager) queryQuota(txn *memdb.Txn, req *Request) (Quota, error) {
 	return nil, nil
 }
 
+// QueryResolveRoleQuotas checks if there's a quota for the request mount path
+// which requires ResolveRoleOperation.
+func (m *Manager) QueryResolveRoleQuotas(req *Request) (bool, error) {
+	m.dbAndCacheLock.RLock()
+	defer m.dbAndCacheLock.RUnlock()
+
+	txn := m.db.Txn(false)
+
+	// ns would have been made non-empty during insertion. Use non-empty
+	// value during query as well.
+	if req.NamespacePath == "" {
+		req.NamespacePath = "root"
+	}
+
+	// Check for any role-based quotas on the request namespaces/mount path.
+	for _, qType := range quotaTypes() {
+		// Use the namespace and mount as indexes and find all matches with a
+		// set Role field (see: 'true' as the last argument). We can't use
+		// indexNamespaceMountRole for this, because Role is a StringFieldIndex,
+		// which won't match on an empty string.
+		quota, err := txn.First(qType, indexNamespaceMount, req.NamespacePath, req.MountPath, false, true)
+		if err != nil {
+			return false, err
+		}
+		if quota != nil {
+			return true, nil
+		}
+	}
+	return false, nil
+}
+
 // DeleteQuota removes a quota rule from the db for a given name
 func (m *Manager) DeleteQuota(ctx context.Context, qType string, name string) error {
 	m.quotaLock.Lock()

vault/quotas/quotas_test.go

@@ -135,3 +135,48 @@ func TestQuotas_Precedence(t *testing.T) {
 	checkQuotaFunc(t, "testns/nested2/nested3/", "", "", "", rateLimitMultiNestedNsInheritableQuota)
 	checkQuotaFunc(t, "testns/nested2/nested3/nested4/nested5", "", "", "", rateLimitMultiNestedNsInheritableQuota)
 }
+
+// TestQuotas_QueryRoleQuotas checks to see if quota creation on a mount
+// requires a call to ResolveRoleOperation.
+func TestQuotas_QueryResolveRole_RateLimitQuotas(t *testing.T) {
+	leaseWalkFunc := func(context.Context, func(request *Request) bool) error {
+		return nil
+	}
+	qm, err := NewManager(logging.NewVaultLogger(log.Trace), leaseWalkFunc, metricsutil.BlackholeSink())
+	require.NoError(t, err)
+
+	rlqReq := &Request{
+		Type:          TypeRateLimit,
+		Path:          "",
+		MountPath:     "mount1/",
+		NamespacePath: "",
+		ClientAddress: "127.0.0.1",
+	}
+	// Check that we have no quotas requiring role resolution on mount1/
+	required, err := qm.QueryResolveRoleQuotas(rlqReq)
+	require.NoError(t, err)
+	require.False(t, required)
+
+	// Create a non-role-based RLQ on mount1/ and make sure it doesn't require role resolution
+	rlq := NewRateLimitQuota("tq", rlqReq.NamespacePath, rlqReq.MountPath, rlqReq.Path, rlqReq.Role, false, 1*time.Minute, 10*time.Second, 10)
+	require.NoError(t, qm.SetQuota(context.Background(), TypeRateLimit.String(), rlq, false))
+
+	required, err = qm.QueryResolveRoleQuotas(rlqReq)
+	require.NoError(t, err)
+	require.False(t, required)
+
+	// Create a role-based RLQ on mount1/ and make sure it requires role resolution
+	rlqReq.Role = "test"
+	rlq = NewRateLimitQuota("tq", rlqReq.NamespacePath, rlqReq.MountPath, rlqReq.Path, rlqReq.Role, false, 1*time.Minute, 10*time.Second, 10)
+	require.NoError(t, qm.SetQuota(context.Background(), TypeRateLimit.String(), rlq, false))
+
+	required, err = qm.QueryResolveRoleQuotas(rlqReq)
+	require.NoError(t, err)
+	require.True(t, required)
+
+	// Check that we have no quotas requiring role resolution on mount2/
+	rlqReq.MountPath = "mount2/"
+	required, err = qm.QueryResolveRoleQuotas(rlqReq)
+	require.NoError(t, err)
+	require.False(t, required)
+}

vault/request_handling.go

@@ -1485,7 +1485,8 @@ func (c *Core) handleLoginRequest(ctx context.Context, req *logical.Request) (re
 
 		// Check for request role in context to role based quotas
 		var role string
-		if reqRole := ctx.Value(logical.CtxKeyRequestRole{}); reqRole != nil {
+		reqRole := ctx.Value(logical.CtxKeyRequestRole{})
+		if reqRole != nil {
 			role = reqRole.(string)
 		}
 
@@ -1686,6 +1687,13 @@ func (c *Core) handleLoginRequest(ctx context.Context, req *logical.Request) (re
 		// Attach the display name, might be used by audit backends
 		req.DisplayName = auth.DisplayName
 
+		// If this is not a role-based quota, we still need to associate the
+		// login role with this lease for later lease-count quotas to be
+		// accurate.
+		if reqRole == nil && resp.Auth.TokenType != logical.TokenTypeBatch {
+			role = c.DetermineRoleFromLoginRequest(ctx, req.MountPoint, req.Data)
+		}
+
 		leaseGen, respTokenCreate, errCreateToken := c.LoginCreateToken(ctx, ns, req.Path, source, role, resp)
 		leaseGenerated = leaseGen
 		if errCreateToken != nil {