customize AES-CFB Keys

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 {