From b47954ea64b99de194b24de095803c511d91caa5 Mon Sep 17 00:00:00 2001
From: Senis <thank243@gmail.com>
Date: Fri, 4 Nov 2022 08:58:41 +0800
Subject: [PATCH] fix: minor bugs

---
 common/limiter/limiter.go | 37 +++++++++++++++++++++++++++++++++----
 1 file changed, 33 insertions(+), 4 deletions(-)

diff --git a/common/limiter/limiter.go b/common/limiter/limiter.go
index 58e5e96..e5fb1fc 100644
--- a/common/limiter/limiter.go
+++ b/common/limiter/limiter.go
@@ -2,8 +2,11 @@
 package limiter
 
 import (
+	"context"
 	"fmt"
+	"strings"
 	"sync"
+	"time"
 
 	"github.com/go-redis/redis/v8"
 	"golang.org/x/time/rate"
@@ -139,16 +142,42 @@ func (l *Limiter) GetUserBucket(tag string, email string, ip string) (limiter *r
 	if value, ok := l.InboundInfo.Load(tag); ok {
 		inboundInfo := value.(*InboundInfo)
 		nodeLimit := inboundInfo.NodeSpeedLimit
-		var userLimit uint64 = 0
-		var deviceLimit int = 0
-		var uid int = 0
+		var (
+			userLimit                           uint64 = 0
+			deviceLimit, uid, globalDeviceLimit int
+		)
+
 		if v, ok := inboundInfo.UserInfo.Load(email); ok {
 			u := v.(UserInfo)
 			uid = u.UID
 			userLimit = u.SpeedLimit
 			deviceLimit = u.DeviceLimit
 		}
-		// Report online device
+
+		// Global device limit
+		if globalDeviceLimit > 0 {
+			ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+			defer cancel()
+
+			trimEmail := strings.Split(email, "|")[1]
+			exist, err := l.r.Exists(ctx, trimEmail).Result()
+			if err != nil {
+				newError(fmt.Sprintf("Redis: %v", err)).AtError().WriteToLog()
+			} else {
+				if exist == 0 {
+					l.r.HSet(ctx, trimEmail, ip, uid)
+					l.r.Expire(ctx, trimEmail, time.Duration(l.g.expiry)*time.Minute)
+				} else {
+					l.r.HSet(ctx, trimEmail, ip, uid)
+				}
+				if l.r.HLen(ctx, trimEmail).Val() > int64(l.g.limit) {
+					l.r.HDel(ctx, trimEmail, ip)
+					return nil, false, true
+				}
+			}
+		}
+
+		// Local device limit
 		ipMap := new(sync.Map)
 		ipMap.Store(ip, uid)
 		// If any device is online