block users

Author: tyler-french

enterprise/server/quota/BUILD

@@ -34,15 +34,19 @@ go_test(
     deps = [
         "//enterprise/server/backends/authdb",
         "//enterprise/server/backends/userdb",
+        "//enterprise/server/experiments",
         "//proto:quota_go_proto",
         "//server/environment",
         "//server/tables",
         "//server/testutil/pubsub",
+        "//server/testutil/testauth",
         "//server/testutil/testenv",
         "//server/util/db",
         "//server/util/query_builder",
         "@com_github_google_go_cmp//cmp",
         "@com_github_google_go_cmp//cmp/cmpopts",
+        "@com_github_open_feature_go_sdk//openfeature",
+        "@com_github_open_feature_go_sdk//openfeature/memprovider",
         "@com_github_stretchr_testify//assert",
         "@com_github_stretchr_testify//require",
         "@org_golang_google_protobuf//testing/protocmp",

enterprise/server/quota/quota_manager.go

@@ -29,9 +29,7 @@ import (
 	qpb "github.com/buildbuddy-io/buildbuddy/proto/quota"
 )
 
-var (
-	quotaManagerEnabled = flag.Bool("app.enable_quota_management", false, "If set, quota management will be enabled")
-)
+var quotaManagerEnabled = flag.Bool("app.enable_quota_management", false, "If set, quota management will be enabled")
 
 const (
 	// The maximum number of attempts to update rate limit data in redis.
@@ -48,6 +46,10 @@ const (
 	// when there is an update.
 	pubSubChannelName = "quota-change-notifications"
 
+	// The names of the flagd experiments for quota management.
+	blockedQuotaExperimentName = "quota.blocked"
+	bucketQuotaExperimentName  = "quota.buckets"
+
 	namespaceSeperator = ":"
 )
 
@@ -84,6 +86,61 @@ func bucketToRow(namespace string, from *qpb.Bucket) *tables.QuotaBucket {
 	return res
 }
 
+func bucketRowFromMap(namespace string, bucketMap map[string]interface{}) (*tables.QuotaBucket, error) {
+	name, ok := bucketMap["name"].(string)
+	if !ok || name == "" {
+		return nil, status.InvalidArgumentError("bucket.name must be a non-empty string")
+	}
+
+	maxRateInterface, ok := bucketMap["max_rate"]
+	if !ok {
+		return nil, status.InvalidArgumentError("bucket.max_rate is required")
+	}
+	maxRateMap, ok := maxRateInterface.(map[string]interface{})
+	if !ok {
+		return nil, status.InvalidArgumentError("bucket.max_rate must be an object")
+	}
+
+	numRequests, err := interfaceToInt64(maxRateMap["num_requests"])
+	if err != nil || numRequests <= 0 {
+		return nil, status.InvalidArgumentErrorf("bucket.max_rate.num_requests must be a positive number: %s", err)
+	}
+
+	period, err := interfaceToInt64(maxRateMap["period_usec"])
+	if err != nil {
+		return nil, status.InvalidArgumentErrorf("bucket.max_rate.period_usec is invalid: %s", err)
+	}
+	if period == 0 {
+		return nil, status.InvalidArgumentError("bucket.max_rate.period cannot be zero")
+	}
+
+	maxBurst, err := interfaceToInt64(bucketMap["max_burst"])
+	if err != nil {
+		return nil, status.InvalidArgumentErrorf("bucket.max_burst must be a number: %s", err)
+	}
+
+	return &tables.QuotaBucket{
+		Namespace:          namespace,
+		Name:               name,
+		NumRequests:        numRequests,
+		PeriodDurationUsec: period,
+		MaxBurst:           maxBurst,
+	}, nil
+}
+
+func interfaceToInt64(v interface{}) (int64, error) {
+	switch val := v.(type) {
+	case int64:
+		return val, nil
+	case int:
+		return int64(val), nil
+	case float64:
+		return int64(val), nil
+	default:
+		return 0, status.InvalidArgumentErrorf("expected number type, got %T", v)
+	}
+}
+
 func namespaceConfigToProto(from *namespaceConfig) *qpb.Namespace {
 	res := &qpb.Namespace{
 		Name: from.name,
@@ -135,46 +192,45 @@ func fetchConfigFromDB(env environment.Env, namespace string) (map[string]*names
 			namespace,
 		)
 	}
-	err := db.ScanEach(
-		rq,
-		func(ctx context.Context, qbg *struct {
-			tables.QuotaBucket
-			*tables.QuotaGroup
-		}) error {
-			qb := &qbg.QuotaBucket
-			ns, ok := config[qb.Namespace]
-			if !ok {
-				ns = &namespaceConfig{
-					name:            qb.Namespace,
-					assignedBuckets: make(map[string]*assignedBucket),
-				}
-				config[qb.Namespace] = ns
-			}
-			bucket, ok := ns.assignedBuckets[qb.Name]
-			if !ok {
-				if err := validateBucket(bucketToProto(qb)); err != nil {
-					return status.InternalErrorf("invalid bucket: %v", qbg.QuotaBucket)
-				}
-				bucket = &assignedBucket{
-					bucket: &qbg.QuotaBucket,
-				}
-				ns.assignedBuckets[qb.Name] = bucket
-			}
 
-			qg := qbg.QuotaGroup
-			if qg == nil {
-				// No Quota Group for this bucket
-				return nil
+	type quotaBucketGroup struct {
+		tables.QuotaBucket
+		*tables.QuotaGroup
+	}
+
+	if err := db.ScanEach(rq, func(ctx context.Context, qbg *quotaBucketGroup) error {
+		qb := &qbg.QuotaBucket
+		ns, ok := config[qb.Namespace]
+		if !ok {
+			ns = &namespaceConfig{
+				name:            qb.Namespace,
+				assignedBuckets: make(map[string]*assignedBucket),
 			}
-			if qg.BucketName == defaultBucketName {
-				log.Warningf("Doesn't need to create QuotaGroup for default bucket in namespace %q", qg.Namespace)
-				return nil
+			config[qb.Namespace] = ns
+		}
+		bucket, ok := ns.assignedBuckets[qb.Name]
+		if !ok {
+			if err := validateBucket(bucketToProto(qb)); err != nil {
+				return status.InternalErrorf("invalid bucket: %v", qbg.QuotaBucket)
 			}
-			bucket.quotaKeys = append(bucket.quotaKeys, qg.QuotaKey)
+			bucket = &assignedBucket{
+				bucket: &qbg.QuotaBucket,
+			}
+			ns.assignedBuckets[qb.Name] = bucket
+		}
+
+		qg := qbg.QuotaGroup
+		if qg == nil {
+			// No Quota Group for this bucket
 			return nil
-		},
-	)
-	if err != nil {
+		}
+		if qg.BucketName == defaultBucketName {
+			log.Warningf("Doesn't need to create QuotaGroup for default bucket in namespace %q", qg.Namespace)
+			return nil
+		}
+		bucket.quotaKeys = append(bucket.quotaKeys, qg.QuotaKey)
+		return nil
+	}); err != nil {
 		return nil, status.InternalErrorf("fetchConfigFromDB query failed: %s", err)
 	}
 	return config, nil
@@ -200,8 +256,6 @@ func validateBucket(bucket *qpb.Bucket) error {
 }
 
 type Bucket interface {
-	// Config returns a copy of the QuotaBucket. Used for testing.
-	Config() tables.QuotaBucket
 	Allow(ctx context.Context, key string, quantity int64) (bool, error)
 }
 
@@ -319,6 +373,65 @@ func (qm *QuotaManager) createNamespace(env environment.Env, name string, config
 	return ns, nil
 }
 
+// blockedBucket always denies requests. Used for blocking/banning users.
+type blockedBucket struct{}
+
+func (b *blockedBucket) Allow(ctx context.Context, key string, quantity int64) (bool, error) {
+	return false, nil
+}
+
+// quota requirements can also be configured in flagd experiment config. user/groups
+// can either be blocked completely, or a bucket can be configured similar to MySQL.
+//
+// quota.buckets must have valid structure, for example:
+//
+//	"rpc:/buildbuddy.service.BuildBuddyService/GetInvocation": {
+//	  "name": "custom",
+//	  "max_rate": {"num_requests": 100, "period": "1m"},
+//	  "max_burst": 10
+//	}
+func (qm *QuotaManager) readFlagdQuota(ctx context.Context, namespace string) Bucket {
+	if qm.env.GetExperimentFlagProvider() == nil {
+		return nil
+	}
+
+	if qm.env.GetExperimentFlagProvider().Boolean(ctx, blockedQuotaExperimentName, false) {
+		return &blockedBucket{}
+	}
+
+	bucketsConfig := qm.env.GetExperimentFlagProvider().Object(ctx, bucketQuotaExperimentName, nil)
+	if bucketsConfig == nil {
+		return nil
+	}
+	namespaceConfig, ok := bucketsConfig[namespace]
+	if !ok {
+		return nil
+	}
+	bucketMap, ok := namespaceConfig.(map[string]interface{})
+	if !ok {
+		log.CtxWarningf(ctx, "Invalid quota.buckets config for namespace %q: expected object, got %T", namespace, namespaceConfig)
+		return nil
+	}
+
+	// Directly create bucket row from map to avoid expensive proto conversion on every request
+	bucketRow, err := bucketRowFromMap(namespace, bucketMap)
+	if err != nil {
+		log.CtxWarningf(ctx, "Failed to parse quota bucket config for namespace %q: %s", namespace, err)
+		return nil
+	}
+
+	// TODO: cache created buckets to avoid recreating on every request. But this will only
+	// happen if the group matches and the namespace matches, so we should only create
+	// the bucket when we need to limit it.
+	bucket, err := qm.bucketCreator(qm.env, bucketRow)
+	if err != nil {
+		log.CtxWarningf(ctx, "Failed to create quota bucket for namespace %q: %s", namespace, err)
+		return nil
+	}
+
+	return bucket
+}
+
 // findBucket finds the bucket given a namespace and key. If the key is found in
 // bucketsByKey map, return the corresponding bucket. Otherwise, return the
 // default bucket. Returns nil if the namespace is not found or the default bucket
@@ -339,13 +452,31 @@ func (qm *QuotaManager) findBucket(nsName string, key string) Bucket {
 
 func (qm *QuotaManager) Allow(ctx context.Context, namespace string, quantity int64) error {
 	key, err := quota.GetKey(ctx, qm.env)
-
 	if err != nil {
 		metrics.QuotaKeyEmptyCount.With(prometheus.Labels{
 			metrics.QuotaNamespace: namespace,
 		}).Inc()
 		return nil
 	}
+
+	flagdBucket := qm.readFlagdQuota(ctx, namespace)
+	if flagdBucket != nil {
+		allow, err := flagdBucket.Allow(ctx, key, quantity)
+		if err != nil {
+			log.CtxWarningf(ctx, "Flagd quota check for %q failed: %s", namespace, err)
+			// There is some error when determining whether the request should be
+			// allowed. Do not block the traffic when the quota system has issues.
+		} else if !allow {
+			metrics.QuotaExceeded.With(prometheus.Labels{
+				metrics.QuotaNamespace: namespace,
+				metrics.QuotaKey:       key,
+				metrics.QuotaSource:    "flagd",
+			}).Inc()
+			return status.ResourceExhaustedErrorf("quota exceeded for %q", namespace)
+		}
+	}
+
+	// Also check DB based quota check.
 	b := qm.findBucket(namespace, key)
 	if b == nil {
 		// The bucket is not found, b/c either the namespace or the default bucket
@@ -354,7 +485,7 @@ func (qm *QuotaManager) Allow(ctx context.Context, namespace string, quantity in
 	}
 	allow, err := b.Allow(ctx, key, quantity)
 	if err != nil {
-		log.CtxWarningf(ctx, "Quota check for %q failed: %s", namespace, err)
+		log.CtxWarningf(ctx, "MySQL quota check for %q failed: %s", namespace, err)
 		// There is some error when determining whether the request should be
 		// allowed. Do not block the traffic when the quota system has issues.
 		return nil
@@ -365,6 +496,7 @@ func (qm *QuotaManager) Allow(ctx context.Context, namespace string, quantity in
 		metrics.QuotaExceeded.With(prometheus.Labels{
 			metrics.QuotaNamespace: namespace,
 			metrics.QuotaKey:       key,
+			metrics.QuotaSource:    "mysql",
 		}).Inc()
 		return status.ResourceExhaustedErrorf("quota exceeded for %q", namespace)
 	}
@@ -570,7 +702,6 @@ func (qm *QuotaManager) reloadNamespaces() error {
 		ns, err := qm.createNamespace(qm.env, nsName, nsConfig)
 		if err != nil {
 			return err
-
 		}
 		qm.namespaces.Store(nsName, ns)
 	}