init

gold/link/link.go

@@ -45,11 +45,11 @@ func (l *Link) Read() *head.Packet {
 }
 
 func (l *Link) Write(p *head.Packet) (n int, err error) {
-	var d []byte
-	d, err = p.Mashal(me.String(), l.peerip.String())
-	logrus.Debugln("[link] write data", string(d))
+	p.Data, err = l.Encode(p.Data)
 	if err == nil {
-		d, err = l.Encode(d)
+		var d []byte
+		d, err = p.Mashal(me.String(), l.peerip.String())
+		logrus.Debugln("[link] write data", string(d))
 		if err == nil {
 			n, err = myconn.WriteToUDP(d, l.endpoint)
 		}

gold/link/local.go

@@ -45,33 +45,30 @@ func listen() (conn *net.UDPConn, err error) {
 				n, addr, err := conn.ReadFromUDP(lbf)
 				if err == nil {
 					lbf = lbf[:n]
-					p, ok := IsEndpointInPeer(addr.String())
-					logrus.Infoln("[link] recv from endpoint", addr)
-					logrus.Debugln("[link] recv:", string(lbf))
-					if ok {
-						packet := head.Packet{}
-						d, err := p.Decode(lbf)
-						if err == nil {
-							packet.UnMashal(d)
-							r := packet.DataSZ - uint32(len(packet.Data))
-							if r > 0 {
-								i := 0
-								n := 0
-								remain := make([]byte, r)
-								for r > 0 {
-									n, _, err = conn.ReadFromUDP(remain[i:])
-									if err == nil {
-										i += n
-										r -= uint32(n)
-									} else {
-										logrus.Errorln("[link.listen]", err)
-										return
-									}
-								}
+					packet := head.Packet{}
+					err = packet.UnMashal(lbf)
+					if err == nil {
+						r := int(packet.DataSZ) - len(packet.Data)
+						if r > 0 {
+							remain, err := readAll(conn, r)
+							if err == nil {
 								packet.Data = append(packet.Data, remain...)
 							}
-							logrus.Infoln("[link] deliver to", p.peerip)
-							p.pipe <- &packet
+						}
+						p, ok := IsInPeer(packet.Src)
+						logrus.Infoln("[link] recv from endpoint", addr, "src", packet.Src, "dst", packet.Dst)
+						logrus.Debugln("[link] recv:", string(lbf))
+						if ok {
+							packet.Data, err = p.Decode(packet.Data)
+							if err == nil {
+								logrus.Infoln("[link] deliver to", p.peerip)
+								if p.EndPoint == "" {
+									logrus.Infoln("[link] set endpoint of peer", p.peerip, "to", addr.String())
+									p.endpoint = addr
+									p.EndPoint = addr.String()
+								}
+								p.pipe <- &packet
+							}
 						}
 					}
 				}
@@ -80,3 +77,22 @@ func listen() (conn *net.UDPConn, err error) {
 	}
 	return
 }
+
+func readAll(conn *net.UDPConn, sz int) ([]byte, error) {
+	i := 0
+	n := 0
+	r := sz
+	var err error
+	remain := make([]byte, r)
+	for sz > 0 {
+		n, _, err = conn.ReadFromUDP(remain[i:])
+		if err == nil {
+			i += n
+			r -= n
+		} else {
+			logrus.Errorln("[link] read all err:", err)
+			return nil, err
+		}
+	}
+	return remain, nil
+}

gold/link/peer.go

@@ -2,16 +2,10 @@ package link
 
 import (
 	"net"
-	"sync"
 
 	"github.com/fumiama/WireGold/gold/head"
 )
 
-var (
-	eps  = make(map[string]*Link)
-	epmu sync.RWMutex
-)
-
 func AddPeer(peerip string, pubicKey [32]byte, endPoint string, keepAlive int64) (l *Link) {
 	peerip = net.ParseIP(peerip).String()
 	var ok bool
@@ -19,24 +13,23 @@ func AddPeer(peerip string, pubicKey [32]byte, endPoint string, keepAlive int64)
 	if ok {
 		return
 	}
-	e, err := net.ResolveUDPAddr("udp", endPoint)
-	if err != nil {
-		panic(err)
-	}
 	l = &Link{
 		PubicKey:  pubicKey,
-		EndPoint:  endPoint,
 		KeepAlive: keepAlive,
 		pipe:      make(chan *head.Packet, 32),
 		peerip:    net.ParseIP(peerip),
-		endpoint:  e,
+	}
+	if endPoint != "" {
+		e, err := net.ResolveUDPAddr("udp", endPoint)
+		if err != nil {
+			panic(err)
+		}
+		l.EndPoint = endPoint
+		l.endpoint = e
 	}
 	connmapmu.Lock()
-	epmu.Lock()
 	connections[peerip] = l
-	eps[endPoint] = l
 	connmapmu.Unlock()
-	epmu.Unlock()
 	return
 }
 
@@ -46,10 +39,3 @@ func IsInPeer(peer string) (p *Link, ok bool) {
 	connmapmu.RUnlock()
 	return
 }
-
-func IsEndpointInPeer(ep string) (p *Link, ok bool) {
-	epmu.RLock()
-	p, ok = eps[ep]
-	epmu.RUnlock()
-	return
-}