Introduce pools to obfuscated2

mtglib/internal/obfuscated2/client_handshake.go

@@ -5,12 +5,20 @@ import (
 	"crypto/subtle"
 	"encoding/hex"
 	"fmt"
+	"io"
 )
 
 // Connection Type secure. We support only fake tls.
 var clientHandshakeMagic = []byte{0xdd, 0xdd, 0xdd, 0xdd}
 
-func ClientHandshake(secret []byte, handshakeFrame *HandhakeFrame) (int16, cipher.Stream, cipher.Stream, error) {
+func ClientHandshake(secret []byte, reader io.Reader) (int16, cipher.Stream, cipher.Stream, error) {
+	handshakeFrame := acquireHandshakeFrame()
+	defer releaseHandshakeFrame(handshakeFrame)
+
+	if _, err := io.ReadFull(reader, handshakeFrame.data[:]); err != nil {
+		return 0, nil, nil, fmt.Errorf("cannot read frame: %w", err)
+	}
+
 	decHasher := acquireSha256Hasher()
 	defer releaseSha256Hasher(decHasher)
 
@@ -21,17 +29,22 @@ func ClientHandshake(secret []byte, handshakeFrame *HandhakeFrame) (int16, ciphe
 	encHasher := acquireSha256Hasher()
 	defer releaseSha256Hasher(encHasher)
 
-	invertedFrame := handshakeFrame.invert()
+	invertedFrame := acquireHandshakeFrame()
+	defer releaseHandshakeFrame(invertedFrame)
+
+	for i, v := range handshakeFrame.data {
+		invertedFrame.data[handshakeFrameLen-1-i] = v
+	}
+
 	encHasher.Write(invertedFrame.key()) // nolint: errcheck
-	encHasher.Write(secret)               // nolint: errcheck
+	encHasher.Write(secret)              // nolint: errcheck
 	encryptor := makeAesCtr(encHasher.Sum(nil), invertedFrame.iv())
 
-	decryptedFrame := HandhakeFrame{}
+	decryptor.XORKeyStream(handshakeFrame.data[:], handshakeFrame.data[:])
-	decryptor.XORKeyStream(decryptedFrame.data[:], handshakeFrame.data[:])
 
-	if magic := decryptedFrame.magic(); subtle.ConstantTimeCompare(clientHandshakeMagic, magic) != 1 {
+	if magic := handshakeFrame.magic(); subtle.ConstantTimeCompare(clientHandshakeMagic, magic) != 1 {
 		return 0, nil, nil, fmt.Errorf("unsupported connection type: %s", hex.EncodeToString(magic))
 	}
 
-	return decryptedFrame.dc(), encryptor, decryptor, nil
+	return handshakeFrame.dc(), encryptor, decryptor, nil
 }

mtglib/internal/obfuscated2/frame.go

@@ -1,10 +1,6 @@
 package obfuscated2
 
-import (
+import "encoding/binary"
-	"encoding/binary"
-	"fmt"
-	"io"
-)
 
 const (
 	handshakeFrameLen = 64
@@ -32,48 +28,24 @@ const (
 //    - 4 bytes of 'magic' - this has some settings like a connection type
 //    - 2 bytes of 'DC'. DC is little endian int16
 //    - 2 bytes of noise
-type HandhakeFrame struct {
+type handshakeFrame struct {
 	data [handshakeFrameLen]byte
 }
 
-func (f *HandhakeFrame) Fingerprint() []byte {
+func (h *handshakeFrame) dc() int16 {
-	return f.data[handshakeFrameOffsetStart:handshakeFrameOffsetEnd]
+	data := h.data[handshakeFrameOffsetDC:handshakeFrameOffsetEnd]
-}
-
-func (f *HandhakeFrame) dc() int16 {
-	data := f.data[handshakeFrameOffsetDC:handshakeFrameOffsetEnd]
 
 	return int16(binary.LittleEndian.Uint16(data))
 }
 
-func (f *HandhakeFrame) key() []byte {
+func (h *handshakeFrame) key() []byte {
-	return f.data[handshakeFrameLenKey:handshakeFrameOffsetIV]
+	return h.data[handshakeFrameLenKey:handshakeFrameOffsetIV]
 }
 
-func (f *HandhakeFrame) iv() []byte {
+func (h *handshakeFrame) iv() []byte {
-	return f.data[handshakeFrameOffsetIV:handshakeFrameOffsetMagic]
+	return h.data[handshakeFrameOffsetIV:handshakeFrameOffsetMagic]
 }
 
-func (f *HandhakeFrame) magic() []byte {
+func (h *handshakeFrame) magic() []byte {
-	return f.data[handshakeFrameOffsetMagic:handshakeFrameOffsetDC]
+	return h.data[handshakeFrameOffsetMagic:handshakeFrameOffsetDC]
-}
-
-func (f *HandhakeFrame) invert() *HandhakeFrame {
-	newFrame := &HandhakeFrame{}
-
-	for i, v := range f.data {
-		newFrame.data[handshakeFrameLen-1-i] = v
-	}
-
-	return newFrame
-}
-
-func ReadHandshakeFrame(reader io.Reader) (*HandhakeFrame, error) {
-	frame := &HandhakeFrame{}
-
-	if _, err := io.ReadFull(reader, frame.data[:]); err != nil {
-		return nil, fmt.Errorf("cannot read frame data: %w", err)
-	}
-
-	return frame, nil
 }

mtglib/internal/obfuscated2/pools.go

@@ -6,11 +6,18 @@ import (
 	"sync"
 )
 
-var sha256HasherPool = sync.Pool{
+var (
-	New: func() interface{} {
+	sha256HasherPool = sync.Pool{
-		return sha256.New()
+		New: func() interface{} {
-	},
+			return sha256.New()
-}
+		},
+	}
+	handshakeFramePool = sync.Pool{
+		New: func() interface{} {
+			return &handshakeFrame{}
+		},
+	}
+)
 
 func acquireSha256Hasher() hash.Hash {
 	return sha256HasherPool.Get().(hash.Hash)
@@ -20,3 +27,11 @@ func releaseSha256Hasher(h hash.Hash) {
 	h.Reset()
 	sha256HasherPool.Put(h)
 }
+
+func acquireHandshakeFrame() *handshakeFrame {
+	return handshakeFramePool.Get().(*handshakeFrame)
+}
+
+func releaseHandshakeFrame(h *handshakeFrame) {
+	handshakeFramePool.Put(h)
+}

mtglib/proxy.go

@@ -95,12 +95,7 @@ func (p *Proxy) Shutdown() {
 }
 
 func (p *Proxy) doObfuscated2Handshake(ctx *streamContext) error {
-	handshakeFrame, err := obfuscated2.ReadHandshakeFrame(ctx.clientConn)
+	dc, encryptor, decryptor, err := obfuscated2.ClientHandshake(p.secret.Key[:], ctx.clientConn)
-	if err != nil {
-		return fmt.Errorf("cannot read handshake frame: %w", err)
-	}
-
-	dc, encryptor, decryptor, err := obfuscated2.ClientHandshake(p.secret.Key[:], handshakeFrame)
 	if err != nil {
 		return fmt.Errorf("cannot process client handshake: %w", err)
 	}