// Package notify is a package which provide common tcp/udp/unix socket service package notify import ( "context" "errors" "fmt" "math/rand" "net" "strings" "sync" "time" "b612.me/starcrypto" "b612.me/starnet" ) var aesKey = []byte{0x19, 0x96, 0x11, 0x27, 228, 187, 187, 231, 142, 137, 230, 179, 189, 229, 184, 133} func encodeFunc(data []byte) []byte { return starcrypto.AesEncryptCFB(data, aesKey) } func decodeFunc(data []byte) []byte { return starcrypto.AesDecryptCFB(data, aesKey) } // StarNotifyS 为Server端 type StarNotifyS struct { // Queue 是用来处理收发信息的简单消息队列 Queue *starnet.StarQueue // FuncLists 记录了被通知项所记录的函数 aesKey []byte FuncLists map[string]func(SMsg) string funcMu sync.Mutex defaultFunc func(SMsg) string Connected func(SMsg) nickName map[string]string stopSign context.Context cancel context.CancelFunc connPool sync.Map connMu sync.Mutex lockPool map[string]SMsg lockMu sync.Mutex udpPool map[string]*net.UDPAddr listener net.Listener isUDP bool Sync bool // 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 queue *starnet.StarQueue } func (star *StarNotifyS) SetAesKey(key []byte) { star.aesKey = key star.Queue.EncodeFunc = func(data []byte) []byte { return starcrypto.AesEncryptCFB(data, key) } star.Queue.DecodeFunc = func(data []byte) []byte { return starcrypto.AesDecryptCFB(data, key) } } func (star *StarNotifyS) GetAesKey() []byte { if len(star.aesKey) == 0 { return aesKey } return star.aesKey } func (star *StarNotifyS) getName(conn string) string { for k, v := range star.nickName { if v == conn { return k } } return "" } func (star *StarNotifyS) Stoped() <-chan struct{} { return star.stopSign.Done() } // GetConnPool 获取所有Client端信息 func (star *StarNotifyS) GetConnPool() []SMsg { var result []SMsg star.connPool.Range(func(k, val interface{}) bool { v := val.(net.Conn) result = append(result, SMsg{Conn: v, mode: "pa", nickName: star.setNickName, getName: star.getName, queue: star.Queue}) return true }) for _, v := range star.udpPool { result = append(result, SMsg{UDP: v, Uconn: star.UDPConn, mode: "pa0", nickName: star.setNickName, getName: star.getName, queue: star.Queue}) } return result } // GetClient 获取所有Client端信息 func (star *StarNotifyS) GetClient(name string) (SMsg, error) { if str, ok := star.nickName[name]; ok { if tmp, ok := star.connPool.Load(str); ok { conn := tmp.(net.Conn) return SMsg{Conn: conn, mode: "pa", nickName: star.setNickName, getName: star.getName, queue: star.Queue}, nil } if conn, ok := star.udpPool[str]; ok { return SMsg{UDP: conn, Uconn: star.UDPConn, mode: "pa0", nickName: star.setNickName, getName: star.getName, queue: star.Queue}, 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(nmsg.queue.BuildMessage([]byte(nmsg.mode + "||" + nmsg.addSlash(nmsg.Key) + "||" + msg))) } else { err = WriteToUDP(nmsg.Uconn, nmsg.UDP, nmsg.queue.BuildMessage([]byte(nmsg.mode+"||"+nmsg.addSlash(nmsg.Key)+"||"+msg))) } 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(nmsg.queue.BuildMessage([]byte("pa||" + nmsg.addSlash(key) + "||" + value))) } else { err = WriteToUDP(nmsg.Uconn, nmsg.UDP, nmsg.queue.BuildMessage([]byte("pa||"+nmsg.addSlash(key)+"||"+value))) } 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("%d%06d", time.Now().UnixNano(), rand.Intn(999999)) if source.Uconn == nil { _, err = source.Conn.Write(star.Queue.BuildMessage([]byte(mode + "||" + source.addSlash(key) + "||" + value))) } else { err = WriteToUDP(source.Uconn, source.UDP, star.Queue.BuildMessage([]byte(mode+"||"+source.addSlash(key)+"||"+value))) } if err != nil { return SMsg{}, err } if int64(tmout) > 0 { tmceed = time.After(tmout) } source.wait = make(chan int, 2) star.lockMu.Lock() star.lockPool[mode] = source star.lockMu.Unlock() select { case <-source.wait: star.lockMu.Lock() res := star.lockPool[mode] delete(star.lockPool, mode) star.lockMu.Unlock() return res, nil case <-tmceed: return SMsg{}, errors.New("Time Exceed") } } func (star *StarNotifyS) starinits() { builder := starnet.NewQueue() builder.EncodeFunc = encodeFunc builder.DecodeFunc = decodeFunc builder.Encode = true star.stopSign, star.cancel = context.WithCancel(context.Background()) star.Queue = builder star.udpPool = make(map[string]*net.UDPAddr) star.FuncLists = make(map[string]func(SMsg) string) star.nickName = make(map[string]string) star.lockPool = make(map[string]SMsg) star.Online = false star.Queue.RestoreDuration(time.Millisecond * 50) } // 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 { WriteToUDP(star.UDPConn, v, star.Queue.BuildMessage([]byte("b612ryzstop"))) star.connMu.Lock() delete(star.udpPool, k) star.connMu.Unlock() 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, 81920) 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, queue: star.Queue}) } } star.connMu.Lock() star.udpPool[addr.String()] = addr star.connMu.Unlock() } 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() star.connPool.Range(func(a, b interface{}) bool { k := a.(string) v := b.(net.Conn) v.Close() star.connPool.Delete(a) for k2, v2 := range star.nickName { if v2 == k { star.funcMu.Lock() delete(star.nickName, k2) star.funcMu.Unlock() } } return true }) 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() star.connPool.Delete(fmt.Sprint(conn)) for k, v := range star.nickName { if v == fmt.Sprint(conn) { delete(star.nickName, k) } } break } } }(conn) star.connPool.Store(fmt.Sprint(conn), conn) if star.Connected != nil { go star.Connected(SMsg{Conn: conn, nickName: star.setNickName, getName: star.getName, queue: star.Queue}) } } }() 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.Load(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.funcMu.Lock() star.nickName[name] = conn star.funcMu.Unlock() return nil } // SetNotify 用于设置通知关键词的调用函数 func (star *StarNotifyS) SetNotify(name string, data func(SMsg) string) { star.funcMu.Lock() defer star.funcMu.Unlock() if data == nil { if _, ok := star.FuncLists[name]; ok { delete(star.FuncLists, name) } return } 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 * 500) continue } mode, key, value := star.analyseData(string(data.Msg)) if mode == key && mode == value && mode == "" { continue } var rmsg SMsg if !star.isUDP { rmsg = SMsg{data.Conn.(net.Conn), key, value, nil, nil, mode, nil, star.setNickName, star.getName, star.Queue} } else { rmsg = SMsg{nil, key, value, data.Conn.(*net.UDPAddr), star.UDPConn, mode, nil, star.setNickName, star.getName, star.Queue} if key == "b612ryzstop" { star.connMu.Lock() delete(star.udpPool, rmsg.UDP.String()) star.connMu.Unlock() for k, v := range star.nickName { if v == rmsg.UDP.String() { delete(star.nickName, k) } } continue } } replyFunc := func(key string, rmsg SMsg) { 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) } } } if mode[0:2] != "sr" { if !star.Sync { go replyFunc(key, rmsg) } else { replyFunc(key, rmsg) } } else { if sa, ok := star.lockPool[mode]; ok { rmsg.wait = sa.wait star.lockMu.Lock() star.lockPool[mode] = rmsg star.lockPool[mode].wait <- 1 star.lockMu.Unlock() } else { if !star.Sync { go replyFunc(key, rmsg) } else { replyFunc(key, rmsg) } } } } } func (star *StarNotifyS) analyseData(msg string) (mode, key, value string) { slice := strings.SplitN(msg, "||", 3) if len(slice) < 3 { return "", "", "" } return slice[0], star.trim(slice[1]), slice[2] } // ServerStop 用于终止Server端运行 func (star *StarNotifyS) ServerStop() { star.cancel() }