-
-
Notifications
You must be signed in to change notification settings - Fork 11
/
limiter.go
197 lines (178 loc) · 5.13 KB
/
limiter.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
// Package limiter is to control the links that go into the dispatcher
package limiter
import (
"fmt"
"golang.org/x/time/rate"
"sync"
"github.com/SSPanel-NeXT/NeXT-Server/api"
)
type UserInfo struct {
UID int
SpeedLimit uint64
DeviceLimit int
}
type InboundInfo struct {
Tag string
NodeSpeedLimit uint64
UserInfo *sync.Map // Key: Email value: UserInfo
BucketHub *sync.Map // key: Email, value: *rate.Limiter
UserOnlineIP *sync.Map // Key: Email, value: {Key: IP, value: UID}
}
type Limiter struct {
InboundInfo *sync.Map // Key: Tag, Value: *InboundInfo
}
func New() *Limiter {
return &Limiter{
InboundInfo: new(sync.Map),
}
}
func (l *Limiter) AddInboundLimiter(tag string, nodeSpeedLimit uint64, userList *[]api.UserInfo) error {
inboundInfo := &InboundInfo{
Tag: tag,
NodeSpeedLimit: nodeSpeedLimit,
BucketHub: new(sync.Map),
UserOnlineIP: new(sync.Map),
}
userMap := new(sync.Map)
for _, u := range *userList {
userMap.Store(fmt.Sprintf("%s|%s|%d", tag, u.Email, u.UID), UserInfo{
UID: u.UID,
SpeedLimit: u.SpeedLimit,
DeviceLimit: u.DeviceLimit,
})
}
inboundInfo.UserInfo = userMap
l.InboundInfo.Store(tag, inboundInfo) // Replace the old inbound info
return nil
}
func (l *Limiter) UpdateInboundLimiter(tag string, updatedUserList *[]api.UserInfo) error {
if value, ok := l.InboundInfo.Load(tag); ok {
inboundInfo := value.(*InboundInfo)
// Update User info
for _, u := range *updatedUserList {
inboundInfo.UserInfo.Store(fmt.Sprintf("%s|%s|%d", tag, u.Email, u.UID), UserInfo{
UID: u.UID,
SpeedLimit: u.SpeedLimit,
DeviceLimit: u.DeviceLimit,
})
// Update old limiter bucket
limit := determineRate(inboundInfo.NodeSpeedLimit, u.SpeedLimit)
if limit > 0 {
if bucket, ok := inboundInfo.BucketHub.Load(fmt.Sprintf("%s|%s|%d", tag, u.Email, u.UID)); ok {
limiter := bucket.(*rate.Limiter)
limiter.SetLimit(rate.Limit(limit))
limiter.SetBurst(int(limit))
}
} else {
inboundInfo.BucketHub.Delete(fmt.Sprintf("%s|%s|%d", tag, u.Email, u.UID))
}
}
} else {
return fmt.Errorf("no such inbound in limiter: %s", tag)
}
return nil
}
func (l *Limiter) DeleteInboundLimiter(tag string) error {
l.InboundInfo.Delete(tag)
return nil
}
func (l *Limiter) GetOnlineDevice(tag string) (*[]api.OnlineUser, error) {
var onlineUser []api.OnlineUser
if value, ok := l.InboundInfo.Load(tag); ok {
inboundInfo := value.(*InboundInfo)
// Clear Speed Limiter bucket for users who are not online
inboundInfo.BucketHub.Range(func(key, value interface{}) bool {
email := key.(string)
if _, exists := inboundInfo.UserOnlineIP.Load(email); !exists {
inboundInfo.BucketHub.Delete(email)
}
return true
})
inboundInfo.UserOnlineIP.Range(func(key, value interface{}) bool {
email := key.(string)
ipMap := value.(*sync.Map)
ipMap.Range(func(key, value interface{}) bool {
uid := value.(int)
ip := key.(string)
onlineUser = append(onlineUser, api.OnlineUser{UID: uid, IP: ip})
return true
})
inboundInfo.UserOnlineIP.Delete(email) // Reset online device
return true
})
} else {
return nil, fmt.Errorf("no such inbound in limiter: %s", tag)
}
return &onlineUser, nil
}
func (l *Limiter) GetUserBucket(tag string, email string, ip string) (limiter *rate.Limiter, SpeedLimit bool, Reject bool) {
if value, ok := l.InboundInfo.Load(tag); ok {
var (
userLimit uint64 = 0
deviceLimit, uid int
)
inboundInfo := value.(*InboundInfo)
nodeLimit := inboundInfo.NodeSpeedLimit
if v, ok := inboundInfo.UserInfo.Load(email); ok {
u := v.(UserInfo)
uid = u.UID
userLimit = u.SpeedLimit
deviceLimit = u.DeviceLimit
}
// Local device limit
ipMap := new(sync.Map)
ipMap.Store(ip, uid)
// If any device is online
if v, ok := inboundInfo.UserOnlineIP.LoadOrStore(email, ipMap); ok {
ipMap := v.(*sync.Map)
// If this is a new ip
if _, ok := ipMap.LoadOrStore(ip, uid); !ok {
counter := 0
ipMap.Range(func(key, value interface{}) bool {
counter++
return true
})
if counter > deviceLimit && deviceLimit > 0 {
ipMap.Delete(ip)
return nil, false, true
}
}
}
// Speed limit
limit := determineRate(nodeLimit, userLimit) // Determine the speed limit rate
if limit > 0 {
limiter := rate.NewLimiter(rate.Limit(limit), int(limit)) // Byte/s
if v, ok := inboundInfo.BucketHub.LoadOrStore(email, limiter); ok {
bucket := v.(*rate.Limiter)
return bucket, true, false
} else {
return limiter, true, false
}
} else {
return nil, false, false
}
} else {
newError("Get Inbound Limiter information failed").AtDebug().WriteToLog()
return nil, false, false
}
}
// determineRate returns the minimum non-zero rate
func determineRate(nodeLimit, userLimit uint64) (limit uint64) {
if nodeLimit == 0 || userLimit == 0 {
if nodeLimit > userLimit {
return nodeLimit
} else if nodeLimit < userLimit {
return userLimit
} else {
return 0
}
} else {
if nodeLimit > userLimit {
return userLimit
} else if nodeLimit < userLimit {
return nodeLimit
} else {
return nodeLimit
}
}
}