diff --git a/config/cfg.go b/config/cfg.go index ceec970..c62f61c 100644 --- a/config/cfg.go +++ b/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"` } diff --git a/go.mod b/go.mod index 000ff33..1bf8e52 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index b25de4a..bcb05cc 100644 --- a/go.sum +++ b/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= diff --git a/gold/link/link.go b/gold/link/link.go index 51a06c3..0221bd2 100644 --- a/gold/link/link.go +++ b/gold/link/link.go @@ -37,6 +37,8 @@ type Link struct { status int8 // 是否允许转发 allowtrans bool + // 是否对数据进行 zstd 压缩 + usezstd bool // udp 数据包的最大大小 mtu uint16 // 随机放缩 mtu 范围 (只减不增) diff --git a/gold/link/listen.go b/gold/link/listen.go index 26d59e8..5020382 100644 --- a/gold/link/listen.go +++ b/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() diff --git a/gold/link/peer.go b/gold/link/peer.go index 326a079..19ff734 100644 --- a/gold/link/peer.go +++ b/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), diff --git a/gold/link/send.go b/gold/link/send.go index e2daf15..298054d 100644 --- a/gold/link/send.go +++ b/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 diff --git a/upper/services/wg/wg.go b/upper/services/wg/wg.go index c842244..b81ae3b 100644 --- a/upper/services/wg/wg.go +++ b/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, }) } }