[management] Auto update geolite (#2297)

introduces helper functions to fetch and verify database versions, downloads new files if outdated, and deletes old ones. It also refactors filename handling to improve clarity and consistency, adding options to disable auto-updating via a flag. The changes aim to simplify GeoLite database management for admins.

.github/workflows/test-infrastructure-files.yml

@@ -209,8 +209,8 @@ jobs:
         working-directory: infrastructure_files/artifacts
         run: |
           sleep 30
-          docker compose exec management ls -l /var/lib/netbird/ | grep -i GeoLite2-City.mmdb
-          docker compose exec management ls -l /var/lib/netbird/ | grep -i geonames.db
+          docker compose exec management ls -l /var/lib/netbird/ | grep -i GeoLite2-City_[0-9]*.mmdb
+          docker compose exec management ls -l /var/lib/netbird/ | grep -i geonames_[0-9]*.db
 
   test-getting-started-script:
     runs-on: ubuntu-latest
@@ -237,7 +237,7 @@ jobs:
         run: test -f management.json
 
       - name: test turnserver.conf file gen postgres
-        run: | 
+        run: |
           set -x
           test -f turnserver.conf
           grep external-ip turnserver.conf
@@ -278,7 +278,7 @@ jobs:
         run: test -f management.json
 
       - name: test turnserver.conf file gen CockroachDB
-        run: | 
+        run: |
           set -x
           test -f turnserver.conf
           grep external-ip turnserver.conf
@@ -291,21 +291,3 @@ jobs:
 
       - name: test relay.env file gen CockroachDB
         run: test -f relay.env
-
-  test-download-geolite2-script:
-    runs-on: ubuntu-latest
-    steps:
-      - name: Install jq
-        run: sudo apt-get update && sudo apt-get install -y unzip sqlite3
-
-      - name: Checkout code
-        uses: actions/checkout@v3
-
-      - name: test script
-        run: bash -x infrastructure_files/download-geolite2.sh
-
-      - name: test mmdb file exists
-        run: test -f GeoLite2-City.mmdb
-
-      - name: test geonames file exists
-        run: test -f geonames.db

.gitignore

@@ -29,4 +29,3 @@ infrastructure_files/setup.env
 infrastructure_files/setup-*.env
 .vscode
 .DS_Store
-GeoLite2-City*

management/cmd/management.go

@@ -123,6 +123,8 @@ var (
 			return nil
 		},
 		RunE: func(cmd *cobra.Command, args []string) error {
+			flag.Parse()
+
 			ctx, cancel := context.WithCancel(cmd.Context())
 			defer cancel()
 			//nolint
@@ -178,11 +180,11 @@ var (
 				}
 			}
 
-			geo, err := geolocation.NewGeolocation(ctx, config.Datadir)
+			geo, err := geolocation.NewGeolocation(ctx, config.Datadir, !disableGeoliteUpdate)
 			if err != nil {
-				log.WithContext(ctx).Warnf("could not initialize geo location service: %v, we proceed without geo support", err)
+				log.WithContext(ctx).Warnf("could not initialize geolocation service. proceeding without geolocation support: %v", err)
 			} else {
-				log.WithContext(ctx).Infof("geo location service has been initialized from %s", config.Datadir)
+				log.WithContext(ctx).Infof("geolocation service has been initialized from %s", config.Datadir)
 			}
 
 			integratedPeerValidator, err := integrations.NewIntegratedValidator(ctx, eventStore)

management/cmd/root.go

@@ -24,6 +24,7 @@ var (
 	logFile                  string
 	disableMetrics           bool
 	disableSingleAccMode     bool
+	disableGeoliteUpdate     bool
 	idpSignKeyRefreshEnabled bool
 	userDeleteFromIDPEnabled bool
 
@@ -65,6 +66,7 @@ func init() {
 	mgmtCmd.Flags().StringVar(&dnsDomain, "dns-domain", defaultSingleAccModeDomain, fmt.Sprintf("Domain used for peer resolution. This is appended to the peer's name, e.g. pi-server. %s. Max length is 192 characters to allow appending to a peer name with up to 63 characters.", defaultSingleAccModeDomain))
 	mgmtCmd.Flags().BoolVar(&idpSignKeyRefreshEnabled, idpSignKeyRefreshEnabledFlagName, false, "Enable cache headers evaluation to determine signing key rotation period. This will refresh the signing key upon expiry.")
 	mgmtCmd.Flags().BoolVar(&userDeleteFromIDPEnabled, "user-delete-from-idp", false, "Allows to delete user from IDP when user is deleted from account")
+	mgmtCmd.Flags().BoolVar(&disableGeoliteUpdate, "disable-geolite-update", true, "disables automatic updates to the Geolite2 geolocation databases")
 	rootCmd.MarkFlagRequired("config") //nolint
 
 	rootCmd.PersistentFlags().StringVar(&logLevel, "log-level", "info", "")

management/server/geolocation/database.go

@@ -1,10 +1,9 @@
 package geolocation
 
 import (
+	"context"
 	"encoding/csv"
-	"fmt"
 	"io"
-	"net/url"
 	"os"
 	"path"
 	"strconv"
@@ -20,26 +19,27 @@ const (
 	geoLiteCityZipURL       = "https://pkgs.netbird.io/geolocation-dbs/GeoLite2-City-CSV/download?suffix=zip"
 	geoLiteCitySha256TarURL = "https://pkgs.netbird.io/geolocation-dbs/GeoLite2-City/download?suffix=tar.gz.sha256"
 	geoLiteCitySha256ZipURL = "https://pkgs.netbird.io/geolocation-dbs/GeoLite2-City-CSV/download?suffix=zip.sha256"
+	geoLiteCityMMDB         = "GeoLite2-City.mmdb"
+	geoLiteCityCSV          = "GeoLite2-City-Locations-en.csv"
 )
 
 // loadGeolocationDatabases loads the MaxMind databases.
-func loadGeolocationDatabases(dataDir string) error {
-	files := []string{MMDBFileName, GeoSqliteDBFile}
-	for _, file := range files {
+func loadGeolocationDatabases(ctx context.Context, dataDir string, mmdbFile string, geonamesdbFile string) error {
+	for _, file := range []string{mmdbFile, geonamesdbFile} {
 		exists, _ := fileExists(path.Join(dataDir, file))
 		if exists {
 			continue
 		}
 
-		log.Infof("geo location file %s not found , file will be downloaded", file)
+		log.WithContext(ctx).Infof("Geolocation database file %s not found, file will be downloaded", file)
 
 		switch file {
-		case MMDBFileName:
+		case mmdbFile:
 			extractFunc := func(src string, dst string) error {
 				if err := decompressTarGzFile(src, dst); err != nil {
 					return err
 				}
-				return copyFile(path.Join(dst, MMDBFileName), path.Join(dataDir, MMDBFileName))
+				return copyFile(path.Join(dst, geoLiteCityMMDB), path.Join(dataDir, mmdbFile))
 			}
 			if err := loadDatabase(
 				geoLiteCitySha256TarURL,
@@ -49,13 +49,13 @@ func loadGeolocationDatabases(dataDir string) error {
 				return err
 			}
 
-		case GeoSqliteDBFile:
+		case geonamesdbFile:
 			extractFunc := func(src string, dst string) error {
 				if err := decompressZipFile(src, dst); err != nil {
 					return err
 				}
-				extractedCsvFile := path.Join(dst, "GeoLite2-City-Locations-en.csv")
-				return importCsvToSqlite(dataDir, extractedCsvFile)
+				extractedCsvFile := path.Join(dst, geoLiteCityCSV)
+				return importCsvToSqlite(dataDir, extractedCsvFile, geonamesdbFile)
 			}
 
 			if err := loadDatabase(
@@ -79,7 +79,12 @@ func loadDatabase(checksumURL string, fileURL string, extractFunc func(src strin
 	}
 	defer os.RemoveAll(temp)
 
-	checksumFile := path.Join(temp, getDatabaseFileName(checksumURL))
+	checksumFilename, err := getFilenameFromURL(checksumURL)
+	if err != nil {
+		return err
+	}
+	checksumFile := path.Join(temp, checksumFilename)
+
 	err = downloadFile(checksumURL, checksumFile)
 	if err != nil {
 		return err
@@ -90,7 +95,12 @@ func loadDatabase(checksumURL string, fileURL string, extractFunc func(src strin
 		return err
 	}
 
-	dbFile := path.Join(temp, getDatabaseFileName(fileURL))
+	dbFilename, err := getFilenameFromURL(fileURL)
+	if err != nil {
+		return err
+	}
+	dbFile := path.Join(temp, dbFilename)
+
 	err = downloadFile(fileURL, dbFile)
 	if err != nil {
 		return err
@@ -104,13 +114,13 @@ func loadDatabase(checksumURL string, fileURL string, extractFunc func(src strin
 }
 
 // importCsvToSqlite imports a CSV file into a SQLite database.
-func importCsvToSqlite(dataDir string, csvFile string) error {
+func importCsvToSqlite(dataDir string, csvFile string, geonamesdbFile string) error {
 	geonames, err := loadGeonamesCsv(csvFile)
 	if err != nil {
 		return err
 	}
 
-	db, err := gorm.Open(sqlite.Open(path.Join(dataDir, GeoSqliteDBFile)), &gorm.Config{
+	db, err := gorm.Open(sqlite.Open(path.Join(dataDir, geonamesdbFile)), &gorm.Config{
 		Logger:          logger.Default.LogMode(logger.Silent),
 		CreateBatchSize: 1000,
 		PrepareStmt:     true,
@@ -178,18 +188,6 @@ func loadGeonamesCsv(filepath string) ([]GeoNames, error) {
 	return geoNames, nil
 }
 
-// getDatabaseFileName extracts the file name from a given URL string.
-func getDatabaseFileName(urlStr string) string {
-	u, err := url.Parse(urlStr)
-	if err != nil {
-		panic(err)
-	}
-
-	ext := u.Query().Get("suffix")
-	fileName := fmt.Sprintf("%s.%s", path.Base(u.Path), ext)
-	return fileName
-}
-
 // copyFile performs a file copy operation from the source file to the destination.
 func copyFile(src string, dst string) error {
 	srcFile, err := os.Open(src)