only record tcp connecting ip on real time

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
+	})
+}