block users

Author: tyler-french

enterprise/server/backends/authdb/authdb.go

@@ -367,7 +367,7 @@ func (d *AuthDB) fillChildGroupIDs(ctx context.Context, akg *apiKeyGroup) error
 		return nil
 	}
 	rq := d.h.NewQuery(ctx, "authdb_get_child_group_ids").Raw(`
-		SELECT child.group_id 
+		SELECT child.group_id
 		FROM "Groups" parent
 		JOIN "Groups" child on parent.saml_idp_metadata_url = child.saml_idp_metadata_url
 		WHERE parent.group_id = ?
@@ -428,7 +428,6 @@ func (d *AuthDB) GetAPIKeyGroupFromAPIKey(ctx context.Context, apiKey string) (i
 		).Raw(
 			q, args...,
 		).Take(akg)
-
 		if err != nil {
 			if db.IsRecordNotFound(err) {
 				if d.apiKeyGroupCache != nil {
@@ -475,7 +474,6 @@ func (d *AuthDB) GetAPIKeyGroupFromAPIKeyID(ctx context.Context, apiKeyID string
 		q,
 		args...,
 	).Take(akg)
-
 	if err != nil {
 		if db.IsRecordNotFound(err) {
 			return nil, status.UnauthenticatedErrorf("Invalid API key ID %q", redactInvalidAPIKey(apiKeyID))
@@ -495,7 +493,7 @@ func (d *AuthDB) LookupUserFromSubID(ctx context.Context, subID string) (*tables
 	rq := d.h.NewQueryWithOpts(ctx, "authdb_lookup_user_groups", db.Opts().WithStaleReads()).Raw(`
 		SELECT u.*, g.*, ug.*
 		FROM (
-			SELECT * FROM "Users" 
+			SELECT * FROM "Users"
 			WHERE sub_id = ?
 			ORDER BY user_id ASC
 			LIMIT 1

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,
@@ -140,7 +197,8 @@ func fetchConfigFromDB(env environment.Env, namespace string) (map[string]*names
 		func(ctx context.Context, qbg *struct {
 			tables.QuotaBucket
 			*tables.QuotaGroup
-		}) error {
+		},
+		) error {
 			qb := &qbg.QuotaBucket
 			ns, ok := config[qb.Namespace]
 			if !ok {
@@ -200,8 +258,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 +375,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 +454,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 +487,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 +498,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 +704,6 @@ func (qm *QuotaManager) reloadNamespaces() error {
 		ns, err := qm.createNamespace(qm.env, nsName, nsConfig)
 		if err != nil {
 			return err
-
 		}
 		qm.namespaces.Store(nsName, ns)
 	}