// Package notify is a package which provide common tcp/udp/unix socket service package notify import ( "context" "errors" "fmt" "math/rand" "net" "strings" "time" "b612.me/starainrt" ) var builder *starainrt.StarQueue func init() { builder = starainrt.NewQueue() } // StarNotifyS 为Server端 type StarNotifyS struct { // Queue 是用来处理收发信息的简单消息队列 Queue *starainrt.StarQueue // FuncLists 记录了被通知项所记录的函数 FuncLists map[string]func(SMsg) string defaultFunc 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 // UDPConn UDP监听 UDPConn *net.UDPConn // Online 当前链接是否处于活跃状态 Online bool // ReadDeadline tcp/unix中读超时设置,udp请直接调用UDPConn ReadDeadline time.Time // WriteDeadline tcp/unix中写超时设置,udp请直接调用UDPConn WriteDeadline time.Time // Deadline tcp/unix中超时设置,udp请直接调用UDPConn Deadline time.Time } // SMsg 指明当前服务端被通知的关键字 type SMsg struct { 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", nickName: star.setNickName, getName: star.getName}) } for _, v := range star.udpPool { 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) { if v == byte(124) || v == byte(92) { key = append(key, byte(92)) } key = append(key, v) } 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 if nmsg.uconn == nil { _, err = nmsg.Conn.Write(builder.BuildMessage([]byte(nmsg.mode + "||" + nmsg.addSlash(nmsg.Key) + "||" + msg))) } else { _, err = nmsg.uconn.WriteToUDP(builder.BuildMessage([]byte(nmsg.mode+"||"+nmsg.addSlash(nmsg.Key)+"||"+msg)), nmsg.UDP) } return err } // Send 用于向client端发送key-value数据 func (nmsg *SMsg) Send(key, value string) error { var err error if nmsg.uconn == nil { _, err = nmsg.Conn.Write(builder.BuildMessage([]byte("pa||" + nmsg.addSlash(key) + "||" + value))) } else { _, err = nmsg.uconn.WriteToUDP(builder.BuildMessage([]byte("pa||"+nmsg.addSlash(key)+"||"+value)), nmsg.UDP) } 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 var tmceed <-chan time.Time rand.Seed(time.Now().UnixNano()) mode := "sr" + fmt.Sprintf("%05d", rand.Intn(99999)) if source.uconn == nil { _, err = source.Conn.Write(builder.BuildMessage([]byte(mode + "||" + source.addSlash(key) + "||" + value))) } else { _, err = source.uconn.WriteToUDP(builder.BuildMessage([]byte(mode+"||"+source.addSlash(key)+"||"+value)), source.UDP) } if err != nil { return SMsg{}, err } if int64(tmout) > 0 { tmceed = time.After(tmout) } source.wait = make(chan int, 2) star.lockPool[mode] = source select { case <-source.wait: res := star.lockPool[mode] delete(star.lockPool, mode) return res, nil case <-tmceed: return SMsg{}, errors.New("Time Exceed") } } func (star *StarNotifyS) starinits() { star.stopSign, star.cancel = context.WithCancel(context.Background()) star.Queue = starainrt.NewQueue() 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 star.Queue.RestoreDuration(time.Second * 2) } // NewNotifyS 开启一个新的Server端通知 func NewNotifyS(netype, value string) (*StarNotifyS, error) { if netype[0:3] != "udp" { return notudps(netype, value) } return doudps(netype, value) } func doudps(netype, value string) (*StarNotifyS, error) { var star StarNotifyS star.starinits() star.isUDP = true udpaddr, err := net.ResolveUDPAddr(netype, value) if err != nil { return nil, err } star.UDPConn, err = net.ListenUDP(netype, udpaddr) if err != nil { return nil, err } go star.notify() go func() { <-star.stopSign.Done() for k, v := range star.udpPool { star.UDPConn.WriteToUDP(star.Queue.BuildMessage([]byte("b612ryzstop")), v) delete(star.udpPool, k) for k2, v2 := range star.nickName { if v2 == k { delete(star.nickName, k2) } } } star.UDPConn.Close() star.Online = false return }() go func() { for { buf := make([]byte, 8192) n, addr, err := star.UDPConn.ReadFromUDP(buf) if n != 0 { 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, nickName: star.setNickName, getName: star.getName}) } } star.udpPool[addr.String()] = addr } if err != nil { continue } } }() star.Online = true return &star, nil } func notudps(netype, value string) (*StarNotifyS, error) { var err error var star StarNotifyS star.starinits() star.isUDP = false star.listener, err = net.Listen(netype, value) if err != nil { return nil, err } go star.notify() go func() { <-star.stopSign.Done() 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) } } } star.listener.Close() star.Online = false return }() go func() { for { conn, err := star.listener.Accept() if err != nil { select { case <-star.stopSign.Done(): star.listener.Close() return default: continue } } if !star.ReadDeadline.IsZero() { conn.SetReadDeadline(star.ReadDeadline) } if !star.WriteDeadline.IsZero() { conn.SetWriteDeadline(star.WriteDeadline) } if !star.Deadline.IsZero() { conn.SetDeadline(star.Deadline) } go func(conn net.Conn) { for { buf := make([]byte, 8192) n, err := conn.Read(buf) if n != 0 { star.Queue.ParseMessage(buf[0:n], conn) } if err != nil { conn.Close() 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[fmt.Sprint(conn)] = conn if star.Connected != nil { go star.Connected(SMsg{Conn: conn, nickName: star.setNickName, getName: star.getName}) } } }() star.Online = true 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 } // SetDefaultNotify 用于设置默认关键词的调用函数 func (star *StarNotifyS) SetDefaultNotify(data func(SMsg) string) { star.defaultFunc = data } func (star *StarNotifyS) trim(name string) string { var slash bool = false var key []byte for _, v := range []byte(name) { if v == byte(92) && !slash { slash = true continue } slash = false key = append(key, v) } return string(key) } func (star *StarNotifyS) notify() { for { select { case <-star.stopSign.Done(): return default: } data, err := star.Queue.RestoreOne() if err != nil { time.Sleep(time.Millisecond * 20) continue } 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, star.setNickName, star.getName} } else { 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 } } if mode[0:2] != "sr" { go func() { if msg, ok := star.FuncLists[key]; ok { sdata := msg(rmsg) if sdata == "" { return } rmsg.Reply(sdata) } else { if star.defaultFunc != nil { sdata := star.defaultFunc(rmsg) if sdata == "" { return } rmsg.Reply(sdata) } } }() } else { if sa, ok := star.lockPool[mode]; ok { rmsg.wait = sa.wait star.lockPool[mode] = rmsg star.lockPool[mode].wait <- 1 } else { go func() { if msg, ok := star.FuncLists[key]; ok { sdata := msg(rmsg) if sdata == "" { return } rmsg.Reply(sdata) } else { if star.defaultFunc != nil { sdata := star.defaultFunc(rmsg) if sdata == "" { return } rmsg.Reply(sdata) } } }() } } } } func (star *StarNotifyS) analyseData(msg string) (mode, key, value string) { slice := strings.SplitN(msg, "||", 3) return slice[0], star.trim(slice[1]), slice[2] } // ServerStop 用于终止Server端运行 func (star *StarNotifyS) ServerStop() { star.cancel() star.Stop <- 1 star.Stop <- 1 star.Stop <- 1 }