diff --git a/client.go b/client.go index ce3b3bf..7d3de6e 100644 --- a/client.go +++ b/client.go @@ -10,6 +10,7 @@ import ( "sync" "time" + "b612.me/starcrypto" "b612.me/starnet" ) @@ -33,6 +34,7 @@ type StarNotifyC struct { // Online 当前链接是否处于活跃状态 Online bool lockPool map[string]CMsg + aesKey []byte } // CMsg 指明当前客户端被通知的关键字 @@ -44,11 +46,12 @@ type CMsg struct { } func (star *StarNotifyC) starinitc() { + builder := starnet.NewQueue() + builder.EncodeFunc = encodeFunc + builder.DecodeFunc = decodeFunc + builder.Encode = true star.stopSign, star.cancel = context.WithCancel(context.Background()) - star.Queue = starnet.NewQueue() - star.Queue.EncodeFunc = encodeFunc - star.Queue.DecodeFunc = decodeFunc - star.Queue.Encode = true + star.Queue = builder star.FuncLists = make(map[string]func(CMsg)) star.UseChannel = false star.clientSign = make(map[string]chan string) @@ -57,6 +60,23 @@ func (star *StarNotifyC) starinitc() { star.Queue.RestoreDuration(time.Millisecond * 50) } +func (star *StarNotifyC) 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 *StarNotifyC) GetAesKey() []byte { + if len(star.aesKey) == 0 { + return aesKey + } + return star.aesKey +} + // Notify 用于获取一个通知 func (star *StarNotifyC) Notify(key string) chan string { if _, ok := star.clientSign[key]; !ok { diff --git a/client_test.go b/client_test.go index b21e0b7..99c16c1 100644 --- a/client_test.go +++ b/client_test.go @@ -34,7 +34,7 @@ func Test_usechannel(t *testing.T) { txt = <-client.Notify("nihao") fmt.Println("client", txt) server.ServerStop() - <-client.Stop + <-client.Stoped() client.ClientStop() time.Sleep(time.Second * 3) } @@ -70,7 +70,7 @@ func Test_nochannel(t *testing.T) { client.SendValue("nihao", "lalala") time.Sleep(time.Second * 3) server.ServerStop() - <-client.Stop + <-client.Stoped() client.ClientStop() time.Sleep(time.Second * 3) } @@ -81,6 +81,7 @@ func Test_pipec(t *testing.T) { fmt.Println(err) return } + server.SetAesKey([]byte("abcdefg123456789")) server.SetNotify("ni\\||hao", func(data SMsg) string { fmt.Println("name-get", data.GetName()) fmt.Println("name-set", data.SetName("iiiis")) @@ -97,6 +98,7 @@ func Test_pipec(t *testing.T) { fmt.Println(err) return } + client.SetAesKey([]byte("abcdefg123456789")) client.UseChannel = false sa, err := client.SendValueWait("ni\\||hao", "lalaeee", time.Second*10) if err != nil { @@ -113,7 +115,7 @@ func Test_pipec(t *testing.T) { fmt.Println("sukidesu") time.Sleep(time.Second * 3) server.ServerStop() - <-client.Stop + <-client.Stoped() client.ClientStop() time.Sleep(time.Second * 2) } @@ -153,7 +155,7 @@ func Test_pips(t *testing.T) { fmt.Println(server.SendWait(testmsg, "nihao", "wozuinb", time.Second*20)) fmt.Println("sakura") server.ServerStop() - <-client.Stop + <-client.Stoped() client.ClientStop() time.Sleep(time.Second * 3) } diff --git a/server.go b/server.go index 3bea8a7..55bd6c9 100644 --- a/server.go +++ b/server.go @@ -16,7 +16,6 @@ import ( "b612.me/starnet" ) -var builder *starnet.StarQueue var aesKey = []byte{0x19, 0x96, 0x11, 0x27, 228, 187, 187, 231, 142, 137, 230, 179, 189, 229, 184, 133} func encodeFunc(data []byte) []byte { @@ -27,18 +26,12 @@ func decodeFunc(data []byte) []byte { return starcrypto.AesDecryptCFB(data, aesKey) } -func init() { - builder = starnet.NewQueue() - builder.EncodeFunc = encodeFunc - builder.DecodeFunc = decodeFunc - builder.Encode = true -} - // 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 @@ -78,6 +71,24 @@ type SMsg struct { 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 { @@ -97,11 +108,11 @@ 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}) + 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}) + result = append(result, SMsg{UDP: v, uconn: star.UDPConn, mode: "pa0", nickName: star.setNickName, getName: star.getName, queue: star.Queue}) } return result } @@ -111,10 +122,10 @@ 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}, nil + 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}, nil + 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") @@ -157,9 +168,9 @@ func (nmsg *SMsg) ReplyRaw(msg interface{}) error { 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))) + _, err = nmsg.Conn.Write(nmsg.queue.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) + _, err = nmsg.uconn.WriteToUDP(nmsg.queue.BuildMessage([]byte(nmsg.mode+"||"+nmsg.addSlash(nmsg.Key)+"||"+msg)), nmsg.UDP) } return err } @@ -168,9 +179,9 @@ func (nmsg *SMsg) Reply(msg string) error { 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))) + _, err = nmsg.Conn.Write(nmsg.queue.BuildMessage([]byte("pa||" + nmsg.addSlash(key) + "||" + value))) } else { - _, err = nmsg.uconn.WriteToUDP(builder.BuildMessage([]byte("pa||"+nmsg.addSlash(key)+"||"+value)), nmsg.UDP) + _, err = nmsg.uconn.WriteToUDP(nmsg.queue.BuildMessage([]byte("pa||"+nmsg.addSlash(key)+"||"+value)), nmsg.UDP) } return err } @@ -198,9 +209,9 @@ func (star *StarNotifyS) SendWait(source SMsg, key, value string, tmout time.Dur 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(builder.BuildMessage([]byte(mode + "||" + source.addSlash(key) + "||" + value))) + _, err = source.Conn.Write(star.Queue.BuildMessage([]byte(mode + "||" + source.addSlash(key) + "||" + value))) } else { - _, err = source.uconn.WriteToUDP(builder.BuildMessage([]byte(mode+"||"+source.addSlash(key)+"||"+value)), source.UDP) + _, err = source.uconn.WriteToUDP(star.Queue.BuildMessage([]byte(mode+"||"+source.addSlash(key)+"||"+value)), source.UDP) } if err != nil { return SMsg{}, err @@ -225,11 +236,12 @@ func (star *StarNotifyS) SendWait(source SMsg, key, value string, tmout time.Dur } 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 = starnet.NewQueue() - star.Queue.EncodeFunc = encodeFunc - star.Queue.DecodeFunc = decodeFunc - star.Queue.Encode = true + star.Queue = builder star.udpPool = make(map[string]*net.UDPAddr) star.FuncLists = make(map[string]func(SMsg) string) star.nickName = make(map[string]string) @@ -284,7 +296,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, nickName: star.setNickName, getName: star.getName}) + go star.Connected(SMsg{UDP: addr, uconn: star.UDPConn, nickName: star.setNickName, getName: star.getName, queue: star.Queue}) } } star.connMu.Lock() @@ -372,7 +384,7 @@ func notudps(netype, value string) (*StarNotifyS, error) { }(conn) star.connPool.Store(fmt.Sprint(conn), conn) if star.Connected != nil { - go star.Connected(SMsg{Conn: conn, nickName: star.setNickName, getName: star.getName}) + go star.Connected(SMsg{Conn: conn, nickName: star.setNickName, getName: star.getName, queue: star.Queue}) } } }() @@ -452,9 +464,9 @@ func (star *StarNotifyS) notify() { } var rmsg SMsg if !star.isUDP { - rmsg = SMsg{data.Conn.(net.Conn), key, value, nil, nil, mode, nil, star.setNickName, star.getName} + 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} + 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())