优化

internal/service/security_session.go

@@ -2,6 +2,7 @@ package service
 
 import (
 	"context"
+	"strconv"
 	"sync"
 	"time"
 
@@ -22,22 +23,82 @@ func NewSecurityService(uc *biz.SecurityUsecase, runtime *conf.Runtime, cache bi
 	return &SecurityService{uc: uc, runtime: runtime, cache: cache, tokens: tokens}
 }
 
-func (s *SecurityService) CacheGet(ctx context.Context, key string) (string, bool, error) {
+func (s *SecurityService) cacheGet(ctx context.Context, key string) (string, bool, error) {
 	return s.cache.Get(ctx, key)
 }
 
-func (s *SecurityService) CacheSet(ctx context.Context, key, value string, expiration time.Duration) error {
+func (s *SecurityService) cacheSet(ctx context.Context, key, value string, expiration time.Duration) error {
 	return s.cache.Set(ctx, key, value, expiration)
 }
 
-func (s *SecurityService) CacheDelete(ctx context.Context, key string) error {
+func (s *SecurityService) cacheDelete(ctx context.Context, key string) error {
 	return s.cache.Delete(ctx, key)
 }
 
-func (s *SecurityService) CacheIncrement(ctx context.Context, key string, expiration time.Duration) (int64, error) {
+func (s *SecurityService) cacheIncrement(ctx context.Context, key string, expiration time.Duration) (int64, error) {
 	return s.cache.Increment(ctx, key, expiration)
 }
 
+func (s *SecurityService) ActiveToken(ctx context.Context, username string) (string, error) {
+	value, _, err := s.cacheGet(ctx, activeTokenKey(username))
+	return value, err
+}
+
+func (s *SecurityService) LoginLocked(ctx context.Context, username string) (bool, error) {
+	_, locked, err := s.cacheGet(ctx, "login:lock:"+username)
+	return locked, err
+}
+
+func (s *SecurityService) IncrementLoginFailure(ctx context.Context, username string, expiration time.Duration) (int64, error) {
+	return s.cacheIncrement(ctx, "login:fail:"+username, expiration)
+}
+
+func (s *SecurityService) LockLogin(ctx context.Context, username string, expiration time.Duration) error {
+	return s.cacheSet(ctx, "login:lock:"+username, "1", expiration)
+}
+
+func (s *SecurityService) ClearLoginState(ctx context.Context, username string) {
+	_ = s.cacheDelete(ctx, "login:fail:"+username)
+	_ = s.cacheDelete(ctx, "login:lock:"+username)
+}
+
+func (s *SecurityService) EnsureLoginIPCounter(ctx context.Context, ip string, expiration time.Duration) (int, error) {
+	value, exists, err := s.cacheGet(ctx, ip)
+	if err != nil {
+		return 0, err
+	}
+	if exists {
+		return strconv.Atoi(value)
+	}
+	if expiration <= 0 {
+		expiration = time.Hour
+	}
+	if err = s.cacheSet(ctx, ip, "1", expiration); err != nil {
+		return 0, err
+	}
+	return 1, nil
+}
+
+func (s *SecurityService) IncrementLoginIP(ctx context.Context, ip string, expiration time.Duration) (int64, error) {
+	return s.cacheIncrement(ctx, ip, expiration)
+}
+
+func (s *SecurityService) IncrementRateLimit(ctx context.Context, key string, expiration time.Duration) (int64, error) {
+	return s.cacheIncrement(ctx, key, expiration)
+}
+
+func (s *SecurityService) SetCaptcha(ctx context.Context, id, value string, expiration time.Duration) error {
+	return s.cacheSet(ctx, "captcha:"+id, value, expiration)
+}
+
+func (s *SecurityService) GetCaptcha(ctx context.Context, id string) (string, bool, error) {
+	return s.cacheGet(ctx, "captcha:"+id)
+}
+
+func (s *SecurityService) DeleteCaptcha(ctx context.Context, id string) error {
+	return s.cacheDelete(ctx, "captcha:"+id)
+}
+
 func (s *SecurityService) UseMultipoint() bool {
 	config := s.runtime.Admin()
 	return config != nil && config.System != nil && config.System.UseMultipoint
@@ -49,7 +110,7 @@ func (s *SecurityService) ActiveTokenMatches(ctx context.Context, username, toke
 	if !s.UseMultipoint() {
 		return true, nil
 	}
-	active, ok, err := s.CacheGet(ctx, activeTokenKey(username))
+	active, ok, err := s.cacheGet(ctx, activeTokenKey(username))
 	return ok && active == token, err
 }
 
@@ -62,5 +123,5 @@ func (s *SecurityService) RotateActiveToken(ctx context.Context, username, oldTo
 			return err
 		}
 	}
-	return s.CacheSet(ctx, activeTokenKey(username), newToken, expiration)
+	return s.cacheSet(ctx, activeTokenKey(username), newToken, expiration)
 }