customize AES-CFB Keys

master
兔子 4 years ago
parent 164c412c24
commit bdcfcd05db

@ -10,6 +10,7 @@ import (
"sync" "sync"
"time" "time"
"b612.me/starcrypto"
"b612.me/starnet" "b612.me/starnet"
) )
@ -33,6 +34,7 @@ type StarNotifyC struct {
// Online 当前链接是否处于活跃状态 // Online 当前链接是否处于活跃状态
Online bool Online bool
lockPool map[string]CMsg lockPool map[string]CMsg
aesKey []byte
} }
// CMsg 指明当前客户端被通知的关键字 // CMsg 指明当前客户端被通知的关键字
@ -44,11 +46,12 @@ type CMsg struct {
} }
func (star *StarNotifyC) starinitc() { 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.stopSign, star.cancel = context.WithCancel(context.Background())
star.Queue = starnet.NewQueue() star.Queue = builder
star.Queue.EncodeFunc = encodeFunc
star.Queue.DecodeFunc = decodeFunc
star.Queue.Encode = true
star.FuncLists = make(map[string]func(CMsg)) star.FuncLists = make(map[string]func(CMsg))
star.UseChannel = false star.UseChannel = false
star.clientSign = make(map[string]chan string) star.clientSign = make(map[string]chan string)
@ -57,6 +60,23 @@ func (star *StarNotifyC) starinitc() {
star.Queue.RestoreDuration(time.Millisecond * 50) 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 用于获取一个通知 // Notify 用于获取一个通知
func (star *StarNotifyC) Notify(key string) chan string { func (star *StarNotifyC) Notify(key string) chan string {
if _, ok := star.clientSign[key]; !ok { if _, ok := star.clientSign[key]; !ok {

@ -34,7 +34,7 @@ func Test_usechannel(t *testing.T) {
txt = <-client.Notify("nihao") txt = <-client.Notify("nihao")
fmt.Println("client", txt) fmt.Println("client", txt)
server.ServerStop() server.ServerStop()
<-client.Stop <-client.Stoped()
client.ClientStop() client.ClientStop()
time.Sleep(time.Second * 3) time.Sleep(time.Second * 3)
} }
@ -70,7 +70,7 @@ func Test_nochannel(t *testing.T) {
client.SendValue("nihao", "lalala") client.SendValue("nihao", "lalala")
time.Sleep(time.Second * 3) time.Sleep(time.Second * 3)
server.ServerStop() server.ServerStop()
<-client.Stop <-client.Stoped()
client.ClientStop() client.ClientStop()
time.Sleep(time.Second * 3) time.Sleep(time.Second * 3)
} }
@ -81,6 +81,7 @@ func Test_pipec(t *testing.T) {
fmt.Println(err) fmt.Println(err)
return return
} }
server.SetAesKey([]byte("abcdefg123456789"))
server.SetNotify("ni\\||hao", func(data SMsg) string { server.SetNotify("ni\\||hao", func(data SMsg) string {
fmt.Println("name-get", data.GetName()) fmt.Println("name-get", data.GetName())
fmt.Println("name-set", data.SetName("iiiis")) fmt.Println("name-set", data.SetName("iiiis"))
@ -97,6 +98,7 @@ func Test_pipec(t *testing.T) {
fmt.Println(err) fmt.Println(err)
return return
} }
client.SetAesKey([]byte("abcdefg123456789"))
client.UseChannel = false client.UseChannel = false
sa, err := client.SendValueWait("ni\\||hao", "lalaeee", time.Second*10) sa, err := client.SendValueWait("ni\\||hao", "lalaeee", time.Second*10)
if err != nil { if err != nil {
@ -113,7 +115,7 @@ func Test_pipec(t *testing.T) {
fmt.Println("sukidesu") fmt.Println("sukidesu")
time.Sleep(time.Second * 3) time.Sleep(time.Second * 3)
server.ServerStop() server.ServerStop()
<-client.Stop <-client.Stoped()
client.ClientStop() client.ClientStop()
time.Sleep(time.Second * 2) 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(server.SendWait(testmsg, "nihao", "wozuinb", time.Second*20))
fmt.Println("sakura") fmt.Println("sakura")
server.ServerStop() server.ServerStop()
<-client.Stop <-client.Stoped()
client.ClientStop() client.ClientStop()
time.Sleep(time.Second * 3) time.Sleep(time.Second * 3)
} }

@ -16,7 +16,6 @@ import (
"b612.me/starnet" "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} var aesKey = []byte{0x19, 0x96, 0x11, 0x27, 228, 187, 187, 231, 142, 137, 230, 179, 189, 229, 184, 133}
func encodeFunc(data []byte) []byte { func encodeFunc(data []byte) []byte {
@ -27,18 +26,12 @@ func decodeFunc(data []byte) []byte {
return starcrypto.AesDecryptCFB(data, aesKey) return starcrypto.AesDecryptCFB(data, aesKey)
} }
func init() {
builder = starnet.NewQueue()
builder.EncodeFunc = encodeFunc
builder.DecodeFunc = decodeFunc
builder.Encode = true
}
// StarNotifyS 为Server端 // StarNotifyS 为Server端
type StarNotifyS struct { type StarNotifyS struct {
// Queue 是用来处理收发信息的简单消息队列 // Queue 是用来处理收发信息的简单消息队列
Queue *starnet.StarQueue Queue *starnet.StarQueue
// FuncLists 记录了被通知项所记录的函数 // FuncLists 记录了被通知项所记录的函数
aesKey []byte
FuncLists map[string]func(SMsg) string FuncLists map[string]func(SMsg) string
funcMu sync.Mutex funcMu sync.Mutex
defaultFunc func(SMsg) string defaultFunc func(SMsg) string
@ -78,6 +71,24 @@ type SMsg struct {
wait chan int wait chan int
nickName func(string, string) error nickName func(string, string) error
getName func(string) string 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 { func (star *StarNotifyS) getName(conn string) string {
@ -97,11 +108,11 @@ func (star *StarNotifyS) GetConnPool() []SMsg {
var result []SMsg var result []SMsg
star.connPool.Range(func(k, val interface{}) bool { star.connPool.Range(func(k, val interface{}) bool {
v := val.(net.Conn) 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 return true
}) })
for _, v := range star.udpPool { 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 return result
} }
@ -111,10 +122,10 @@ func (star *StarNotifyS) GetClient(name string) (SMsg, error) {
if str, ok := star.nickName[name]; ok { if str, ok := star.nickName[name]; ok {
if tmp, ok := star.connPool.Load(str); ok { if tmp, ok := star.connPool.Load(str); ok {
conn := tmp.(net.Conn) 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 { 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") return SMsg{}, errors.New("Not Found")
@ -157,9 +168,9 @@ func (nmsg *SMsg) ReplyRaw(msg interface{}) error {
func (nmsg *SMsg) Reply(msg string) error { func (nmsg *SMsg) Reply(msg string) error {
var err error var err error
if nmsg.uconn == nil { 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 { } 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 return err
} }
@ -168,9 +179,9 @@ func (nmsg *SMsg) Reply(msg string) error {
func (nmsg *SMsg) Send(key, value string) error { func (nmsg *SMsg) Send(key, value string) error {
var err error var err error
if nmsg.uconn == nil { 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 { } 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 return err
} }
@ -198,9 +209,9 @@ func (star *StarNotifyS) SendWait(source SMsg, key, value string, tmout time.Dur
rand.Seed(time.Now().UnixNano()) rand.Seed(time.Now().UnixNano())
mode := "sr" + fmt.Sprintf("%d%06d", time.Now().UnixNano(), rand.Intn(999999)) mode := "sr" + fmt.Sprintf("%d%06d", time.Now().UnixNano(), rand.Intn(999999))
if source.uconn == nil { 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 { } 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 { if err != nil {
return SMsg{}, err return SMsg{}, err
@ -225,11 +236,12 @@ func (star *StarNotifyS) SendWait(source SMsg, key, value string, tmout time.Dur
} }
func (star *StarNotifyS) starinits() { 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.stopSign, star.cancel = context.WithCancel(context.Background())
star.Queue = starnet.NewQueue() star.Queue = builder
star.Queue.EncodeFunc = encodeFunc
star.Queue.DecodeFunc = decodeFunc
star.Queue.Encode = true
star.udpPool = make(map[string]*net.UDPAddr) star.udpPool = make(map[string]*net.UDPAddr)
star.FuncLists = make(map[string]func(SMsg) string) star.FuncLists = make(map[string]func(SMsg) string)
star.nickName = make(map[string]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) star.Queue.ParseMessage(buf[0:n], addr)
if _, ok := star.udpPool[addr.String()]; !ok { if _, ok := star.udpPool[addr.String()]; !ok {
if star.Connected != nil { 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() star.connMu.Lock()
@ -372,7 +384,7 @@ func notudps(netype, value string) (*StarNotifyS, error) {
}(conn) }(conn)
star.connPool.Store(fmt.Sprint(conn), conn) star.connPool.Store(fmt.Sprint(conn), conn)
if star.Connected != nil { 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 var rmsg SMsg
if !star.isUDP { 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 { } 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" { if key == "b612ryzstop" {
star.connMu.Lock() star.connMu.Lock()
delete(star.udpPool, rmsg.UDP.String()) delete(star.udpPool, rmsg.UDP.String())

Loading…
Cancel
Save