From 94d46d2c6588701f4eac49dcf5558c49ffd3e757 Mon Sep 17 00:00:00 2001 From: 9seconds Date: Mon, 23 Feb 2026 22:36:49 +0100 Subject: [PATCH] Add fetching of addresses from proxyGetConfig endpoint --- mtglib/internal/dc/addr_test.go | 16 --- mtglib/internal/dc/init.go | 33 +++-- mtglib/internal/dc/init_test.go | 43 ++++++ mtglib/internal/dc/public_config_updater.go | 93 ++++++++++++ .../internal/dc/public_config_updater_test.go | 135 ++++++++++++++++++ mtglib/internal/dc/telegram.go | 43 ++---- mtglib/internal/dc/updater.go | 39 +++++ mtglib/internal/dc/updater_test.go | 55 +++++++ mtglib/internal/dc/view.go | 11 +- mtglib/internal/dc/view_test.go | 16 +-- mtglib/proxy.go | 15 +- 11 files changed, 419 insertions(+), 80 deletions(-) delete mode 100644 mtglib/internal/dc/addr_test.go create mode 100644 mtglib/internal/dc/init_test.go create mode 100644 mtglib/internal/dc/public_config_updater.go create mode 100644 mtglib/internal/dc/public_config_updater_test.go create mode 100644 mtglib/internal/dc/updater.go create mode 100644 mtglib/internal/dc/updater_test.go diff --git a/mtglib/internal/dc/addr_test.go b/mtglib/internal/dc/addr_test.go deleted file mode 100644 index c288c0d..0000000 --- a/mtglib/internal/dc/addr_test.go +++ /dev/null @@ -1,16 +0,0 @@ -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()) -} diff --git a/mtglib/internal/dc/init.go b/mtglib/internal/dc/init.go index c253f6e..2c0f42c 100644 --- a/mtglib/internal/dc/init.go +++ b/mtglib/internal/dc/init.go @@ -1,5 +1,10 @@ package dc +import ( + "context" + "time" +) + type preferIP uint8 const ( @@ -10,7 +15,18 @@ const ( ) const ( + // Default DC to connect to if not sure. DefaultDC = 2 + + // How often should we request updates from + // https://core.telegram.org/getProxyConfig + PublicConfigUpdateEach = time.Hour + PublicConfigUpdateURLv4 = "https://core.telegram.org/getProxyConfig" + PublicConfigUpdateURLv6 = "https://core.telegram.org/getProxyConfigV6" + + // How often should we extract hosts from Telegram using help.getConfig + // method. + OwnConfigUpdateEach = time.Hour ) type Logger interface { @@ -18,6 +34,10 @@ type Logger interface { WarningError(msg string, err error) } +type Updater interface { + Run(ctx context.Context) +} + var ( // https://github.com/telegramdesktop/tdesktop/blob/master/Telegram/SourceFiles/mtproto/mtproto_dc_options.cpp#L30 defaultDCAddrSet = dcAddrSet{ @@ -57,17 +77,4 @@ var ( }, }, } - - 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"}, - }, - }, - } ) diff --git a/mtglib/internal/dc/init_test.go b/mtglib/internal/dc/init_test.go new file mode 100644 index 0000000..426a547 --- /dev/null +++ b/mtglib/internal/dc/init_test.go @@ -0,0 +1,43 @@ +package dc + +import ( + "context" + + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/suite" +) + +type LoggerMock struct { + mock.Mock +} + +func (m *LoggerMock) Info(msg string) { + m.Called(msg) +} + +func (m *LoggerMock) WarningError(msg string, err error) { + m.Called(msg, err) +} + +type UpdaterTestSuiteBase struct { + suite.Suite + + ctx context.Context + ctxCancel context.CancelFunc + loggerMock *LoggerMock +} + +func (s *UpdaterTestSuiteBase) SetupTest() { + ctx, cancel := context.WithCancel(context.Background()) + + s.loggerMock = &LoggerMock{} + s.loggerMock.On("Info", mock.AnythingOfType("string")) + s.loggerMock.On("WarningError", mock.AnythingOfType("string"), mock.Anything) + + s.ctx = ctx + s.ctxCancel = cancel +} + +func (s *UpdaterTestSuiteBase) TearDownTest() { + s.ctxCancel() +} diff --git a/mtglib/internal/dc/public_config_updater.go b/mtglib/internal/dc/public_config_updater.go new file mode 100644 index 0000000..40ecbe7 --- /dev/null +++ b/mtglib/internal/dc/public_config_updater.go @@ -0,0 +1,93 @@ +package dc + +import ( + "bufio" + "context" + "fmt" + "io" + "net/http" + "regexp" + "strconv" +) + +var publicConfigRe = regexp.MustCompile(`^\s*proxy_for\s+(\d+)\s+(\S+?)?;\s*$`) + +type PublicConfigUpdater struct { + updater + + http *http.Client + tg *Telegram +} + +func (p PublicConfigUpdater) Run(ctx context.Context, url, network string) { + p.run(ctx, func() error { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + panic(err) + } + + resp, err := p.http.Do(req) + if err != nil { + if resp != nil { + io.Copy(io.Discard, resp.Body) + resp.Body.Close() + } + return fmt.Errorf("cannot fetch url %s: %w", url, err) + } + + if resp.StatusCode >= http.StatusBadRequest { + return fmt.Errorf("unexpected status code from %s: %d", url, resp.StatusCode) + } + + scanner := bufio.NewScanner(resp.Body) + addrs := map[int][]Addr{} + + for scanner.Scan() { + matches := publicConfigRe.FindStringSubmatch(scanner.Text()) + if len(matches) != 3 { + continue + } + + dc, err := strconv.Atoi(matches[1]) + if err != nil { + continue + } + + switch dc { + // this is a list of DC we currently support. Other are ignored. + case 203: // CDN DC + p.logger.Info(fmt.Sprintf("found %s address for DC %d", matches[2], dc)) + addrs[dc] = append(addrs[dc], Addr{ + Network: network, + Address: matches[2], + }) + } + } + + if err := scanner.Err(); err != nil { + return fmt.Errorf("cannot read response body from %s: %w", url, err) + } + + p.tg.lock.Lock() + defer p.tg.lock.Unlock() + + if network == "tcp4" { + p.tg.view.publicConfigs.v4 = addrs + } else { + p.tg.view.publicConfigs.v6 = addrs + } + + return nil + }) +} + +func NewPublicConfigUpdater(tg *Telegram, logger Logger, client *http.Client) PublicConfigUpdater { + return PublicConfigUpdater{ + updater: updater{ + logger: logger, + period: PublicConfigUpdateEach, + }, + http: client, + tg: tg, + } +} diff --git a/mtglib/internal/dc/public_config_updater_test.go b/mtglib/internal/dc/public_config_updater_test.go new file mode 100644 index 0000000..4da1773 --- /dev/null +++ b/mtglib/internal/dc/public_config_updater_test.go @@ -0,0 +1,135 @@ +package dc + +import ( + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" + "github.com/stretchr/testify/suite" +) + +type PublicConfigUpdaterTestSuite struct { + UpdaterTestSuiteBase + + u PublicConfigUpdater + lock sync.Mutex + srv *httptest.Server + responseHandler func(w http.ResponseWriter) +} + +func (s *PublicConfigUpdaterTestSuite) SetupSuite() { + s.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + s.lock.Lock() + s.responseHandler(w) + s.lock.Unlock() + })) +} + +func (s *PublicConfigUpdaterTestSuite) TearDownSuite() { + s.srv.Close() +} + +func (s *PublicConfigUpdaterTestSuite) SetupTest() { + s.UpdaterTestSuiteBase.SetupTest() + + tg, err := New("prefer-ipv4") + require.NoError(s.T(), err) + + s.u = NewPublicConfigUpdater(tg, s.loggerMock, s.srv.Client()) +} + +func (s *PublicConfigUpdaterTestSuite) Test502StatusCode() { + done := false + + s.responseHandler = func(w http.ResponseWriter) { + w.WriteHeader(http.StatusBadGateway) + done = true + } + go s.u.Run(s.ctx, s.srv.URL, "tcp4") + + s.Eventually(func() bool { + s.lock.Lock() + defer s.lock.Unlock() + + return done + }, time.Second, 10*time.Millisecond) + + s.Len(s.u.tg.view.publicConfigs.v4, 0) +} + +func (s *PublicConfigUpdaterTestSuite) TestEmptyFile() { + done := false + + s.responseHandler = func(w http.ResponseWriter) { + done = true + w.WriteHeader(http.StatusOK) + } + go s.u.Run(s.ctx, s.srv.URL, "tcp4") + + s.Eventually(func() bool { + s.lock.Lock() + defer s.lock.Unlock() + + return done + }, time.Second, 10*time.Millisecond) + + s.Len(s.u.tg.view.publicConfigs.v4, 0) +} + +func (s *PublicConfigUpdaterTestSuite) TestGarbage() { + result := ` +proxy_for -1 -1; +proxy_for 100 100.10.0.0:3333; +lala 0 0 +` + done := false + + s.responseHandler = func(w http.ResponseWriter) { + done = true + w.WriteHeader(http.StatusOK) + w.Write([]byte(result)) + } + go s.u.Run(s.ctx, s.srv.URL, "tcp4") + + s.Eventually(func() bool { + s.lock.Lock() + defer s.lock.Unlock() + + return done + }, time.Second, 10*time.Millisecond) + + s.Len(s.u.tg.view.publicConfigs.v4, 0) +} + +func (s *PublicConfigUpdaterTestSuite) TestOk() { + result := ` +proxy_for 203 100.10.0.0:3333; +proxy_for -100 101.10.0.0:3333; +` + done := false + + s.responseHandler = func(w http.ResponseWriter) { + done = true + w.WriteHeader(http.StatusOK) + w.Write([]byte(result)) + } + go s.u.Run(s.ctx, s.srv.URL, "tcp4") + + s.Eventually(func() bool { + s.lock.Lock() + defer s.lock.Unlock() + + return done + }, time.Second, 10*time.Millisecond) + + s.Len(s.u.tg.view.publicConfigs.v4, 1) + s.Len(s.u.tg.view.publicConfigs.v4[203], 1) + s.Equal("100.10.0.0:3333", s.u.tg.view.publicConfigs.v4[203][0].Address) +} + +func TestPublicConfigUpdater(t *testing.T) { + suite.Run(t, &PublicConfigUpdaterTestSuite{}) +} diff --git a/mtglib/internal/dc/telegram.go b/mtglib/internal/dc/telegram.go index e9f261f..c9a85d4 100644 --- a/mtglib/internal/dc/telegram.go +++ b/mtglib/internal/dc/telegram.go @@ -1,17 +1,23 @@ package dc import ( + "context" "fmt" - "net" "strings" + "sync" ) type Telegram struct { + ctx context.Context + lock sync.RWMutex view dcView preferIP preferIP } func (t *Telegram) GetAddresses(dc int) []Addr { + t.lock.RLock() + defer t.lock.RUnlock() + switch t.preferIP { case preferIPOnlyIPv4: return t.view.getV4(dc) @@ -24,7 +30,7 @@ func (t *Telegram) GetAddresses(dc int) []Addr { return append(t.view.getV6(dc), t.view.getV4(dc)...) } -func New(ipPreference string, userOverrides map[int][]string) (*Telegram, error) { +func New(ipPreference string) (*Telegram, error) { var pref preferIP switch strings.ToLower(ipPreference) { @@ -40,40 +46,7 @@ func New(ipPreference string, userOverrides map[int][]string) (*Telegram, error) return nil, fmt.Errorf("unknown ip preference %s", ipPreference) } - overrides := dcAddrSet{ - v4: map[int][]Addr{}, - v6: map[int][]Addr{}, - } - 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, - }, preferIP: pref, }, nil } diff --git a/mtglib/internal/dc/updater.go b/mtglib/internal/dc/updater.go new file mode 100644 index 0000000..5d57af4 --- /dev/null +++ b/mtglib/internal/dc/updater.go @@ -0,0 +1,39 @@ +package dc + +import ( + "context" + "time" +) + +type updater struct { + logger Logger + period time.Duration +} + +func (u updater) run(ctx context.Context, callback func() error) { + ticker := time.NewTicker(u.period) + + defer func() { + ticker.Stop() + + select { + case <-ticker.C: + default: + } + }() + + for { + u.logger.Info("start update") + if err := callback(); err != nil { + u.logger.WarningError("cannot update: %w", err) + } + u.logger.Info("updated") + + select { + case <-ctx.Done(): + u.logger.Info("stop updating") + return + case <-ticker.C: + } + } +} diff --git a/mtglib/internal/dc/updater_test.go b/mtglib/internal/dc/updater_test.go new file mode 100644 index 0000000..f6916a4 --- /dev/null +++ b/mtglib/internal/dc/updater_test.go @@ -0,0 +1,55 @@ +package dc + +import ( + "sync" + "testing" + "time" + + "github.com/stretchr/testify/suite" +) + +type UpdaterTestSuite struct { + UpdaterTestSuiteBase + + u updater +} + +func (s *UpdaterTestSuite) SetupTest() { + s.UpdaterTestSuiteBase.SetupTest() + s.u = updater{ + logger: s.loggerMock, + period: 100 * time.Millisecond, + } +} + +func (s *UpdaterTestSuite) TestPeriodicUpdates() { + ticker := time.NewTicker(10 * time.Millisecond) + defer ticker.Stop() + + lock := &sync.Mutex{} + collected := []time.Time{} + + go s.u.run(s.ctx, func() error { + select { + case <-s.ctx.Done(): + case value := <-ticker.C: + lock.Lock() + collected = append(collected, value) + lock.Unlock() + } + + return nil + }) + + s.Eventually(func() bool { + lock.Lock() + defer lock.Unlock() + + return len(collected) == 3 + }, time.Second, 10*time.Millisecond) +} + +func TestUpdater(t *testing.T) { + t.Parallel() + suite.Run(t, &UpdaterTestSuite{}) +} diff --git a/mtglib/internal/dc/view.go b/mtglib/internal/dc/view.go index efee1c0..de06499 100644 --- a/mtglib/internal/dc/view.go +++ b/mtglib/internal/dc/view.go @@ -1,20 +1,21 @@ package dc type dcView struct { - overrides dcAddrSet + publicConfigs dcAddrSet + ownConfigs dcAddrSet } func (d dcView) getV4(dc int) []Addr { - addrs := d.overrides.getV4(dc) - addrs = append(addrs, defaultDCOverridesAddrSet.getV4(dc)...) + addrs := d.publicConfigs.getV4(dc) + addrs = append(addrs, d.ownConfigs.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 := d.publicConfigs.getV6(dc) + addrs = append(addrs, d.ownConfigs.getV6(dc)...) addrs = append(addrs, defaultDCAddrSet.getV6(dc)...) return addrs diff --git a/mtglib/internal/dc/view_test.go b/mtglib/internal/dc/view_test.go index 74d46cb..90742c0 100644 --- a/mtglib/internal/dc/view_test.go +++ b/mtglib/internal/dc/view_test.go @@ -16,7 +16,7 @@ type ViewTestSuite struct { func (suite *ViewTestSuite) SetupSuite() { suite.view = dcView{ - overrides: dcAddrSet{ + publicConfigs: dcAddrSet{ v4: map[int][]Addr{ 111: { {Network: "tcp4", Address: "127.0.0.1:443"}, @@ -37,15 +37,14 @@ func (suite *ViewTestSuite) SetupSuite() { func (suite *ViewTestSuite) TestGetV4() { testData := map[int][]Addr{ 111: { - {"tcp4", "127.0.0.1:443"}, + {Network: "tcp4", Address: "127.0.0.1:443"}, }, 203: { - {"tcp4", "127.0.0.2:443"}, - {"tcp4", "91.105.192.100:443"}, + {Network: "tcp4", Address: "127.0.0.2:443"}, }, 2: { - {"tcp4", "149.154.167.51:443"}, - {"tcp4", "95.161.76.100:443"}, + {Network: "tcp4", Address: "149.154.167.51:443"}, + {Network: "tcp4", Address: "95.161.76.100:443"}, }, } @@ -60,11 +59,10 @@ func (suite *ViewTestSuite) TestGetV6() { testData := map[int][]Addr{ 111: {}, 203: { - {"tcp6", "xxx"}, - {"tcp6", "[2a0a:f280:0203:000a:5000:0000:0000:0100]:443"}, + {Network: "tcp6", Address: "xxx"}, }, 1: { - {"tcp6", "[2001:b28:f23d:f001::a]:443"}, + {Network: "tcp6", Address: "[2001:b28:f23d:f001::a]:443"}, }, } diff --git a/mtglib/proxy.go b/mtglib/proxy.go index fe69862..78f239d 100644 --- a/mtglib/proxy.go +++ b/mtglib/proxy.go @@ -297,12 +297,15 @@ func NewProxy(opts ProxyOpts) (*Proxy, error) { return nil, fmt.Errorf("invalid settings: %w", err) } - tg, err := dc.New(opts.getPreferIP(), opts.DCOverrides) + tg, err := dc.New(opts.getPreferIP()) if err != nil { return nil, fmt.Errorf("cannot build telegram dc fetcher: %w", err) } ctx, cancel := context.WithCancel(context.Background()) + logger := opts.getLogger("proxy") + updatersLogger := logger.Named("telegram-updaters") + proxy := &Proxy{ ctx: ctx, ctxCancel: cancel, @@ -312,7 +315,7 @@ func NewProxy(opts ProxyOpts) (*Proxy, error) { blocklist: opts.IPBlocklist, allowlist: opts.IPAllowlist, eventStream: opts.EventStream, - logger: opts.getLogger("proxy"), + logger: logger, domainFrontingPort: opts.getDomainFrontingPort(), tolerateTimeSkewness: opts.getTolerateTimeSkewness(), allowFallbackOnUnknownDC: opts.AllowFallbackOnUnknownDC, @@ -322,6 +325,14 @@ func NewProxy(opts ProxyOpts) (*Proxy, error) { }, } + publicConfigUpdater := dc.NewPublicConfigUpdater( + tg, + updatersLogger.Named("public-config"), + opts.Network.MakeHTTPClient(nil), + ) + go publicConfigUpdater.Run(ctx, dc.PublicConfigUpdateURLv4, "tcp4") + go publicConfigUpdater.Run(ctx, dc.PublicConfigUpdateURLv6, "tcp6") + pool, err := ants.NewPoolWithFunc(opts.getConcurrency(), func(arg any) { proxy.ServeConn(arg.(essentials.Conn)) //nolint: forcetypeassert