Files
gemini-banlancer/internal/service/security_service.go
2025-11-22 14:20:05 +08:00

84 lines
2.4 KiB
Go
Raw Blame History

This file contains invisible Unicode characters
This file contains invisible Unicode characters that are indistinguishable to humans but may be processed differently by a computer. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
// Filename: internal/service/security_service.go
package service
import (
"context"
"crypto/sha256" // [NEW] Import crypto library for hashing
"encoding/hex" // [NEW] Import hex encoding
"fmt"
"gemini-balancer/internal/models"
"gemini-balancer/internal/repository" // [NEW] Import repository
"gemini-balancer/internal/settings"
"gemini-balancer/internal/store"
"github.com/sirupsen/logrus"
"gorm.io/gorm"
)
const loginAttemptsKey = "security:login_attempts"
type SecurityService struct {
repo repository.AuthTokenRepository
store store.Store
SettingsManager *settings.SettingsManager
logger *logrus.Entry
}
// NewSecurityService signature updated to accept the repository.
func NewSecurityService(repo repository.AuthTokenRepository, store store.Store, settingsManager *settings.SettingsManager, logger *logrus.Logger) *SecurityService {
return &SecurityService{
repo: repo,
store: store,
SettingsManager: settingsManager,
logger: logger.WithField("component", "SecurityService🛡"),
}
}
// AuthenticateToken is now secure and efficient.
func (s *SecurityService) AuthenticateToken(tokenValue string) (*models.AuthToken, error) {
if tokenValue == "" {
return nil, gorm.ErrRecordNotFound
}
// [REFACTORED]
// 1. Hash the incoming plaintext token.
hash := sha256.Sum256([]byte(tokenValue))
tokenHash := hex.EncodeToString(hash[:])
// 2. Delegate the lookup to the repository using the hash.
return s.repo.GetTokenByHashedValue(tokenHash)
}
// IsIPBanned
func (s *SecurityService) IsIPBanned(ctx context.Context, ip string) (bool, error) {
banKey := fmt.Sprintf("banned_ip:%s", ip)
return s.store.Exists(ctx, banKey)
}
// RecordFailedLoginAttempt
func (s *SecurityService) RecordFailedLoginAttempt(ctx context.Context, ip string) error {
if !s.SettingsManager.IsIPBanEnabled() {
return nil
}
count, err := s.store.HIncrBy(ctx, loginAttemptsKey, ip, 1)
if err != nil {
return err
}
maxAttempts := s.SettingsManager.GetMaxLoginAttempts()
if count >= int64(maxAttempts) {
banDuration := s.SettingsManager.GetIPBanDuration()
banKey := fmt.Sprintf("banned_ip:%s", ip)
if err := s.store.Set(ctx, banKey, []byte("1"), banDuration); err != nil {
return err
}
s.logger.Warnf("IP BANNED: IP [%s] has been banned for %v due to excessive failed login attempts.", ip, banDuration)
s.store.HDel(ctx, loginAttemptsKey, ip)
}
return nil
}