diff --git a/README.md b/README.md index a24076ee5..c07b05ed1 100644 --- a/README.md +++ b/README.md @@ -194,11 +194,14 @@ supported environment variables: | `MTG_STATSD_PREFIX` | `--statsd-prefix` | `mtg` | Which bucket prefix we should use. For example, if you set `mtg`, then metric `traffic.ingress` would be send as `mtg.traffic.ingress`. | | `MTG_STATSD_TAGS_FORMAT` | `--statsd-tags-format` | | Which tags format we should use. By default, we are using default vanilla statsd tags format but if you want to send directly to InfluxDB or Datadog, please specify it there. Possible options are `influxdb` and `datadog`. | | `MTG_STATSD_TAGS` | `--statsd-tags` | | Which tags should we send to statsd with our metrics. Please specify them as `key=value` pairs. | -| `MTG_BUFFER_WRITE` | `-w`, `--write-buffer` | `64KB` | The size of TCP write buffer in bytes. Write buffer is the buffer for messages which are going from client to Telegram. | -| `MTG_BUFFER_READ` | `-r`, `--read-buffer` | `128KB` | The size of TCP read buffer in bytes. Read buffer is the buffer for messages from Telegram to client. | +| `MTG_BUFFER_WRITE` | `-w`, `--write-buffer` | `32KB` | The size of TCP write buffer in bytes. Write buffer is the buffer for messages which are going from client to Telegram. | +| `MTG_BUFFER_READ` | `-r`, `--read-buffer` | `32KB` | The size of TCP read buffer in bytes. Read buffer is the buffer for messages from Telegram to client. | | `MTG_ANTIREPLAY_MAXSIZE` | `--anti-replay-max-size` | `128MB` | Max size of antireplay cache. | | `MTG_CLOAK_PORT` | `--cloak-port` | `443` | Which port we should use to connect to cloaked host in FakeTLS mode. | | `MTG_MULTIPLEX_PERCONNECTION` | `--multiplex-per-connection` | `50` | How many client connections can share a single Telegram connection in adtag mode | +| `MTG_NTP_SERVERS` | `--ntp-server` | default pool | A list of NTP servers to use. | +| `MTG_PREFER_DIRECT_IP` | `--prefer-ip` | `ipv6` | Which IP protocol to prefer if possible. Works mostly in direct mode. | + Usually you want to modify only read/write buffer sizes. If you feel that proxy is slow, try to increase both sizes giving more priority to diff --git a/faketls/client_protocol.go b/faketls/client_protocol.go index 1f6a20750..135c42256 100644 --- a/faketls/client_protocol.go +++ b/faketls/client_protocol.go @@ -63,7 +63,12 @@ func (c *ClientProtocol) tlsHandshake(conn io.ReadWriter) error { return fmt.Errorf("cannot read initial record: %w", err) } - clientHello, err := tlstypes.ParseClientHello(helloRecord.Data.Bytes()) + buf := acquireBytesBuffer() + defer releaseBytesBuffer(buf) + + helloRecord.Data.WriteBytes(buf) + + clientHello, err := tlstypes.ParseClientHello(buf.Bytes()) if err != nil { return fmt.Errorf("cannot parse client hello: %w", err) } diff --git a/faketls/cloak.go b/faketls/cloak.go index c33d06f0b..d9a2cc2ca 100644 --- a/faketls/cloak.go +++ b/faketls/cloak.go @@ -28,15 +28,9 @@ func cloak(one, another io.ReadWriteCloser) { wg.Add(2) - go func() { - defer wg.Done() - io.Copy(one, another) // nolint: errcheck - }() + go cloakPipe(one, another, wg) - go func() { - defer wg.Done() - io.Copy(another, one) // nolint: errcheck - }() + go cloakPipe(another, one, wg) go func() { wg.Wait() @@ -69,3 +63,12 @@ func cloak(one, another io.ReadWriteCloser) { <-ctx.Done() } + +func cloakPipe(one io.Writer, another io.Reader, wg *sync.WaitGroup) { + defer wg.Done() + + buf := acquireCloakBuffer() + defer releaseCloakBuffer(buf) + + io.CopyBuffer(one, another, *buf) // nolint: errcheck +} diff --git a/faketls/pools.go b/faketls/pools.go new file mode 100644 index 000000000..ca7b6c4d1 --- /dev/null +++ b/faketls/pools.go @@ -0,0 +1,39 @@ +package faketls + +import ( + "bytes" + "sync" +) + +const cloakBufferSize = 1024 + +var ( + poolBytesBuffer = sync.Pool{ + New: func() interface{} { + return &bytes.Buffer{} + }, + } + poolCloakBuffer = sync.Pool{ + New: func() interface{} { + rv := make([]byte, cloakBufferSize) + return &rv + }, + } +) + +func acquireBytesBuffer() *bytes.Buffer { + return poolBytesBuffer.Get().(*bytes.Buffer) +} + +func acquireCloakBuffer() *[]byte { + return poolCloakBuffer.Get().(*[]byte) +} + +func releaseBytesBuffer(buf *bytes.Buffer) { + buf.Reset() + poolBytesBuffer.Put(buf) +} + +func releaseCloakBuffer(buf *[]byte) { + poolCloakBuffer.Put(buf) +} diff --git a/go.mod b/go.mod index de09dfba3..a80997bee 100644 --- a/go.mod +++ b/go.mod @@ -11,10 +11,10 @@ require ( github.com/prometheus/procfs v0.0.11 // indirect github.com/smira/go-statsd v1.3.1 go.uber.org/zap v1.14.1 - golang.org/x/crypto v0.0.0-20200317142112-1b76d66859c6 + golang.org/x/crypto v0.0.0-20200323165209-0ec3e9974c59 golang.org/x/lint v0.0.0-20200302205851-738671d3881b // indirect - golang.org/x/net v0.0.0-20200319234117-63522dbf7eec // indirect - golang.org/x/sys v0.0.0-20200317113312-5766fd39f98d + golang.org/x/net v0.0.0-20200324143707-d3edc9973b7e // indirect + golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd golang.org/x/tools v0.0.0-20200319210407-521f4a0cd458 // indirect gopkg.in/alecthomas/kingpin.v2 v2.2.6 honnef.co/go/tools v0.0.1-2020.1.3 // indirect diff --git a/go.sum b/go.sum index 0054be28d..2775b1704 100644 --- a/go.sum +++ b/go.sum @@ -123,8 +123,8 @@ golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2 h1:VklqNMn3ovrHsnt90Pveol golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= golang.org/x/crypto v0.0.0-20190510104115-cbcb75029529/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= -golang.org/x/crypto v0.0.0-20200317142112-1b76d66859c6 h1:TjszyFsQsyZNHwdVdZ5m7bjmreu0znc2kRYsEml9/Ww= -golang.org/x/crypto v0.0.0-20200317142112-1b76d66859c6/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= +golang.org/x/crypto v0.0.0-20200323165209-0ec3e9974c59 h1:3zb4D3T4G8jdExgVU/95+vQXfpEPiMdCaZgmGVxjNHM= +golang.org/x/crypto v0.0.0-20200323165209-0ec3e9974c59/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= golang.org/x/lint v0.0.0-20190930215403-16217165b5de h1:5hukYrvBGR8/eNkX5mdUezrA6JiaEZDtJb9Ei+1LlBs= golang.org/x/lint v0.0.0-20190930215403-16217165b5de/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc= golang.org/x/lint v0.0.0-20200302205851-738671d3881b h1:Wh+f8QHJXR411sJR8/vRBTZ7YapZaRvUcLFFJhusH0k= @@ -140,8 +140,8 @@ golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn golang.org/x/net v0.0.0-20190613194153-d28f0bde5980/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= -golang.org/x/net v0.0.0-20200319234117-63522dbf7eec h1:w0SItUiQ4sBiXBAwWNkyu8Fu2Qpn/dtDIcoPkPDqjRw= -golang.org/x/net v0.0.0-20200319234117-63522dbf7eec/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= +golang.org/x/net v0.0.0-20200324143707-d3edc9973b7e h1:3G+cUijn7XD+S4eJFddp53Pv7+slrESplyjG25HgL+k= +golang.org/x/net v0.0.0-20200324143707-d3edc9973b7e/go.mod h1:qpuaurCH72eLCgpAm/N6yyVIVM9cpaDIP3A8BGJEC5A= golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20181221193216-37e7f081c4d4 h1:YUO/7uOKsKeq9UokNS62b8FYywz3ker1l1vDZRCRefw= golang.org/x/sync v0.0.0-20181221193216-37e7f081c4d4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= @@ -154,8 +154,8 @@ golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7w golang.org/x/sys v0.0.0-20190422165155-953cdadca894/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20200106162015-b016eb3dc98e/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20200122134326-e047566fdf82/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= -golang.org/x/sys v0.0.0-20200317113312-5766fd39f98d h1:62ap6LNOjDU6uGmKXHJbSfciMoV+FeI1sRXx/pLDL44= -golang.org/x/sys v0.0.0-20200317113312-5766fd39f98d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd h1:xhmwyvizuTgC2qz7ZlMluP20uW+C3Rm0FD/WLDX8884= +golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs= golang.org/x/tools v0.0.0-20190621195816-6e04913cbbac/go.mod h1:/rFqwRUd4F7ZHNgwSSTFct+R/Kf4OFW1sUzUTQQTgfc= diff --git a/main.go b/main.go index 5be6c30a3..3e87f14bb 100644 --- a/main.go +++ b/main.go @@ -92,13 +92,13 @@ var ( "Write buffer size. You can think about it as a buffer from client to Telegram."). Short('w'). Envar("MTG_BUFFER_WRITE"). - Default("64KB"). + Default("32KB"). Bytes() runReadBufferSize = runCommand.Flag("read-buffer", "Read buffer size. You can think about it as a buffer from Telegram to client."). Short('r'). Envar("MTG_BUFFER_READ"). - Default("128KB"). + Default("32KB"). Bytes() runTLSCloakPort = runCommand.Flag("cloak-port", "Port which should be used for host cloaking."). diff --git a/proxy/direct.go b/proxy/direct.go index 7067cbb4c..2d13a13be 100644 --- a/proxy/direct.go +++ b/proxy/direct.go @@ -11,7 +11,7 @@ import ( "github.com/9seconds/mtg/protocol" ) -const directPipeBufferSize = 1024 * 1024 +const directPipeBufferSize = 1024 func directConnection(request *protocol.TelegramRequest) error { telegramConnRaw, err := obfuscated2.TelegramProtocol(request) @@ -42,8 +42,9 @@ func directPipe(dst io.WriteCloser, src io.ReadCloser, wg *sync.WaitGroup, logge wg.Done() }() - buf := make([]byte, directPipeBufferSize) - if _, err := io.CopyBuffer(dst, src, buf); err != nil { + buf := [directPipeBufferSize]byte{} + + if _, err := io.CopyBuffer(dst, src, buf[:]); err != nil { logger.Debugw("Cannot pump sockets", "error", err) } } diff --git a/tlstypes/client_hello.go b/tlstypes/client_hello.go index dbdcde4ba..ffb6e8748 100644 --- a/tlstypes/client_hello.go +++ b/tlstypes/client_hello.go @@ -25,7 +25,7 @@ func (c ClientHello) Digest() []byte { } mac := hmac.New(sha256.New, config.C.Secret) - mac.Write(rec.Bytes()) // nolint: errcheck + rec.WriteBytes(mac) computedDigest := mac.Sum(nil) for i := range computedDigest { diff --git a/tlstypes/consts.go b/tlstypes/consts.go index 72e6935d6..6e4156a28 100644 --- a/tlstypes/consts.go +++ b/tlstypes/consts.go @@ -1,5 +1,7 @@ package tlstypes +import "io" + type RecordType uint8 const ( @@ -69,11 +71,16 @@ var ( ) type Byter interface { - Bytes() []byte + WriteBytes(io.Writer) + Len() int } type RawBytes []byte -func (r RawBytes) Bytes() []byte { - return []byte(r) +func (r RawBytes) WriteBytes(writer io.Writer) { + writer.Write(r) // nolint: errcheck +} + +func (r RawBytes) Len() int { + return len(r) } diff --git a/tlstypes/handshake.go b/tlstypes/handshake.go index ec0accfdf..f70fc923f 100644 --- a/tlstypes/handshake.go +++ b/tlstypes/handshake.go @@ -1,7 +1,7 @@ package tlstypes import ( - "bytes" + "io" "github.com/9seconds/mtg/utils" ) @@ -14,24 +14,31 @@ type Handshake struct { Tail Byter } -func (h *Handshake) Bytes() []byte { - buf := bytes.Buffer{} - packetBuf := bytes.Buffer{} +func (h *Handshake) WriteBytes(writer io.Writer) { + packetBuf := acquireBytesBuffer() + defer releaseBytesBuffer(packetBuf) - buf.WriteByte(byte(h.Type)) + writer.Write([]byte{byte(h.Type)}) // nolint: errcheck packetBuf.Write(h.Version.Bytes()) packetBuf.Write(h.Random[:]) packetBuf.WriteByte(byte(len(h.SessionID))) packetBuf.Write(h.SessionID) - packetBuf.Write(h.Tail.Bytes()) + h.Tail.WriteBytes(packetBuf) sizeUint24 := utils.ToUint24(uint32(packetBuf.Len())) sizeUint24Bytes := sizeUint24[:] sizeUint24Bytes[0], sizeUint24Bytes[2] = sizeUint24Bytes[2], sizeUint24Bytes[0] - buf.Write(sizeUint24Bytes) - packetBuf.WriteTo(&buf) // nolint: errcheck + writer.Write(sizeUint24Bytes) // nolint: errcheck + packetBuf.WriteTo(writer) // nolint: errcheck +} + +func (h *Handshake) Len() int { + buf := acquireBytesBuffer() + defer releaseBytesBuffer(buf) + + h.WriteBytes(buf) - return buf.Bytes() + return buf.Len() } diff --git a/tlstypes/pools.go b/tlstypes/pools.go new file mode 100644 index 000000000..f51aae72b --- /dev/null +++ b/tlstypes/pools.go @@ -0,0 +1,23 @@ +package tlstypes + +import ( + "bytes" + "sync" +) + +var ( + poolBytesBuffer = sync.Pool{ + New: func() interface{} { + return &bytes.Buffer{} + }, + } +) + +func acquireBytesBuffer() *bytes.Buffer { + return poolBytesBuffer.Get().(*bytes.Buffer) +} + +func releaseBytesBuffer(buf *bytes.Buffer) { + buf.Reset() + poolBytesBuffer.Put(buf) +} diff --git a/tlstypes/record.go b/tlstypes/record.go index d6a71ddf0..5dae77e54 100644 --- a/tlstypes/record.go +++ b/tlstypes/record.go @@ -15,16 +15,15 @@ type Record struct { Data Byter } -func (r Record) Bytes() []byte { - buf := bytes.Buffer{} - data := r.Data.Bytes() - - buf.WriteByte(byte(r.Type)) - buf.Write(r.Version.Bytes()) - binary.Write(&buf, binary.BigEndian, uint16(len(data))) // nolint: errcheck - buf.Write(data) +func (r Record) WriteBytes(writer io.Writer) { + writer.Write([]byte{byte(r.Type)}) // nolint: errcheck + writer.Write(r.Version.Bytes()) // nolint: errcheck + binary.Write(writer, binary.BigEndian, uint16(r.Data.Len())) // nolint: errcheck + r.Data.WriteBytes(writer) +} - return buf.Bytes() +func (r Record) Len() int { + return 1 + 2 + 2 + r.Data.Len() } func ReadRecord(reader io.Reader) (Record, error) { diff --git a/tlstypes/server_hello.go b/tlstypes/server_hello.go index 3c9249c60..a92325ca3 100644 --- a/tlstypes/server_hello.go +++ b/tlstypes/server_hello.go @@ -20,20 +20,22 @@ type ServerHello struct { } func (s ServerHello) WelcomePacket() []byte { + buf := &bytes.Buffer{} + s.Random = [32]byte{} rec := Record{ Type: RecordTypeHandshake, Version: Version12, Data: &s, } - buf := bytes.NewBuffer(rec.Bytes()) + rec.WriteBytes(buf) recChangeCipher := Record{ Type: RecordTypeChangeCipherSpec, Version: Version12, Data: RawBytes([]byte{0x01}), } - buf.Write(recChangeCipher.Bytes()) + recChangeCipher.WriteBytes(buf) hostCert := make([]byte, 1024+mrand.Intn(3092)) rand.Read(hostCert) // nolint: errcheck @@ -43,7 +45,8 @@ func (s ServerHello) WelcomePacket() []byte { Version: Version12, Data: RawBytes(hostCert), } - buf.Write(recData.Bytes()) + recData.WriteBytes(buf) + packet := buf.Bytes() mac := hmac.New(sha256.New, config.C.Secret) diff --git a/utils/init_tcp.go b/utils/init_tcp.go index 3bf90cad8..d9166c108 100644 --- a/utils/init_tcp.go +++ b/utils/init_tcp.go @@ -3,10 +3,13 @@ package utils import ( "fmt" "net" + "time" "github.com/9seconds/mtg/config" ) +const tcpKeepAlivePingPeriod = 2 * time.Second + func InitTCP(conn net.Conn) error { tcpConn := conn.(*net.TCPConn) @@ -22,5 +25,13 @@ func InitTCP(conn net.Conn) error { return fmt.Errorf("cannot set write buffer size: %w", err) } + if err := tcpConn.SetKeepAlive(true); err != nil { + return fmt.Errorf("cannot enable keep-alive: %w", err) + } + + if err := tcpConn.SetKeepAlivePeriod(tcpKeepAlivePingPeriod); err != nil { + return fmt.Errorf("cannot set keep-alive period: %w", err) + } + return nil } diff --git a/wrappers/packet/mtproto_frame.go b/wrappers/packet/mtproto_frame.go index 6a1cd1aea..6190d39fa 100644 --- a/wrappers/packet/mtproto_frame.go +++ b/wrappers/packet/mtproto_frame.go @@ -42,7 +42,9 @@ type wrapperMtprotoFrame struct { } func (w *wrapperMtprotoFrame) Read() (conntypes.Packet, error) { // nolint: funlen - buf := &bytes.Buffer{} + buf := acquireMtprotoFrameBytesBuffer() + defer releaseMtprotoFrameBytesBuffer(buf) + sum := crc32.NewIEEE() writer := io.MultiWriter(buf, sum) @@ -71,7 +73,6 @@ func (w *wrapperMtprotoFrame) Read() (conntypes.Packet, error) { // nolint: funl } buf.Reset() - buf.Grow(int(messageLength) - 4 - 4) if _, err := io.CopyN(writer, w.parent, int64(messageLength)-4-4); err != nil { return nil, fmt.Errorf("cannot read the message frame: %w", err) @@ -113,8 +114,8 @@ func (w *wrapperMtprotoFrame) Write(p conntypes.Packet) error { messageLength := 4 + 4 + len(p) + 4 paddingLength := (aes.BlockSize - messageLength%aes.BlockSize) % aes.BlockSize - buf := &bytes.Buffer{} - buf.Grow(messageLength + paddingLength) + buf := acquireMtprotoFrameBytesBuffer() + defer releaseMtprotoFrameBytesBuffer(buf) binary.Write(buf, binary.LittleEndian, uint32(messageLength)) // nolint: errcheck binary.Write(buf, binary.LittleEndian, w.writeSeqNo) // nolint: errcheck diff --git a/wrappers/packet/pools.go b/wrappers/packet/pools.go new file mode 100644 index 000000000..c27e30ecd --- /dev/null +++ b/wrappers/packet/pools.go @@ -0,0 +1,23 @@ +package packet + +import ( + "bytes" + "sync" +) + +var ( + poolMtprotoFrameBytesBuffer = sync.Pool{ + New: func() interface{} { + return &bytes.Buffer{} + }, + } +) + +func acquireMtprotoFrameBytesBuffer() *bytes.Buffer { + return poolMtprotoFrameBytesBuffer.Get().(*bytes.Buffer) +} + +func releaseMtprotoFrameBytesBuffer(buf *bytes.Buffer) { + buf.Reset() + poolMtprotoFrameBytesBuffer.Put(buf) +} diff --git a/wrappers/packetack/client_abridged.go b/wrappers/packetack/client_abridged.go index 1b8aa9aeb..556a4f2a8 100644 --- a/wrappers/packetack/client_abridged.go +++ b/wrappers/packetack/client_abridged.go @@ -88,7 +88,9 @@ func (w *wrapperClientAbridged) Write(packet conntypes.Packet, acks *conntypes.C return nil case packetLength < clientAbridgedLargePacketLength: length24 := utils.ToUint24(uint32(packetLength)) - buf := bytes.Buffer{} + + buf := acquireClientBytesBuffer() + defer releaseClientBytesBuffer(buf) buf.WriteByte(byte(clientAbridgedSmallPacketLength)) buf.Write(length24[:]) diff --git a/wrappers/packetack/client_intermediate_secure.go b/wrappers/packetack/client_intermediate_secure.go index 153e77957..9c3c448f0 100644 --- a/wrappers/packetack/client_intermediate_secure.go +++ b/wrappers/packetack/client_intermediate_secure.go @@ -1,7 +1,6 @@ package packetack import ( - "bytes" "encoding/binary" "fmt" "math/rand" @@ -35,11 +34,13 @@ func (w *wrapperClientIntermediateSecure) Write(packet conntypes.Packet, acks *c return nil } - buf := bytes.Buffer{} + buf := acquireClientBytesBuffer() + defer releaseClientBytesBuffer(buf) + paddingLength := rand.Intn(4) buf.Grow(4 + len(packet) + paddingLength) - binary.Write(&buf, binary.LittleEndian, uint32(len(packet)+paddingLength)) // nolint: errcheck + binary.Write(buf, binary.LittleEndian, uint32(len(packet)+paddingLength)) // nolint: errcheck buf.Write(packet) buf.Write(make([]byte, paddingLength)) diff --git a/wrappers/packetack/pools.go b/wrappers/packetack/pools.go new file mode 100644 index 000000000..30a0d7d39 --- /dev/null +++ b/wrappers/packetack/pools.go @@ -0,0 +1,23 @@ +package packetack + +import ( + "bytes" + "sync" +) + +var ( + poolClientBytesBuffer = sync.Pool{ + New: func() interface{} { + return &bytes.Buffer{} + }, + } +) + +func acquireClientBytesBuffer() *bytes.Buffer { + return poolClientBytesBuffer.Get().(*bytes.Buffer) +} + +func releaseClientBytesBuffer(buf *bytes.Buffer) { + buf.Reset() + poolClientBytesBuffer.Put(buf) +} diff --git a/wrappers/packetack/proxy.go b/wrappers/packetack/proxy.go index 62ac7d2de..bd98edfaa 100644 --- a/wrappers/packetack/proxy.go +++ b/wrappers/packetack/proxy.go @@ -23,8 +23,8 @@ type wrapperProxy struct { func (w *wrapperProxy) Write(packet conntypes.Packet, acks *conntypes.ConnectionAcks) error { buf := bytes.Buffer{} - flags := w.flags + if acks.Quick { flags |= rpc.ProxyRequestFlagsQuickAck } @@ -43,6 +43,7 @@ func (w *wrapperProxy) Write(packet conntypes.Packet, acks *conntypes.Connection buf.WriteByte(byte(len(config.C.AdTag))) buf.Write(config.C.AdTag) buf.Write(make([]byte, (4-buf.Len()%4)%4)) + buf.Grow(len(packet)) buf.Write(packet) return w.proxy.Write(buf.Bytes()) diff --git a/wrappers/stream/faketls.go b/wrappers/stream/faketls.go index 3db845b78..d074dcbb1 100644 --- a/wrappers/stream/faketls.go +++ b/wrappers/stream/faketls.go @@ -1,6 +1,7 @@ package stream import ( + "bytes" "errors" "fmt" "net" @@ -39,13 +40,19 @@ func (w *wrapperFakeTLS) WriteTimeout(p []byte, timeout time.Duration) (int, err func (w *wrapperFakeTLS) write(p []byte, writeFunc func([]byte) (int, error)) (int, error) { sum := 0 + buf := acquireBytesBuffer() + defer releaseBytesBuffer(buf) + for _, v := range tlstypes.MakeRecords(p) { - _, err := writeFunc(v.Bytes()) + buf.Reset() + v.WriteBytes(buf) + + _, err := writeFunc(buf.Bytes()) if err != nil { return sum, err } - sum += len(v.Data.Bytes()) + sum += v.Data.Len() } return sum, nil @@ -86,7 +93,10 @@ func NewFakeTLS(socket conntypes.StreamReadWriteCloser) conntypes.StreamReadWrit switch rec.Type { case tlstypes.RecordTypeChangeCipherSpec: case tlstypes.RecordTypeApplicationData: - return rec.Data.Bytes(), nil + buf := &bytes.Buffer{} + rec.Data.WriteBytes(buf) + + return buf.Bytes(), nil default: return nil, fmt.Errorf("unsupported record type %v", rec.Type) } diff --git a/wrappers/stream/mtproto_cipher.go b/wrappers/stream/mtproto_cipher.go index a46b52816..fd9916944 100644 --- a/wrappers/stream/mtproto_cipher.go +++ b/wrappers/stream/mtproto_cipher.go @@ -1,7 +1,6 @@ package stream import ( - "bytes" "crypto/aes" "crypto/cipher" "crypto/md5" // nolint: gosec @@ -54,7 +53,9 @@ func mtprotoDeriveKeys(purpose mtprotoCipherPurpose, resp *rpc.NonceResponse, client, remote *net.TCPAddr, secret []byte) ([]byte, []byte) { - message := bytes.Buffer{} + message := acquireBytesBuffer() + defer releaseBytesBuffer(message) + message.Write(resp.Nonce) // nolint: gosec message.Write(req.Nonce) // nolint: gosec message.Write(req.CryptoTS) // nolint: gosec diff --git a/wrappers/stream/obfuscated2.go b/wrappers/stream/obfuscated2.go index d94b245b3..e005c846d 100644 --- a/wrappers/stream/obfuscated2.go +++ b/wrappers/stream/obfuscated2.go @@ -40,16 +40,26 @@ func (w *wrapperObfuscated2) Read(p []byte) (int, error) { } func (w *wrapperObfuscated2) WriteTimeout(p []byte, timeout time.Duration) (int, error) { - buf := make([]byte, len(p)) - copy(buf, p) + buffer := acquireBytesBuffer() + defer releaseBytesBuffer(buffer) + + buffer.Write(p) + + buf := buffer.Bytes() + w.encryptor.XORKeyStream(buf, buf) return w.parent.WriteTimeout(buf, timeout) } func (w *wrapperObfuscated2) Write(p []byte) (int, error) { - buf := make([]byte, len(p)) - copy(buf, p) + buffer := acquireBytesBuffer() + defer releaseBytesBuffer(buffer) + + buffer.Write(p) + + buf := buffer.Bytes() + w.encryptor.XORKeyStream(buf, buf) return w.parent.Write(buf) diff --git a/wrappers/stream/pools.go b/wrappers/stream/pools.go new file mode 100644 index 000000000..3fd990703 --- /dev/null +++ b/wrappers/stream/pools.go @@ -0,0 +1,23 @@ +package stream + +import ( + "bytes" + "sync" +) + +var ( + poolBytesBuffer = sync.Pool{ + New: func() interface{} { + return &bytes.Buffer{} + }, + } +) + +func acquireBytesBuffer() *bytes.Buffer { + return poolBytesBuffer.Get().(*bytes.Buffer) +} + +func releaseBytesBuffer(buf *bytes.Buffer) { + buf.Reset() + poolBytesBuffer.Put(buf) +}