优化代码结构

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