Merge pull request #183 from 9seconds/test-stage

Add option to use telegram test dcs

mtglib/internal/telegram/address_pool.go

@@ -0,0 +1,33 @@
+package telegram
+
+import "math/rand"
+
+type addressPool struct {
+	v4 [][]tgAddr
+	v6 [][]tgAddr
+}
+
+func (a addressPool) getV4(dc int) []tgAddr {
+	return a.get(a.v4, dc-1)
+}
+
+func (a addressPool) getV6(dc int) []tgAddr {
+	return a.get(a.v6, dc-1)
+}
+
+func (a addressPool) get(addresses [][]tgAddr, dc int) []tgAddr {
+	if dc < 0 || dc >= len(addresses) {
+		return nil
+	}
+
+	rv := make([]tgAddr, len(addresses[dc]))
+	copy(rv, addresses[dc])
+
+	if len(rv) > 1 {
+		rand.Shuffle(len(rv), func(i, j int) {
+			rv[i], rv[j] = rv[j], rv[i]
+		})
+	}
+
+	return rv
+}

mtglib/internal/telegram/init.go

@@ -21,30 +21,63 @@ type tgAddr struct {
 
 // https://github.com/telegramdesktop/tdesktop/blob/master/Telegram/SourceFiles/mtproto/mtproto_dc_options.cpp#L30
 var (
-	v4Addresses = [5][]tgAddr{
+	productionV4Addresses = [][]tgAddr{
-		{
+		{ // dc1
 			{network: "tcp4", address: "149.154.175.50:443"},
 		},
-		{
+		{ // dc2
 			{network: "tcp4", address: "149.154.167.51:443"},
 			{network: "tcp4", address: "95.161.76.100:443"},
 		},
-		{
+		{ // dc3
 			{network: "tcp4", address: "149.154.175.100:443"},
 		},
-		{
+		{ // dc4
 			{network: "tcp4", address: "149.154.167.91:443"},
 		},
-		{
+		{ // dc5
 			{network: "tcp4", address: "149.154.171.5:443"},
 		},
 	}
-	v6Addresses = [5]tgAddr{
+	productionV6Addresses = [][]tgAddr{
+		{ // dc1
 			{network: "tcp6", address: "[2001:b28:f23d:f001::a]:443"},
+		},
+		{ // dc2
 			{network: "tcp6", address: "[2001:67c:04e8:f002::a]:443"},
+		},
+		{ // dc3
 			{network: "tcp6", address: "[2001:b28:f23d:f003::a]:443"},
+		},
+		{ // dc4
 			{network: "tcp6", address: "[2001:67c:04e8:f004::a]:443"},
+		},
+		{ // dc5
 			{network: "tcp6", address: "[2001:b28:f23f:f005::a]:443"},
+		},
+	}
+
+	testV4Addresses = [][]tgAddr{
+		{ // dc1
+			{network: "tcp4", address: "149.154.175.10:443"},
+		},
+		{ // dc2
+			{network: "tcp4", address: "149.154.167.40:443"},
+		},
+		{ // dc3
+			{network: "tcp4", address: "149.154.175.117:443"},
+		},
+	}
+	testV6Addresses = [][]tgAddr{
+		{ // dc1
+			{network: "tcp6", address: "[2001:b28:f23d:f001::e]:443"},
+		},
+		{ // dc2
+			{network: "tcp6", address: "[2001:67c:04e8:f002::e]:443"},
+		},
+		{ // dc3
+			{network: "tcp6", address: "[2001:b28:f23d:f003::e]:443"},
+		},
 	}
 )
 

mtglib/internal/telegram/telegram.go

@@ -3,7 +3,6 @@ package telegram
 import (
 	"context"
 	"fmt"
-	"math/rand"
 	"net"
 	"strings"
 )
@@ -11,30 +10,21 @@ import (
 type Telegram struct {
 	dialer   Dialer
 	preferIP preferIP
+	pool     addressPool
 }
 
 func (t Telegram) Dial(ctx context.Context, dc int) (net.Conn, error) {
-	if dc <= 0 || dc > 5 {
-		return nil, fmt.Errorf("do not know how to dial to %d", dc)
-	}
-
 	var addresses []tgAddr
 
-	if t.preferIP == preferIPOnlyIPv6 {
-		addresses = []tgAddr{v6Addresses[dc-1]}
-	} else {
-		addresses = append(addresses, v4Addresses[dc-1]...)
-		rand.Shuffle(len(addresses), func(i, j int) {
-			addresses[i], addresses[j] = addresses[j], addresses[i]
-		})
-	}
-
 	switch t.preferIP {
+	case preferIPOnlyIPv4:
+		addresses = t.pool.getV4(dc)
+	case preferIPOnlyIPv6:
+		addresses = t.pool.getV6(dc)
 	case preferIPPreferIPv4:
-		addresses = append(addresses, v6Addresses[dc-1])
+		addresses = append(t.pool.getV4(dc), t.pool.getV6(dc)...)
 	case preferIPPreferIPv6:
-		addresses = append([]tgAddr{v6Addresses[dc-1]}, addresses...)
+		addresses = append(t.pool.getV6(dc), t.pool.getV4(dc)...)
-	case preferIPOnlyIPv4, preferIPOnlyIPv6:
 	}
 
 	var (
@@ -52,7 +42,7 @@ func (t Telegram) Dial(ctx context.Context, dc int) (net.Conn, error) {
 	return nil, fmt.Errorf("cannot dial to %d dc: %w", dc, err)
 }
 
-func New(dialer Dialer, ipPreference string) (*Telegram, error) {
+func New(dialer Dialer, ipPreference string, useTestDCs bool) (*Telegram, error) {
 	var pref preferIP
 
 	switch strings.ToLower(ipPreference) {
@@ -68,8 +58,18 @@ func New(dialer Dialer, ipPreference string) (*Telegram, error) {
 		return nil, fmt.Errorf("unknown ip preference %s", ipPreference)
 	}
 
+	pool := addressPool{
+		v4: productionV4Addresses,
+		v6: productionV6Addresses,
+	}
+	if useTestDCs {
+		pool.v4 = testV4Addresses
+		pool.v6 = testV6Addresses
+	}
+
 	return &Telegram{
 		dialer:   dialer,
 		preferIP: pref,
+		pool:     pool,
 	}, nil
 }

mtglib/internal/telegram/telegram_internal_test.go

@@ -23,7 +23,7 @@ type TelegramTestSuite struct {
 
 func (suite *TelegramTestSuite) SetupTest() {
 	suite.dialerMock = &testlib.MtglibNetworkMock{}
-	suite.t, _ = New(suite.dialerMock, "prefer-ipv4")
+	suite.t, _ = New(suite.dialerMock, "prefer-ipv4", false)
 }
 
 func (suite *TelegramTestSuite) TearDownTest() {
@@ -53,8 +53,8 @@ func (suite *TelegramTestSuite) TestDialToCorrectIPs() {
 
 	for i := 1; i <= 5; i++ {
 		testData[i] = []tgAddr{}
-		testData[i] = append(testData[i], v4Addresses[i-1]...)
+		testData[i] = append(testData[i], productionV4Addresses[i-1]...)
-		testData[i] = append(testData[i], v6Addresses[i-1])
+		testData[i] = append(testData[i], productionV6Addresses[i-1]...)
 	}
 
 	for i, v := range testData {
@@ -77,10 +77,10 @@ func (suite *TelegramTestSuite) TestDialToCorrectIPs() {
 
 func (suite *TelegramTestSuite) TestDialPreferIPRange() {
 	testData := map[string][]tgAddr{
-		"prefer-ipv4": {v4Addresses[0][0], v6Addresses[0]},
+		"prefer-ipv4": {testV4Addresses[0][0], testV6Addresses[0][0]},
-		"prefer-ipv6": {v6Addresses[0], v4Addresses[0][0]},
+		"prefer-ipv6": {testV6Addresses[0][0], testV4Addresses[0][0]},
-		"only-ipv4":   {v4Addresses[0][0]},
+		"only-ipv4":   {testV4Addresses[0][0]},
-		"only-ipv6":   {v6Addresses[0]},
+		"only-ipv6":   {testV6Addresses[0][0]},
 	}
 
 	for k, v := range testData {
@@ -95,7 +95,7 @@ func (suite *TelegramTestSuite) TestDialPreferIPRange() {
 					Return((*net.TCPConn)(nil), io.EOF)
 			}
 
-			tg, _ := New(suite.dialerMock, name)
+			tg, _ := New(suite.dialerMock, name, true)
 			_, err := tg.Dial(context.Background(), 1)
 
 			assert.True(t, errors.Is(err, io.EOF))
@@ -105,8 +105,8 @@ func (suite *TelegramTestSuite) TestDialPreferIPRange() {
 
 func (suite *TelegramTestSuite) TestDialPreferIPPriority() {
 	testData := map[string]tgAddr{
-		"prefer-ipv4": v4Addresses[0][0],
+		"prefer-ipv4": productionV4Addresses[0][0],
-		"prefer-ipv6": v6Addresses[0],
+		"prefer-ipv6": productionV6Addresses[0][0],
 	}
 
 	for k, v := range testData {
@@ -121,7 +121,7 @@ func (suite *TelegramTestSuite) TestDialPreferIPPriority() {
 				Once().
 				Return(conn, nil)
 
-			tg, _ := New(suite.dialerMock, name)
+			tg, _ := New(suite.dialerMock, name, false)
 
 			res, err := tg.Dial(context.Background(), 1)
 			assert.NoError(t, err)
@@ -131,7 +131,7 @@ func (suite *TelegramTestSuite) TestDialPreferIPPriority() {
 }
 
 func (suite *TelegramTestSuite) TestUnknownPreferIP() {
-	_, err := New(suite.dialerMock, "xxx")
+	_, err := New(suite.dialerMock, "xxx", false)
 	suite.Error(err)
 }
 

mtglib/proxy.go

@@ -270,7 +270,7 @@ func NewProxy(opts ProxyOpts) (*Proxy, error) {
 		return nil, fmt.Errorf("invalid settings: %w", err)
 	}
 
-	tg, err := telegram.New(opts.Network, opts.getPreferIP())
+	tg, err := telegram.New(opts.Network, opts.getPreferIP(), opts.UseTestDCs)
 	if err != nil {
 		return nil, fmt.Errorf("cannot build telegram dialer: %w", err)
 	}

mtglib/proxy_opts.go

@@ -89,6 +89,15 @@ type ProxyOpts struct {
 	//
 	// This is an optional setting.
 	PreferIP string
+
+	// UseTestDCs defines if we have to connect to production or to staging
+	// DCs of Telegram.
+	//
+	// This is required if you use mtglib as an integration library for
+	// your Telegram-related projects.
+	//
+	// This is an optional setting.
+	UseTestDCs bool
 }
 
 func (p ProxyOpts) valid() error {