Support different client connection types

obfuscated2/frame.go

@@ -7,6 +7,8 @@ import (
 	"io"
 
 	"github.com/juju/errors"
+
+	"github.com/9seconds/mtg/mtproto"
 )
 
 // [frameOffsetFirst:frameOffsetKey:frameOffsetIV:frameOffsetMagic:frameOffsetDC:frameOffsetEnd]
@@ -22,13 +24,9 @@ const (
 	frameOffsetMagic = frameOffsetIV + frameLenMagic
 	frameOffsetDC    = frameOffsetMagic + frameLenDC
 
-	tgMagicByte = byte(239)
-
 	FrameLen = 64
 )
 
-var tgMagicBytes = []byte{tgMagicByte, tgMagicByte, tgMagicByte, tgMagicByte}
-
 // Frame represents handshake frame. Telegram sends 64 bytes of obfuscated2
 // initialization data first.
 // https://blog.susanka.eu/how-telegram-obfuscates-its-mtproto-traffic/
@@ -61,9 +59,9 @@ func (f Frame) DC() (n int16) {
 	return
 }
 
-// Valid checks that *decrypted* frame is valid. Only magic bytes are checked.
-func (f Frame) Valid() bool {
-	return bytes.Equal(f.Magic(), tgMagicBytes)
+// ConnectionType identifies connection type of the handshake frame.
+func (f Frame) ConnectionType() (mtproto.ConnectionType, error) {
+	return mtproto.ConnectionTagFromHandshake(f.Magic())
 }
 
 // Invert inverts frame for extracting encryption keys. Pkease check that link:
@@ -94,7 +92,7 @@ func ExtractFrame(conn io.Reader) (*Frame, error) {
 	return frame, nil
 }
 
-func generateFrame() *Frame {
+func generateFrame(connectionType mtproto.ConnectionType) *Frame {
 	frame := MakeFrame()
 	data := *frame
 
@@ -116,7 +114,9 @@ func generateFrame() *Frame {
 			continue
 		}
 
-		copy(data.Magic(), tgMagicBytes)
+		// error has to be checked before calling this function
+		tag, _ := connectionType.Tag() // nolint: errcheck
+		copy(data.Magic(), tag)
 
 		return frame
 	}

obfuscated2/frame_test.go

@@ -2,9 +2,12 @@ package obfuscated2
 
 import (
 	"bytes"
+	"strconv"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+
+	"github.com/9seconds/mtg/mtproto"
 )
 
 func TestFrameKey(t *testing.T) {
@@ -28,7 +31,7 @@ func TestFrameIV(t *testing.T) {
 func TestFrameMagic(t *testing.T) {
 	toCompare := make([]byte, 4)
 	for i := 0; i < 4; i++ {
-		toCompare[i] = tgMagicByte
+		toCompare[i] = 0xee
 	}
 
 	assert.Equal(t, toCompare, makeFrame().Magic())
@@ -40,10 +43,13 @@ func TestFrameDC(t *testing.T) {
 
 func TestFrameValid(t *testing.T) {
 	frame := makeFrame()
-	assert.True(t, frame.Valid())
+	connType, err := frame.ConnectionType()
+	assert.Nil(t, err)
+	assert.Equal(t, connType, mtproto.ConnectionTypeIntermediate)
 
 	frame[8+32+16+2] = byte(3)
-	assert.False(t, frame.Valid())
+	_, err = frame.ConnectionType()
+	assert.NotNil(t, err)
 }
 
 func TestFrameDoubleInvert(t *testing.T) {
@@ -66,7 +72,18 @@ func TestFrameInvert(t *testing.T) {
 }
 
 func TestFrameGenerateValid(t *testing.T) {
-	assert.True(t, generateFrame().Valid())
+	validTests := []mtproto.ConnectionType{
+		mtproto.ConnectionTypeIntermediate,
+		mtproto.ConnectionTypeAbridged,
+	}
+	for _, test := range validTests {
+		t.Run(strconv.Itoa(int(test)), func(tt *testing.T) {
+			frame := generateFrame(test)
+			conType, err := frame.ConnectionType()
+			assert.Nil(t, err)
+			assert.Equal(t, conType, test)
+		})
+	}
 }
 
 func makeFrame() Frame {
@@ -79,7 +96,7 @@ func makeFrame() Frame {
 		f[i] = byte(2)
 	}
 	for i := (8 + 32 + 16); i < (8 + 32 + 16 + 4); i++ {
-		f[i] = tgMagicByte
+		f[i] = 0xee
 	}
 	for i := (8 + 32 + 16 + 4); i < (8 + 32 + 16 + 4 + 2); i++ {
 		f[i] = byte(3)

obfuscated2/obfuscated2.go

@@ -6,6 +6,8 @@ import (
 	"crypto/sha256"
 
 	"github.com/juju/errors"
+
+	"github.com/9seconds/mtg/mtproto"
 )
 
 // Obfuscated2 contains AES CTR encryption and decryption streams
@@ -19,7 +21,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, *mtproto.ConnectionOpts, error) {
 	decHasher := sha256.New()
 	decHasher.Write(frame.Key()) // nolint: errcheck
 	decHasher.Write(secret)      // nolint: errcheck
@@ -34,23 +36,28 @@ func ParseObfuscated2ClientFrame(secret []byte, frame *Frame) (*Obfuscated2, int
 	decryptedFrame := MakeFrame()
 	defer ReturnFrame(decryptedFrame)
 	decryptor.XORKeyStream(*decryptedFrame, *frame)
-	if !decryptedFrame.Valid() {
-		return nil, 0, errors.New("Unknown protocol")
+	connType, err := decryptedFrame.ConnectionType()
+	if err != nil {
+		return nil, nil, errors.Annotate(err, "Unknown protocol")
 	}
 
 	obfs := &Obfuscated2{
 		Decryptor: decryptor,
 		Encryptor: encryptor,
 	}
+	connOpts := &mtproto.ConnectionOpts{
+		DC:             decryptedFrame.DC(),
+		ConnectionType: connType,
+	}
 
-	return obfs, decryptedFrame.DC(), nil
+	return obfs, connOpts, nil
 }
 
 // MakeTelegramObfuscated2Frame creates new handshake frame to send to
 // Telegram.
 // https://blog.susanka.eu/how-telegram-obfuscates-its-mtproto-traffic/
-func MakeTelegramObfuscated2Frame() (*Obfuscated2, *Frame) {
-	frame := generateFrame()
+func MakeTelegramObfuscated2Frame(opts *mtproto.ConnectionOpts) (*Obfuscated2, *Frame) {
+	frame := generateFrame(opts.ConnectionType)
 
 	encryptor := makeStreamCipher(frame.Key(), frame.IV())
 	decryptorFrame := frame.Invert()

obfuscated2/obfuscated2_test.go

@@ -5,20 +5,31 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+
+	"github.com/9seconds/mtg/mtproto"
 )
 
 func TestObfs2TelegramFrameDecrypt(t *testing.T) {
-	_, frame := MakeTelegramObfuscated2Frame()
+	connOpts := &mtproto.ConnectionOpts{
+		DC:             1,
+		ConnectionType: mtproto.ConnectionTypeIntermediate,
+	}
+	_, frame := MakeTelegramObfuscated2Frame(connOpts)
 	decryptor := makeStreamCipher(frame.Key(), frame.IV())
 
 	decrypted := make(Frame, FrameLen)
 	decryptor.XORKeyStream(decrypted, *frame)
 
-	assert.True(t, decrypted.Valid())
+	_, err := decrypted.ConnectionType()
+	assert.Nil(t, err)
 }
 
 func TestObfs2TelegramDecryptEncryptDecrypt(t *testing.T) {
-	obfs2, frame := MakeTelegramObfuscated2Frame()
+	connOpts := &mtproto.ConnectionOpts{
+		DC:             1,
+		ConnectionType: mtproto.ConnectionTypeIntermediate,
+	}
+	obfs2, frame := MakeTelegramObfuscated2Frame(connOpts)
 	inverted := frame.Invert()
 	encryptor := makeStreamCipher(inverted.Key(), inverted.IV())
 
@@ -34,7 +45,7 @@ func TestObfs2TelegramDecryptEncryptDecrypt(t *testing.T) {
 func TestObfs2Full(t *testing.T) {
 	secret := []byte{1, 2, 3, 4, 5}
 
-	clientFrame := generateFrame()
+	clientFrame := generateFrame(mtproto.ConnectionTypeIntermediate)
 	clientHasher := sha256.New()
 	clientHasher.Write(clientFrame.Key())
 	clientHasher.Write(secret)
@@ -55,11 +66,16 @@ func TestObfs2Full(t *testing.T) {
 	clientObfs, _, err := ParseObfuscated2ClientFrame(secret, &encrypted)
 	assert.Nil(t, err)
 
-	tgObfs, tgFrame := MakeTelegramObfuscated2Frame()
+	connOpts := &mtproto.ConnectionOpts{
+		DC:             1,
+		ConnectionType: mtproto.ConnectionTypeIntermediate,
+	}
+	tgObfs, tgFrame := MakeTelegramObfuscated2Frame(connOpts)
 	tgDecryptor := makeStreamCipher(tgFrame.Key(), tgFrame.IV())
 	decrypted := make(Frame, FrameLen)
 	tgDecryptor.XORKeyStream(decrypted, *tgFrame)
-	assert.True(t, decrypted.Valid())
+	_, err = decrypted.ConnectionType()
+	assert.Nil(t, err)
 
 	tgInvertedFrame := tgFrame.Invert()
 	tgEncryptor := makeStreamCipher(tgInvertedFrame.Key(), tgInvertedFrame.IV())