sm-user-plugin/classes/throttle.go

package classes

import (
	"context"
	"errors"
	"time"

	"git.golem15.com/golem15/sm-user-plugin/models"
	"git.golem15.com/golem15/summercms/modules/backpack"
	"gorm.io/gorm"
)

var (
	errBanned    = errors.New("user is banned")
	errSuspended = errors.New("user is suspended")
)

type throttleSettings struct {
	use     bool
	limit   int
	minutes int
}

type throttleSettingsKey struct{}

func withThrottle(ctx context.Context, s throttleSettings) context.Context {
	if ctx == nil {
		ctx = context.Background()
	}
	return context.WithValue(ctx, throttleSettingsKey{}, s)
}

// ContextWithThrottle copies golem15.user.throttle.* onto ctx for CheckAndRecordLogin.
func ContextWithThrottle(ctx context.Context, app *backpack.App) context.Context {
	s := throttleSettings{use: true, limit: 5, minutes: 15}
	if app != nil && app.Config != nil {
		if app.Config.Has("golem15.user.throttle.use_throttle") {
			s.use = app.Config.Bool("golem15.user.throttle.use_throttle")
		}
		if n := app.Config.Int("golem15.user.throttle.attempt_limit"); n > 0 {
			s.limit = n
		}
		if n := app.Config.Int("golem15.user.throttle.suspension_minutes"); n > 0 {
			s.minutes = n
		}
	}
	return withThrottle(ctx, s)
}

func throttleFrom(ctx context.Context) throttleSettings {
	if ctx != nil {
		if s, ok := ctx.Value(throttleSettingsKey{}).(throttleSettings); ok {
			return s
		}
	}
	return throttleSettings{use: true, limit: 5, minutes: 15}
}

// IsLoginRejected reports a ban or an active suspension. Callers map it to the
// generic invalid-login response so the cause is not distinguishable.
func IsLoginRejected(err error) bool {
	return errors.Is(err, errBanned) || errors.Is(err, errSuspended)
}

// RejectIfThrottled enforces a ban or an active suspension without recording
// an attempt. Login calls it before the password check; a missing row is allowed.
func RejectIfThrottled(ctx context.Context, db *gorm.DB, user *models.User, ip string) error {
	if user == nil || user.ID == 0 || db == nil {
		return nil
	}
	settings := throttleFrom(ctx)
	if !settings.use {
		return nil
	}
	if settings.minutes <= 0 {
		settings.minutes = 15
	}
	var row models.Throttle
	err := db.WithContext(ctx).
		Where("user_id = ? AND (ip_address = ? OR ip_address IS NULL)", user.ID, ip).
		Order("ip_address NULLS LAST").
		Take(&row).Error
	if errors.Is(err, gorm.ErrRecordNotFound) {
		return nil
	}
	if err != nil {
		return err
	}
	return throttleBlocked(row, settings)
}

func throttleBlocked(row models.Throttle, settings throttleSettings) error {
	if row.IsBanned {
		return errBanned
	}
	window := time.Duration(settings.minutes) * time.Minute
	if row.IsSuspended && row.SuspendedAt != nil && time.Now().Before(row.SuspendedAt.Add(window)) {
		return errSuspended
	}
	return nil
}

// CheckAndRecordLogin applies Winter's failed-login throttle.
// Login calls RejectIfThrottled before the password check, then this function
// once with the outcome. A nil user is ignored so an unknown email never
// creates a row.
func CheckAndRecordLogin(ctx context.Context, db *gorm.DB, user *models.User, ip string, ok bool) error {
	if user == nil || user.ID == 0 || db == nil {
		return nil
	}
	settings := throttleFrom(ctx)
	if !settings.use {
		return nil
	}
	if settings.limit <= 0 {
		settings.limit = 5
	}
	if settings.minutes <= 0 {
		settings.minutes = 15
	}
	var row models.Throttle
	err := db.WithContext(ctx).
		Where("user_id = ? AND (ip_address = ? OR ip_address IS NULL)", user.ID, ip).
		Order("ip_address NULLS LAST").
		Take(&row).Error
	if errors.Is(err, gorm.ErrRecordNotFound) {
		row = models.Throttle{UserID: user.ID, IPAddress: &ip}
		if err := db.WithContext(ctx).Create(&row).Error; err != nil {
			return err
		}
	} else if err != nil {
		return err
	}
	if err := throttleBlocked(row, settings); err != nil {
		return err
	}
	if ok {
		return db.WithContext(ctx).Model(&row).Updates(map[string]any{
			"attempts":     0,
			"is_suspended": false,
			"suspended_at": nil,
		}).Error
	}
	attempts := row.Attempts + 1
	updates := map[string]any{"attempts": attempts}
	if attempts >= settings.limit {
		updates["is_suspended"] = true
		updates["suspended_at"] = time.Now()
	}
	return db.WithContext(ctx).Model(&row).Updates(updates).Error
}