update new function

server.go

@@ -26,12 +26,14 @@ type StarNotifyS struct {
 	// FuncLists 记录了被通知项所记录的函数
 	FuncLists   map[string]func(SMsg) string
 	defaultFunc func(SMsg) string
-	Connected   func(SMsg) string
+	Connected   func(SMsg)
+	nickName    map[string]string
 	stopSign    context.Context
 	cancel      context.CancelFunc
 	connPool    map[string]net.Conn
 	lockPool    map[string]SMsg
 	udpPool     map[string]*net.UDPAddr
+	listener    net.Listener
 	isUDP       bool
 	// Stop 停止信 号
 	Stop chan int
@@ -50,27 +52,65 @@ type StarNotifyS struct {
 
 // SMsg 指明当前服务端被通知的关键字
 type SMsg struct {
-	Conn  net.Conn
-	Key   string
-	Value string
-	UDP   *net.UDPAddr
-	uconn *net.UDPConn
-	mode  string
-	wait  chan int
+	Conn     net.Conn
+	Key      string
+	Value    string
+	UDP      *net.UDPAddr
+	uconn    *net.UDPConn
+	mode     string
+	wait     chan int
+	nickName func(string, string) error
+	getName  func(string) string
+}
+
+func (star *StarNotifyS) getName(conn string) string {
+	for k, v := range star.nickName {
+		if v == conn {
+			return k
+		}
+	}
+	return ""
 }
 
 // GetConnPool 获取所有Client端信息
 func (star *StarNotifyS) GetConnPool() []SMsg {
 	var result []SMsg
 	for _, v := range star.connPool {
-		result = append(result, SMsg{Conn: v, mode: "pa"})
+		result = append(result, SMsg{Conn: v, mode: "pa", nickName: star.setNickName, getName: star.getName})
 	}
 	for _, v := range star.udpPool {
-		result = append(result, SMsg{UDP: v, uconn: star.UDPConn, mode: "pa0"})
+		result = append(result, SMsg{UDP: v, uconn: star.UDPConn, mode: "pa0", nickName: star.setNickName, getName: star.getName})
 	}
 	return result
 }
 
+// GetConnPool 获取所有Client端信息
+func (star *StarNotifyS) GetClient(name string) (SMsg, error) {
+	if str, ok := star.nickName[name]; ok {
+		if conn, ok := star.connPool[str]; ok {
+			return SMsg{Conn: conn, mode: "pa", nickName: star.setNickName, getName: star.getName}, nil
+		}
+		if conn, ok := star.udpPool[str]; ok {
+			return SMsg{UDP: conn, uconn: star.UDPConn, mode: "pa0", nickName: star.setNickName, getName: star.getName}, nil
+		}
+	}
+	return SMsg{}, errors.New("Not Found")
+}
+
+func (nmsg *SMsg) GetName() string {
+	if nmsg.uconn != nil {
+		return nmsg.getName(nmsg.UDP.String())
+	}
+	return nmsg.getName(fmt.Sprint(nmsg.Conn))
+}
+
+func (nmsg *SMsg) SetName(name string) error {
+	if nmsg.uconn != nil {
+		return nmsg.nickName(name, nmsg.UDP.String())
+	}
+	return nmsg.nickName(name, fmt.Sprint(nmsg.Conn))
+}
+
 func (nmsg *SMsg) addSlash(name string) string {
 	var key []byte
 	for _, v := range []byte(name) {
@@ -82,6 +122,14 @@ func (nmsg *SMsg) addSlash(name string) string {
 	return string(key)
 }
 
+func (nmsg *SMsg) ReplyRaw(msg interface{}) error {
+	encodeData, err := encode(msg)
+	if err != nil {
+		return err
+	}
+	return nmsg.Reply(string(encodeData))
+}
+
 // Reply 用于向client端回复数据
 func (nmsg *SMsg) Reply(msg string) error {
 	var err error
@@ -104,6 +152,22 @@ func (nmsg *SMsg) Send(key, value string) error {
 	return err
 }
 
+func (nmsg *SMsg) SendRaw(key string, msg interface{}) error {
+	encodeData, err := encode(msg)
+	if err != nil {
+		return err
+	}
+	return nmsg.Send(key, string(encodeData))
+}
+
+func (star *StarNotifyS) SendWaitRaw(source SMsg, key string, msg interface{}, tmout time.Duration) (SMsg, error) {
+	encodeData, err := encode(msg)
+	if err != nil {
+		return SMsg{}, err
+	}
+	return star.SendWait(source, key, string(encodeData), tmout)
+}
+
 // SendWait 用于向client端发送key-value数据，并等待
 func (star *StarNotifyS) SendWait(source SMsg, key, value string, tmout time.Duration) (SMsg, error) {
 	var err error
@@ -139,6 +203,7 @@ func (star *StarNotifyS) starinits() {
 	star.udpPool = make(map[string]*net.UDPAddr)
 	star.FuncLists = make(map[string]func(SMsg) string)
 	star.connPool = make(map[string]net.Conn)
+	star.nickName = make(map[string]string)
 	star.lockPool = make(map[string]SMsg)
 	star.Stop = make(chan int, 5)
 	star.Online = false
@@ -170,7 +235,12 @@ func doudps(netype, value string) (*StarNotifyS, error) {
 		<-star.stopSign.Done()
 		for k, v := range star.udpPool {
 			star.UDPConn.WriteToUDP(star.Queue.BuildMessage([]byte("b612ryzstop")), v)
-			delete(star.connPool, k)
+			delete(star.udpPool, k)
+			for k2, v2 := range star.nickName {
+				if v2 == k {
+					delete(star.nickName, k2)
+				}
+			}
 		}
 		star.UDPConn.Close()
 		star.Online = false
@@ -184,7 +254,7 @@ func doudps(netype, value string) (*StarNotifyS, error) {
 				star.Queue.ParseMessage(buf[0:n], addr)
 				if _, ok := star.udpPool[addr.String()]; !ok {
 					if star.Connected != nil {
-						go star.Connected(SMsg{UDP: addr, uconn: star.UDPConn})
+						go star.Connected(SMsg{UDP: addr, uconn: star.UDPConn, nickName: star.setNickName, getName: star.getName})
 					}
 				}
 				star.udpPool[addr.String()] = addr
@@ -199,10 +269,11 @@ func doudps(netype, value string) (*StarNotifyS, error) {
 }
 
 func notudps(netype, value string) (*StarNotifyS, error) {
+	var err error
 	var star StarNotifyS
 	star.starinits()
 	star.isUDP = false
-	listener, err := net.Listen(netype, value)
+	star.listener, err = net.Listen(netype, value)
 	if err != nil {
 		return nil, err
 	}
@@ -212,18 +283,23 @@ func notudps(netype, value string) (*StarNotifyS, error) {
 		for k, v := range star.connPool {
 			v.Close()
 			delete(star.connPool, k)
+			for k2, v2 := range star.nickName {
+				if v2 == k {
+					delete(star.nickName, k2)
+				}
+			}
 		}
-		listener.Close()
+		star.listener.Close()
 		star.Online = false
 		return
 	}()
 	go func() {
 		for {
-			conn, err := listener.Accept()
+			conn, err := star.listener.Accept()
 			if err != nil {
 				select {
 				case <-star.stopSign.Done():
-					listener.Close()
+					star.listener.Close()
 					return
 				default:
 					continue
@@ -247,14 +323,19 @@ func notudps(netype, value string) (*StarNotifyS, error) {
 					}
 					if err != nil {
 						conn.Close()
-						delete(star.connPool, conn.RemoteAddr().String())
+						delete(star.connPool, fmt.Sprint(conn))
+						for k, v := range star.nickName {
+							if v == fmt.Sprint(conn) {
+								delete(star.nickName, k)
+							}
+						}
 						break
 					}
 				}
 			}(conn)
-			star.connPool[conn.RemoteAddr().String()] = conn
+			star.connPool[fmt.Sprint(conn)] = conn
 			if star.Connected != nil {
-				go star.Connected(SMsg{Conn: conn})
+				go star.Connected(SMsg{Conn: conn, nickName: star.setNickName, getName: star.getName})
 			}
 		}
 	}()
@@ -262,6 +343,26 @@ func notudps(netype, value string) (*StarNotifyS, error) {
 	return &star, nil
 }
 
+func (star *StarNotifyS) GetListenerInfo() net.Listener {
+	return star.listener
+}
+
+// SetNotify 用于设置通知关键词的调用函数
+func (star *StarNotifyS) setNickName(name string, conn string) error {
+	if _, ok := star.connPool[conn]; !ok {
+		if _, ok := star.udpPool[conn]; !ok {
+			return errors.New("Conn Not Found")
+		}
+	}
+	for k, v := range star.nickName {
+		if v == conn {
+			delete(star.nickName, k)
+		}
+	}
+	star.nickName[name] = conn
+	return nil
+}
+
 // SetNotify 用于设置通知关键词的调用函数
 func (star *StarNotifyS) SetNotify(name string, data func(SMsg) string) {
 	star.FuncLists[name] = data
@@ -301,11 +402,16 @@ func (star *StarNotifyS) notify() {
 		mode, key, value := star.analyseData(string(data.Msg))
 		var rmsg SMsg
 		if !star.isUDP {
-			rmsg = SMsg{data.Conn.(net.Conn), key, value, nil, nil, mode, nil}
+			rmsg = SMsg{data.Conn.(net.Conn), key, value, nil, nil, mode, nil, star.setNickName, star.getName}
 		} else {
-			rmsg = SMsg{nil, key, value, data.Conn.(*net.UDPAddr), star.UDPConn, mode, nil}
+			rmsg = SMsg{nil, key, value, data.Conn.(*net.UDPAddr), star.UDPConn, mode, nil, star.setNickName, star.getName}
 			if key == "b612ryzstop" {
 				delete(star.udpPool, rmsg.UDP.String())
+				for k, v := range star.nickName {
+					if v == rmsg.UDP.String() {
+						delete(star.nickName, k)
+					}
+				}
 				continue
 			}
 		}