From 4a49c4db96cee54c676c030a61d29a288a6157d5 Mon Sep 17 00:00:00 2001
From: Eugene Burkov <e.burkov@adguard.com>
Date: Tue, 26 Nov 2024 20:35:16 +0300
Subject: [PATCH] Pull request 2307: AGDNS-2556 Custom updater URL

Squashed commit of the following:

commit 73f946138ccb4f89141f192b6cb1a21887604ab4
Merge: c58847bfb d578c713f
Author: Eugene Burkov <E.Burkov@AdGuard.COM>
Date:   Tue Nov 26 17:42:29 2024 +0300

    Merge branch 'master' into AGDNS-2556-custom-update-url

commit c58847bfb08131263e1cff4813eb4a466f613d91
Author: Eugene Burkov <E.Burkov@AdGuard.COM>
Date:   Tue Nov 26 17:34:11 2024 +0300

    home: imp logging

commit 0d451621d76fdf2c363d223eb29c4442d8f36dc8
Author: Eugene Burkov <E.Burkov@AdGuard.COM>
Date:   Tue Nov 26 15:12:04 2024 +0300

    home: rename config field

commit c7f3822929e9199f8f411f1a0ad072c643feb42f
Author: Eugene Burkov <E.Burkov@AdGuard.COM>
Date:   Tue Nov 26 15:07:09 2024 +0300

    all: enable updater for some cases

commit 872cd3a18c876076ea643624336cfc0a4296a81d
Author: Eugene Burkov <E.Burkov@AdGuard.COM>
Date:   Fri Nov 22 19:09:18 2024 +0300

    updater: imp test

commit c9efb412e7411b769df54b7247fe168047fb9799
Merge: c989eef71 abb738013
Author: Eugene Burkov <E.Burkov@AdGuard.COM>
Date:   Fri Nov 22 17:51:46 2024 +0300

    Merge branch 'master' into AGDNS-2556-custom-update-url

commit c989eef715ae7edd98d7b2d5df06fd3d04153209
Author: Eugene Burkov <E.Burkov@AdGuard.COM>
Date:   Fri Nov 22 17:46:34 2024 +0300

    all: imp code

commit 0452d8b356e6d0b73b097d43b97b7027fcca752d
Author: Eugene Burkov <E.Burkov@AdGuard.COM>
Date:   Fri Nov 22 15:37:21 2024 +0300

    all: add custom url to updater
---
 internal/home/config.go                   |   6 +
 internal/home/controlupdate.go            |  33 ++---
 internal/home/home.go                     | 144 +++++++++++++++-------
 internal/home/web.go                      |   2 +
 internal/updater/check.go                 |   8 +-
 internal/updater/check_test.go            |   6 +-
 internal/updater/updater.go               |  31 +++--
 internal/updater/updater_internal_test.go |  50 ++++----
 internal/updater/updater_test.go          |   5 +-
 9 files changed, 183 insertions(+), 102 deletions(-)

diff --git a/internal/home/config.go b/internal/home/config.go
index 810ec24e..5242dbc6 100644
--- a/internal/home/config.go
+++ b/internal/home/config.go
@@ -162,6 +162,12 @@ type configuration struct {
 	// SchemaVersion is the version of the configuration schema.  See
 	// [configmigrate.LastSchemaVersion].
 	SchemaVersion uint `yaml:"schema_version"`
+
+	// UnsafeUseCustomUpdateIndexURL is the URL to the custom update index.
+	//
+	// NOTE: It's only exists for testing purposes and should not be used in
+	// release.
+	UnsafeUseCustomUpdateIndexURL bool `yaml:"unsafe_use_custom_update_index_url,omitempty"`
 }
 
 // httpConfig is a block with HTTP configuration params.
diff --git a/internal/home/controlupdate.go b/internal/home/controlupdate.go
index 50a1a6f3..aeab8810 100644
--- a/internal/home/controlupdate.go
+++ b/internal/home/controlupdate.go
@@ -75,30 +75,31 @@ func (web *webAPI) handleVersionJSON(w http.ResponseWriter, r *http.Request) {
 // update server.
 func (web *webAPI) requestVersionInfo(resp *versionResponse, recheck bool) (err error) {
 	updater := web.conf.updater
-	for i := 0; i != 3; i++ {
+	for range 3 {
 		resp.VersionInfo, err = updater.VersionInfo(recheck)
-		if err != nil {
-			var terr temporaryError
-			if errors.As(err, &terr) && terr.Temporary() {
-				// Temporary network error.  This case may happen while we're
-				// restarting our DNS server.  Log and sleep for some time.
-				//
-				// See https://github.com/AdguardTeam/AdGuardHome/issues/934.
-				d := time.Duration(i) * time.Second
-				log.Info("update: temp net error: %q; sleeping for %s and retrying", err, d)
-				time.Sleep(d)
+		if err == nil {
+			return nil
+		}
 
-				continue
-			}
+		var terr temporaryError
+		if errors.As(err, &terr) && terr.Temporary() {
+			// Temporary network error.  This case may happen while we're
+			// restarting our DNS server.  Log and sleep for some time.
+			//
+			// See https://github.com/AdguardTeam/AdGuardHome/issues/934.
+			const sleepTime = 2 * time.Second
+
+			log.Info("update: temp net error: %v; sleeping for %s and retrying", err, sleepTime)
+			time.Sleep(sleepTime)
+
+			continue
 		}
 
 		break
 	}
 
 	if err != nil {
-		vcu := updater.VersionCheckURL()
-
-		return fmt.Errorf("getting version info from %s: %w", vcu, err)
+		return fmt.Errorf("getting version info: %w", err)
 	}
 
 	return nil
diff --git a/internal/home/home.go b/internal/home/home.go
index 34620e2c..da146a6d 100644
--- a/internal/home/home.go
+++ b/internal/home/home.go
@@ -12,7 +12,6 @@ import (
 	"net/url"
 	"os"
 	"os/signal"
-	"path"
 	"path/filepath"
 	"runtime"
 	"slices"
@@ -495,11 +494,42 @@ func checkPorts() (err error) {
 	return nil
 }
 
+// isUpdateEnabled returns true if the update is enabled for current
+// configuration.  It also logs the decision.  customURL should be true if the
+// updater is using a custom URL.
+func isUpdateEnabled(ctx context.Context, l *slog.Logger, opts *options, customURL bool) (ok bool) {
+	if opts.disableUpdate {
+		l.DebugContext(ctx, "updates are disabled by command-line option")
+
+		return false
+	}
+
+	switch version.Channel() {
+	case
+		version.ChannelDevelopment,
+		version.ChannelCandidate:
+		if customURL {
+			l.DebugContext(ctx, "updates are enabled because custom url is used")
+		} else {
+			l.DebugContext(ctx, "updates are disabled for development and candidate builds")
+		}
+
+		return customURL
+	default:
+		l.DebugContext(ctx, "updates are enabled")
+
+		return true
+	}
+}
+
+// initWeb initializes the web module.
 func initWeb(
+	ctx context.Context,
 	opts options,
 	clientBuildFS fs.FS,
 	upd *updater.Updater,
 	l *slog.Logger,
+	customURL bool,
 ) (web *webAPI, err error) {
 	var clientFS fs.FS
 	if opts.localFrontend {
@@ -513,17 +543,7 @@ func initWeb(
 		}
 	}
 
-	disableUpdate := opts.disableUpdate
-	switch version.Channel() {
-	case
-		version.ChannelDevelopment,
-		version.ChannelCandidate:
-		disableUpdate = true
-	}
-
-	if disableUpdate {
-		log.Info("AdGuard Home updates are disabled")
-	}
+	disableUpdate := !isUpdateEnabled(ctx, l, &opts, customURL)
 
 	webConf := &webConfig{
 		updater: upd,
@@ -544,7 +564,7 @@ func initWeb(
 
 	web = newWebAPI(webConf, l)
 	if web == nil {
-		return nil, fmt.Errorf("initializing web: %w", err)
+		return nil, errors.Error("can not initialize web")
 	}
 
 	return web, nil
@@ -557,6 +577,8 @@ func fatalOnError(err error) {
 }
 
 // run configures and starts AdGuard Home.
+//
+// TODO(e.burkov):  Make opts a pointer.
 func run(opts options, clientBuildFS fs.FS, done chan struct{}) {
 	// Configure working dir.
 	err := initWorkingDir(opts)
@@ -604,33 +626,13 @@ func run(opts options, clientBuildFS fs.FS, done chan struct{}) {
 	execPath, err := os.Executable()
 	fatalOnError(errors.Annotate(err, "getting executable path: %w"))
 
-	u := &url.URL{
-		Scheme: urlutil.SchemeHTTPS,
-		// TODO(a.garipov): Make configurable.
-		Host: "static.adtidy.org",
-		Path: path.Join("adguardhome", version.Channel(), "version.json"),
-	}
-
 	confPath := configFilePath()
-	log.Debug("using config path %q for updater", confPath)
 
-	upd := updater.NewUpdater(&updater.Config{
-		Client:          config.Filtering.HTTPClient,
-		Version:         version.Version(),
-		Channel:         version.Channel(),
-		GOARCH:          runtime.GOARCH,
-		GOOS:            runtime.GOOS,
-		GOARM:           version.GOARM(),
-		GOMIPS:          version.GOMIPS(),
-		WorkDir:         Context.workDir,
-		ConfName:        confPath,
-		ExecPath:        execPath,
-		VersionCheckURL: u.String(),
-	})
+	upd, customURL := newUpdater(ctx, slogLogger, Context.workDir, confPath, execPath, config)
 
 	// TODO(e.burkov): This could be made earlier, probably as the option's
 	// effect.
-	cmdlineUpdate(opts, upd, slogLogger)
+	cmdlineUpdate(ctx, slogLogger, opts, upd)
 
 	if !Context.firstRun {
 		// Save the updated config.
@@ -658,7 +660,7 @@ func run(opts options, clientBuildFS fs.FS, done chan struct{}) {
 		onConfigModified()
 	}
 
-	Context.web, err = initWeb(opts, clientBuildFS, upd, slogLogger)
+	Context.web, err = initWeb(ctx, opts, clientBuildFS, upd, slogLogger, customURL)
 	fatalOnError(err)
 
 	statsDir, querylogDir, err := checkStatsAndQuerylogDirs(&Context, config)
@@ -696,6 +698,57 @@ func run(opts options, clientBuildFS fs.FS, done chan struct{}) {
 	<-done
 }
 
+// newUpdater creates a new AdGuard Home updater.  customURL is true if the user
+// has specified a custom version announcement URL.
+func newUpdater(
+	ctx context.Context,
+	l *slog.Logger,
+	workDir string,
+	confPath string,
+	execPath string,
+	config *configuration,
+) (upd *updater.Updater, customURL bool) {
+	// envName is the name of the environment variable that can be used to
+	// override the default version check URL.
+	const envName = "ADGUARD_HOME_TEST_UPDATE_VERSION_URL"
+
+	customURLStr := os.Getenv(envName)
+
+	var versionURL *url.URL
+	switch {
+	case version.Channel() == version.ChannelRelease:
+		// Only enable custom version URL for development builds.
+		l.DebugContext(ctx, "custom version url is disabled for release builds")
+	case !config.UnsafeUseCustomUpdateIndexURL:
+		l.DebugContext(ctx, "custom version url is disabled in config")
+	default:
+		versionURL, _ = url.Parse(customURLStr)
+	}
+
+	err := urlutil.ValidateHTTPURL(versionURL)
+	if customURL = err == nil; customURL {
+		l.DebugContext(ctx, "parsing custom version url", slogutil.KeyError, err)
+
+		versionURL = updater.DefaultVersionURL()
+	}
+
+	l.DebugContext(ctx, "creating updater", "config_path", confPath)
+
+	return updater.NewUpdater(&updater.Config{
+		Client:          config.Filtering.HTTPClient,
+		Version:         version.Version(),
+		Channel:         version.Channel(),
+		GOARCH:          runtime.GOARCH,
+		GOOS:            runtime.GOOS,
+		GOARM:           version.GOARM(),
+		GOMIPS:          version.GOMIPS(),
+		WorkDir:         workDir,
+		ConfName:        confPath,
+		ExecPath:        execPath,
+		VersionCheckURL: versionURL,
+	}), customURL
+}
+
 // checkPermissions checks and migrates permissions of the files and directories
 // used by AdGuard Home, if needed.
 func checkPermissions(workDir, confPath, dataDir, statsDir, querylogDir string) {
@@ -1010,7 +1063,7 @@ type jsonError struct {
 }
 
 // cmdlineUpdate updates current application and exits.  l must not be nil.
-func cmdlineUpdate(opts options, upd *updater.Updater, l *slog.Logger) {
+func cmdlineUpdate(ctx context.Context, l *slog.Logger, opts options, upd *updater.Updater) {
 	if !opts.performUpdate {
 		return
 	}
@@ -1023,20 +1076,19 @@ func cmdlineUpdate(opts options, upd *updater.Updater, l *slog.Logger) {
 	err := initDNSServer(nil, nil, nil, nil, nil, nil, &tlsConfigSettings{}, l)
 	fatalOnError(err)
 
-	log.Info("cmdline update: performing update")
+	l.InfoContext(ctx, "performing update via cli")
 
 	info, err := upd.VersionInfo(true)
 	if err != nil {
-		vcu := upd.VersionCheckURL()
-		log.Error("getting version info from %s: %s", vcu, err)
+		l.ErrorContext(ctx, "getting version info", slogutil.KeyError, err)
 
-		os.Exit(1)
+		os.Exit(osutil.ExitCodeFailure)
 	}
 
 	if info.NewVersion == version.Version() {
-		log.Info("no updates available")
+		l.InfoContext(ctx, "no updates available")
 
-		os.Exit(0)
+		os.Exit(osutil.ExitCodeSuccess)
 	}
 
 	err = upd.Update(Context.firstRun)
@@ -1044,10 +1096,10 @@ func cmdlineUpdate(opts options, upd *updater.Updater, l *slog.Logger) {
 
 	err = restartService()
 	if err != nil {
-		log.Debug("restarting service: %s", err)
-		log.Info("AdGuard Home was not installed as a service. " +
+		l.DebugContext(ctx, "restarting service", slogutil.KeyError, err)
+		l.InfoContext(ctx, "AdGuard Home was not installed as a service. "+
 			"Please restart running instances of AdGuardHome manually.")
 	}
 
-	os.Exit(0)
+	os.Exit(osutil.ExitCodeSuccess)
 }
diff --git a/internal/home/web.go b/internal/home/web.go
index 099f9aeb..1909720b 100644
--- a/internal/home/web.go
+++ b/internal/home/web.go
@@ -101,6 +101,8 @@ type webAPI struct {
 
 // newWebAPI creates a new instance of the web UI and API server.  l must not be
 // nil.
+//
+// TODO(a.garipov):  Return a proper error.
 func newWebAPI(conf *webConfig, l *slog.Logger) (w *webAPI) {
 	log.Info("web: initializing")
 
diff --git a/internal/updater/check.go b/internal/updater/check.go
index 84da6281..2a3e2cfe 100644
--- a/internal/updater/check.go
+++ b/internal/updater/check.go
@@ -13,6 +13,7 @@ import (
 	"github.com/AdguardTeam/golibs/errors"
 	"github.com/AdguardTeam/golibs/ioutil"
 	"github.com/AdguardTeam/golibs/log"
+	"github.com/c2h5oh/datasize"
 )
 
 // TODO(a.garipov): Make configurable.
@@ -28,8 +29,9 @@ type VersionInfo struct {
 	CanAutoUpdate aghalg.NullBool `json:"can_autoupdate,omitempty"`
 }
 
-// MaxResponseSize is responses on server's requests maximum length in bytes.
-const MaxResponseSize = 64 * 1024
+// maxVersionRespSize is the maximum length in bytes for version information
+// response.
+const maxVersionRespSize datasize.ByteSize = 64 * datasize.KB
 
 // VersionInfo downloads the latest version information.  If forceRecheck is
 // false and there are cached results, those results are returned.
@@ -51,7 +53,7 @@ func (u *Updater) VersionInfo(forceRecheck bool) (vi VersionInfo, err error) {
 	}
 	defer func() { err = errors.WithDeferred(err, resp.Body.Close()) }()
 
-	r := ioutil.LimitReader(resp.Body, MaxResponseSize)
+	r := ioutil.LimitReader(resp.Body, maxVersionRespSize.Bytes())
 
 	// This use of ReadAll is safe, because we just limited the appropriate
 	// ReadCloser.
diff --git a/internal/updater/check_test.go b/internal/updater/check_test.go
index e61ba443..5a7c0f5d 100644
--- a/internal/updater/check_test.go
+++ b/internal/updater/check_test.go
@@ -51,9 +51,11 @@ func TestUpdater_VersionInfo(t *testing.T) {
 	}))
 	t.Cleanup(srv.Close)
 
-	fakeURL, err := url.JoinPath(srv.URL, "adguardhome", version.ChannelBeta, "version.json")
+	srvURL, err := url.Parse(srv.URL)
 	require.NoError(t, err)
 
+	fakeURL := srvURL.JoinPath("adguardhome", version.ChannelBeta, "version.json")
+
 	u := updater.NewUpdater(&updater.Config{
 		Client:          srv.Client(),
 		Version:         "v0.103.0-beta.1",
@@ -134,7 +136,7 @@ func TestUpdater_VersionInfo_others(t *testing.T) {
 			GOARCH:          tc.arch,
 			GOARM:           tc.arm,
 			GOMIPS:          tc.mips,
-			VersionCheckURL: fakeURL.String(),
+			VersionCheckURL: fakeURL,
 		})
 
 		info, err := u.VersionInfo(false)
diff --git a/internal/updater/updater.go b/internal/updater/updater.go
index 85cb69b0..fe1e7727 100644
--- a/internal/updater/updater.go
+++ b/internal/updater/updater.go
@@ -9,8 +9,10 @@ import (
 	"io"
 	"io/fs"
 	"net/http"
+	"net/url"
 	"os"
 	"os/exec"
+	"path"
 	"path/filepath"
 	"strings"
 	"sync"
@@ -21,6 +23,7 @@ import (
 	"github.com/AdguardTeam/golibs/errors"
 	"github.com/AdguardTeam/golibs/ioutil"
 	"github.com/AdguardTeam/golibs/log"
+	"github.com/AdguardTeam/golibs/netutil/urlutil"
 )
 
 // Updater is the AdGuard Home updater.
@@ -61,10 +64,23 @@ type Updater struct {
 	prevCheckResult VersionInfo
 }
 
+// DefaultVersionURL returns the default URL for the version announcement.
+func DefaultVersionURL() *url.URL {
+	return &url.URL{
+		Scheme: urlutil.SchemeHTTPS,
+		Host:   "static.adtidy.org",
+		Path:   path.Join("adguardhome", version.Channel(), "version.json"),
+	}
+}
+
 // Config is the AdGuard Home updater configuration.
 type Config struct {
 	Client *http.Client
 
+	// VersionCheckURL is URL to the latest version announcement.  It must not
+	// be nil, see [DefaultVersionURL].
+	VersionCheckURL *url.URL
+
 	Version string
 	Channel string
 	GOARCH  string
@@ -81,12 +97,9 @@ type Config struct {
 
 	// ExecPath is path to the executable file.
 	ExecPath string
-
-	// VersionCheckURL is url to the latest version announcement.
-	VersionCheckURL string
 }
 
-// NewUpdater creates a new Updater.
+// NewUpdater creates a new Updater.  conf must not be nil.
 func NewUpdater(conf *Config) *Updater {
 	return &Updater{
 		client: conf.Client,
@@ -101,7 +114,7 @@ func NewUpdater(conf *Config) *Updater {
 		confName:        conf.ConfName,
 		workDir:         conf.WorkDir,
 		execPath:        conf.ExecPath,
-		versionCheckURL: conf.VersionCheckURL,
+		versionCheckURL: conf.VersionCheckURL.String(),
 
 		mu: &sync.RWMutex{},
 	}
@@ -167,14 +180,6 @@ func (u *Updater) NewVersion() (nv string) {
 	return u.newVersion
 }
 
-// VersionCheckURL returns the version check URL.
-func (u *Updater) VersionCheckURL() (vcu string) {
-	u.mu.RLock()
-	defer u.mu.RUnlock()
-
-	return u.versionCheckURL
-}
-
 // prepare fills all necessary fields in Updater object.
 func (u *Updater) prepare() (err error) {
 	u.updateDir = filepath.Join(u.workDir, fmt.Sprintf("agh-update-%s", u.newVersion))
diff --git a/internal/updater/updater_internal_test.go b/internal/updater/updater_internal_test.go
index e233db3d..67c16dc1 100644
--- a/internal/updater/updater_internal_test.go
+++ b/internal/updater/updater_internal_test.go
@@ -1,6 +1,7 @@
 package updater
 
 import (
+	"net/url"
 	"os"
 	"path/filepath"
 	"testing"
@@ -45,7 +46,7 @@ func TestUpdater_internal(t *testing.T) {
 	for _, tc := range testCases {
 		exePath := filepath.Join(wd, tc.exeName)
 
-		// start server for returning package file
+		// Start server for returning package file.
 		pkgData, err := os.ReadFile(filepath.Join("testdata", tc.archiveName))
 		require.NoError(t, err)
 
@@ -59,6 +60,9 @@ func TestUpdater_internal(t *testing.T) {
 			ExecPath: exePath,
 			WorkDir:  wd,
 			ConfName: yamlPath,
+			// TODO(e.burkov):  Rewrite the test to use a fake version check
+			// URL with a fake URLs for the package files.
+			VersionCheckURL: &url.URL{},
 		})
 
 		u.newVersion = "v0.103.1"
@@ -72,36 +76,40 @@ func TestUpdater_internal(t *testing.T) {
 
 		u.clean()
 
-		// check backup files
-		d, err := os.ReadFile(filepath.Join(wd, "agh-backup", "AdGuardHome.yaml"))
-		require.NoError(t, err)
+		require.True(t, t.Run("backup", func(t *testing.T) {
+			var d []byte
+			d, err = os.ReadFile(filepath.Join(wd, "agh-backup", "AdGuardHome.yaml"))
+			require.NoError(t, err)
 
-		assert.Equal(t, "AdGuardHome.yaml", string(d))
+			assert.Equal(t, "AdGuardHome.yaml", string(d))
 
-		d, err = os.ReadFile(filepath.Join(wd, "agh-backup", tc.exeName))
-		require.NoError(t, err)
+			d, err = os.ReadFile(filepath.Join(wd, "agh-backup", tc.exeName))
+			require.NoError(t, err)
 
-		assert.Equal(t, tc.exeName, string(d))
+			assert.Equal(t, tc.exeName, string(d))
+		}))
 
-		// check updated files
-		d, err = os.ReadFile(exePath)
-		require.NoError(t, err)
+		require.True(t, t.Run("updated", func(t *testing.T) {
+			var d []byte
+			d, err = os.ReadFile(exePath)
+			require.NoError(t, err)
 
-		assert.Equal(t, "1", string(d))
+			assert.Equal(t, "1", string(d))
 
-		d, err = os.ReadFile(readmePath)
-		require.NoError(t, err)
+			d, err = os.ReadFile(readmePath)
+			require.NoError(t, err)
 
-		assert.Equal(t, "2", string(d))
+			assert.Equal(t, "2", string(d))
 
-		d, err = os.ReadFile(licensePath)
-		require.NoError(t, err)
+			d, err = os.ReadFile(licensePath)
+			require.NoError(t, err)
 
-		assert.Equal(t, "3", string(d))
+			assert.Equal(t, "3", string(d))
 
-		d, err = os.ReadFile(yamlPath)
-		require.NoError(t, err)
+			d, err = os.ReadFile(yamlPath)
+			require.NoError(t, err)
 
-		assert.Equal(t, "AdGuardHome.yaml", string(d))
+			assert.Equal(t, "AdGuardHome.yaml", string(d))
+		}))
 	}
 }
diff --git a/internal/updater/updater_test.go b/internal/updater/updater_test.go
index 4af567c0..735d9c99 100644
--- a/internal/updater/updater_test.go
+++ b/internal/updater/updater_test.go
@@ -65,7 +65,10 @@ func TestUpdater_Update(t *testing.T) {
 	srv := httptest.NewServer(mux)
 	t.Cleanup(srv.Close)
 
-	versionCheckURL, err := url.JoinPath(srv.URL, versionPath)
+	srvURL, err := url.Parse(srv.URL)
+	require.NoError(t, err)
+
+	versionCheckURL := srvURL.JoinPath(versionPath)
 	require.NoError(t, err)
 
 	u := updater.NewUpdater(&updater.Config{