fix: Add deviceLimit to the prefix of email string

Solve different node and different device limit the less limit node cannot work normally.
This commit is contained in:
Senis John 2022-12-07 18:36:01 +08:00
parent 85d73408c3
commit 9a06f78653
No known key found for this signature in database
GPG Key ID: 845E9E4727C3E1A4

View File

@ -4,6 +4,7 @@ package limiter
import ( import (
"context" "context"
"fmt" "fmt"
"strconv"
"strings" "strings"
"sync" "sync"
"time" "time"
@ -63,16 +64,17 @@ func (l *Limiter) AddInboundLimiter(tag string, nodeSpeedLimit uint64, userList
gs := goCacheStore.NewGoCache(goCache.New(time.Duration(globalLimit.Expiry)*time.Second, 10*time.Minute)) gs := goCacheStore.NewGoCache(goCache.New(time.Duration(globalLimit.Expiry)*time.Second, 10*time.Minute))
// init redis store // init redis store
redisClient := redis.NewClient(&redis.Options{ rs := redisStore.NewRedis(redis.NewClient(
&redis.Options{
Addr: globalLimit.RedisAddr, Addr: globalLimit.RedisAddr,
Password: globalLimit.RedisPassword, Password: globalLimit.RedisPassword,
DB: globalLimit.RedisDB, DB: globalLimit.RedisDB,
}) }),
rs := redisStore.NewRedis(redisClient, store.WithExpiration(time.Duration(globalLimit.Expiry)*time.Second)) store.WithExpiration(time.Duration(globalLimit.Expiry)*time.Second))
// init chained cache. First use local go-cache, if go-cache is nil, then use redis cache // init chained cache. First use local go-cache, if go-cache is nil, then use redis cache
cacheManager := cache.NewChain[any]( cacheManager := cache.NewChain[any](
cache.New[any](gs), cache.New[any](gs), // go-cache is priority
cache.New[any](rs), cache.New[any](rs),
) )
inboundInfo.GlobalLimit.globalOnlineIP = marshaler.New(cacheManager) inboundInfo.GlobalLimit.globalOnlineIP = marshaler.New(cacheManager)
@ -223,29 +225,30 @@ func globalLimit(inboundInfo *InboundInfo, email string, uid int, ip string, dev
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(inboundInfo.GlobalLimit.config.Timeout)*time.Second) ctx, cancel := context.WithTimeout(context.Background(), time.Duration(inboundInfo.GlobalLimit.config.Timeout)*time.Second)
defer cancel() defer cancel()
email := email[strings.Index(email, "|")+1:] // reformat email for unique key // reformat email for unique key
email := strings.Replace(email, inboundInfo.Tag, strconv.Itoa(deviceLimit), 1)
v, err := inboundInfo.GlobalLimit.globalOnlineIP.Get(ctx, email, new(map[string]int)) v, err := inboundInfo.GlobalLimit.globalOnlineIP.Get(ctx, email, new(map[string]int))
if err != nil { if err != nil {
switch err.(type) { if _, ok := err.(*store.NotFound); ok {
case *store.NotFound:
// If the email is a new device // If the email is a new device
if v == nil { go pushIP(inboundInfo, email, &map[string]int{ip: uid})
go pushIP(inboundInfo, email, ip, uid) } else {
}
default:
newError("cache service").Base(err).AtError().WriteToLog() newError("cache service").Base(err).AtError().WriteToLog()
} }
return false return false
} }
ipMap := *v.(*map[string]int) ipMap := v.(*map[string]int)
// Reject device reach limit directly // Reject device reach limit directly
if deviceLimit > 0 && len(ipMap) > deviceLimit { if deviceLimit > 0 && len(*ipMap) > deviceLimit {
return true return true
} }
// If the ip is not in cache // If the ip is not in cache
if _, ok := ipMap[ip]; !ok { if _, ok := (*ipMap)[ip]; !ok {
go pushIP(inboundInfo, email, ip, uid) (*ipMap)[ip] = uid
go pushIP(inboundInfo, email, ipMap)
} }
} }
@ -253,11 +256,11 @@ func globalLimit(inboundInfo *InboundInfo, email string, uid int, ip string, dev
} }
// push the ip to cache // push the ip to cache
func pushIP(inboundInfo *InboundInfo, email string, ip string, uid int) { func pushIP(inboundInfo *InboundInfo, email string, ipMap *map[string]int) {
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(inboundInfo.GlobalLimit.config.Timeout)*time.Second) ctx, cancel := context.WithTimeout(context.Background(), time.Duration(inboundInfo.GlobalLimit.config.Timeout)*time.Second)
defer cancel() defer cancel()
if err := inboundInfo.GlobalLimit.globalOnlineIP.Set(ctx, email, map[string]int{ip: uid}); err != nil { if err := inboundInfo.GlobalLimit.globalOnlineIP.Set(ctx, email, ipMap); err != nil {
newError("cache service").Base(err).AtError().WriteToLog() newError("cache service").Base(err).AtError().WriteToLog()
} }
} }