Add enforcement mechanism for realm quotas

cmd/adminapi/main.go

@@ -150,7 +150,7 @@ func realMain(ctx context.Context) error {
 		// Install the APIKey Auth Middleware
 		sub.Use(requireAPIKey)
 
-		issueapiController, err := issueapi.New(ctx, config, db, h)
+		issueapiController, err := issueapi.New(ctx, config, db, limiterStore, h)
 		if err != nil {
 			return fmt.Errorf("issueapi.New: %w", err)
 		}

cmd/server/main.go

@@ -252,7 +252,7 @@ func realMain(ctx context.Context) error {
 		sub.Handle("", homeController.HandleHome()).Methods("GET")
 
 		// API for creating new verification codes. Called via AJAX.
-		issueapiController, err := issueapi.New(ctx, config, db, h)
+		issueapiController, err := issueapi.New(ctx, config, db, limiterStore, h)
 		if err != nil {
 			return fmt.Errorf("issueapi.New: %w", err)
 		}

pkg/config/admin_server_config.go

@@ -47,6 +47,7 @@ type AdminAPIServerConfig struct {
 
 	CollisionRetryCount uint          `env:"COLLISION_RETRY_COUNT,default=6"`
 	AllowedSymptomAge   time.Duration `env:"ALLOWED_PAST_SYMPTOM_DAYS,default=336h"` // 336h is 14 days.
+	EnforceRealmQuotas  bool          `env:"ENFORCE_REALM_QUOTAS, default=false"`
 }
 
 // NewAdminAPIServerConfig returns the environment config for the Admin API server.
@@ -85,6 +86,14 @@ func (c *AdminAPIServerConfig) GetAllowedSymptomAge() time.Duration {
 	return c.AllowedSymptomAge
 }
 
+func (c *AdminAPIServerConfig) GetEnforceRealmQuotas() bool {
+	return c.EnforceRealmQuotas
+}
+
+func (c *AdminAPIServerConfig) GetRateLimitConfig() *ratelimit.Config {
+	return &c.RateLimit
+}
+
 func (c *AdminAPIServerConfig) ObservabilityExporterConfig() *observability.Config {
 	return &c.Observability
 }

pkg/config/issue.go

@@ -14,11 +14,17 @@
 
 package config
 
-import "time"
+import (
+	"time"
+
+	"github.com/google/exposure-notifications-verification-server/pkg/ratelimit"
+)
 
 // IssueAPIConfig is an interface that represents what is needed of the verification
 // code issue API.
 type IssueAPIConfig interface {
 	GetCollisionRetryCount() uint
 	GetAllowedSymptomAge() time.Duration
+	GetEnforceRealmQuotas() bool
+	GetRateLimitConfig() *ratelimit.Config
 }

pkg/config/server_config.go

@@ -59,6 +59,7 @@ type ServerConfig struct {
 	ServerName          string        `env:"SERVER_NAME,default=Diagnosis Verification Server"`
 	CollisionRetryCount uint          `env:"COLLISION_RETRY_COUNT,default=6"`
 	AllowedSymptomAge   time.Duration `env:"ALLOWED_PAST_SYMPTOM_DAYS,default=336h"` // 336h is 14 days.
+	EnforceRealmQuotas  bool          `env:"ENFORCE_REALM_QUOTAS, default=false"`
 
 	AssetsPath string `env:"ASSETS_PATH,default=./cmd/server/assets"`
 
@@ -111,6 +112,14 @@ func (c *ServerConfig) GetAllowedSymptomAge() time.Duration {
 	return c.AllowedSymptomAge
 }
 
+func (c *ServerConfig) GetEnforceRealmQuotas() bool {
+	return c.EnforceRealmQuotas
+}
+
+func (c *ServerConfig) GetRateLimitConfig() *ratelimit.Config {
+	return &c.RateLimit
+}
+
 func (c *ServerConfig) ObservabilityExporterConfig() *observability.Config {
 	return &c.Observability
 }

pkg/controller/issueapi/issue.go

@@ -24,6 +24,7 @@ import (
 	"github.com/google/exposure-notifications-verification-server/pkg/api"
 	"github.com/google/exposure-notifications-verification-server/pkg/controller"
 	"github.com/google/exposure-notifications-verification-server/pkg/database"
+	"github.com/google/exposure-notifications-verification-server/pkg/digest"
 	"github.com/google/exposure-notifications-verification-server/pkg/observability"
 	"github.com/google/exposure-notifications-verification-server/pkg/otp"
 	"github.com/google/exposure-notifications-verification-server/pkg/sms"
@@ -194,6 +195,33 @@ func (c *Controller) HandleIssue() http.Handler {
 			}
 		}
 
+		// If we got this far, we're about to issue a code.
+		dig, err := digest.HMACUint(realm.ID, c.config.GetRateLimitConfig().HMACKey)
+		if err != nil {
+			controller.InternalError(w, r, c.h, err)
+			return
+		}
+		key := fmt.Sprintf("realm:quota:%s", dig)
+		limit, _, reset, ok, err := c.limiter.Take(ctx, key)
+		if err != nil {
+			logger.Errorw("failed to take from limiter", "error", err)
+			stats.Record(ctx, c.metrics.QuotaErrors.M(1))
+			c.h.RenderJSON(w, http.StatusInternalServerError, api.Errorf("failed to verify realm stats, please try again"))
+			return
+		}
+		if !ok {
+			logger.Warnw("realm has exceeded daily quota",
+				"realm", realm.ID,
+				"limit", limit,
+				"reset", reset)
+			stats.Record(ctx, c.metrics.QuotaExceeded.M(1))
+
+			if c.config.GetEnforceRealmQuotas() {
+				c.h.RenderJSON(w, http.StatusTooManyRequests, api.Errorf("exceeded realm quota"))
+				return
+			}
+		}
+
 		now := time.Now().UTC()
 		expiryTime := now.Add(realm.CodeDuration.Duration)
 		longExpiryTime := now.Add(realm.LongCodeDuration.Duration)

pkg/controller/issueapi/issueapi.go

@@ -27,32 +27,35 @@ import (
 
 	"github.com/google/exposure-notifications-server/pkg/logging"
 
+	"github.com/sethvargo/go-limiter"
 	"go.uber.org/zap"
 )
 
 type Controller struct {
-	config config.IssueAPIConfig
-	db     *database.Database
-	h      *render.Renderer
-	logger *zap.SugaredLogger
+	config  config.IssueAPIConfig
+	db      *database.Database
+	h       *render.Renderer
+	limiter limiter.Store
+	logger  *zap.SugaredLogger
 
 	validTestType map[string]struct{}
 
 	metrics *Metrics
 }
 
 // New creates a new IssueAPI controller.
-func New(ctx context.Context, config config.IssueAPIConfig, db *database.Database, h *render.Renderer) (*Controller, error) {
+func New(ctx context.Context, config config.IssueAPIConfig, db *database.Database, limiter limiter.Store, h *render.Renderer) (*Controller, error) {
 	metrics, err := registerMetrics()
 	if err != nil {
 		return nil, err
 	}
 
 	return &Controller{
-		config: config,
-		db:     db,
-		h:      h,
-		logger: logging.FromContext(ctx),
+		config:  config,
+		db:      db,
+		h:       h,
+		limiter: limiter,
+		logger:  logging.FromContext(ctx),
 		validTestType: map[string]struct{}{
 			api.TestTypeConfirmed: {},
 			api.TestTypeLikely:    {},

pkg/controller/issueapi/metrics.go

@@ -30,6 +30,8 @@ var (
 
 type Metrics struct {
 	IssueAttempts   *stats.Int64Measure
+	QuotaErrors     *stats.Int64Measure
+	QuotaExceeded   *stats.Int64Measure
 	CodesIssued     *stats.Int64Measure
 	CodeIssueErrors *stats.Int64Measure
 	SMSSent         *stats.Int64Measure
@@ -48,6 +50,28 @@ func registerMetrics() (*Metrics, error) {
 		return nil, fmt.Errorf("stat view registration failure: %w", err)
 	}
 
+	mQuotaErrors := stats.Int64(MetricPrefix+"/quota_errors", "The number of errors when taking from the limiter", stats.UnitDimensionless)
+	if err := view.Register(&view.View{
+		Name:        MetricPrefix + "/quota_errors_count",
+		Measure:     mQuotaErrors,
+		Description: "The count of the number of errors to the limiter",
+		TagKeys:     []tag.Key{observability.RealmTagKey},
+		Aggregation: view.Count(),
+	}); err != nil {
+		return nil, fmt.Errorf("stat view registration failure: %w", err)
+	}
+
+	mQuotaExceeded := stats.Int64(MetricPrefix+"/quota_exceeded", "The number of times quota has been exceeded", stats.UnitDimensionless)
+	if err := view.Register(&view.View{
+		Name:        MetricPrefix + "/quota_exceeded_count",
+		Measure:     mQuotaExceeded,
+		Description: "The count of the number of times quota has been exceeded",
+		TagKeys:     []tag.Key{observability.RealmTagKey},
+		Aggregation: view.Count(),
+	}); err != nil {
+		return nil, fmt.Errorf("stat view registration failure: %w", err)
+	}
+
 	mCodesIssued := stats.Int64(MetricPrefix+"/codes_issued", "The number of verification codes issued", stats.UnitDimensionless)
 	if err := view.Register(&view.View{
 		Name:        MetricPrefix + "/codes_issued_count",
@@ -94,6 +118,8 @@ func registerMetrics() (*Metrics, error) {
 
 	return &Metrics{
 		IssueAttempts:   mIssueAttempts,
+		QuotaErrors:     mQuotaErrors,
+		QuotaExceeded:   mQuotaExceeded,
 		CodesIssued:     mCodesIssued,
 		CodeIssueErrors: mCodesIssued,
 		SMSSent:         mSMSSent,