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/errors" "git.sendico.io/tech/sxchnge/shared/logging" "go.mongodb.org/mongo-driver/v2/mongo" "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 *mongo.Database, 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.CreateMongoRepository(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 }