Add fetching of addresses from proxyGetConfig endpoint

mtglib/internal/dc/addr_test.go

@@ -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())
-}

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"},
-			},
-		},
-	}
 )

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()
+}

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,
+	}
+}

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{})
+}

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
 }

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:
+		}
+	}
+}

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{})
+}

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

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"},
 		},
 	}
 

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