diff --git a/internal/conns/map.go b/internal/conns/map.go index f02f3c8..cc7e51a 100644 --- a/internal/conns/map.go +++ b/internal/conns/map.go @@ -3,6 +3,7 @@ package conns import ( + "github.com/iwind/TeaGo/types" "net" "sync" ) @@ -10,14 +11,14 @@ import ( var SharedMap = NewMap() type Map struct { - m map[string]map[int]net.Conn // ip => { port => Conn } + m map[string]map[string]net.Conn // ip => { network_port => Conn } locker sync.RWMutex } func NewMap() *Map { return &Map{ - m: map[string]map[int]net.Conn{}, + m: map[string]map[string]net.Conn{}, } } @@ -25,21 +26,19 @@ func (this *Map) Add(conn net.Conn) { if conn == nil { return } - tcpAddr, ok := conn.RemoteAddr().(*net.TCPAddr) + + key, ip, ok := this.connAddr(conn) if !ok { return } - var ip = tcpAddr.IP.String() - var port = tcpAddr.Port - this.locker.Lock() defer this.locker.Unlock() connMap, ok := this.m[ip] if !ok { - this.m[ip] = map[int]net.Conn{port: conn} + this.m[ip] = map[string]net.Conn{key: conn} } else { - connMap[port] = conn + connMap[key] = conn } } @@ -47,14 +46,11 @@ func (this *Map) Remove(conn net.Conn) { if conn == nil { return } - tcpAddr, ok := conn.RemoteAddr().(*net.TCPAddr) + key, ip, ok := this.connAddr(conn) if !ok { return } - var ip = tcpAddr.IP.String() - var port = tcpAddr.Port - this.locker.Lock() defer this.locker.Unlock() @@ -62,7 +58,7 @@ func (this *Map) Remove(conn net.Conn) { if !ok { return } - delete(connMap, port) + delete(connMap, key) if len(connMap) == 0 { delete(this.m, ip) @@ -121,3 +117,24 @@ func (this *Map) AllConns() []net.Conn { return result } + +func (this *Map) connAddr(conn net.Conn) (key string, ip string, ok bool) { + if conn == nil { + return + } + + var addr = conn.RemoteAddr() + switch realAddr := addr.(type) { + case *net.TCPAddr: + return addr.Network() + types.String(realAddr.Port), realAddr.IP.String(), true + case *net.UDPAddr: + return addr.Network() + types.String(realAddr.Port), realAddr.IP.String(), true + default: + var s = addr.String() + host, port, err := net.SplitHostPort(s) + if err != nil { + return + } + return addr.Network() + port, host, true + } +} diff --git a/internal/nodes/node.go b/internal/nodes/node.go index ac4e3fe..f7b3605 100644 --- a/internal/nodes/node.go +++ b/internal/nodes/node.go @@ -693,6 +693,7 @@ func (this *Node) listenSock() error { var lastReadAt int64 var lastWriteAt int64 var lastErrString = "" + var protocol = "tcp" clientConn, ok := conn.(*ClientConn) if ok { createdAt = clientConn.CreatedAt() @@ -703,6 +704,8 @@ func (this *Node) listenSock() error { if lastErr != nil { lastErrString = lastErr.Error() } + } else { + protocol = "udp" } var age int64 = -1 var lastReadAge int64 = -1 @@ -719,6 +722,7 @@ func (this *Node) listenSock() error { } connMaps = append(connMaps, maps.Map{ + "protocol": protocol, "addr": conn.RemoteAddr().String(), "age": age, "readAge": lastReadAge,