Просмотр исходного кода

Merge pull request #210 from 9seconds/fallback-to-random-dc

Fallback to another DC if given is unknown
tags/v2.1.2^2
Sergey Arkhipov 4 лет назад
Родитель
Сommit
f742066c54
Аккаунт пользователя с таким Email не найден

+ 8
- 0
mtglib/internal/telegram/address_pool.go Просмотреть файл

7
 	v6 [][]tgAddr
7
 	v6 [][]tgAddr
8
 }
8
 }
9
 
9
 
10
+func (a addressPool) isValidDC(dc int) bool {
11
+	return dc > 0 && dc <= len(a.v4) && dc <= len(a.v6)
12
+}
13
+
14
+func (a addressPool) getRandomDC() int {
15
+	return 1 + rand.Intn(len(a.v4))
16
+}
17
+
10
 func (a addressPool) getV4(dc int) []tgAddr {
18
 func (a addressPool) getV4(dc int) []tgAddr {
11
 	return a.get(a.v4, dc-1)
19
 	return a.get(a.v4, dc-1)
12
 }
20
 }

+ 8
- 0
mtglib/internal/telegram/telegram.go Просмотреть файл

42
 	return nil, fmt.Errorf("cannot dial to %d dc: %w", dc, err)
42
 	return nil, fmt.Errorf("cannot dial to %d dc: %w", dc, err)
43
 }
43
 }
44
 
44
 
45
+func (t Telegram) IsKnownDC(dc int) bool {
46
+	return t.pool.isValidDC(dc)
47
+}
48
+
49
+func (t Telegram) GetFallbackDC() int {
50
+	return t.pool.getRandomDC()
51
+}
52
+
45
 func New(dialer Dialer, ipPreference string, useTestDCs bool) (*Telegram, error) {
53
 func New(dialer Dialer, ipPreference string, useTestDCs bool) (*Telegram, error) {
46
 	var pref preferIP
54
 	var pref preferIP
47
 
55
 

+ 18
- 0
mtglib/internal/telegram/telegram_internal_test.go Просмотреть файл

44
 		suite.T().Run(strconv.Itoa(value), func(t *testing.T) {
44
 		suite.T().Run(strconv.Itoa(value), func(t *testing.T) {
45
 			_, err := suite.t.Dial(context.Background(), value)
45
 			_, err := suite.t.Dial(context.Background(), value)
46
 			assert.Error(t, err)
46
 			assert.Error(t, err)
47
+			assert.False(t, suite.t.IsKnownDC(value))
47
 		})
48
 		})
48
 	}
49
 	}
49
 }
50
 }
71
 
72
 
72
 			_, err := suite.t.Dial(context.Background(), idx)
73
 			_, err := suite.t.Dial(context.Background(), idx)
73
 			assert.True(t, errors.Is(err, io.EOF))
74
 			assert.True(t, errors.Is(err, io.EOF))
75
+			assert.True(t, suite.t.IsKnownDC(idx))
74
 		})
76
 		})
75
 	}
77
 	}
76
 }
78
 }
135
 	suite.Error(err)
137
 	suite.Error(err)
136
 }
138
 }
137
 
139
 
140
+func (suite *TelegramTestSuite) TestFallbackDC() {
141
+	dcs := make([]int, 10)
142
+
143
+	for i := 0; i < len(dcs); i++ {
144
+		dcs[i] = suite.t.GetFallbackDC()
145
+	}
146
+
147
+	for _, v := range dcs {
148
+		value := v
149
+
150
+		suite.T().Run(strconv.Itoa(value), func(t *testing.T) {
151
+			assert.True(t, suite.t.IsKnownDC(value))
152
+		})
153
+	}
154
+}
155
+
138
 func TestTelegram(t *testing.T) {
156
 func TestTelegram(t *testing.T) {
139
 	t.Parallel()
157
 	t.Parallel()
140
 	suite.Run(t, &TelegramTestSuite{})
158
 	suite.Run(t, &TelegramTestSuite{})

+ 10
- 1
mtglib/proxy.go Просмотреть файл

207
 }
207
 }
208
 
208
 
209
 func (p *Proxy) doTelegramCall(ctx *streamContext) error {
209
 func (p *Proxy) doTelegramCall(ctx *streamContext) error {
210
-	conn, err := p.telegram.Dial(ctx, ctx.dc)
210
+	dc := ctx.dc
211
+
212
+	if !p.telegram.IsKnownDC(dc) {
213
+		dc = p.telegram.GetFallbackDC()
214
+		ctx.logger = ctx.logger.BindInt("fallback_dc", dc)
215
+
216
+		ctx.logger.Warning("unknown DC, fallbacks")
217
+	}
218
+
219
+	conn, err := p.telegram.Dial(ctx, dc)
211
 	if err != nil {
220
 	if err != nil {
212
 		return fmt.Errorf("cannot dial to Telegram: %w", err)
221
 		return fmt.Errorf("cannot dial to Telegram: %w", err)
213
 	}
222
 	}

Загрузка…
Отмена
Сохранить