sxchnge/backend/shared/ratelimit/store/db.go

package ratelimits

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"strings"
	"time"

	"git.sendico.io/tech/sxchnge/shared/database/record"
	"git.sendico.io/tech/sxchnge/shared/database/repository"
	ri "git.sendico.io/tech/sxchnge/shared/database/repository/index"
	"git.sendico.io/tech/sxchnge/shared/database/source"
	"git.sendico.io/tech/sxchnge/shared/errors"
	"git.sendico.io/tech/sxchnge/shared/logging"
	"go.uber.org/zap"
)

const (
	defaultCollection = "rate_limits"

	fieldScopeKind   = "scopeKind"
	fieldScopeHash   = "scopeHash"
	fieldWindowStart = "windowStart"
	fieldCount       = "count"
	fieldExpiresAt   = "expiresAt"
)

// Store provides atomic, fixed-window abuse accounting shared by every BFF
// instance connected to the same database.
type Store interface {
	Current(context.Context, string, string, time.Time) (int, error)
	Increment(context.Context, string, string, time.Time, time.Time) (int, error)
}

func HashScope(kind, value string) string {
	sum := sha256.Sum256([]byte(strings.TrimSpace(kind) + "\x00" + strings.TrimSpace(value)))
	return hex.EncodeToString(sum[:])
}

type record struct {
	storable.Base `bson:",inline"`
	ScopeKind     string    `bson:"scopeKind"`
	ScopeHash     string    `bson:"scopeHash"`
	WindowStart   time.Time `bson:"windowStart"`
	Count         int       `bson:"count"`
	ExpiresAt     time.Time `bson:"expiresAt"`
}

func (r *record) Collection() string { return defaultCollection }

type DB struct {
	repo   repository.Repository
	logger mlogger.Logger
}

func Create(logger mlogger.Logger, database source.Source, collection string) (*DB, error) {
	if database == nil {
		return nil, merrors.InvalidArgument("database is required", "database")
	}
	collection = strings.TrimSpace(collection)
	if collection == "" {
		collection = defaultCollection
	}
	db := &DB{
		repo:   repository.CreateRepository(database, collection),
		logger: logger.Named("rate_limits"),
	}
	if err := db.ensureIndexes(); err != nil {
		db.logger.Error("Failed to create rate-limit indexes", zap.Error(err), zap.String("collection", collection))
		return nil, err
	}
	return db, nil
}

func (db *DB) Current(
	ctx context.Context,
	scopeKind, scopeHash string,
	windowStart time.Time,
) (int, error) {
	filter, err := scopeFilter(scopeKind, scopeHash, windowStart)
	if err != nil {
		return 0, err
	}
	current := &record{}
	if err := db.repo.FindOneByFilter(ctx, filter, current); err != nil {
		if errors.Is(err, merrors.ErrNoData) {
			return 0, nil
		}
		return 0, err
	}
	return current.Count, nil
}

func (db *DB) Increment(
	ctx context.Context,
	scopeKind, scopeHash string,
	windowStart, expiresAt time.Time,
) (int, error) {
	filter, err := scopeFilter(scopeKind, scopeHash, windowStart)
	if err != nil {
		return 0, err
	}
	if expiresAt.IsZero() || !expiresAt.After(windowStart) {
		return 0, merrors.InvalidArgument("expiry must be after window start", "expiresAt")
	}

	current := &record{}
	if err := db.repo.FindOneByFilter(ctx, filter, current); err == nil {
		updated := &record{}
		matched, patchErr := db.repo.PatchOneByFilter(ctx, filter, repository.Patch().
			Inc(repository.Field(fieldCount), 1).
			Set(repository.Field(fieldExpiresAt), expiresAt.UTC()), updated)
		if patchErr != nil {
			return 0, patchErr
		}
		if !matched {
			return db.Increment(ctx, scopeKind, scopeHash, windowStart, expiresAt)
		}
		return updated.Count, nil
	} else if !errors.Is(err, merrors.ErrNoData) {
		return 0, err
	}

	created := &record{
		ScopeKind:   strings.TrimSpace(scopeKind),
		ScopeHash:   strings.TrimSpace(scopeHash),
		WindowStart: windowStart.UTC(),
		Count:       1,
		ExpiresAt:   expiresAt.UTC(),
	}
	if err := db.repo.Insert(ctx, created, nil); err != nil {
		if errors.Is(err, merrors.ErrDataConflict) {
			return db.Increment(ctx, scopeKind, scopeHash, windowStart, expiresAt)
		}
		return 0, err
	}
	return created.Count, nil
}

func (db *DB) ensureIndexes() error {
	ttl := int32(0)
	indexes := []*ri.Definition{
		{
			Keys: []ri.Key{
				{Field: fieldScopeKind, Sort: ri.Asc},
				{Field: fieldScopeHash, Sort: ri.Asc},
				{Field: fieldWindowStart, Sort: ri.Asc},
			},
			Unique: true,
			Name:   "unique_rate_limit_window",
		},
		{
			Keys: []ri.Key{{Field: fieldExpiresAt, Sort: ri.Asc}},
			TTL:  &ttl,
			Name: "ttl_rate_limit_expiry",
		},
	}
	for _, index := range indexes {
		if err := db.repo.CreateIndex(index); err != nil {
			return err
		}
	}
	return nil
}

func scopeFilter(scopeKind, scopeHash string, windowStart time.Time) (repository.FilterQuery, error) {
	scopeKind = strings.TrimSpace(scopeKind)
	if scopeKind == "" {
		return nil, merrors.InvalidArgument("scope kind is required", "scopeKind")
	}
	scopeHash = strings.TrimSpace(scopeHash)
	if scopeHash == "" {
		return nil, merrors.InvalidArgument("scope hash is required", "scopeHash")
	}
	if windowStart.IsZero() {
		return nil, merrors.InvalidArgument("window start is required", "windowStart")
	}
	return repository.Query().
		Filter(repository.Field(fieldScopeKind), scopeKind).
		Filter(repository.Field(fieldScopeHash), scopeHash).
		Filter(repository.Field(fieldWindowStart), windowStart.UTC()), nil
}