Fix tests

mtglib/internal/dc/init.go

@@ -38,9 +38,8 @@ type Updater interface {
 	Run(ctx context.Context)
 }
 
-var (
+// https://github.com/telegramdesktop/tdesktop/blob/master/Telegram/SourceFiles/mtproto/mtproto_dc_options.cpp#L30
-	// https://github.com/telegramdesktop/tdesktop/blob/master/Telegram/SourceFiles/mtproto/mtproto_dc_options.cpp#L30
+var defaultDCAddrSet = dcAddrSet{
-	defaultDCAddrSet = dcAddrSet{
 	v4: map[int][]Addr{
 		1: {
 			{Network: "tcp4", Address: "149.154.175.50:443"},
@@ -76,5 +75,4 @@ var (
 			{Network: "tcp6", Address: "[2001:b28:f23f:f005::a]:443"},
 		},
 	},
-	}
+}
-)

mtglib/internal/dc/public_config_updater.go

@@ -19,7 +19,7 @@ type PublicConfigUpdater struct {
 	tg   *Telegram
 }
 
-func (p PublicConfigUpdater) Run(ctx context.Context, url, network string) {
+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 {
@@ -81,8 +81,8 @@ func (p PublicConfigUpdater) Run(ctx context.Context, url, network string) {
 	})
 }
 
-func NewPublicConfigUpdater(tg *Telegram, logger Logger, client *http.Client) PublicConfigUpdater {
+func NewPublicConfigUpdater(tg *Telegram, logger Logger, client *http.Client) *PublicConfigUpdater {
-	return PublicConfigUpdater{
+	return &PublicConfigUpdater{
 		updater: updater{
 			logger: logger,
 			period: PublicConfigUpdateEach,

mtglib/internal/dc/public_config_updater_test.go

@@ -14,7 +14,7 @@ import (
 type PublicConfigUpdaterTestSuite struct {
 	UpdaterTestSuiteBase
 
-	u               PublicConfigUpdater
+	u               *PublicConfigUpdater
 	lock            sync.Mutex
 	srv             *httptest.Server
 	responseHandler func(w http.ResponseWriter)
@@ -42,39 +42,27 @@ func (s *PublicConfigUpdaterTestSuite) SetupTest() {
 }
 
 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.u.Run(s.ctx, s.srv.URL, "tcp4")
 
-	s.Eventually(func() bool {
+	time.Sleep(100 * time.Millisecond)
-		s.lock.Lock()
+	s.ctxCancel()
-		defer s.lock.Unlock()
+	s.u.Wait()
-
-		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.u.Run(s.ctx, s.srv.URL, "tcp4")
 
-	s.Eventually(func() bool {
+	time.Sleep(100 * time.Millisecond)
-		s.lock.Lock()
+	s.ctxCancel()
-		defer s.lock.Unlock()
+	s.u.Wait()
-
-		return done
-	}, time.Second, 10*time.Millisecond)
 
 	s.Len(s.u.tg.view.publicConfigs.v4, 0)
 }
@@ -85,21 +73,16 @@ 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)) //nolint: errcheck
 	}
-	go s.u.Run(s.ctx, s.srv.URL, "tcp4")
+	s.u.Run(s.ctx, s.srv.URL, "tcp4")
 
-	s.Eventually(func() bool {
+	time.Sleep(100 * time.Millisecond)
-		s.lock.Lock()
+	s.ctxCancel()
-		defer s.lock.Unlock()
+	s.u.Wait()
-
-		return done
-	}, time.Second, 10*time.Millisecond)
 
 	s.Len(s.u.tg.view.publicConfigs.v4, 0)
 }
@@ -109,21 +92,16 @@ func (s *PublicConfigUpdaterTestSuite) TestOk() {
 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)) //nolint: errcheck
 	}
-	go s.u.Run(s.ctx, s.srv.URL, "tcp4")
+	s.u.Run(s.ctx, s.srv.URL, "tcp4")
 
-	s.Eventually(func() bool {
+	time.Sleep(100 * time.Millisecond)
-		s.lock.Lock()
+	s.ctxCancel()
-		defer s.lock.Unlock()
+	s.u.Wait()
-
-		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)

mtglib/internal/dc/updater.go

@@ -2,15 +2,22 @@ package dc
 
 import (
 	"context"
+	"sync"
 	"time"
 )
 
 type updater struct {
+	wg     sync.WaitGroup
 	logger Logger
 	period time.Duration
 }
 
-func (u updater) run(ctx context.Context, callback func() error) {
+func (u *updater) Wait() {
+	u.wg.Wait()
+}
+
+func (u *updater) run(ctx context.Context, callback func() error) {
+	u.wg.Go(func() {
 		ticker := time.NewTicker(u.period)
 
 		defer func() {
@@ -36,4 +43,5 @@ func (u updater) run(ctx context.Context, callback func() error) {
 			case <-ticker.C:
 			}
 		}
+	})
 }

mtglib/proxy.go

@@ -30,6 +30,7 @@ type Proxy struct {
 	domainFrontingIP         string
 	workerPool               *ants.PoolWithFunc
 	telegram                 *dc.Telegram
+	configUpdater            *dc.PublicConfigUpdater
 	clientObfuscatror        obfuscation.Obfuscator
 
 	secret          Secret
@@ -152,6 +153,7 @@ func (p *Proxy) Shutdown() {
 	p.ctxCancel()
 	p.streamWaitGroup.Wait()
 	p.workerPool.Release()
+	p.configUpdater.Wait()
 
 	p.allowlist.Shutdown()
 	p.blocklist.Shutdown()
@@ -328,18 +330,18 @@ func NewProxy(opts ProxyOpts) (*Proxy, error) {
 		tolerateTimeSkewness:     opts.getTolerateTimeSkewness(),
 		allowFallbackOnUnknownDC: opts.AllowFallbackOnUnknownDC,
 		telegram:                 tg,
+		configUpdater: dc.NewPublicConfigUpdater(
+			tg,
+			updatersLogger.Named("public-config"),
+			opts.Network.MakeHTTPClient(nil),
+		),
 		clientObfuscatror: obfuscation.Obfuscator{
 			Secret: opts.Secret.Key[:],
 		},
 	}
 
-	publicConfigUpdater := dc.NewPublicConfigUpdater(
+	proxy.configUpdater.Run(ctx, dc.PublicConfigUpdateURLv4, "tcp4")
-		tg,
+	proxy.configUpdater.Run(ctx, dc.PublicConfigUpdateURLv6, "tcp6")
-		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) {