From bc14e406c9895634190e3571e74fde7aa6808262 Mon Sep 17 00:00:00 2001 From: yuzuki999 Date: Wed, 17 May 2023 00:56:14 +0800 Subject: [PATCH] only record tcp connecting ip on real time --- core/app/dispatcher/default.go | 6 ++-- limiter/clear.go | 14 ++++++++ limiter/conn.go | 65 +++++++++++++++++++++++++--------- limiter/conn_test.go | 21 +++++++---- limiter/limiter.go | 15 ++++++-- limiter/task.go | 1 - 6 files changed, 95 insertions(+), 27 deletions(-) create mode 100644 limiter/clear.go delete mode 100644 limiter/task.go diff --git a/core/app/dispatcher/default.go b/core/app/dispatcher/default.go index 7ef40c5..294bdb5 100644 --- a/core/app/dispatcher/default.go +++ b/core/app/dispatcher/default.go @@ -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()) }() diff --git a/limiter/clear.go b/limiter/clear.go new file mode 100644 index 0000000..14a49c7 --- /dev/null +++ b/limiter/clear.go @@ -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 +} diff --git a/limiter/conn.go b/limiter/conn.go index 4c853fe..3d3cb6a 100644 --- a/limiter/conn.go +++ b/limiter/conn.go @@ -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 + }) +} diff --git a/limiter/conn_test.go b/limiter/conn_test.go index c685910..ad3bacd 100644 --- a/limiter/conn_test.go +++ b/limiter/conn_test.go @@ -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() }() diff --git a/limiter/limiter.go b/limiter/limiter.go index df3ba34..5dd55a4 100644 --- a/limiter/limiter.go +++ b/limiter/limiter.go @@ -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 diff --git a/limiter/task.go b/limiter/task.go deleted file mode 100644 index 6a935d2..0000000 --- a/limiter/task.go +++ /dev/null @@ -1 +0,0 @@ -package limiter