优化代码结构，router增加ttl

2026-06-27 14:20:27 +08:00 · 2022-04-15 15:41:09 +08:00
parent dcaff382bf
commit 826c89c678
9 changed files with 152 additions and 70 deletions
--- a/go.mod
+++ b/go.mod
@@ -1,6 +1,6 @@
 module github.com/fumiama/WireGold

-go 1.16
+go 1.18

 require (
 	github.com/fumiama/blake2b-simd v0.0.0-20220412092318-c99573b3b2b1
@@ -9,5 +9,13 @@ require (
 	github.com/fumiama/gofastTEA v0.0.9
 	github.com/fumiama/water v0.0.0-20211231134027-da391938d6ac
 	github.com/sirupsen/logrus v1.8.1
+	github.com/wdvxdr1123/ZeroBot v1.5.0-mid.0.20220415071800-9e52436ab5c0
 	gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b
 )
+
+require (
+	github.com/fumiama/wintun v0.0.0-20211229152851-8bc97c8034c0 // indirect
+	golang.org/x/crypto v0.0.0-20210921155107-089bfa567519 // indirect
+	golang.org/x/sys v0.0.0-20220412211240-33da011f77ad // indirect
+	golang.org/x/text v0.3.7 // indirect
+)
--- a/go.sum
+++ b/go.sum
@@ -17,18 +17,18 @@ github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZN
 github.com/sirupsen/logrus v1.8.1 h1:dJKuHgqk1NNQlqoA6BTlM1Wf9DOH3NBjQyu0h9+AZZE=
 github.com/sirupsen/logrus v1.8.1/go.mod h1:yWOB1SBYBC5VeMP7gHvWumXLIWorT60ONWic61uBYv0=
 github.com/songgao/water v0.0.0-20200317203138-2b4b6d7c09d8 h1:TG/diQgUe0pntT/2D9tmUCz4VNwm9MfrtPr0SU2qSX8=
-github.com/songgao/water v0.0.0-20200317203138-2b4b6d7c09d8/go.mod h1:P5HUIBuIWKbyjl083/loAegFkfbFNx5i2qEP4CNbm7E=
-github.com/stretchr/testify v1.2.2 h1:bSDNvY7ZPG5RlJ8otE/7V6gMiyenm9RtJ7IUVIAoJ1w=
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
+github.com/stretchr/testify v1.7.1 h1:5TQK59W5E3v0r2duFAb7P95B6hEeOyEnHRa8MjYSMTY=
+github.com/wdvxdr1123/ZeroBot v1.5.0-mid.0.20220415071800-9e52436ab5c0 h1:SoNu5MZZsbleecVCBth/s610wLXWKIr18KADhZURCDw=
+github.com/wdvxdr1123/ZeroBot v1.5.0-mid.0.20220415071800-9e52436ab5c0/go.mod h1:K2vu0mslV8s4qhIAu/a03Z7YW24qjM0j3imIR+k21KI=
 golang.org/x/crypto v0.0.0-20210921155107-089bfa567519 h1:7I4JAnoQBe7ZtJcBaYHi5UtiO8tQHbUSXxL+pnGRANg=
 golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
 golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
 golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.0.0-20211103235746-7861aae1554b/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.0.0-20211216021012-1d35b9e2eb4e h1:fLOSk5Q00efkSvAm+4xcoXD+RRmLmmulPn5I3Y9F2EM=
-golang.org/x/sys v0.0.0-20211216021012-1d35b9e2eb4e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20220412211240-33da011f77ad h1:ntjMns5wyP/fN65tdBD4g8J5w8n015+iIIs9rtjXkY0=
+golang.org/x/sys v0.0.0-20220412211240-33da011f77ad/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.7 h1:olpwvP2KacW1ZWvsR7uQhoyTYvKAupfQrRGBFM352Gk=
--- a/gold/link/link.go
+++ b/gold/link/link.go
@@ -24,14 +24,14 @@ type Link struct {
 	endpoint *net.UDPAddr
 	// 本机允许接收/发送的 ip 网段
 	allowedips []*net.IPNet
-	// 是否允许转发
-	allowtrans bool
-	// 连接的状态，详见下方 const
-	status int
 	// 连接所用对称加密密钥
 	key []tea.TEA
 	// 本机信息
 	me *Me
+	// 连接的状态，详见下方 const
+	status int
+	// 是否允许转发
+	allowtrans bool
 }

 const (
--- a/gold/link/me.go
+++ b/gold/link/me.go
@@ -6,12 +6,14 @@ import (
 	"net"
 	"strconv"
 	"sync"
+	"time"

 	"github.com/fumiama/WireGold/gold/head"
 	"github.com/fumiama/WireGold/helper"
 	"github.com/fumiama/WireGold/lower"
 	"github.com/fumiama/water/waterutil"
 	"github.com/sirupsen/logrus"
+	"github.com/wdvxdr1123/ZeroBot/extension/ttl"
 )

 // Me 是本机的抽象
@@ -51,15 +53,23 @@ type Me struct {
 	srcport, dstport, mtu uint16
 }

+type MyConfig struct {
+	MyIPwithMask          string
+	MyEndpoint            string
+	PrivateKey            *[32]byte
+	NIC                   lower.NICIO
+	SrcPort, DstPort, MTU uint16
+}
+
 // NewMe 设置本机参数
-func NewMe(privateKey *[32]byte, myipwithmask string, myEndpoint string, nic lower.NICIO, srcport, dstport, mtu uint16) (m Me) {
-	m.privKey = *privateKey
+func NewMe(cfg *MyConfig) (m Me) {
+	m.privKey = *cfg.PrivateKey
 	var err error
-	m.myend, err = net.ResolveUDPAddr("udp", myEndpoint)
+	m.myend, err = net.ResolveUDPAddr("udp", cfg.MyEndpoint)
 	if err != nil {
 		panic(err)
 	}
-	ip, cidr, err := net.ParseCIDR(myipwithmask)
+	ip, cidr, err := net.ParseCIDR(cfg.MyIPwithMask)
 	if err != nil {
 		panic(err)
 	}
@@ -70,18 +80,23 @@ func NewMe(privateKey *[32]byte, myipwithmask string, myEndpoint string, nic low
 		panic(err)
 	}
 	m.connections = make(map[string]*Link)
-	m.nic = nic
+	m.nic = cfg.NIC
 	m.router = &Router{
 		list:  make([]*net.IPNet, 1, 16),
 		table: make(map[string]*Link, 16),
+		cache: ttl.NewCache[string, *Link](time.Minute),
 	}
 	m.router.SetDefault(nil)
-	m.loop = m.AddPeer(m.me.String(), nil, "127.0.0.1:56789", []string{myipwithmask}, nil, 0, 0, false, nic != nil)
-	m.srcport = srcport
-	m.dstport = dstport
-	m.mtu = mtu & 0xfff8
-	m.writer = helper.SelectWriter()
-	go m.initrecvpool()
+	m.loop = m.AddPeer(&PeerConfig{
+		PeerIP:     m.me.String(),
+		EndPoint:   "127.0.0.1:56789",
+		AllowedIPs: []string{cfg.MyIPwithMask},
+		NoPipe:     cfg.NIC != nil,
+	})
+	m.srcport = cfg.SrcPort
+	m.dstport = cfg.DstPort
+	m.mtu = cfg.MTU & 0xfff8
+	m.initrecvpool()
 	return
 }

@@ -125,13 +140,13 @@ func (m *Me) ListenFromNIC() (written int64, err error) {
 	return io.Copy(m, m.nic)
 }

-type PacketID [2]byte
+type packetID [2]byte

-func newpacketid(packet []byte) PacketID {
+func newpacketid(packet []byte) packetID {
 	return waterutil.IPv4Identification(packet)
 }

-func (p PacketID) issame(packet []byte) bool {
+func (p packetID) issame(packet []byte) bool {
 	return p == waterutil.IPv4Identification(packet)
 }

--- a/gold/link/peer.go
+++ b/gold/link/peer.go
@@ -10,27 +10,36 @@ import (
 	"github.com/sirupsen/logrus"
 )

+type PeerConfig struct {
+	PeerIP                  string
+	EndPoint                string
+	AllowedIPs, Querys      []string
+	PubicKey                *[32]byte
+	KeepAliveDur, QueryTick int64
+	AllowTrans, NoPipe      bool
+}
+
 // AddPeer 添加一个 peer
-func (m *Me) AddPeer(peerip string, pubicKey *[32]byte, endPoint string, allowedIPs, querys []string, keepAliveDur, queryTick int64, allowTrans, nopipe bool) (l *Link) {
-	peerip = net.ParseIP(peerip).String()
+func (m *Me) AddPeer(cfg *PeerConfig) (l *Link) {
+	cfg.PeerIP = net.ParseIP(cfg.PeerIP).String()
 	var ok bool
-	l, ok = m.IsInPeer(peerip)
+	l, ok = m.IsInPeer(cfg.PeerIP)
 	if ok {
 		return
 	}
 	l = &Link{
-		pubk:       pubicKey,
-		peerip:     net.ParseIP(peerip),
-		allowtrans: allowTrans,
+		pubk:       cfg.PubicKey,
+		peerip:     net.ParseIP(cfg.PeerIP),
+		allowtrans: cfg.AllowTrans,
 		me:         m,
 	}

-	if !nopipe {
+	if !cfg.NoPipe {
 		l.pipe = make(chan *head.Packet, 32)
 	}
-	if pubicKey != nil {
+	if cfg.PubicKey != nil {
 		c := curve.Get(m.privKey[:])
-		k, err := c.Shared(pubicKey)
+		k, err := c.Shared(cfg.PubicKey)
 		if err == nil {
 			l.key = make([]tea.TEA, 16)
 			for i := range l.key {
@@ -38,31 +47,31 @@ func (m *Me) AddPeer(peerip string, pubicKey *[32]byte, endPoint string, allowed
 			}
 		}
 	}
-	if endPoint != "" {
-		e, err := net.ResolveUDPAddr("udp", endPoint)
+	if cfg.EndPoint != "" {
+		e, err := net.ResolveUDPAddr("udp", cfg.EndPoint)
 		if err != nil {
 			panic(err)
 		}
 		l.endpoint = e
 	}
-	if allowedIPs != nil {
-		l.allowedips = make([]*net.IPNet, 0, len(allowedIPs))
-		for _, ipnet := range allowedIPs {
+	if cfg.AllowedIPs != nil {
+		l.allowedips = make([]*net.IPNet, 0, len(cfg.AllowedIPs))
+		for _, ipnet := range cfg.AllowedIPs {
 			_, cidr, err := net.ParseCIDR(ipnet)
 			if err == nil {
 				l.allowedips = append(l.allowedips, cidr)
 				l.me.router.SetItem(cidr, l)
 				l.me.connmapmu.Lock()
-				l.me.connections[peerip] = l
+				l.me.connections[cfg.PeerIP] = l
 				l.me.connmapmu.Unlock()
 			} else {
 				panic(err)
 			}
 		}
 	}
-	logrus.Infoln("[peer] add peer:", peerip, "allow:", allowedIPs)
-	go l.keepAlive(keepAliveDur)
-	go l.sendquery(time.Second*time.Duration(queryTick), querys...)
+	logrus.Infoln("[peer] add peer:", cfg.PeerIP, "allow:", cfg.AllowedIPs)
+	go l.keepAlive(cfg.KeepAliveDur)
+	go l.sendquery(time.Second*time.Duration(cfg.QueryTick), cfg.Querys...)
 	return
 }

--- a/gold/link/recv.go
+++ b/gold/link/recv.go
@@ -7,6 +7,7 @@ import (
 	"unsafe"

 	"github.com/fumiama/WireGold/gold/head"
+	"github.com/fumiama/WireGold/helper"
 	"github.com/sirupsen/logrus"
 )

@@ -16,6 +17,9 @@ func (l *Link) Read() *head.Packet {
 }

 func (m *Me) initrecvpool() {
+	if m.writer == nil {
+		m.writer = helper.SelectWriter()
+	}
 	if m.recving == nil {
 		m.recving = make(map[[32]byte]*head.Packet, 128)
 	}
@@ -23,23 +27,25 @@ func (m *Me) initrecvpool() {
 	m.clock = make(map[*head.Packet]uint8, 128)
 	var delhs []*head.Packet
 	t := time.NewTicker(time.Second)
-	for range t.C {
-		m.recvmu.Lock()
-		for k, v := range m.clock {
-			if v > 10 { // 10s
-				delete(m.recving, k.Hash)
-				delhs = append(delhs, k)
-			} else {
-				m.clock[k]++
+	go func() {
+		for range t.C {
+			m.recvmu.Lock()
+			for k, v := range m.clock {
+				if v > 10 { // 10s
+					delete(m.recving, k.Hash)
+					delhs = append(delhs, k)
+				} else {
+					m.clock[k]++
+				}
 			}
+			for _, k := range delhs {
+				delete(m.clock, k)
+				logrus.Warnln("[recv] drop timeout packet from", k.Src)
+			}
+			delhs = delhs[:0]
+			m.recvmu.Unlock()
 		}
-		for _, k := range delhs {
-			delete(m.clock, k)
-			logrus.Warnln("[recv] drop timeout packet from", k.Src)
-		}
-		delhs = delhs[:0]
-		m.recvmu.Unlock()
-	}
+	}()
 }

 func (m *Me) wait(data []byte) *head.Packet {
--- a/gold/link/router.go
+++ b/gold/link/router.go
@@ -6,6 +6,7 @@ import (
 	"sync"

 	"github.com/sirupsen/logrus"
+	"github.com/wdvxdr1123/ZeroBot/extension/ttl"
 )

 type Router struct {
@@ -13,6 +14,7 @@ type Router struct {
 	table map[string]*Link
 	mu    sync.RWMutex
 	list  []*net.IPNet
+	cache *ttl.Cache[string, *Link]
 }

 // Accept 判断是否应当接受 ip 发来的包
@@ -41,6 +43,11 @@ func (r *Router) SetDefault(l *Link) {

 // NextHop 得到前往 ip 的下一跳的 link
 func (r *Router) NextHop(ip string) (l *Link) {
+	l = r.cache.Get(ip)
+	if l != nil {
+		logrus.Debugln("[router] get cached nexthop to", ip, "link", l)
+		return
+	}
 	ipb := net.ParseIP(ip)
 	if ipb == nil {
 		logrus.Errorln("[router] nil ip")
@@ -56,6 +63,7 @@ func (r *Router) NextHop(ip string) (l *Link) {
 		if c.Contains(ipb) {
 			l = r.table[c.String()]
 			logrus.Debugln("[router] get nexthop to", ipb, "-->", c, "link", l)
+			r.cache.Set(ip, l)
 			return l
 		}
 	}
--- a/upper/services/tunnel/tunnel_test.go
+++ b/upper/services/tunnel/tunnel_test.go
@@ -27,10 +27,34 @@ func TestTunnel(t *testing.T) {
 	t.Log("peer priv key:", hex.EncodeToString(peerpk.Private()[:]))
 	t.Log("peer publ key:", hex.EncodeToString(peerpk.Public()[:]))

-	m := link.NewMe(selfpk.Private(), "192.168.1.2/32", "127.0.0.1:1236", nil, 1, 1, 4096)
-	m.AddPeer("192.168.1.3", peerpk.Public(), "127.0.0.1:1237", []string{"192.168.1.3/32"}, nil, 0, 0, false, false)
-	p := link.NewMe(peerpk.Private(), "192.168.1.3/32", "127.0.0.1:1237", nil, 1, 1, 4096)
-	p.AddPeer("192.168.1.2", selfpk.Public(), "127.0.0.1:1236", []string{"192.168.1.2/32"}, nil, 0, 0, false, false)
+	m := link.NewMe(&link.MyConfig{
+		MyIPwithMask: "192.168.1.2/32",
+		MyEndpoint:   "127.0.0.1:1236",
+		PrivateKey:   selfpk.Private(),
+		SrcPort:      1,
+		DstPort:      1,
+		MTU:          4096,
+	})
+	m.AddPeer(&link.PeerConfig{
+		PeerIP:     "192.168.1.3",
+		EndPoint:   "127.0.0.1:1237",
+		AllowedIPs: []string{"192.168.1.3/32"},
+		PubicKey:   peerpk.Public(),
+	})
+	p := link.NewMe(&link.MyConfig{
+		MyIPwithMask: "192.168.1.3/32",
+		MyEndpoint:   "127.0.0.1:1237",
+		PrivateKey:   peerpk.Private(),
+		SrcPort:      1,
+		DstPort:      1,
+		MTU:          4096,
+	})
+	p.AddPeer(&link.PeerConfig{
+		PeerIP:     "192.168.1.2",
+		EndPoint:   "127.0.0.1:1236",
+		AllowedIPs: []string{"192.168.1.2/32"},
+		PubicKey:   selfpk.Public(),
+	})
 	tunnme, err := Create(&m, "192.168.1.3")
 	if err != nil {
 		t.Fatal(err)
--- a/upper/services/wg/wg.go
+++ b/upper/services/wg/wg.go
@@ -65,7 +65,7 @@ func (wg *WG) Stop() {
 	_ = wg.me.Close()
 }

-func (wg *WG) init(srcport, destport, mtu uint16) {
+func (wg *WG) init(srcport, dstport, mtu uint16) {
 	cidrsmap := make(map[string]bool, 32)
 	_, mysubnet, err := net.ParseCIDR(wg.c.SubNet)
 	if err != nil {
@@ -89,13 +89,15 @@ func (wg *WG) init(srcport, destport, mtu uint16) {
 		i++
 	}

-	wg.me = link.NewMe(
-		&wg.key,
-		wg.c.IP+"/32",
-		wg.c.EndPoint,
-		lower.NewNIC(wg.c.IP, wg.c.SubNet, fmt.Sprintf("%d", mtu), cidrs...),
-		srcport, destport, mtu,
-	)
+	wg.me = link.NewMe(&link.MyConfig{
+		MyIPwithMask: wg.c.IP + "/32",
+		MyEndpoint:   wg.c.EndPoint,
+		PrivateKey:   &wg.key,
+		NIC:          lower.NewNIC(wg.c.IP, wg.c.SubNet, fmt.Sprintf("%d", mtu), cidrs...),
+		SrcPort:      srcport,
+		DstPort:      dstport,
+		MTU:          mtu,
+	})

 	for _, peer := range wg.c.Peers {
 		var peerkey [32]byte
@@ -107,6 +109,16 @@ func (wg *WG) init(srcport, destport, mtu uint16) {
 		if n != 32 {
 			panic("peer public key length is not 32")
 		}
-		wg.me.AddPeer(peer.IP, &peerkey, peer.EndPoint, peer.AllowedIPs, peer.QueryList, peer.KeepAliveSeconds, peer.QuerySeconds, peer.AllowTrans, true)
+		wg.me.AddPeer(&link.PeerConfig{
+			PeerIP:       peer.IP,
+			EndPoint:     peer.EndPoint,
+			AllowedIPs:   peer.AllowedIPs,
+			Querys:       peer.QueryList,
+			PubicKey:     &peerkey,
+			KeepAliveDur: peer.KeepAliveSeconds,
+			QueryTick:    peer.QuerySeconds,
+			AllowTrans:   peer.AllowTrans,
+			NoPipe:       true,
+		})
 	}
 }