Merge pull request #6 from 9seconds/syncpool

Use syncpool

Gopkg.lock

 [[projects]]
   name = "github.com/stretchr/testify"
   packages = ["assert"]
-  revision = "12b6f73e6084dad08a7c6e575284b177ecafbc71"
-  version = "v1.2.1"
+  revision = "f35b8ab0b5a2cef36673838d662e249dd9c94686"
+  version = "v1.2.2"
 
 [[projects]]
   name = "go.uber.org/atomic"

client/direct.go

 	if err != nil {
 		return 0, nil, errors.Annotate(err, "Cannot extract frame")
 	}
+	defer obfuscated2.ReturnFrame(frame)
 
 	obfs2, dc, err := obfuscated2.ParseObfuscated2ClientFrame(conf.Secret, frame)
 	if err != nil {

main.go

 import (
 	"encoding/json"
 	"io"
+	"math/rand"
 	"os"
+	"time"
 
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
 )
 
 func main() {
+	rand.Seed(time.Now().UTC().UnixNano())
+
 	app.Version(version)
 	kingpin.MustParse(app.Parse(os.Args[1:]))
 

obfuscated2/frame.go

 
 // Invert inverts frame for extracting encryption keys. Pkease check that link:
 // https://blog.susanka.eu/how-telegram-obfuscates-its-mtproto-traffic/
-func (f Frame) Invert() Frame {
-	reversed := make(Frame, FrameLen)
-	copy(reversed, f)
+func (f Frame) Invert() *Frame {
+	reversed := MakeFrame()
+	copy(*reversed, f)
 
 	for i := 0; i < frameLenKey+frameLenIV; i++ {
-		reversed[frameOffsetFirst+i] = f[frameOffsetIV-1-i]
+		(*reversed)[frameOffsetFirst+i] = f[frameOffsetIV-1-i]
 	}
 
 	return reversed
 }
 
 // ExtractFrame extracts exact obfuscated2 handshake frame from given reader.
-func ExtractFrame(conn io.Reader) (Frame, error) {
-	buf := &bytes.Buffer{}
+func ExtractFrame(conn io.Reader) (*Frame, error) {
+	frame := MakeFrame()
+	buf := bytes.NewBuffer(*frame)
+	buf.Reset()
+
 	if _, err := io.CopyN(buf, conn, FrameLen); err != nil {
+		ReturnFrame(frame)
 		return nil, errors.Annotate(err, "Cannot extract obfuscated header")
 	}
+	copy(*frame, buf.Bytes())
 
-	return Frame(buf.Bytes()), nil
+	return frame, nil
 }
 
-func generateFrame() Frame {
-	data := make(Frame, FrameLen)
+func generateFrame() *Frame {
+	frame := MakeFrame()
+	data := *frame
 
 	for {
 		if _, err := rand.Read(data); err != nil {
 
 		copy(data.Magic(), tgMagicBytes)
 
-		return data
+		return frame
 	}
 }

obfuscated2/frame_pool.go

+package obfuscated2
+
+import "sync"
+
+var framePool sync.Pool
+
+// MakeFrame returns new pointer to the handshake frame.
+func MakeFrame() *Frame {
+	return framePool.Get().(*Frame)
+}
+
+// ReturnFrame returns pointer to the handshake frame back to the pool.
+func ReturnFrame(f *Frame) {
+	framePool.Put(f)
+}
+
+func init() {
+	framePool = sync.Pool{
+		New: func() interface{} {
+			data := make(Frame, FrameLen)
+			return &data
+		},
+	}
+}

obfuscated2/frame_test.go

 package obfuscated2
 
 import (
+	"bytes"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 
 func TestFrameDoubleInvert(t *testing.T) {
 	frame := makeFrame()
-	assert.Equal(t, frame, frame.Invert().Invert())
+	assert.True(t, bytes.Equal(frame, *frame.Invert().Invert()))
 }
 
 func TestFrameInvert(t *testing.T) {
 	frame := makeFrame()
 	reversed := frame.Invert()
 
-	assert.Exactly(t, frame[:8], reversed[:8])
-	assert.Exactly(t, frame[56:], reversed[56:])
+	assert.Exactly(t, frame[:8], (*reversed)[:8])
+	assert.Exactly(t, frame[56:], (*reversed)[56:])
 
 	toCompare := make([]byte, 48)
 	for i := 0; i < 48; i++ {
 		toCompare[i] = frame[55-i]
 	}
-	assert.Equal(t, []byte(reversed[8:56]), toCompare)
+	assert.Equal(t, []byte((*reversed)[8:56]), toCompare)
 }
 
 func TestFrameGenerateValid(t *testing.T) {

obfuscated2/obfuscated2.go

 // details: http://telegra.ph/telegram-blocks-wtf-05-26
 //
 // Beware, link above is in russian.
-func ParseObfuscated2ClientFrame(secret []byte, frame Frame) (*Obfuscated2, int16, error) {
+func ParseObfuscated2ClientFrame(secret []byte, frame *Frame) (*Obfuscated2, int16, error) {
 	decHasher := sha256.New()
 	decHasher.Write(frame.Key()) // nolint: errcheck
 	decHasher.Write(secret)      // nolint: errcheck
 	encHasher.Write(secret)              // nolint: errcheck
 	encryptor := makeStreamCipher(encHasher.Sum(nil), invertedFrame.IV())
 
-	decryptedFrame := make(Frame, FrameLen)
-	decryptor.XORKeyStream(decryptedFrame, frame)
+	decryptedFrame := MakeFrame()
+	defer ReturnFrame(decryptedFrame)
+	decryptor.XORKeyStream(*decryptedFrame, *frame)
 	if !decryptedFrame.Valid() {
 		return nil, 0, errors.New("Unknown protocol")
 	}
 // MakeTelegramObfuscated2Frame creates new handshake frame to send to
 // Telegram.
 // https://blog.susanka.eu/how-telegram-obfuscates-its-mtproto-traffic/
-func MakeTelegramObfuscated2Frame() (*Obfuscated2, Frame) {
+func MakeTelegramObfuscated2Frame() (*Obfuscated2, *Frame) {
 	frame := generateFrame()
 
 	encryptor := makeStreamCipher(frame.Key(), frame.IV())
 	decryptorFrame := frame.Invert()
 	decryptor := makeStreamCipher(decryptorFrame.Key(), decryptorFrame.IV())
 
-	copyFrame := make(Frame, frameOffsetIV)
-	copy(copyFrame, frame)
-	encryptor.XORKeyStream(frame, frame)
-	copy(frame, copyFrame)
+	copyFrame := MakeFrame()
+	defer ReturnFrame(copyFrame)
+	copy((*copyFrame)[:frameOffsetIV], (*frame)[:frameOffsetIV])
+	encryptor.XORKeyStream(*frame, *frame)
+	copy((*frame)[:frameOffsetIV], (*copyFrame)[:frameOffsetIV])
 
 	obfs := &Obfuscated2{
 		Decryptor: decryptor,

obfuscated2/obfuscated2_test.go

 	decryptor := makeStreamCipher(frame.Key(), frame.IV())
 
 	decrypted := make(Frame, FrameLen)
-	decryptor.XORKeyStream(decrypted, frame)
+	decryptor.XORKeyStream(decrypted, *frame)
 
 	assert.True(t, decrypted.Valid())
 }
 
 	encryptor := makeStreamCipher(clientKey, clientFrame.IV())
 	encrypted := make(Frame, FrameLen)
-	encryptor.XORKeyStream(encrypted, clientFrame)
-	copy(encrypted[:56], clientFrame[:56])
+	encryptor.XORKeyStream(encrypted, *clientFrame)
+	copy(encrypted[:56], (*clientFrame)[:56])
 
 	invertedClientFrame := clientFrame.Invert()
 	clientHasher = sha256.New()
 	invertedClientKey := clientHasher.Sum(nil)
 	clientDecryptor := makeStreamCipher(invertedClientKey, invertedClientFrame.IV())
 
-	clientObfs, _, err := ParseObfuscated2ClientFrame(secret, encrypted)
+	clientObfs, _, err := ParseObfuscated2ClientFrame(secret, &encrypted)
 	assert.Nil(t, err)
 
 	tgObfs, tgFrame := MakeTelegramObfuscated2Frame()
 	tgDecryptor := makeStreamCipher(tgFrame.Key(), tgFrame.IV())
 	decrypted := make(Frame, FrameLen)
-	tgDecryptor.XORKeyStream(decrypted, tgFrame)
+	tgDecryptor.XORKeyStream(decrypted, *tgFrame)
 	assert.True(t, decrypted.Valid())
 
 	tgInvertedFrame := tgFrame.Invert()

proxy/copy_pool.go

+package proxy
+
+import "sync"
+
+const copyBufferSize = 30 * 1024
+
+var copyPool sync.Pool
+
+func init() {
+	copyPool = sync.Pool{
+		New: func() interface{} {
+			data := make([]byte, copyBufferSize)
+			return &data
+		},
+	}
+}

proxy/server.go

 
 	wait := &sync.WaitGroup{}
 	wait.Add(2)
-	go func() {
-		defer wait.Done()
-		io.Copy(clientConn, tgConn) // nolint: errcheck
-	}()
-	go func() {
-		defer wait.Done()
-		io.Copy(tgConn, clientConn) // nolint: errcheck
-	}()
+
+	go s.pipe(clientConn, tgConn, wait)
+	go s.pipe(tgConn, clientConn, wait)
+
 	<-ctx.Done()
 	wait.Wait()
 
 	return conn, nil
 }
 
+func (s *Server) pipe(dst io.Writer, src io.Reader, wait *sync.WaitGroup) {
+	defer wait.Done()
+
+	buf := copyPool.Get().(*[]byte)
+	defer copyPool.Put(buf)
+
+	io.CopyBuffer(dst, src, *buf) // nolint: errcheck
+}
+
 // NewServer creates new instance of MTPROTO proxy.
 func NewServer(conf *config.Config, logger *zap.SugaredLogger, stat *Stats) *Server {
 	return &Server{

run-mtg.sh

   chmod 0400 "$SECRET_PATH"
 )
 
-# docker pull "$IMAGE_NAME"
+docker pull "$IMAGE_NAME"
 docker ps --filter "Name=$CONTAINER_NAME" -aq | xargs -r docker rm -fv
 docker run \
     -d \

telegram/direct.go

 
 func (t *directTelegram) Init(conn io.ReadWriteCloser) (io.ReadWriteCloser, error) {
 	obfs2, frame := obfuscated2.MakeTelegramObfuscated2Frame()
-	if n, err := conn.Write(frame); err != nil || n != len(frame) {
+	defer obfuscated2.ReturnFrame(frame)
+
+	if n, err := conn.Write(*frame); err != nil || n != len(*frame) {
 		return nil, errors.Annotate(err, "Cannot write hadnshake frame")
 	}
 

wrappers/buffer_pool.go

+package wrappers
+
+import (
+	"bytes"
+	"sync"
+)
+
+var bufPool sync.Pool
+
+func getBuffer() *bytes.Buffer {
+	buf := bufPool.Get().(*bytes.Buffer)
+	buf.Reset()
+
+	return buf
+}
+
+func putBuffer(buf *bytes.Buffer) {
+	bufPool.Put(buf)
+}
+
+func init() {
+	bufPool = sync.Pool{
+		New: func() interface{} {
+			return &bytes.Buffer{}
+		},
+	}
+}

wrappers/streamcipherrwc.go

 
 // Write writes into connection.
 func (c *StreamCipherReadWriteCloser) Write(p []byte) (int, error) {
-	encrypted := make([]byte, len(p))
+	// This is to decrease an amount of allocations. Unfortunately, escape
+	// analysis in (at least Golang 1.10) is absolutely not perfect. For
+	// example, it understands that we want to have a slice locally, right?
+	// But since slice is effectively 2 ints + uintptr to [number]byte, the
+	// most heavyweight part is placed in heap.
+	buf := getBuffer()
+	defer putBuffer(buf)
+	buf.Grow(len(p))
+	buf.Write(p)
+
+	encrypted := buf.Bytes()
 	c.encryptor.XORKeyStream(encrypted, p)
-	allWritten := 0
-
-	for len(encrypted) > 0 {
-		n, err := c.conn.Write(encrypted)
-		allWritten += n
-		if err != nil {
-			return allWritten, err
-		}
-		encrypted = encrypted[n:]
-	}
 
-	return allWritten, nil
+	return c.conn.Write(encrypted)
 }
 
 // Close closes underlying connection.