Propagate connection protocol

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)
 

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

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}

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) {

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{

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

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 {