Support different client connection types

client/client.go

 	"net"
 
 	"github.com/9seconds/mtg/config"
+	"github.com/9seconds/mtg/mtproto"
 )
 
 // Init has to initialize client connection based on given config.
-type Init func(net.Conn, *config.Config) (int16, io.ReadWriteCloser, error)
+type Init func(net.Conn, *config.Config) (*mtproto.ConnectionOpts, io.ReadWriteCloser, error)

client/direct.go

 	"github.com/juju/errors"
 
 	"github.com/9seconds/mtg/config"
+	"github.com/9seconds/mtg/mtproto"
 	"github.com/9seconds/mtg/obfuscated2"
 	"github.com/9seconds/mtg/wrappers"
 )
 
 // DirectInit initializes client to access Telegram bypassing middleproxies.
-func DirectInit(conn net.Conn, conf *config.Config) (int16, io.ReadWriteCloser, error) {
+func DirectInit(conn net.Conn, conf *config.Config) (*mtproto.ConnectionOpts, io.ReadWriteCloser, error) {
 	socket := wrappers.NewTimeoutRWC(conn, conf.TimeoutRead, conf.TimeoutWrite)
 	frame, err := obfuscated2.ExtractFrame(socket)
 	if err != nil {
-		return 0, nil, errors.Annotate(err, "Cannot extract frame")
+		return nil, nil, errors.Annotate(err, "Cannot extract frame")
 	}
 	defer obfuscated2.ReturnFrame(frame)
 
-	obfs2, dc, err := obfuscated2.ParseObfuscated2ClientFrame(conf.Secret, frame)
+	obfs2, connOpts, err := obfuscated2.ParseObfuscated2ClientFrame(conf.Secret, frame)
 	if err != nil {
-		return 0, nil, errors.Annotate(err, "Cannot parse obfuscated frame")
+		return nil, nil, errors.Annotate(err, "Cannot parse obfuscated frame")
 	}
 
 	socket = wrappers.NewStreamCipherRWC(socket, obfs2.Encryptor, obfs2.Decryptor)
 
-	return dc, socket, nil
+	return connOpts, socket, nil
 }

mtproto/connection_options.go

+package mtproto
+
+import (
+	"bytes"
+
+	"github.com/juju/errors"
+)
+
+// ConnectionType is a type of obfuscated2/mtproto connection requested
+// by the user.
+type ConnectionType uint8
+
+// ConnectionOpts presents an options, metadata on connection requested
+// by the user on handshake.
+type ConnectionOpts struct {
+	DC             int16
+	ConnectionType ConnectionType
+}
+
+// Different connection types which user requests from Telegram.
+const (
+	ConnectionTypeUnknown ConnectionType = iota
+	ConnectionTypeAbridged
+	ConnectionTypeIntermediate
+)
+
+// Connection tags for mtproto handshakes.
+var (
+	ConnectionTagAbridged     = []byte{0xef, 0xef, 0xef, 0xef}
+	ConnectionTagIntermediate = []byte{0xee, 0xee, 0xee, 0xee}
+)
+
+// Tag maps connection type to the corresponding handshake tag.
+func (t ConnectionType) Tag() ([]byte, error) {
+	switch t {
+	case ConnectionTypeAbridged:
+		return ConnectionTagAbridged, nil
+	case ConnectionTypeIntermediate:
+		return ConnectionTagIntermediate, nil
+	default:
+		return nil, errors.Errorf("Unknown connection type %d", t)
+	}
+}
+
+// ConnectionTagFromHandshake maps magic bytes to the connection type.
+func ConnectionTagFromHandshake(magic []byte) (ConnectionType, error) {
+	if bytes.Equal(magic, ConnectionTagIntermediate) {
+		return ConnectionTypeIntermediate, nil
+	}
+	if bytes.Equal(magic, ConnectionTagAbridged) {
+		return ConnectionTypeAbridged, nil
+	}
+
+	return ConnectionTypeUnknown, errors.New("Unknown handshake protocol")
+}

obfuscated2/frame.go

 	"io"
 
 	"github.com/juju/errors"
+
+	"github.com/9seconds/mtg/mtproto"
 )
 
 // [frameOffsetFirst:frameOffsetKey:frameOffsetIV:frameOffsetMagic:frameOffsetDC:frameOffsetEnd]
 	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/
 	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:
 	return frame, nil
 }
 
-func generateFrame() *Frame {
+func generateFrame(connectionType mtproto.ConnectionType) *Frame {
 	frame := MakeFrame()
 	data := *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

 
 import (
 	"bytes"
+	"strconv"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+
+	"github.com/9seconds/mtg/mtproto"
 )
 
 func TestFrameKey(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())
 
 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) {
 }
 
 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 {
 		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

 	"crypto/sha256"
 
 	"github.com/juju/errors"
+
+	"github.com/9seconds/mtg/mtproto"
 )
 
 // Obfuscated2 contains AES CTR encryption and decryption streams
 // 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
 	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

 	"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())
 
 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)
 	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())

proxy/server.go

 
 	"github.com/9seconds/mtg/client"
 	"github.com/9seconds/mtg/config"
+	"github.com/9seconds/mtg/mtproto"
 	"github.com/9seconds/mtg/telegram"
 	"github.com/9seconds/mtg/wrappers"
 )
 		"socketid", socketID,
 	)
 
-	dc, clientConn, err := s.getClientStream(ctx, cancel, conn, socketID)
+	connOpts, clientConn, err := s.getClientStream(ctx, cancel, conn, socketID)
 	if err != nil {
 		s.logger.Warnw("Cannot initialize client connection",
 			"addr", conn.RemoteAddr().String(),
 	}
 	defer clientConn.Close() // nolint: errcheck
 
-	tgConn, err := s.getTelegramStream(ctx, cancel, dc, socketID)
+	tgConn, err := s.getTelegramStream(ctx, cancel, connOpts, socketID)
 	if err != nil {
 		s.logger.Warnw("Cannot initialize Telegram connection",
 			"socketid", socketID,
 	)
 }
 
-func (s *Server) getClientStream(ctx context.Context, cancel context.CancelFunc, conn net.Conn, socketID string) (int16, io.ReadWriteCloser, error) {
-	dc, socket, err := s.clientInit(conn, s.conf)
+func (s *Server) getClientStream(ctx context.Context, cancel context.CancelFunc, conn net.Conn, socketID string) (*mtproto.ConnectionOpts, io.ReadWriteCloser, error) {
+	connOpts, socket, err := s.clientInit(conn, s.conf)
 	if err != nil {
-		return 0, nil, errors.Annotate(err, "Cannot init client connection")
+		return nil, nil, errors.Annotate(err, "Cannot init client connection")
 	}
 
 	socket = wrappers.NewTrafficRWC(socket, s.stats.addIncomingTraffic, s.stats.addOutgoingTraffic)
 	socket = wrappers.NewLogRWC(socket, s.logger, socketID, "client")
 	socket = wrappers.NewCtxRWC(ctx, cancel, socket)
 
-	return dc, socket, nil
+	return connOpts, socket, nil
 }
 
-func (s *Server) getTelegramStream(ctx context.Context, cancel context.CancelFunc, dc int16, socketID string) (io.ReadWriteCloser, error) {
-	conn, err := s.tg.Dial(dc)
+func (s *Server) getTelegramStream(ctx context.Context, cancel context.CancelFunc, connOpts *mtproto.ConnectionOpts, socketID string) (io.ReadWriteCloser, error) {
+	conn, err := s.tg.Dial(connOpts)
 	if err != nil {
 		return nil, errors.Annotate(err, "Cannot connect to Telegram")
 	}
 
 	conn = wrappers.NewTrafficRWC(conn, s.stats.addIncomingTraffic, s.stats.addOutgoingTraffic)
-	conn, err = s.tg.Init(conn)
+	conn, err = s.tg.Init(connOpts, conn)
 	if err != nil {
 		return nil, errors.Annotate(err, "Cannot handshake Telegram")
 	}

telegram/dialer.go

 	return wrappers.NewTimeoutRWC(conn, t.conf.TimeoutRead, t.conf.TimeoutWrite), nil
 }
 
-func newDialer(conf *config.Config) *tgDialer {
-	return &tgDialer{
+func newDialer(conf *config.Config) tgDialer {
+	return tgDialer{
 		Dialer: net.Dialer{Timeout: conf.TimeoutRead},
 		conf:   conf,
 	}

telegram/direct.go

 	"github.com/juju/errors"
 
 	"github.com/9seconds/mtg/config"
+	"github.com/9seconds/mtg/mtproto"
 	"github.com/9seconds/mtg/obfuscated2"
 	"github.com/9seconds/mtg/wrappers"
 )
 	baseTelegram
 }
 
-func (t *directTelegram) Dial(dcIdx int16) (io.ReadWriteCloser, error) {
-	if dcIdx < 0 {
-		dcIdx = -dcIdx
-	} else if dcIdx == 0 {
-		dcIdx = 1
+func (t *directTelegram) Dial(connOpts *mtproto.ConnectionOpts) (io.ReadWriteCloser, error) {
+	dc := connOpts.DC
+	if dc < 0 {
+		dc = -dc
+	} else if dc == 0 {
+		dc = 1
 	}
 
-	return t.baseTelegram.Dial(dcIdx - 1)
+	return t.baseTelegram.dial(dc - 1)
 }
 
-func (t *directTelegram) Init(conn io.ReadWriteCloser) (io.ReadWriteCloser, error) {
-	obfs2, frame := obfuscated2.MakeTelegramObfuscated2Frame()
+func (t *directTelegram) Init(connOpts *mtproto.ConnectionOpts, conn io.ReadWriteCloser) (io.ReadWriteCloser, error) {
+	obfs2, frame := obfuscated2.MakeTelegramObfuscated2Frame(connOpts)
 	defer obfuscated2.ReturnFrame(frame)
 
 	if n, err := conn.Write(*frame); err != nil || n != len(*frame) {

telegram/telegram.go

 	"math/rand"
 
 	"github.com/juju/errors"
+
+	"github.com/9seconds/mtg/mtproto"
 )
 
 // Telegram defines an interface to connect to Telegram. This
 // encapsulates logic of working with middleproxies or direct
 // connections.
 type Telegram interface {
-	Dial(int16) (io.ReadWriteCloser, error)
-	Init(io.ReadWriteCloser) (io.ReadWriteCloser, error)
+	Dial(*mtproto.ConnectionOpts) (io.ReadWriteCloser, error)
+	Init(*mtproto.ConnectionOpts, io.ReadWriteCloser) (io.ReadWriteCloser, error)
 }
 
 type baseTelegram struct {
-	dialer *tgDialer
+	dialer tgDialer
 
 	v4Addresses map[int16][]string
 	v6Addresses map[int16][]string
 }
 
-func (b *baseTelegram) Dial(dcIdx int16) (io.ReadWriteCloser, error) {
+func (b *baseTelegram) dial(dcIdx int16) (io.ReadWriteCloser, error) {
 	addrs := make([]string, 2)
 	if addr, ok := b.v6Addresses[dcIdx]; ok && len(addr) > 0 {
 		addrs = append(addrs, addr[rand.Intn(len(addr))])