Create internal DC package

mtglib/internal/dc/addr.go

@@ -0,0 +1,10 @@
+package dc
+
+type Addr struct {
+	Network string
+	Address string
+}
+
+func (d Addr) String() string {
+	return d.Address
+}

mtglib/internal/dc/addr_set.go

@@ -0,0 +1,33 @@
+package dc
+
+import "math/rand/v2"
+
+type dcAddrSet struct {
+	v4 map[int][]Addr
+	v6 map[int][]Addr
+}
+
+func (d dcAddrSet) getV4(dc int) []Addr {
+	if d.v4 == nil {
+		return nil
+	}
+	return d.get(d.v4[dc])
+}
+
+func (d dcAddrSet) getV6(dc int) []Addr {
+	if d.v6 == nil {
+		return nil
+	}
+	return d.get(d.v6[dc])
+}
+
+func (d dcAddrSet) get(addrs []Addr) []Addr {
+	otherSet := make([]Addr, 0, len(addrs))
+	otherSet = append(otherSet, addrs...)
+
+	rand.Shuffle(len(otherSet), func(i, j int) {
+		otherSet[i], otherSet[j] = otherSet[j], otherSet[i]
+	})
+
+	return otherSet
+}

mtglib/internal/dc/addr_test.go

@@ -0,0 +1,16 @@
+package dc_test
+
+import (
+	"testing"
+
+	"github.com/9seconds/mtg/v2/mtglib/internal/dc"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestAddr(t *testing.T) {
+	t.Parallel()
+
+	addr := dc.Addr{Network: "tcp4", Address: "127.0.0.1:443"}
+
+	assert.Equal(t, "127.0.0.1:443", addr.String())
+}

mtglib/internal/dc/init.go

@@ -0,0 +1,79 @@
+package dc
+
+import "time"
+
+type preferIP uint8
+
+const (
+	preferIPOnlyIPv4 preferIP = iota
+	preferIPOnlyIPv6
+	preferIPPreferIPv4
+	preferIPPreferIPv6
+)
+
+const (
+	DefaultDC                    = 2
+	DefaultUpdateDCAddressesEach = time.Hour
+
+	defaultAppID   = 123456
+	defaultAppHash = ""
+)
+
+type Logger interface {
+	Info(msg string)
+	WarningError(msg string, err error)
+}
+
+var (
+	// https://github.com/telegramdesktop/tdesktop/blob/master/Telegram/SourceFiles/mtproto/mtproto_dc_options.cpp#L30
+	defaultDCAddrSet = dcAddrSet{
+		v4: map[int][]Addr{
+			1: {
+				{Network: "tcp4", Address: "149.154.175.50:443"},
+			},
+			2: {
+				{Network: "tcp4", Address: "149.154.167.51:443"},
+				{Network: "tcp4", Address: "95.161.76.100:443"},
+			},
+			3: {
+				{Network: "tcp4", Address: "149.154.175.100:443"},
+			},
+			4: {
+				{Network: "tcp4", Address: "149.154.167.91:443"},
+			},
+			5: {
+				{Network: "tcp4", Address: "149.154.171.5:443"},
+			},
+		},
+		v6: map[int][]Addr{
+			1: {
+				{Network: "tcp6", Address: "[2001:b28:f23d:f001::a]:443"},
+			},
+			2: {
+				{Network: "tcp6", Address: "[2001:67c:04e8:f002::a]:443"},
+			},
+			3: {
+				{Network: "tcp6", Address: "[2001:b28:f23d:f003::a]:443"},
+			},
+			4: {
+				{Network: "tcp6", Address: "[2001:67c:04e8:f004::a]:443"},
+			},
+			5: {
+				{Network: "tcp6", Address: "[2001:b28:f23f:f005::a]:443"},
+			},
+		},
+	}
+
+	defaultDCOverridesAddrSet = dcAddrSet{
+		v4: map[int][]Addr{
+			203: {
+				{Network: "tcp4", Address: "91.105.192.100:443"},
+			},
+		},
+		v6: map[int][]Addr{
+			203: {
+				{Network: "tcp6", Address: "[2a0a:f280:0203:000a:5000:0000:0000:0100]:443"},
+			},
+		},
+	}
+)

mtglib/internal/dc/telegram.go

@@ -0,0 +1,160 @@
+package dc
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/gotd/td/telegram"
+)
+
+type Telegram struct {
+	logger   Logger
+	lock     sync.RWMutex
+	view     dcView
+	preferIP preferIP
+	client   *telegram.Client
+}
+
+func (t *Telegram) GetAddresses(dc int) []Addr {
+	t.lock.RLock()
+	defer t.lock.RUnlock()
+
+	switch t.preferIP {
+	case preferIPOnlyIPv4:
+		return t.view.getV4(dc)
+	case preferIPOnlyIPv6:
+		return t.view.getV4(dc)
+	case preferIPPreferIPv4:
+		return append(t.view.getV4(dc), t.view.getV6(dc)...)
+	}
+
+	return append(t.view.getV6(dc), t.view.getV4(dc)...)
+}
+
+func (t *Telegram) Run(ctx context.Context, updateEach time.Duration) {
+	if updateEach == 0 {
+		updateEach = DefaultUpdateDCAddressesEach
+	}
+
+	t.update(ctx)
+
+	ticker := time.NewTicker(updateEach)
+	defer func() {
+		ticker.Stop()
+
+		select {
+		case <-ctx.Done():
+		case <-ticker.C:
+		default:
+		}
+	}()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			t.update(ctx)
+		}
+	}
+}
+
+func (t *Telegram) update(ctx context.Context) {
+	collected := dcAddrSet{
+		v4: map[int][]Addr{},
+		v6: map[int][]Addr{},
+	}
+
+	err := t.client.Run(ctx, func(tgctx context.Context) error {
+		conf, err := t.client.API().HelpGetConfig(tgctx)
+		if err != nil {
+			return err
+		}
+
+		for _, opt := range conf.DCOptions {
+			addr := net.JoinHostPort(opt.IPAddress, strconv.Itoa(opt.Port))
+
+			if opt.Ipv6 {
+				collected.v6[opt.ID] = append(collected.v6[opt.ID], Addr{
+					Network: "tcp6",
+					Address: addr,
+				})
+			} else {
+				collected.v4[opt.ID] = append(collected.v4[opt.ID], Addr{
+					Network: "tcp4",
+					Address: addr,
+				})
+			}
+		}
+
+		return nil
+	})
+	if err != nil {
+		t.logger.WarningError("update has failed", err)
+		return
+	}
+
+	t.lock.Lock()
+	t.view.collected = collected
+	t.lock.Unlock()
+
+	t.logger.Info(fmt.Sprintf("updated DC list: %v", collected))
+}
+
+func New(logger Logger, ipPreference string, userOverrides map[int][]string) (*Telegram, error) {
+	var pref preferIP
+
+	switch strings.ToLower(ipPreference) {
+	case "prefer-ipv4":
+		pref = preferIPPreferIPv4
+	case "prefer-ipv6":
+		pref = preferIPPreferIPv6
+	case "only-ipv4":
+		pref = preferIPOnlyIPv4
+	case "only-ipv6":
+		pref = preferIPOnlyIPv6
+	default:
+		return nil, fmt.Errorf("unknown ip preference %s", ipPreference)
+	}
+
+	overrides := dcAddrSet{}
+	for dc, addrs := range userOverrides {
+		for _, addr := range addrs {
+			host, _, err := net.SplitHostPort(addr)
+			if err != nil {
+				return nil, fmt.Errorf("incorrect host %s: %w", addr, err)
+			}
+
+			parsed := net.ParseIP(host)
+			if parsed == nil {
+				return nil, fmt.Errorf("incorrect host %s", addr)
+			}
+
+			if parsed.To4() != nil {
+				overrides.v4[dc] = append(overrides.v4[dc], Addr{
+					Network: "tcp4",
+					Address: addr,
+				})
+			} else {
+				overrides.v6[dc] = append(overrides.v6[dc], Addr{
+					Network: "tcp6",
+					Address: addr,
+				})
+			}
+		}
+	}
+
+	return &Telegram{
+		view: dcView{
+			overrides: overrides,
+		},
+		logger:   logger,
+		client:   telegram.NewClient(defaultAppID, defaultAppHash, telegram.Options{}),
+		preferIP: pref,
+	}, nil
+}

mtglib/internal/dc/view.go

@@ -0,0 +1,24 @@
+package dc
+
+type dcView struct {
+	overrides dcAddrSet
+	collected dcAddrSet
+}
+
+func (d dcView) getV4(dc int) []Addr {
+	addrs := d.overrides.getV4(dc)
+	addrs = append(addrs, defaultDCOverridesAddrSet.getV4(dc)...)
+	addrs = append(addrs, d.collected.getV4(dc)...)
+	addrs = append(addrs, defaultDCAddrSet.getV4(dc)...)
+
+	return addrs
+}
+
+func (d dcView) getV6(dc int) []Addr {
+	addrs := d.overrides.getV6(dc)
+	addrs = append(addrs, defaultDCOverridesAddrSet.getV6(dc)...)
+	addrs = append(addrs, d.collected.getV6(dc)...)
+	addrs = append(addrs, defaultDCAddrSet.getV6(dc)...)
+
+	return addrs
+}

mtglib/internal/dc/view_test.go

@@ -0,0 +1,92 @@
+package dc
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/suite"
+)
+
+type ViewTestSuite struct {
+	suite.Suite
+
+	view dcView
+}
+
+func (suite *ViewTestSuite) SetupSuite() {
+	suite.view = dcView{
+		overrides: dcAddrSet{
+			v4: map[int][]Addr{
+				111: {
+					{Network: "tcp4", Address: "127.0.0.1:443"},
+				},
+				203: {
+					{Network: "tcp4", Address: "127.0.0.2:443"},
+				},
+			},
+			v6: map[int][]Addr{
+				203: {
+					{Network: "tcp6", Address: "xxx"},
+				},
+			},
+		},
+		collected: dcAddrSet{
+			v4: map[int][]Addr{
+				1: {
+					{Network: "tcp4", Address: "127.1.0.1:443"},
+				},
+			},
+		},
+	}
+}
+
+func (suite *ViewTestSuite) TestGetV4() {
+	testData := map[int][]Addr{
+		111: {
+			{"tcp4", "127.0.0.1:443"},
+		},
+		203: {
+			{"tcp4", "127.0.0.2:443"},
+			{"tcp4", "91.105.192.100:443"},
+		},
+		2: {
+			{"tcp4", "149.154.167.51:443"},
+			{"tcp4", "95.161.76.100:443"},
+		},
+		1: {
+			{"tcp4", "127.1.0.1:443"},
+			{"tcp4", "149.154.175.50:443"},
+		},
+	}
+
+	for dc, addresses := range testData {
+		suite.T().Run(fmt.Sprintf("dc%d", dc), func(t *testing.T) {
+			assert.ElementsMatch(t, addresses, suite.view.getV4(dc))
+		})
+	}
+}
+
+func (suite *ViewTestSuite) TestGetV6() {
+	testData := map[int][]Addr{
+		111: {},
+		203: {
+			{"tcp6", "xxx"},
+			{"tcp6", "[2a0a:f280:0203:000a:5000:0000:0000:0100]:443"},
+		},
+		1: {
+			{"tcp6", "[2001:b28:f23d:f001::a]:443"},
+		},
+	}
+
+	for dc, addresses := range testData {
+		suite.T().Run(fmt.Sprintf("dc%d", dc), func(t *testing.T) {
+			assert.ElementsMatch(t, addresses, suite.view.getV6(dc))
+		})
+	}
+}
+
+func TestView(t *testing.T) {
+	t.Parallel()
+	suite.Run(t, &ViewTestSuite{})
+}

mtglib/internal/telegram/dc_addresses.go

@@ -1,19 +0,0 @@
-package telegram
-
-type dcAddresses struct {
-	v4 map[int][]tgAddr
-	v6 map[int][]tgAddr
-}
-
-func (a dcAddresses) getV4(dc int) []tgAddr {
-	return a.v4[dc]
-}
-
-func (a dcAddresses) getV6(dc int) []tgAddr {
-	return a.v6[dc]
-}
-
-func (a dcAddresses) isValidDC(dc int) bool {
-	_, ok := a.v4[dc]
-	return ok
-}

mtglib/internal/telegram/init.go

@@ -1,80 +0,0 @@
-package telegram
-
-import (
-	"context"
-	"errors"
-	"time"
-
-	"github.com/9seconds/mtg/v2/essentials"
-)
-
-var errNoAddresses = errors.New("no addresses")
-
-type preferIP uint8
-
-const (
-	preferIPOnlyIPv4 preferIP = iota
-	preferIPOnlyIPv6
-	preferIPPreferIPv4
-	preferIPPreferIPv6
-)
-
-const (
-	defaultDC                    = 2
-	defaultUpdateDCAddressesEach = time.Hour
-	defaultAppID                 = 123456
-	defaultAppHash               = ""
-)
-
-type loggerInterface interface {
-	Info(msg string)
-	WarningError(msg string, err error)
-}
-
-type tgAddr struct {
-	network string
-	address string
-}
-
-// https://github.com/telegramdesktop/tdesktop/blob/master/Telegram/SourceFiles/mtproto/mtproto_dc_options.cpp#L30
-var (
-	defaultV4Addresses = map[int][]tgAddr{
-		1: {
-			{network: "tcp4", address: "149.154.175.50:443"},
-		},
-		2: {
-			{network: "tcp4", address: "149.154.167.51:443"},
-			{network: "tcp4", address: "95.161.76.100:443"},
-		},
-		3: {
-			{network: "tcp4", address: "149.154.175.100:443"},
-		},
-		4: {
-			{network: "tcp4", address: "149.154.167.91:443"},
-		},
-		5: {
-			{network: "tcp4", address: "149.154.171.5:443"},
-		},
-	}
-	defaultV6Addresses = map[int][]tgAddr{
-		1: {
-			{network: "tcp6", address: "[2001:b28:f23d:f001::a]:443"},
-		},
-		2: {
-			{network: "tcp6", address: "[2001:67c:04e8:f002::a]:443"},
-		},
-		3: {
-			{network: "tcp6", address: "[2001:b28:f23d:f003::a]:443"},
-		},
-		4: {
-			{network: "tcp6", address: "[2001:67c:04e8:f004::a]:443"},
-		},
-		5: {
-			{network: "tcp6", address: "[2001:b28:f23f:f005::a]:443"},
-		},
-	}
-)
-
-type Dialer interface {
-	DialContext(ctx context.Context, network, address string) (essentials.Conn, error)
-}

mtglib/internal/telegram/rpc_client.go

@@ -1,42 +0,0 @@
-package telegram
-
-import (
-	"context"
-	"net"
-	"strconv"
-
-	"github.com/gotd/td/telegram"
-)
-
-type rpcClient struct {
-	*telegram.Client
-}
-
-func (r rpcClient) getDCAddresses(logger loggerInterface, ctx context.Context) (dcAddresses, error) {
-	addrs := dcAddresses{
-		v4: map[int][]tgAddr{},
-		v6: map[int][]tgAddr{},
-	}
-
-	err := r.Client.Run(ctx, func(_ context.Context) error {
-		for _, opt := range r.Client.Config().DCOptions {
-			addr := net.JoinHostPort(opt.IPAddress, strconv.Itoa(opt.Port))
-
-			if opt.Ipv6 {
-				addrs.v6[opt.ID] = append(addrs.v6[opt.ID], tgAddr{
-					network: "tcp6",
-					address: addr,
-				})
-			} else {
-				addrs.v4[opt.ID] = append(addrs.v4[opt.ID], tgAddr{
-					network: "tcp4",
-					address: addr,
-				})
-			}
-		}
-
-		return nil
-	})
-
-	return addrs, err
-}

mtglib/internal/telegram/telegram.go

@@ -1,139 +0,0 @@
-package telegram
-
-import (
-	"context"
-	"fmt"
-	"strings"
-	"sync"
-	"time"
-
-	"github.com/9seconds/mtg/v2/essentials"
-	"github.com/gotd/td/telegram"
-)
-
-type Telegram struct {
-	ctx       context.Context
-	ctxCancel context.CancelFunc
-	lock      sync.RWMutex
-
-	dialer    Dialer
-	preferIP  preferIP
-	addresses dcAddresses
-	rpc       rpcClient
-}
-
-func (t *Telegram) Dial(ctx context.Context, dc int) (essentials.Conn, error) {
-	var addresses []tgAddr
-
-	t.lock.RLock()
-	switch t.preferIP {
-	case preferIPOnlyIPv4:
-		addresses = t.addresses.getV4(dc)
-	case preferIPOnlyIPv6:
-		addresses = t.addresses.getV6(dc)
-	case preferIPPreferIPv4:
-		addresses = append(t.addresses.getV4(dc), t.addresses.getV6(dc)...)
-	case preferIPPreferIPv6:
-		addresses = append(t.addresses.getV6(dc), t.addresses.getV4(dc)...)
-	}
-	t.lock.RUnlock()
-
-	var conn essentials.Conn
-
-	err := errNoAddresses
-
-	for _, v := range addresses {
-		conn, err = t.dialer.DialContext(ctx, v.network, v.address)
-		if err == nil {
-			return conn, nil
-		}
-	}
-
-	return nil, fmt.Errorf("cannot dial to %d dc: %w", dc, err)
-}
-
-func (t *Telegram) IsKnownDC(dc int) bool {
-	return t.addresses.isValidDC(dc)
-}
-
-func (t *Telegram) GetFallbackDC() int {
-	return defaultDC
-}
-
-func (t *Telegram) Shutdown() {
-	t.ctxCancel()
-}
-
-func (t *Telegram) Run(logger loggerInterface, updateEach time.Duration) {
-	if updateEach == 0 {
-		updateEach = defaultUpdateDCAddressesEach
-	}
-
-	t.update(logger)
-
-	ticker := time.NewTicker(updateEach)
-	defer func() {
-		ticker.Stop()
-
-		select {
-		case <-ticker.C:
-		default:
-		}
-	}()
-
-	for {
-		select {
-		case <-t.ctx.Done():
-			return
-		case <-ticker.C:
-			t.update(logger)
-		}
-	}
-}
-
-func (t *Telegram) update(logger loggerInterface) {
-	otherAddresses, err := t.rpc.getDCAddresses(logger, t.ctx)
-	if err != nil {
-		logger.WarningError("Cannot update DC list", err)
-		return
-	}
-
-	t.lock.Lock()
-	t.addresses = otherAddresses
-	t.lock.Unlock()
-
-	logger.Info(fmt.Sprintf("DC are updated: %v", t.addresses))
-}
-
-func New(dialer Dialer, ipPreference string) (*Telegram, error) {
-	var pref preferIP
-
-	switch strings.ToLower(ipPreference) {
-	case "prefer-ipv4":
-		pref = preferIPPreferIPv4
-	case "prefer-ipv6":
-		pref = preferIPPreferIPv6
-	case "only-ipv4":
-		pref = preferIPOnlyIPv4
-	case "only-ipv6":
-		pref = preferIPOnlyIPv6
-	default:
-		return nil, fmt.Errorf("unknown ip preference %s", ipPreference)
-	}
-
-	ctx, cancel := context.WithCancel(context.Background())
-
-	return &Telegram{
-		ctx:       ctx,
-		ctxCancel: cancel,
-		dialer:    dialer,
-		preferIP:  pref,
-		addresses: dcAddresses{
-			v4: defaultV4Addresses,
-			v6: defaultV6Addresses,
-		},
-		rpc: rpcClient{
-			Client: telegram.NewClient(defaultAppID, defaultAppHash, telegram.Options{}),
-		},
-	}, nil
-}

mtglib/internal/telegram/telegram_internal_test.go

@@ -1,159 +0,0 @@
-package telegram
-
-import (
-	"context"
-	"errors"
-	"io"
-	"net"
-	"strconv"
-	"testing"
-
-	"github.com/9seconds/mtg/v2/internal/testlib"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/mock"
-	"github.com/stretchr/testify/suite"
-)
-
-type TelegramTestSuite struct {
-	suite.Suite
-
-	dialerMock *testlib.MtglibNetworkMock
-	t          *Telegram
-}
-
-func (suite *TelegramTestSuite) SetupTest() {
-	suite.dialerMock = &testlib.MtglibNetworkMock{}
-	suite.t, _ = New(suite.dialerMock, "prefer-ipv4", false)
-}
-
-func (suite *TelegramTestSuite) TearDownTest() {
-	suite.dialerMock.AssertExpectations(suite.T())
-}
-
-func (suite *TelegramTestSuite) TestUnknownDC() {
-	testData := []int{
-		-1,
-		0,
-		6,
-		100,
-	}
-
-	for _, v := range testData {
-		value := v
-
-		suite.T().Run(strconv.Itoa(value), func(t *testing.T) {
-			_, err := suite.t.Dial(context.Background(), value)
-			assert.Error(t, err)
-			assert.False(t, suite.t.IsKnownDC(value))
-		})
-	}
-}
-
-func (suite *TelegramTestSuite) TestDialToCorrectIPs() {
-	testData := map[int][]tgAddr{}
-
-	for i := 1; i <= 5; i++ {
-		testData[i] = []tgAddr{}
-		testData[i] = append(testData[i], productionV4Addresses[i-1]...)
-		testData[i] = append(testData[i], productionV6Addresses[i-1]...)
-	}
-
-	for i, v := range testData {
-		idx := i
-		addresses := v
-
-		suite.T().Run(strconv.Itoa(idx), func(t *testing.T) {
-			for _, addr := range addresses {
-				suite.dialerMock.
-					On("DialContext", mock.Anything, addr.network, addr.address).
-					Once().
-					Return((*net.TCPConn)(nil), io.EOF)
-			}
-
-			_, err := suite.t.Dial(context.Background(), idx)
-			assert.True(t, errors.Is(err, io.EOF))
-			assert.True(t, suite.t.IsKnownDC(idx))
-		})
-	}
-}
-
-func (suite *TelegramTestSuite) TestDialPreferIPRange() {
-	testData := map[string][]tgAddr{
-		"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 {
-		name := k
-		addresses := v
-
-		suite.T().Run(name, func(t *testing.T) {
-			for _, addr := range addresses {
-				suite.dialerMock.
-					On("DialContext", mock.Anything, addr.network, addr.address).
-					Once().
-					Return((*net.TCPConn)(nil), io.EOF)
-			}
-
-			tg, _ := New(suite.dialerMock, name, true)
-			_, err := tg.Dial(context.Background(), 1)
-
-			assert.True(t, errors.Is(err, io.EOF))
-		})
-	}
-}
-
-func (suite *TelegramTestSuite) TestDialPreferIPPriority() {
-	testData := map[string]tgAddr{
-		"prefer-ipv4": productionV4Addresses[0][0],
-		"prefer-ipv6": productionV6Addresses[0][0],
-	}
-
-	for k, v := range testData {
-		name := k
-		addr := v
-
-		suite.T().Run(name, func(t *testing.T) {
-			conn := &net.TCPConn{}
-
-			suite.dialerMock.
-				On("DialContext", mock.Anything, addr.network, addr.address).
-				Once().
-				Return(conn, nil)
-
-			tg, _ := New(suite.dialerMock, name, false)
-
-			res, err := tg.Dial(context.Background(), 1)
-			assert.NoError(t, err)
-			assert.Equal(t, conn, res)
-		})
-	}
-}
-
-func (suite *TelegramTestSuite) TestUnknownPreferIP() {
-	_, err := New(suite.dialerMock, "xxx", false)
-	suite.Error(err)
-}
-
-func (suite *TelegramTestSuite) TestFallbackDC() {
-	dcs := make([]int, 10)
-
-	for i := 0; i < len(dcs); i++ {
-		dcs[i] = suite.t.GetFallbackDC()
-	}
-
-	for _, v := range dcs {
-		value := v
-
-		suite.T().Run(strconv.Itoa(value), func(t *testing.T) {
-			assert.True(t, suite.t.IsKnownDC(value))
-		})
-	}
-}
-
-func TestTelegram(t *testing.T) {
-	t.Parallel()
-	suite.Run(t, &TelegramTestSuite{})
-}

mtglib/proxy.go

@@ -10,11 +10,11 @@ import (
 	"time"
 
 	"github.com/9seconds/mtg/v2/essentials"
+	"github.com/9seconds/mtg/v2/mtglib/internal/dc"
 	"github.com/9seconds/mtg/v2/mtglib/internal/faketls"
 	"github.com/9seconds/mtg/v2/mtglib/internal/faketls/record"
 	"github.com/9seconds/mtg/v2/mtglib/internal/obfuscated2"
 	"github.com/9seconds/mtg/v2/mtglib/internal/relay"
-	"github.com/9seconds/mtg/v2/mtglib/internal/telegram"
 	"github.com/panjf2000/ants/v2"
 )
 
@@ -28,7 +28,7 @@ type Proxy struct {
 	tolerateTimeSkewness     time.Duration
 	domainFrontingPort       int
 	workerPool               *ants.PoolWithFunc
-	telegram                 *telegram.Telegram
+	telegram                 *dc.Telegram
 
 	secret          Secret
 	network         Network
@@ -144,7 +144,6 @@ func (p *Proxy) Shutdown() {
 	p.ctxCancel()
 	p.streamWaitGroup.Wait()
 	p.workerPool.Release()
-	p.telegram.Shutdown()
 
 	p.allowlist.Shutdown()
 	p.blocklist.Shutdown()
@@ -220,18 +219,26 @@ func (p *Proxy) doObfuscated2Handshake(ctx *streamContext) error {
 }
 
 func (p *Proxy) doTelegramCall(ctx *streamContext) error {
-	dc := ctx.dc
-
-	if p.allowFallbackOnUnknownDC && !p.telegram.IsKnownDC(dc) {
-		dc = p.telegram.GetFallbackDC()
-		ctx.logger = ctx.logger.BindInt("fallback_dc", dc)
+	dcid := ctx.dc
 
+	addresses := p.telegram.GetAddresses(dcid)
+	if len(addresses) == 0 && p.allowFallbackOnUnknownDC {
+		ctx.logger = ctx.logger.BindInt("fallback_dc", dc.DefaultDC)
 		ctx.logger.Warning("unknown DC, fallbacks")
+		addresses = p.telegram.GetAddresses(dc.DefaultDC)
 	}
 
-	conn, err := p.telegram.Dial(ctx, dc)
+	var conn essentials.Conn
+	var err error
+
+	for _, addr := range addresses {
+		conn, err = p.network.Dial(addr.Network, addr.Address)
+		if err == nil {
+			break
+		}
+	}
 	if err != nil {
-		return fmt.Errorf("cannot dial to Telegram: %w", err)
+		return fmt.Errorf("no addresses to call: %w", err)
 	}
 
 	encryptor, decryptor, err := obfuscated2.ServerHandshake(conn)
@@ -293,9 +300,15 @@ func NewProxy(opts ProxyOpts) (*Proxy, error) {
 		return nil, fmt.Errorf("invalid settings: %w", err)
 	}
 
-	tg, err := telegram.New(opts.Network, opts.getPreferIP())
+	logger := opts.getLogger("proxy")
+
+	tg, err := dc.New(
+		logger.Named("telegram"),
+		opts.getPreferIP(),
+		map[int][]string{},
+	) // TODO: propagate value
 	if err != nil {
-		return nil, fmt.Errorf("cannot build telegram dialer: %w", err)
+		return nil, fmt.Errorf("cannot build telegram dc fetcher: %w", err)
 	}
 
 	ctx, cancel := context.WithCancel(context.Background())
@@ -315,7 +328,7 @@ func NewProxy(opts ProxyOpts) (*Proxy, error) {
 		telegram:                 tg,
 	}
 
-	go tg.Run(proxy.logger.Named("telegram"), 0)
+	go tg.Run(ctx, 0) // TODO: propagate value
 
 	pool, err := ants.NewPoolWithFunc(opts.getConcurrency(),
 		func(arg interface{}) {