ROX-23260: Add Rate Limiter Service

emailsender/config/config.go

@@ -10,12 +10,20 @@ import (
 
 // Config contains this application's runtime configuration.
 type Config struct {
-	ClusterID      string `env:"CLUSTER_ID"`
-	ServerAddress  string `env:"SERVER_ADDRESS" envDefault:":8080"`
-	EnableHTTPS    bool   `env:"ENABLE_HTTPS" envDefault:"false"`
-	HTTPSCertFile  string `env:"HTTPS_CERT_FILE" envDefault:""`
-	HTTPSKeyFile   string `env:"HTTPS_KEY_FILE" envDefault:""`
-	MetricsAddress string `env:"METRICS_ADDRESS" envDefault:":9090"`
+	ClusterID                 string `env:"CLUSTER_ID"`
+	ServerAddress             string `env:"SERVER_ADDRESS" envDefault:":8080"`
+	EnableHTTPS               bool   `env:"ENABLE_HTTPS" envDefault:"false"`
+	HTTPSCertFile             string `env:"HTTPS_CERT_FILE" envDefault:""`
+	HTTPSKeyFile              string `env:"HTTPS_KEY_FILE" envDefault:""`
+	MetricsAddress            string `env:"METRICS_ADDRESS" envDefault:":9090"`
+	DatabaseHost              string `env:"DATABASE_HOST" envDefault:"localhost"`
+	DatabasePort              int    `env:"DATABASE_PORT" envDefault:"5432"`
+	DatabaseName              string `env:"DATABASE_NAME" envDefault:"postgres"`
+	DatabaseUser              string `env:"DATABASE_USER" envDefault:"postgres"`
+	DatabasePassword          string `env:"DATABASE_PASSWORD" envDefault:"postgres"`
+	DatabaseSSLMode           string `env:"DATABASE_SSL_MODE" envDefault:"disable"`
+	LimitEmailPerSecond       int    `env:"LIMIT_EMAIL_PER_SECOND" envDefault:"14"`
+	LimitEmailPerDayPerTenant int    `env:"LIMIT_EMAIL_PER_DAY_PER_TENANT" envDefault:"250"`
 }
 
 // GetConfig retrieves the current runtime configuration from the environment and returns it.

emailsender/pkg/db/connect.go

@@ -0,0 +1,80 @@
+package db
+
+import (
+	"fmt"
+	"gorm.io/gorm"
+	"time"
+
+	commonDB "github.com/stackrox/acs-fleet-manager/pkg/db"
+)
+
+// DatabaseClient defines methods for fetching or updating models in DB
+type DatabaseClient interface {
+	GetEmailSentByTenant(tenantID string, date time.Time) (*EmailSentByTenant, error)
+	UpdateEmailSentByTenant(tenantID string, date time.Time, amount int) error
+	GetEmailSentPerSecond() (*EmailSentPerSecond, error)
+	UpdateEmailSentPerSecond(amount int) error
+}
+
+// DatabaseConnection contains dependency for communicating with DB
+type DatabaseConnection struct {
+	DB *gorm.DB
+}
+
+// NewDatabaseConnection creates a new DB connection
+func NewDatabaseConnection(host string, port int, user, password, database, SSLMode string) (*DatabaseConnection, error) {
+	dbConfig := &commonDB.DatabaseConfig{
+		Host:     host,
+		Port:     port,
+		Username: user,
+		Password: password, // pragma: allowlist secret
+		Name:     database,
+		SSLMode:  SSLMode,
+	}
+	connection, _ := commonDB.NewConnectionFactory(dbConfig)
+	return &DatabaseConnection{DB: connection.DB}, nil
+}
+
+// Migrate automatically migrates listed models in the database
+// Documentation: https://gorm.io/docs/migration.html#Auto-Migration
+func (d *DatabaseConnection) Migrate() error {
+	return d.DB.AutoMigrate(&EmailSentPerSecond{}, &EmailSentByTenant{})
+}
+
+// GetEmailSentByTenant returns an instance of EmailSentByTenant representing how many emails tenant sent for provided date
+// Note: date uses only days
+func (d *DatabaseConnection) GetEmailSentByTenant(tenantID string, date time.Time) (*EmailSentByTenant, error) {
+	var emailSentByTenant EmailSentByTenant
+	d.DB.FirstOrCreate(&emailSentByTenant, &EmailSentByTenant{TenantID: tenantID, Date: onlyDate(date)})
+	return &emailSentByTenant, nil
+}
+
+// UpdateEmailSentByTenant updates how many emails sent by tenant for provided date
+// Note: date uses only days
+func (d *DatabaseConnection) UpdateEmailSentByTenant(tenantID string, date time.Time, amount int) error {
+	if result := d.DB.Model(&EmailSentByTenant{}).
+		Where("tenant_id = ? and date = ?", tenantID, onlyDate(date)).
+		Update("amount", amount); result.Error != nil {
+		return fmt.Errorf("failed updating email_sent_by_tenant table: %v", result.Error)
+	}
+	return nil
+}
+
+// GetEmailSentPerSecond returns how many emails sent for the last second
+func (d *DatabaseConnection) GetEmailSentPerSecond() (*EmailSentPerSecond, error) {
+	var emailSentPerSecond EmailSentPerSecond
+	d.DB.FirstOrCreate(&emailSentPerSecond, &EmailSentPerSecond{})
+	return &emailSentPerSecond, nil
+}
+
+// UpdateEmailSentPerSecond updates how many emails sent for the last second
+func (d *DatabaseConnection) UpdateEmailSentPerSecond(amount int) error {
+	if result := d.DB.Save(&EmailSentPerSecond{ID: 1, Amount: amount}); result.Error != nil {
+		return fmt.Errorf("failed updating email_sent_per_second table: %v", result.Error)
+	}
+	return nil
+}
+
+func onlyDate(t time.Time) time.Time {
+	return time.Date(t.Year(), t.Month(), t.Day(), 0, 0, 0, 0, time.UTC)
+}

emailsender/pkg/db/models.go

@@ -0,0 +1,17 @@
+package db
+
+import "time"
+
+// EmailSentPerSecond represent how many emails sent for the last second
+type EmailSentPerSecond struct {
+	ID        uint // primary key
+	UpdatedAt int  // gorm automatically  set to current unix seconds on update. It is equal to zero b default
+	Amount    int
+}
+
+// EmailSentByTenant represents how many emails sent by tenant
+type EmailSentByTenant struct {
+	TenantID string    `gorm:"index"`
+	Date     time.Time `gorm:"index"`
+	Amount   int
+}

emailsender/pkg/email/ratelimiter.go

@@ -0,0 +1,85 @@
+package email
+
+import (
+	"github.com/golang/glog"
+	"github.com/stackrox/acs-fleet-manager/emailsender/pkg/db"
+	"time"
+)
+
+// RateLimiter defines an exact methods for rate limiter
+type RateLimiter interface {
+	Allow(tenantID string) bool
+}
+
+// RateLimiterService contains configuration and dependency for rate limiter
+type RateLimiterService struct {
+	limitPerSecond       int
+	limitPerDayPerTenant int
+	dbConnection         db.DatabaseClient
+}
+
+// NewRateLimiterService creates a new instance of RateLimiterService
+func NewRateLimiterService(dbConnection *db.DatabaseConnection) *RateLimiterService {
+	return &RateLimiterService{
+		dbConnection: dbConnection,
+	}
+}
+
+// Allow calculates whether an email may send now for specific tenant for current timestamp
+func (r *RateLimiterService) Allow(tenantID string) bool {
+	now := time.Now()
+	nowSeconds := now.Unix()
+	allowedPerSecond := r.allowRatePerSecond(tenantID, int(nowSeconds))
+	if !allowedPerSecond {
+		return false
+	}
+	allowedPerTenant := r.allowRatePerTenant(tenantID, now)
+
+	return allowedPerTenant
+}
+
+func (r *RateLimiterService) allowRatePerSecond(tenantID string, now int) bool {
+	emailPerSecond, err := r.dbConnection.GetEmailSentPerSecond()
+	if err != nil {
+		glog.Errorf("Cannot get email sent per second for tenant %s: %v", tenantID, err)
+		return false
+	}
+
+	if emailPerSecond.Amount >= r.limitPerSecond && (now-emailPerSecond.UpdatedAt) < 2 {
+		glog.Warningf("Reached limit for sent emails per second for tenant %s", tenantID)
+		return false
+	}
+	if emailPerSecond.UpdatedAt == 0 || // just created EmailSentPerSecond counter
+		(now-emailPerSecond.UpdatedAt) > 1 || // stale EmailSentPerSecond counter
+		(emailPerSecond.Amount < r.limitPerSecond && (now-emailPerSecond.UpdatedAt) < 2) { // rate is within limit
+		if err = r.dbConnection.UpdateEmailSentPerSecond(emailPerSecond.Amount + 1); err != nil {
+			glog.Errorf("Cannot update email sent per second for tenant %s: %v", tenantID, err)
+		}
+		return true
+	}
+
+	return false
+}
+
+func (r *RateLimiterService) allowRatePerTenant(tenantID string, now time.Time) bool {
+	emailPerTenant, err := r.dbConnection.GetEmailSentByTenant(tenantID, now)
+	if err != nil {
+		glog.Errorf("Cannot get email sent for tenant %s: %v", tenantID, err)
+		return false
+	}
+	hoursDelta := now.Sub(emailPerTenant.Date).Hours()
+
+	if emailPerTenant.Amount >= r.limitPerDayPerTenant && hoursDelta <= 24 {
+		glog.Warningf("Reached limit for sent emails per day per tenant for tenant %s", tenantID)
+		return false
+	}
+
+	if hoursDelta > 24 || (emailPerTenant.Amount < r.limitPerDayPerTenant && hoursDelta <= 24) {
+		if err = r.dbConnection.UpdateEmailSentByTenant(tenantID, now, emailPerTenant.Amount+1); err != nil {
+			glog.Errorf("Cannot update email sent per second for tenant %s: %v", tenantID, err)
+		}
+		return true
+	}
+
+	return false
+}