diff --git a/README.md b/README.md index a24076e..c07b05e 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 1f6a207..135c422 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 c33d06f..d9a2cc2 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 0000000..ca7b6c4 --- /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 de09dfb..a80997b 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 0054be2..2775b17 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 5be6c30..3e87f14 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 7067cbb..2d13a13 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 dbdcde4..ffb6e87 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 72e6935..6e4156a 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 ec0accf..f70fc92 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 +} - return buf.Bytes() +func (h *Handshake) Len() int { + buf := acquireBytesBuffer() + defer releaseBytesBuffer(buf) + + h.WriteBytes(buf) + + return buf.Len() } diff --git a/tlstypes/pools.go b/tlstypes/pools.go new file mode 100644 index 0000000..f51aae7 --- /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 d6a71dd..5dae77e 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() +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) +} - buf.WriteByte(byte(r.Type)) - buf.Write(r.Version.Bytes()) - binary.Write(&buf, binary.BigEndian, uint16(len(data))) // nolint: errcheck - buf.Write(data) - - 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 3c9249c..a92325c 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 3bf90ca..d9166c1 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 6a1cd1a..6190d39 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 0000000..c27e30e --- /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 1b8aa9a..556a4f2 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 153e779..9c3c448 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 0000000..30a0d7d --- /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 62ac7d2..bd98edf 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 3db845b..d074dcb 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 a46b528..fd99169 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 d94b245..e005c84 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 0000000..3fd9907 --- /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) +}