Use pool for frames

Gopkg.lock

@@ -43,8 +43,8 @@
 [[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

@@ -18,6 +18,7 @@ func DirectInit(conn net.Conn, conf *config.Config) (int16, io.ReadWriteCloser,
 	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 {

obfuscated2/frame.go

@@ -68,29 +68,35 @@ func (f Frame) Valid() bool {
 
 // 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 {
@@ -112,6 +118,6 @@ func generateFrame() Frame {
 
 		copy(data.Magic(), tgMagicBytes)
 
-		return data
+		return frame
 	}
 }

obfuscated2/frame_pool.go

@@ -0,0 +1,24 @@
+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

@@ -1,6 +1,7 @@
 package obfuscated2
 
 import (
+	"bytes"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -47,21 +48,21 @@ func TestFrameValid(t *testing.T) {
 
 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

@@ -19,7 +19,7 @@ type Obfuscated2 struct {
 // 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
@@ -31,8 +31,9 @@ func ParseObfuscated2ClientFrame(secret []byte, frame Frame) (*Obfuscated2, int1
 	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")
 	}
@@ -48,17 +49,18 @@ func ParseObfuscated2ClientFrame(secret []byte, frame Frame) (*Obfuscated2, int1
 // 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

@@ -12,7 +12,7 @@ func TestObfs2TelegramFrameDecrypt(t *testing.T) {
 	decryptor := makeStreamCipher(frame.Key(), frame.IV())
 
 	decrypted := make(Frame, FrameLen)
-	decryptor.XORKeyStream(decrypted, frame)
+	decryptor.XORKeyStream(decrypted, *frame)
 
 	assert.True(t, decrypted.Valid())
 }
@@ -42,8 +42,8 @@ func TestObfs2Full(t *testing.T) {
 
 	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()
@@ -52,13 +52,13 @@ func TestObfs2Full(t *testing.T) {
 	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()

telegram/direct.go

@@ -43,7 +43,9 @@ func (t *directTelegram) Dial(dcIdx int16) (io.ReadWriteCloser, error) {
 
 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")
 	}