diff --git a/mtglib/internal/telegram/address_pool.go b/mtglib/internal/telegram/address_pool.go index 0f49e80..2a22371 100644 --- a/mtglib/internal/telegram/address_pool.go +++ b/mtglib/internal/telegram/address_pool.go @@ -7,6 +7,14 @@ type addressPool struct { v6 [][]tgAddr } +func (a addressPool) isValidDC(dc int) bool { + return dc > 0 && dc <= len(a.v4) && dc <= len(a.v6) +} + +func (a addressPool) getRandomDC() int { + return 1 + rand.Intn(len(a.v4)) +} + func (a addressPool) getV4(dc int) []tgAddr { return a.get(a.v4, dc-1) } diff --git a/mtglib/internal/telegram/telegram.go b/mtglib/internal/telegram/telegram.go index 474dfa4..43d4867 100644 --- a/mtglib/internal/telegram/telegram.go +++ b/mtglib/internal/telegram/telegram.go @@ -42,6 +42,14 @@ 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 (t Telegram) IsKnownDC(dc int) bool { + return t.pool.isValidDC(dc) +} + +func (t Telegram) GetFallbackDC() int { + return t.pool.getRandomDC() +} + func New(dialer Dialer, ipPreference string, useTestDCs bool) (*Telegram, error) { var pref preferIP diff --git a/mtglib/internal/telegram/telegram_internal_test.go b/mtglib/internal/telegram/telegram_internal_test.go index 77d5939..794e97a 100644 --- a/mtglib/internal/telegram/telegram_internal_test.go +++ b/mtglib/internal/telegram/telegram_internal_test.go @@ -44,6 +44,7 @@ func (suite *TelegramTestSuite) TestUnknownDC() { 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)) }) } } @@ -71,6 +72,7 @@ func (suite *TelegramTestSuite) TestDialToCorrectIPs() { _, err := suite.t.Dial(context.Background(), idx) assert.True(t, errors.Is(err, io.EOF)) + assert.True(t, suite.t.IsKnownDC(idx)) }) } } @@ -135,6 +137,22 @@ func (suite *TelegramTestSuite) TestUnknownPreferIP() { 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{}) diff --git a/mtglib/proxy.go b/mtglib/proxy.go index b0a21db..5b685ec 100644 --- a/mtglib/proxy.go +++ b/mtglib/proxy.go @@ -207,7 +207,16 @@ func (p *Proxy) doObfuscated2Handshake(ctx *streamContext) error { } func (p *Proxy) doTelegramCall(ctx *streamContext) error { - conn, err := p.telegram.Dial(ctx, ctx.dc) + dc := ctx.dc + + if !p.telegram.IsKnownDC(dc) { + dc = p.telegram.GetFallbackDC() + ctx.logger = ctx.logger.BindInt("fallback_dc", dc) + + ctx.logger.Warning("unknown DC, fallbacks") + } + + conn, err := p.telegram.Dial(ctx, dc) if err != nil { return fmt.Errorf("cannot dial to Telegram: %w", err) }