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 = builder
-	star.Queue.EncodeFunc = encodeFunc
-	star.Queue.DecodeFunc = decodeFunc
-	star.Queue.Encode = true
 	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 {

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

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 = builder
-	star.Queue.EncodeFunc = encodeFunc
-	star.Queue.DecodeFunc = decodeFunc
-	star.Queue.Encode = true
 	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())