customize AES-CFB Keys

master
兔子 4 years ago
parent 164c412c24
commit bdcfcd05db

@ -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 {

@ -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)
}

@ -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())

Loading…
Cancel
Save