From f9a2b0513d00f8fa20a21a71736c6d4407349308 Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Sun, 1 Jul 2018 10:27:36 +0300
Subject: [PATCH] Propagate connection protocol

---
 client/direct.go              |  1 +
 main.go                       |  1 -
 mtproto/connection_options.go | 13 +++++++++++--
 telegram/direct.go            |  2 +-
 telegram/middle.go            |  2 +-
 telegram/middle_caller.go     |  2 +-
 telegram/telegram.go          | 15 ++++++++++-----
 7 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/client/direct.go b/client/direct.go
index 537d12a..f4db89c 100644
--- a/client/direct.go
+++ b/client/direct.go
@@ -35,6 +35,7 @@ func DirectInit(conn net.Conn, conf *config.Config) (*mtproto.ConnectionOpts, io
 	if err != nil {
 		return nil, nil, errors.Annotate(err, "Cannot parse obfuscated frame")
 	}
+	connOpts.ConnectionProto = mtproto.ConnectionProtocolAny
 
 	socket := wrappers.NewStreamCipherRWC(conn, obfs2.Encryptor, obfs2.Decryptor)
 
diff --git a/main.go b/main.go
index 18de1f2..fa63cfc 100644
--- a/main.go
+++ b/main.go
@@ -16,7 +16,6 @@ import (
 
 	"github.com/9seconds/mtg/config"
 	"github.com/9seconds/mtg/proxy"
-	"github.com/9seconds/mtg/telegram"
 	"github.com/juju/errors"
 )
 
diff --git a/mtproto/connection_options.go b/mtproto/connection_options.go
index 8510692..29ee283 100644
--- a/mtproto/connection_options.go
+++ b/mtproto/connection_options.go
@@ -10,11 +10,14 @@ import (
 // by the user.
 type ConnectionType uint8
 
+type ConnectionProtocol uint8
+
 // ConnectionOpts presents an options, metadata on connection requested
 // by the user on handshake.
 type ConnectionOpts struct {
-	DC             int16
-	ConnectionType ConnectionType
+	DC              int16
+	ConnectionType  ConnectionType
+	ConnectionProto ConnectionProtocol
 }
 
 // Different connection types which user requests from Telegram.
@@ -24,6 +27,12 @@ const (
 	ConnectionTypeIntermediate
 )
 
+const (
+	ConnectionProtocolIPv4 ConnectionProtocol = 1
+	ConnectionProtocolIPv6                    = ConnectionProtocolIPv4 << 1
+	ConnectionProtocolAny                     = ConnectionProtocolIPv4 | ConnectionProtocolIPv6
+)
+
 // Connection tags for mtproto handshakes.
 var (
 	ConnectionTagAbridged     = []byte{0xef, 0xef, 0xef, 0xef}
diff --git a/telegram/direct.go b/telegram/direct.go
index 89a7016..f1ee52f 100644
--- a/telegram/direct.go
+++ b/telegram/direct.go
@@ -41,7 +41,7 @@ func (t *directTelegram) Dial(connOpts *mtproto.ConnectionOpts) (io.ReadWriteClo
 		dc = 1
 	}
 
-	return t.baseTelegram.dial(dc - 1)
+	return t.baseTelegram.dial(dc-1, connOpts.ConnectionProto)
 }
 
 func (t *directTelegram) Init(connOpts *mtproto.ConnectionOpts, conn io.ReadWriteCloser) (io.ReadWriteCloser, error) {
diff --git a/telegram/middle.go b/telegram/middle.go
index 6fb8432..01a3b73 100644
--- a/telegram/middle.go
+++ b/telegram/middle.go
@@ -14,7 +14,7 @@ type middleTelegram struct {
 	middleTelegramCaller
 }
 
-func NewMiddleTelegram(conf *config.Config, logger *zap.SugaredLogger) Telegram {
+func NewMiddleTelegram(conf *config.Config, logger *zap.SugaredLogger) *middleTelegram {
 	tg := &middleTelegram{
 		middleTelegramCaller: middleTelegramCaller{
 			baseTelegram: baseTelegram{
diff --git a/telegram/middle_caller.go b/telegram/middle_caller.go
index 27d48fc..2d0d86e 100644
--- a/telegram/middle_caller.go
+++ b/telegram/middle_caller.go
@@ -47,7 +47,7 @@ func (t *middleTelegramCaller) Dial(connOpts *mtproto.ConnectionOpts) (io.ReadWr
 	t.dialerMutex.RLock()
 	defer t.dialerMutex.RUnlock()
 
-	return t.baseTelegram.dial(dc)
+	return t.baseTelegram.dial(dc, connOpts.ConnectionProto)
 }
 
 func (t *middleTelegramCaller) autoUpdate() {
diff --git a/telegram/telegram.go b/telegram/telegram.go
index 0f71efc..6bbb504 100644
--- a/telegram/telegram.go
+++ b/telegram/telegram.go
@@ -24,13 +24,18 @@ type baseTelegram struct {
 	v6Addresses map[int16][]string
 }
 
-func (b *baseTelegram) dial(dcIdx int16) (io.ReadWriteCloser, error) {
+func (b *baseTelegram) dial(dcIdx int16, proto mtproto.ConnectionProtocol) (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))])
+
+	if proto&mtproto.ConnectionProtocolIPv6 != 0 {
+		if addr, ok := b.v6Addresses[dcIdx]; ok && len(addr) > 0 {
+			addrs = append(addrs, addr[rand.Intn(len(addr))])
+		}
 	}
-	if addr, ok := b.v4Addresses[dcIdx]; ok && len(addr) > 0 {
-		addrs = append(addrs, addr[rand.Intn(len(addr))])
+	if proto&mtproto.ConnectionProtocolIPv4 != 0 {
+		if addr, ok := b.v4Addresses[dcIdx]; ok && len(addr) > 0 {
+			addrs = append(addrs, addr[rand.Intn(len(addr))])
+		}
 	}
 
 	for _, addr := range addrs {