From 422ed5b0716c07c17952af5273292aa303d53ac8 Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Tue, 10 Dec 2019 13:59:27 +0300
Subject: [PATCH] Prefer ipv6 protocol on direct connections

---
 telegram/base.go   | 38 ++++++++++++++++++++++++--------------
 telegram/direct.go |  2 +-
 telegram/init.go   |  8 +++++++-
 telegram/middle.go |  2 +-
 4 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/telegram/base.go b/telegram/base.go
index 51cfd42..2f0a693 100644
--- a/telegram/base.go
+++ b/telegram/base.go
@@ -1,10 +1,12 @@
 package telegram
 
 import (
-	"fmt"
+	"errors"
 	"math/rand"
 	"net"
 
+	"go.uber.org/zap"
+
 	"github.com/9seconds/mtg/conntypes"
 	"github.com/9seconds/mtg/utils"
 	"github.com/9seconds/mtg/wrappers/stream"
@@ -12,10 +14,11 @@ import (
 
 type baseTelegram struct {
 	dialer net.Dialer
+	logger *zap.SugaredLogger
 
 	secret      []byte
 	v4DefaultDC conntypes.DC
-	V6DefaultDC conntypes.DC
+	v6DefaultDC conntypes.DC
 	v4Addresses map[conntypes.DC][]string
 	v6Addresses map[conntypes.DC][]string
 }
@@ -26,25 +29,32 @@ func (b *baseTelegram) Secret() []byte {
 
 func (b *baseTelegram) dial(dc conntypes.DC,
 	protocol conntypes.ConnectionProtocol) (conntypes.StreamReadWriteCloser, error) {
-	addr := ""
+	addresses := make([]string, 0, 2)
 
-	switch protocol {
-	case conntypes.ConnectionProtocolIPv4:
-		addr = b.chooseAddress(b.v4Addresses, dc, b.v4DefaultDC)
-	default:
-		addr = b.chooseAddress(b.v6Addresses, dc, b.V6DefaultDC)
+	if protocol&conntypes.ConnectionProtocolIPv6 != 0 {
+		addresses = append(addresses, b.chooseAddress(b.v6Addresses, dc, b.v6DefaultDC))
 	}
 
-	conn, err := b.dialer.Dial("tcp", addr)
-	if err != nil {
-		return nil, fmt.Errorf("dial has failed: %w", err)
+	if protocol&conntypes.ConnectionProtocolIPv4 != 0 {
+		addresses = append(addresses, b.chooseAddress(b.v4Addresses, dc, b.v4DefaultDC))
 	}
 
-	if err := utils.InitTCP(conn); err != nil {
-		return nil, fmt.Errorf("cannot initialize tcp socket: %w", err)
+	for _, addr := range addresses {
+		conn, err := b.dialer.Dial("tcp", addr)
+		if err != nil {
+			b.logger.Infow("Cannot dial to Telegram", "address", addr, "error", err)
+			continue
+		}
+
+		if err := utils.InitTCP(conn); err != nil {
+			b.logger.Infow("Cannot initialize TCP socket", "address", addr, "error", err)
+			continue
+		}
+
+		return stream.NewTelegramConn(dc, conn), nil
 	}
 
-	return stream.NewTelegramConn(dc, conn), nil
+	return nil, errors.New("cannot dial to the chosen DC")
 }
 
 func (b *baseTelegram) chooseAddress(addresses map[conntypes.DC][]string,
diff --git a/telegram/direct.go b/telegram/direct.go
index 3810c9a..5a7183c 100644
--- a/telegram/direct.go
+++ b/telegram/direct.go
@@ -37,5 +37,5 @@ func (d *directTelegram) Dial(dc conntypes.DC,
 		dc = conntypes.DCDefaultIdx
 	}
 
-	return d.baseTelegram.dial(dc-1, protocol)
+	return d.baseTelegram.dial(dc-1, conntypes.ConnectionProtocolAny)
 }
diff --git a/telegram/init.go b/telegram/init.go
index b4a0c19..72da84e 100644
--- a/telegram/init.go
+++ b/telegram/init.go
@@ -4,6 +4,8 @@ import (
 	"net"
 	"sync"
 	"time"
+
+	"go.uber.org/zap"
 )
 
 const telegramDialTimeout = 10 * time.Second
@@ -17,11 +19,14 @@ var (
 
 func Init() {
 	initOnce.Do(func() {
+		logger := zap.S().Named("telegram")
+
 		Direct = &directTelegram{
 			baseTelegram: baseTelegram{
 				dialer:      net.Dialer{Timeout: telegramDialTimeout},
+				logger:      logger.Named("direct"),
 				v4DefaultDC: directV4DefaultIdx,
-				V6DefaultDC: directV6DefaultIdx,
+				v6DefaultDC: directV6DefaultIdx,
 				v4Addresses: directV4Addresses,
 				v6Addresses: directV6Addresses,
 			},
@@ -30,6 +35,7 @@ func Init() {
 		tg := &middleTelegram{
 			baseTelegram: baseTelegram{
 				dialer: net.Dialer{Timeout: telegramDialTimeout},
+				logger: logger.Named("middle"),
 			},
 		}
 		if err := tg.update(); err != nil {
diff --git a/telegram/middle.go b/telegram/middle.go
index 519332e..d0917a9 100644
--- a/telegram/middle.go
+++ b/telegram/middle.go
@@ -45,7 +45,7 @@ func (m *middleTelegram) update() error {
 	m.mutex.Lock()
 	m.secret = secret
 	m.v4DefaultDC = v4DefaultDC
-	m.V6DefaultDC = v6DefaultDC
+	m.v6DefaultDC = v6DefaultDC
 	m.v4Addresses = v4Addresses
 	m.v6Addresses = v6Addresses
 	m.mutex.Unlock()