535 lines
13 KiB
Go
535 lines
13 KiB
Go
// 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()
|
||
}
|