ScuroNeko/mtg
1
0
Fork 0
mirror of https://github.com/ScuroNeko/mtg.git synced 2026-08-31 15:24:01 +03:00
Code Issues Packages Projects Releases Wiki Activity

Add option to use telegram test dcs

Browse Source
This commit is contained in:
9seconds
2021-04-09 16:45:37 +03:00
parent 2b5b325722
commit 115510985a
6 changed files with 118 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mtglib/internal/telegram/address_pool.go
+33
View File
@@ -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
+40 -7
View File
@@ -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
+18 -18
View File
@@ -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
}
mtglib/internal/telegram/telegram_internal_test.go
+12 -12
View File
@@ -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)
}
mtglib/proxy.go
+1 -1
View File
@@ -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
+9
View File
@@ -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 {