From 70c8605cca3b9cba8598bad52d3bb6a36d77d032 Mon Sep 17 00:00:00 2001 From: Dreamacro <8615343+Dreamacro@users.noreply.github.com> Date: Mon, 20 Sep 2021 21:02:18 +0800 Subject: [PATCH] Improve: use one bytes.Buffer pool --- common/pool/buffer.go | 17 ++++++++++++++++ transport/gun/gun.go | 9 ++++----- transport/simple-obfs/tls.go | 3 ++- transport/snell/snell.go | 12 +++++------- transport/ssr/obfs/http_simple.go | 6 ++---- transport/ssr/obfs/tls1.2_ticket_auth.go | 25 ++++++++++-------------- transport/ssr/protocol/packet.go | 8 +++----- transport/ssr/protocol/stream.go | 6 ++---- transport/ssr/tools/bufPool.go | 13 +++--------- transport/trojan/trojan.go | 14 +++++-------- 10 files changed, 53 insertions(+), 60 deletions(-) create mode 100644 common/pool/buffer.go diff --git a/common/pool/buffer.go b/common/pool/buffer.go new file mode 100644 index 0000000..fb328b7 --- /dev/null +++ b/common/pool/buffer.go @@ -0,0 +1,17 @@ +package pool + +import ( + "bytes" + "sync" +) + +var bufferPool = sync.Pool{New: func() interface{} { return &bytes.Buffer{} }} + +func GetBuffer() *bytes.Buffer { + return bufferPool.Get().(*bytes.Buffer) +} + +func PutBuffer(buf *bytes.Buffer) { + buf.Reset() + bufferPool.Put(buf) +} diff --git a/transport/gun/gun.go b/transport/gun/gun.go index 3f97121..e60d13d 100644 --- a/transport/gun/gun.go +++ b/transport/gun/gun.go @@ -5,7 +5,6 @@ package gun import ( "bufio" - "bytes" "crypto/tls" "encoding/binary" "errors" @@ -17,6 +16,8 @@ import ( "sync" "time" + "github.com/Dreamacro/clash/common/pool" + "go.uber.org/atomic" "golang.org/x/net/http2" ) @@ -31,7 +32,6 @@ var ( "content-type": []string{"application/grpc"}, "user-agent": []string{"grpc-go/1.36.0"}, } - bufferPool = sync.Pool{New: func() interface{} { return &bytes.Buffer{} }} ) type DialFn = func(network, addr string) (net.Conn, error) @@ -127,9 +127,8 @@ func (g *Conn) Write(b []byte) (n int, err error) { grpcPayloadLen := uint32(varuintSize + 1 + len(b)) binary.BigEndian.PutUint32(grpcHeader[1:5], grpcPayloadLen) - buf := bufferPool.Get().(*bytes.Buffer) - defer bufferPool.Put(buf) - defer buf.Reset() + buf := pool.GetBuffer() + defer pool.PutBuffer(buf) buf.Write(grpcHeader) buf.Write(protobufHeader[:varuintSize+1]) buf.Write(b) diff --git a/transport/simple-obfs/tls.go b/transport/simple-obfs/tls.go index 6484eb6..914cfe4 100644 --- a/transport/simple-obfs/tls.go +++ b/transport/simple-obfs/tls.go @@ -102,7 +102,8 @@ func (to *TLSObfs) write(b []byte) (int, error) { return len(b), err } - buf := &bytes.Buffer{} + buf := pool.GetBuffer() + defer pool.PutBuffer(buf) buf.Write([]byte{0x17, 0x03, 0x03}) binary.Write(buf, binary.BigEndian, uint16(len(b))) buf.Write(b) diff --git a/transport/snell/snell.go b/transport/snell/snell.go index ecbc90e..8966b4b 100644 --- a/transport/snell/snell.go +++ b/transport/snell/snell.go @@ -1,13 +1,13 @@ package snell import ( - "bytes" "encoding/binary" "errors" "fmt" "io" "net" - "sync" + + "github.com/Dreamacro/clash/common/pool" "github.com/Dreamacro/go-shadowsocks2/shadowaead" ) @@ -31,8 +31,7 @@ const ( ) var ( - bufferPool = sync.Pool{New: func() interface{} { return &bytes.Buffer{} }} - endSignal = []byte{} + endSignal = []byte{} ) type Snell struct { @@ -79,9 +78,8 @@ func (s *Snell) Read(b []byte) (int, error) { } func WriteHeader(conn net.Conn, host string, port uint, version int) error { - buf := bufferPool.Get().(*bytes.Buffer) - buf.Reset() - defer bufferPool.Put(buf) + buf := pool.GetBuffer() + defer pool.PutBuffer(buf) buf.WriteByte(Version) if version == Version2 { buf.WriteByte(CommandConnectV2) diff --git a/transport/ssr/obfs/http_simple.go b/transport/ssr/obfs/http_simple.go index 26fb312..c1ea767 100644 --- a/transport/ssr/obfs/http_simple.go +++ b/transport/ssr/obfs/http_simple.go @@ -10,7 +10,6 @@ import ( "strings" "github.com/Dreamacro/clash/common/pool" - "github.com/Dreamacro/clash/transport/ssr/tools" ) func init() { @@ -102,9 +101,8 @@ func (c *httpConn) Write(b []byte) (int, error) { hosts := strings.Split(host, ",") host = hosts[rand.Intn(len(hosts))] - buf := tools.BufPool.Get().(*bytes.Buffer) - defer tools.BufPool.Put(buf) - defer buf.Reset() + buf := pool.GetBuffer() + defer pool.PutBuffer(buf) if c.post { buf.WriteString("POST /") } else { diff --git a/transport/ssr/obfs/tls1.2_ticket_auth.go b/transport/ssr/obfs/tls1.2_ticket_auth.go index d914112..10f2786 100644 --- a/transport/ssr/obfs/tls1.2_ticket_auth.go +++ b/transport/ssr/obfs/tls1.2_ticket_auth.go @@ -87,9 +87,8 @@ func (c *tls12TicketConn) Read(b []byte) (int, error) { func (c *tls12TicketConn) Write(b []byte) (int, error) { length := len(b) if c.handshakeStatus == 8 { - buf := tools.BufPool.Get().(*bytes.Buffer) - defer tools.BufPool.Put(buf) - defer buf.Reset() + buf := pool.GetBuffer() + defer pool.PutBuffer(buf) for len(b) > 2048 { size := rand.Intn(4096) + 100 if len(b) < size { @@ -115,9 +114,8 @@ func (c *tls12TicketConn) Write(b []byte) (int, error) { if c.handshakeStatus == 0 { c.handshakeStatus = 1 - data := tools.BufPool.Get().(*bytes.Buffer) - defer tools.BufPool.Put(data) - defer data.Reset() + data := pool.GetBuffer() + defer pool.PutBuffer(data) data.Write([]byte{3, 3}) c.packAuthData(data) @@ -126,9 +124,8 @@ func (c *tls12TicketConn) Write(b []byte) (int, error) { data.Write([]byte{0x00, 0x1c, 0xc0, 0x2b, 0xc0, 0x2f, 0xcc, 0xa9, 0xcc, 0xa8, 0xcc, 0x14, 0xcc, 0x13, 0xc0, 0x0a, 0xc0, 0x14, 0xc0, 0x09, 0xc0, 0x13, 0x00, 0x9c, 0x00, 0x35, 0x00, 0x2f, 0x00, 0x0a}) data.Write([]byte{0x1, 0x0}) - ext := tools.BufPool.Get().(*bytes.Buffer) - defer tools.BufPool.Put(ext) - defer ext.Reset() + ext := pool.GetBuffer() + defer pool.PutBuffer(ext) host := c.getHost() ext.Write([]byte{0xff, 0x01, 0x00, 0x01, 0x00}) @@ -145,9 +142,8 @@ func (c *tls12TicketConn) Write(b []byte) (int, error) { binary.Write(data, binary.BigEndian, uint16(ext.Len())) data.ReadFrom(ext) - ret := tools.BufPool.Get().(*bytes.Buffer) - defer tools.BufPool.Put(ret) - defer ret.Reset() + ret := pool.GetBuffer() + defer pool.PutBuffer(ret) ret.Write([]byte{0x16, 3, 1}) binary.Write(ret, binary.BigEndian, uint16(data.Len()+4)) @@ -161,9 +157,8 @@ func (c *tls12TicketConn) Write(b []byte) (int, error) { } return length, nil } else if c.handshakeStatus == 1 && len(b) == 0 { - buf := tools.BufPool.Get().(*bytes.Buffer) - defer tools.BufPool.Put(buf) - defer buf.Reset() + buf := pool.GetBuffer() + defer pool.PutBuffer(buf) buf.Write([]byte{0x14, 3, 3, 0, 1, 1, 0x16, 3, 3, 0, 0x20}) tools.AppendRandBytes(buf, 22) diff --git a/transport/ssr/protocol/packet.go b/transport/ssr/protocol/packet.go index 621dc8c..249db70 100644 --- a/transport/ssr/protocol/packet.go +++ b/transport/ssr/protocol/packet.go @@ -1,10 +1,9 @@ package protocol import ( - "bytes" "net" - "github.com/Dreamacro/clash/transport/ssr/tools" + "github.com/Dreamacro/clash/common/pool" ) type PacketConn struct { @@ -13,9 +12,8 @@ type PacketConn struct { } func (c *PacketConn) WriteTo(b []byte, addr net.Addr) (int, error) { - buf := tools.BufPool.Get().(*bytes.Buffer) - defer tools.BufPool.Put(buf) - defer buf.Reset() + buf := pool.GetBuffer() + defer pool.PutBuffer(buf) err := c.EncodePacket(buf, b) if err != nil { return 0, err diff --git a/transport/ssr/protocol/stream.go b/transport/ssr/protocol/stream.go index 25d46f7..3c84615 100644 --- a/transport/ssr/protocol/stream.go +++ b/transport/ssr/protocol/stream.go @@ -5,7 +5,6 @@ import ( "net" "github.com/Dreamacro/clash/common/pool" - "github.com/Dreamacro/clash/transport/ssr/tools" ) type Conn struct { @@ -37,9 +36,8 @@ func (c *Conn) Read(b []byte) (int, error) { func (c *Conn) Write(b []byte) (int, error) { bLength := len(b) - buf := tools.BufPool.Get().(*bytes.Buffer) - defer tools.BufPool.Put(buf) - defer buf.Reset() + buf := pool.GetBuffer() + defer pool.PutBuffer(buf) err := c.Encode(buf, b) if err != nil { return 0, err diff --git a/transport/ssr/tools/bufPool.go b/transport/ssr/tools/bufPool.go index f3c45c4..ac15c97 100644 --- a/transport/ssr/tools/bufPool.go +++ b/transport/ssr/tools/bufPool.go @@ -2,17 +2,10 @@ package tools import ( "bytes" - "math/rand" - "sync" - - "github.com/Dreamacro/clash/common/pool" + "crypto/rand" + "io" ) -var BufPool = sync.Pool{New: func() interface{} { return &bytes.Buffer{} }} - func AppendRandBytes(b *bytes.Buffer, length int) { - randBytes := pool.Get(length) - defer pool.Put(randBytes) - rand.Read(randBytes) - b.Write(randBytes) + b.ReadFrom(io.LimitReader(rand.Reader, int64(length))) } diff --git a/transport/trojan/trojan.go b/transport/trojan/trojan.go index d62e3f8..d39cbec 100644 --- a/transport/trojan/trojan.go +++ b/transport/trojan/trojan.go @@ -1,7 +1,6 @@ package trojan import ( - "bytes" "crypto/sha256" "crypto/tls" "encoding/binary" @@ -11,6 +10,7 @@ import ( "net" "sync" + "github.com/Dreamacro/clash/common/pool" "github.com/Dreamacro/clash/transport/socks5" ) @@ -22,8 +22,6 @@ const ( var ( defaultALPN = []string{"h2", "http/1.1"} crlf = []byte{'\r', '\n'} - - bufPool = sync.Pool{New: func() interface{} { return &bytes.Buffer{} }} ) type Command = byte @@ -67,9 +65,8 @@ func (t *Trojan) StreamConn(conn net.Conn) (net.Conn, error) { } func (t *Trojan) WriteHeader(w io.Writer, command Command, socks5Addr []byte) error { - buf := bufPool.Get().(*bytes.Buffer) - defer bufPool.Put(buf) - defer buf.Reset() + buf := pool.GetBuffer() + defer pool.PutBuffer(buf) buf.Write(t.hexPassword) buf.Write(crlf) @@ -89,9 +86,8 @@ func (t *Trojan) PacketConn(conn net.Conn) net.PacketConn { } func writePacket(w io.Writer, socks5Addr, payload []byte) (int, error) { - buf := bufPool.Get().(*bytes.Buffer) - defer bufPool.Put(buf) - defer buf.Reset() + buf := pool.GetBuffer() + defer pool.PutBuffer(buf) buf.Write(socks5Addr) binary.Write(buf, binary.BigEndian, uint16(len(payload)))