优化代码结构

gold/head/packet.go

@@ -1,13 +1,16 @@
 package head
 
 import (
-	"crypto/rand"
 	"encoding/json"
+	"unsafe"
+
+	blake2b "github.com/minio/blake2b-simd"
 )
 
 // Packet 是发送和接收的最小单位
 type Packet struct {
 	// DataSZ len(Data)
+	// 不得超过 65507-head 字节
 	DataSZ uint32
 	// Proto 详见 head
 	Proto uint8
@@ -47,11 +50,20 @@ func (p *Packet) UnMashal(data []byte) error {
 }
 
 // Mashal 将自身数据编码为 []byte
-// 同时生成 Hash
 func (p *Packet) Mashal(src string, dst string) ([]byte, error) {
 	p.DataSZ = uint32(len(p.Data))
 	p.Src = src
 	p.Dst = dst
-	rand.Reader.Read(p.Hash[:])
 	return json.Marshal(p)
 }
+
+// FillHash 生成 p.Data 的 Hash
+func (p *Packet) FillHash() {
+	sum := blake2b.New256().Sum(p.Data)
+	p.Hash = *(*[32]byte)(*(*unsafe.Pointer)(unsafe.Pointer(&sum)))
+}
+
+func (p *Packet) IsVaildHash() bool {
+	sum := blake2b.New256().Sum(p.Data)
+	return *(*[32]byte)(*(*unsafe.Pointer)(unsafe.Pointer(&sum))) == p.Hash
+}

gold/link/crypto.go

@@ -50,7 +50,7 @@ func NewMe(privateKey *[32]byte, myIP string, myEndpoint string) (m Me) {
 }
 
 // Encode 使用 TEA 加密
-func (l *Link) Encode(b []byte) (eb []byte, err error) {
+func (l *Link) Encode(b []byte) (eb []byte) {
 	if b == nil {
 		return
 	}
@@ -65,7 +65,7 @@ func (l *Link) Encode(b []byte) (eb []byte, err error) {
 }
 
 // Decode 使用 TEA 解密
-func (l *Link) Decode(b []byte) (db []byte, err error) {
+func (l *Link) Decode(b []byte) (db []byte) {
 	if b == nil {
 		return
 	}

gold/link/link.go

@@ -69,14 +69,13 @@ func (l *Link) Read() *head.Packet {
 
 // Write 向 peer 发包
 func (l *Link) Write(p *head.Packet) (n int, err error) {
-	p.Data, err = l.Encode(p.Data)
+	p.FillHash()
+	p.Data = l.Encode(p.Data)
+	var d []byte
+	d, err = p.Mashal(l.me.me.String(), l.peerip.String())
+	logrus.Debugln("[link] write data", string(d))
 	if err == nil {
-		var d []byte
-		d, err = p.Mashal(l.me.me.String(), l.peerip.String())
-		logrus.Debugln("[link] write data", string(d))
-		if err == nil {
-			n, err = l.me.myconn.WriteToUDP(d, l.NextHop(l.peerip).endpoint)
-		}
+		n, err = l.me.myconn.WriteToUDP(d, l.NextHop(l.peerip).endpoint)
 	}
 	return
 }

gold/link/listen.go

@@ -3,8 +3,9 @@ package link
 import (
 	"net"
 
-	"github.com/fumiama/WireGold/gold/head"
 	"github.com/sirupsen/logrus"
+
+	"github.com/fumiama/WireGold/gold/head"
 )
 
 // 监听本机 endpoint
@@ -38,8 +39,8 @@ func (m *Me) listen() (conn *net.UDPConn, err error) {
 						}
 						if ok {
 							if p.IsToMe(net.ParseIP(packet.Dst)) {
-								packet.Data, err = p.Decode(packet.Data)
-								if err == nil {
+								packet.Data = p.Decode(packet.Data)
+								if packet.IsVaildHash() {
 									switch packet.Proto {
 									case head.ProtoHello:
 										switch p.status {
@@ -64,6 +65,8 @@ func (m *Me) listen() (conn *net.UDPConn, err error) {
 									default:
 										break
 									}
+								} else {
+									logrus.Infoln("[link] drop invalid packet")
 								}
 							} else if p.Accept(net.ParseIP(packet.Dst)) && p.allowtrans {
 								// 转发

gold/link/peer.go

@@ -1,7 +1,6 @@
 package link
 
 import (
-	"fmt"
 	"net"
 	"unsafe"
 
@@ -30,7 +29,6 @@ func (m *Me) AddPeer(peerip string, pubicKey *[32]byte, endPoint string, allowed
 		c := curve.Get(m.privKey[:])
 		k, err := c.Shared(pubicKey)
 		if err == nil {
-			fmt.Println(len(k))
 			l.key = (*[32]byte)(*(*unsafe.Pointer)(unsafe.Pointer(&k)))
 		}
 	}