From 8670a56e39c93dc54a9e19c9509166a6e537e00a Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Thu, 31 May 2018 08:47:42 +0300
Subject: [PATCH] Prefer ipv6 setting

---
 main.go           |  6 +++++-
 proxy/server.go   |  6 ++++--
 proxy/telegram.go | 50 +++++++++++++++++++++++++++++++++++++----------
 3 files changed, 49 insertions(+), 13 deletions(-)

diff --git a/main.go b/main.go
index 2935d89..1a5c0cc 100644
--- a/main.go
+++ b/main.go
@@ -67,6 +67,10 @@ var (
 		Short('s').
 		Envar("MTG_SERVER").
 		String()
+	preferIPv6 = app.Flag("prefer-ipv6", "Use IPv6").
+			Short('6').
+			Envar("MTG_USE_IPV6").
+			Bool()
 
 	secret = app.Arg("secret", "Secret of this proxy.").String()
 )
@@ -117,7 +121,7 @@ func main() {
 	go stat.Serve(*statsIP, *statsPort)
 
 	srv := proxy.NewServer(*bindIP, int(*bindPort), secretBytes, logger,
-		*readTimeout, *writeTimeout, stat)
+		*readTimeout, *writeTimeout, *preferIPv6, stat)
 	if err := srv.Serve(); err != nil {
 		logger.Fatal(err.Error())
 	}
diff --git a/proxy/server.go b/proxy/server.go
index cea4d60..06a00e7 100644
--- a/proxy/server.go
+++ b/proxy/server.go
@@ -26,6 +26,7 @@ type Server struct {
 	readTimeout  time.Duration
 	writeTimeout time.Duration
 	stats        *Stats
+	ipv6         bool
 }
 
 func (s *Server) Serve() error {
@@ -124,7 +125,7 @@ func (s *Server) getClientStream(conn net.Conn, ctx context.Context, cancel cont
 }
 
 func (s *Server) getTelegramStream(dc int16, ctx context.Context, cancel context.CancelFunc, socketID string) (io.ReadWriteCloser, error) {
-	socket, err := dialToTelegram(dc, s.readTimeout)
+	socket, err := dialToTelegram(s.ipv6, dc, s.readTimeout)
 	if err != nil {
 		return nil, errors.Annotate(err, "Cannot dial")
 	}
@@ -149,7 +150,7 @@ func (s *Server) pipe(wait *sync.WaitGroup, reader io.Reader, writer io.Writer)
 }
 
 func NewServer(ip net.IP, port int, secret []byte, logger *zap.SugaredLogger,
-	readTimeout, writeTimeout time.Duration, stat *Stats) *Server {
+	readTimeout, writeTimeout time.Duration, ipv6 bool, stat *Stats) *Server {
 	return &Server{
 		ip:           ip,
 		port:         port,
@@ -159,5 +160,6 @@ func NewServer(ip net.IP, port int, secret []byte, logger *zap.SugaredLogger,
 		readTimeout:  readTimeout,
 		writeTimeout: writeTimeout,
 		stats:        stat,
+		ipv6:         ipv6,
 	}
 }
diff --git a/proxy/telegram.go b/proxy/telegram.go
index 75b42d6..1f5409d 100644
--- a/proxy/telegram.go
+++ b/proxy/telegram.go
@@ -7,24 +7,37 @@ import (
 	"github.com/juju/errors"
 )
 
-var telegramDCIPs = [5]string{
-	"149.154.175.50:443",
-	"149.154.167.51:443",
-	"149.154.175.100:443",
-	"149.154.167.91:443",
-	"149.154.171.5:443",
+type TelegramAddress struct {
+	v4 string
+	v6 string
 }
 
+func (t *TelegramAddress) IPv4() string {
+	return net.JoinHostPort(t.v4, telegramPort)
+}
+
+func (t *TelegramAddress) IPv6() string {
+	return net.JoinHostPort(t.v6, telegramPort)
+}
+
+var TelegramAddresses = []TelegramAddress{
+	TelegramAddress{v4: "149.154.175.50", v6: "2001:b28:f23d:f001::a"},
+	TelegramAddress{v4: "149.154.167.51", v6: "2001:67c:04e8:f002::a"},
+	TelegramAddress{v4: "149.154.175.100", v6: "2001:b28:f23d:f003::a"},
+	TelegramAddress{v4: "149.154.167.91", v6: "2001:67c:04e8:f004::a"},
+	TelegramAddress{v4: "149.154.171.5", v6: "2001:b28:f23f:f005::a"},
+}
+
+const telegramPort = "443"
+
 const telegramKeepAlive = 30 * time.Second
 
-func dialToTelegram(dcIdx int16, timeout time.Duration) (net.Conn, error) {
+func dialToTelegram(ipv6 bool, dcIdx int16, timeout time.Duration) (net.Conn, error) {
 	if dcIdx < 0 || dcIdx >= 5 {
 		return nil, errors.New("Incorrect DC IDX")
 	}
 
-	dialer := net.Dialer{Timeout: timeout}
-	rawConn, err := dialer.Dial("tcp", telegramDCIPs[dcIdx])
-	conn := rawConn.(*net.TCPConn)
+	conn, err := doDial(ipv6, dcIdx, timeout)
 	if err != nil {
 		return nil, errors.Annotate(err, "Cannot dial")
 	}
@@ -38,3 +51,20 @@ func dialToTelegram(dcIdx int16, timeout time.Duration) (net.Conn, error) {
 
 	return conn, nil
 }
+
+func doDial(ipv6 bool, dcIdx int16, timeout time.Duration) (*net.TCPConn, error) {
+	dialer := net.Dialer{Timeout: timeout}
+	addr := TelegramAddresses[dcIdx]
+
+	if ipv6 {
+		if conn, err := dialer.Dial("tcp", addr.IPv6()); err == nil {
+			return conn.(*net.TCPConn), nil
+		}
+	}
+
+	conn, err := dialer.Dial("tcp", addr.IPv4())
+	if err == nil {
+		return conn.(*net.TCPConn), nil
+	}
+	return nil, err
+}