only record tcp connecting ip on real time

This commit is contained in:
yuzuki999 2023-05-17 00:56:14 +08:00
parent 49cdbbbaa4
commit bc14e406c9
6 changed files with 95 additions and 27 deletions

View File

@ -237,7 +237,9 @@ func (d *DefaultDispatcher) getLink(ctx context.Context, network net.Network, sn
return nil, nil, nil, newError("Get limit info error: ", err)
}
// Speed Limit and Device Limit
w, reject := limit.CheckLimit(user.Email, sessionInbound.Source.Address.IP().String())
w, reject := limit.CheckLimit(user.Email,
sessionInbound.Source.Address.IP().String(),
network == net.Network_TCP)
if reject {
newError("Limited ", user.Email, " by conn or ip").AtWarning().WriteToLog()
common.Close(outboundLink.Writer)
@ -473,7 +475,7 @@ func (d *DefaultDispatcher) routedDispatch(ctx context.Context, link *transport.
common.Interrupt(link.Reader)
return
}
} else {
} else if destination.Network == net.Network_TCP {
defer func() {
l.ConnLimiter.DelConnCount(sessionInbound.User.Email, sessionInbound.Source.Address.IP().String())
}()

14
limiter/clear.go Normal file
View File

@ -0,0 +1,14 @@
package limiter
import "log"
func ClearPacketOnlineIP() error {
log.Println("Limiter: Clear packet online ip...")
limitLock.RLock()
for _, l := range limiter {
l.ConnLimiter.ClearPacketOnlineIP()
}
limitLock.RUnlock()
log.Println("Limiter: Clear packet online ip done")
return nil
}

View File

@ -5,10 +5,11 @@ import (
)
type ConnLimiter struct {
realTime bool
ipLimit int
connLimit int
count sync.Map //map[string]int
ip sync.Map //map[string]map[string]*sync.Map
count sync.Map // map[string]int
ip sync.Map // map[string]map[string]int or map[string]struct{}
}
func NewConnLimiter(conn int, ip int) *ConnLimiter {
@ -20,29 +21,45 @@ func NewConnLimiter(conn int, ip int) *ConnLimiter {
}
}
func (c *ConnLimiter) AddConnCount(user string, ip string) (limit bool) {
func (c *ConnLimiter) AddConnCount(user string, ip string, isTcp bool) (limit bool) {
if c.connLimit != 0 {
if v, ok := c.count.Load(user); ok {
if v.(int) >= c.connLimit {
return true
} else {
} else if isTcp { // tcp protocol
c.count.Store(user, v.(int)+1)
}
} else {
} else if isTcp { // tcp protocol
c.count.Store(user, 1)
}
}
if c.ipLimit == 0 {
return false
}
// default user map
ipMap := new(sync.Map)
ipMap.Store(ip, 1)
if isTcp {
ipMap.Store(ip, 2)
} else {
ipMap.Store(ip, 1)
}
// check user online ip
if v, ok := c.ip.LoadOrStore(user, ipMap); ok {
// have online ip
// have user
ips := v.(*sync.Map)
cn := 0
if online, ok := ips.Load(ip); !ok {
ips.Range(func(key, value interface{}) bool {
if online, ok := ips.Load(ip); ok {
// online ip
if isTcp {
// count add
ips.Store(ip, online.(int)+2)
}
} else {
// not online ip for tcp
if _, ok = ips.Load(ip + "o"); ok {
return false
}
ips.Range(func(_, _ interface{}) bool {
cn++
if cn >= c.ipLimit {
limit = true
@ -53,15 +70,17 @@ func (c *ConnLimiter) AddConnCount(user string, ip string) (limit bool) {
if limit {
return
}
ips.Store(ip, 1)
} else {
// have this ip
ips.Store(ip, online.(int)+1)
if isTcp {
ips.Store(ip, 2)
} else {
ips.Store(ip, 1)
}
}
}
return false
return
}
// DelConnCount Delete tcp connection count, no tcp do not use
func (c *ConnLimiter) DelConnCount(user string, ip string) {
if c.connLimit != 0 {
if v, ok := c.count.Load(user); ok {
@ -78,10 +97,10 @@ func (c *ConnLimiter) DelConnCount(user string, ip string) {
if i, ok := c.ip.Load(user); ok {
is := i.(*sync.Map)
if i, ok := is.Load(ip); ok {
if i.(int) == 1 {
if i.(int) == 2 {
is.Delete(ip)
} else {
is.Store(user, i.(int)-1)
is.Store(user, i.(int)-2)
}
notDel := false
c.ip.Range(func(_, _ any) bool {
@ -94,3 +113,17 @@ func (c *ConnLimiter) DelConnCount(user string, ip string) {
}
}
}
// ClearPacketOnlineIP Clear udp,icmp and other packet protocol online ip
func (c *ConnLimiter) ClearPacketOnlineIP() {
c.ip.Range(func(_, v any) bool {
userIp := v.(*sync.Map)
userIp.Range(func(ip, v any) bool {
if v.(int) == 1 {
userIp.Delete(ip)
}
return true
})
return true
})
}

View File

@ -12,15 +12,24 @@ func init() {
}
func TestConnLimiter_AddConnCount(t *testing.T) {
t.Log(c.AddConnCount("1", "1"))
t.Log(c.AddConnCount("1", "2"))
t.Log(c.AddConnCount("1", "1", true))
t.Log(c.AddConnCount("1", "2", true))
}
func TestConnLimiter_DelConnCount(t *testing.T) {
t.Log(c.AddConnCount("1", "1"))
t.Log(c.AddConnCount("1", "2"))
t.Log(c.AddConnCount("1", "1", true))
t.Log(c.AddConnCount("1", "2", true))
c.DelConnCount("1", "1")
t.Log(c.AddConnCount("1", "2"))
t.Log(c.AddConnCount("1", "2", true))
}
func TestConnLimiter_ClearPacketOnlineIP(t *testing.T) {
t.Log(c.AddConnCount("1", "1", false))
t.Log(c.AddConnCount("1", "2", false))
c.ClearPacketOnlineIP()
t.Log(c.AddConnCount("1", "2", true))
c.DelConnCount("1", "2")
t.Log(c.AddConnCount("1", "1", false))
}
func BenchmarkConnLimiter(b *testing.B) {
@ -28,7 +37,7 @@ func BenchmarkConnLimiter(b *testing.B) {
for i := 0; i < b.N; i++ {
wg.Add(1)
go func() {
c.AddConnCount("1", "2")
c.AddConnCount("1", "2", true)
c.DelConnCount("1", "2")
wg.Done()
}()

View File

@ -5,6 +5,8 @@ import (
"fmt"
"github.com/Yuzuki616/V2bX/api/panel"
"github.com/juju/ratelimit"
"github.com/xtls/xray-core/common/task"
"log"
"sync"
"time"
)
@ -14,6 +16,15 @@ var limiter map[string]*Limiter
func Init() {
limiter = map[string]*Limiter{}
c := task.Periodic{
Interval: time.Minute * 2,
Execute: ClearPacketOnlineIP,
}
go func() {
log.Println("Limiter: ClearPacketOnlineIP started")
time.Sleep(time.Minute * 2)
c.Start()
}()
}
type Limiter struct {
@ -104,9 +115,9 @@ func DeleteLimiter(tag string) {
limitLock.Unlock()
}
func (l *Limiter) CheckLimit(email string, ip string) (Bucket *ratelimit.Bucket, Reject bool) {
func (l *Limiter) CheckLimit(email string, ip string, isTcp bool) (Bucket *ratelimit.Bucket, Reject bool) {
// ip and conn limiter
if l.ConnLimiter.AddConnCount(email, ip) {
if l.ConnLimiter.AddConnCount(email, ip, isTcp) {
return nil, true
}
// check and gen speed limit Bucket

View File

@ -1 +0,0 @@
package limiter