feat: add ztsd for data compressing

config/cfg.go

@@ -30,6 +30,7 @@ type Peer struct {
 	QueryList        []string `yaml:"QueryList"`
 	QuerySeconds     int64    `yaml:"QuerySeconds"`
 	AllowTrans       bool     `yaml:"AllowTrans"`
+	UseZstd          bool     `yaml:"UseZstd"`
 	MTU              int64    `yaml:"MTU"`
 	MTURandomRange   int64    `yaml:"MTURandomRange"`
 }

go.mod

@@ -9,6 +9,7 @@ require (
 	github.com/fumiama/go-x25519 v1.0.0
 	github.com/fumiama/gofastTEA v0.0.10
 	github.com/fumiama/water v0.0.0-20211231134027-da391938d6ac
+	github.com/klauspost/compress v1.16.7
 	github.com/sirupsen/logrus v1.9.3
 	golang.org/x/crypto v0.11.1-0.20230731181441-edc325d13aa9
 	gopkg.in/yaml.v3 v3.0.1

go.sum

@@ -15,6 +15,8 @@ github.com/fumiama/water v0.0.0-20211231134027-da391938d6ac h1:A/5A0rODsg+EQHH61
 github.com/fumiama/water v0.0.0-20211231134027-da391938d6ac/go.mod h1:BBnNY9PwK+UUn4trAU+H0qsMEypm7+3Bj1bVFuJItlo=
 github.com/fumiama/wintun v0.0.0-20211229152851-8bc97c8034c0 h1:WfrSFlIlCAtg6Rt2IGna0HhJYSDE45YVHiYqO4wwsEw=
 github.com/fumiama/wintun v0.0.0-20211229152851-8bc97c8034c0/go.mod h1:dPOG7Af/ArO62RgBz2JJTNFByBn/IXWLo/1kZKcLSe8=
+github.com/klauspost/compress v1.16.7 h1:2mk3MPGNzKyxErAw8YaohYh69+pa4sIQSC0fPGCFR9I=
+github.com/klauspost/compress v1.16.7/go.mod h1:ntbaceVETuRiXiv4DpjP66DpAtAGkEQskQzEyD//IeE=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=

gold/link/link.go

@@ -37,6 +37,8 @@ type Link struct {
 	status int8
 	// 是否允许转发
 	allowtrans bool
+	// 是否对数据进行 zstd 压缩
+	usezstd bool
 	// udp 数据包的最大大小
 	mtu uint16
 	// 随机放缩 mtu 范围 (只减不增)

gold/link/listen.go

@@ -1,11 +1,14 @@
 package link
 
 import (
+	"bytes"
+	"io"
 	"net"
 	"runtime"
 	"strconv"
 	"sync"
 
+	"github.com/klauspost/compress/zstd"
 	"github.com/sirupsen/logrus"
 
 	"github.com/fumiama/WireGold/gold/head"
@@ -71,6 +74,16 @@ func (m *Me) listenthread(conn *net.UDPConn, mu *sync.Mutex) {
 					continue
 				}
 			}
+			if p.usezstd {
+				dec, _ := zstd.NewReader(bytes.NewReader(packet.Data))
+				packet.Data, err = io.ReadAll(dec)
+				dec.Close()
+				if err != nil {
+					logrus.Debugln("[listen] drop invalid zstd packet:", err)
+					packet.Put()
+					continue
+				}
+			}
 			if !packet.IsVaildHash() {
 				logrus.Debugln("[listen] drop invalid hash packet")
 				packet.Put()

gold/link/peer.go

@@ -21,6 +21,7 @@ type PeerConfig struct {
 	MTU                     uint16
 	MTURandomRange          uint16
 	AllowTrans, NoPipe      bool
+	UseZstd                 bool
 }
 
 // AddPeer 添加一个 peer
@@ -38,6 +39,7 @@ func (m *Me) AddPeer(cfg *PeerConfig) (l *Link) {
 		pubk:           cfg.PubicKey,
 		peerip:         net.ParseIP(cfg.PeerIP),
 		allowtrans:     cfg.AllowTrans,
+		usezstd:        cfg.UseZstd,
 		me:             m,
 		mtu:            uint16(cfg.MTU),
 		mturandomrange: uint16(cfg.MTURandomRange),

gold/link/send.go

@@ -1,12 +1,16 @@
 package link
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
+	"io"
 	"math/rand"
 	"sync/atomic"
 
 	"github.com/fumiama/WireGold/gold/head"
+	"github.com/fumiama/WireGold/helper"
+	"github.com/klauspost/compress/zstd"
 	"github.com/sirupsen/logrus"
 )
 
@@ -21,21 +25,13 @@ func (l *Link) WriteAndPut(p *head.Packet, istransfer bool) (n int, err error) {
 	logrus.Debugln("[send] mtu:", mtu, ", count:", sndcnt, ", additional data:", uint16(sndcnt))
 	if len(p.Data) <= int(mtu) {
 		if !istransfer {
-			p.FillHash()
-			if l.aead != nil {
-				p.Data = l.EncodePreshared(uint16(sndcnt), p.Data)
-			}
-			p.Data = l.Encode(teatype, p.Data)
+			l.encrypt(p, uint16(sndcnt), teatype)
 		}
 		defer p.Put()
 		return l.write(p, teatype, uint16(sndcnt), mtu, uint32(len(p.Data)), 0, istransfer, false)
 	}
 	if !istransfer {
-		p.FillHash()
-		if l.aead != nil {
-			p.Data = l.EncodePreshared(uint16(sndcnt), p.Data)
-		}
-		p.Data = l.Encode(teatype, p.Data)
+		l.encrypt(p, uint16(sndcnt), teatype)
 	}
 	data := p.Data
 	ttl := p.TTL
@@ -62,6 +58,22 @@ func (l *Link) WriteAndPut(p *head.Packet, istransfer bool) (n int, err error) {
 	return n, err
 }
 
+func (l *Link) encrypt(p *head.Packet, sndcnt uint16, teatype uint8) {
+	p.FillHash()
+	if l.usezstd {
+		w := helper.SelectWriter()
+		defer helper.PutWriter(w)
+		enc, _ := zstd.NewWriter(w, zstd.WithEncoderLevel(zstd.SpeedFastest))
+		defer enc.Close()
+		_, _ = io.Copy(enc, bytes.NewReader(p.Data))
+		p.Data = w.Bytes()
+	}
+	if l.aead != nil {
+		p.Data = l.EncodePreshared(sndcnt, p.Data)
+	}
+	p.Data = l.Encode(teatype, p.Data)
+}
+
 // write 向 peer 发一个包
 func (l *Link) write(p *head.Packet, teatype uint8, additional, mtu uint16, datasz uint32, offset uint16, istransfer, hasmore bool) (n int, err error) {
 	var d []byte

upper/services/wg/wg.go

@@ -140,6 +140,7 @@ func (wg *WG) init(srcport, dstport uint16) {
 			MTURandomRange: uint16(peer.MTURandomRange),
 			AllowTrans:     peer.AllowTrans,
 			NoPipe:         true,
+			UseZstd:        peer.UseZstd,
 		})
 	}
 }