diff --git a/mtglib/internal/telegram/address_pool.go b/mtglib/internal/telegram/address_pool.go new file mode 100644 index 0000000..0f49e80 --- /dev/null +++ b/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 +} diff --git a/mtglib/internal/telegram/init.go b/mtglib/internal/telegram/init.go index af707d3..469478f 100644 --- a/mtglib/internal/telegram/init.go +++ b/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{ - {network: "tcp6", address: "[2001:b28:f23d:f001::a]:443"}, - {network: "tcp6", address: "[2001:67c:04e8:f002::a]:443"}, - {network: "tcp6", address: "[2001:b28:f23d:f003::a]:443"}, - {network: "tcp6", address: "[2001:67c:04e8:f004::a]:443"}, - {network: "tcp6", address: "[2001:b28:f23f:f005::a]:443"}, + 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"}, + }, } ) diff --git a/mtglib/internal/telegram/telegram.go b/mtglib/internal/telegram/telegram.go index 2ee86e2..474dfa4 100644 --- a/mtglib/internal/telegram/telegram.go +++ b/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...) - case preferIPOnlyIPv4, preferIPOnlyIPv6: + addresses = append(t.pool.getV6(dc), t.pool.getV4(dc)...) } 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 } diff --git a/mtglib/internal/telegram/telegram_internal_test.go b/mtglib/internal/telegram/telegram_internal_test.go index fe21ce2..77d5939 100644 --- a/mtglib/internal/telegram/telegram_internal_test.go +++ b/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], v6Addresses[i-1]) + testData[i] = append(testData[i], productionV4Addresses[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-ipv6": {v6Addresses[0], v4Addresses[0][0]}, - "only-ipv4": {v4Addresses[0][0]}, - "only-ipv6": {v6Addresses[0]}, + "prefer-ipv4": {testV4Addresses[0][0], testV6Addresses[0][0]}, + "prefer-ipv6": {testV6Addresses[0][0], testV4Addresses[0][0]}, + "only-ipv4": {testV4Addresses[0][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-ipv6": v6Addresses[0], + "prefer-ipv4": productionV4Addresses[0][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) } diff --git a/mtglib/proxy.go b/mtglib/proxy.go index e07a9fe..9a34fdd 100644 --- a/mtglib/proxy.go +++ b/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) } diff --git a/mtglib/proxy_opts.go b/mtglib/proxy_opts.go index c2fdc69..dc9b57b 100644 --- a/mtglib/proxy_opts.go +++ b/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 {