From f4dbb850cf58edbd51263b7848e6a07056833ece Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Tue, 30 Jan 2024 20:36:39 +0300
Subject: [PATCH 01/19] add locations endpoints

---
 management/server/http/handler.go           |  7 +++++
 management/server/http/locations_handler.go | 31 +++++++++++++++++++++
 2 files changed, 38 insertions(+)
 create mode 100644 management/server/http/locations_handler.go

diff --git a/management/server/http/handler.go b/management/server/http/handler.go
index 305af496c7f..3f64867cf45 100644
--- a/management/server/http/handler.go
+++ b/management/server/http/handler.go
@@ -82,6 +82,7 @@ func APIHandler(accountManager s.AccountManager, jwtValidator jwtclaims.JWTValid
 	api.addDNSSettingEndpoint()
 	api.addEventsEndpoint()
 	api.addPostureCheckEndpoint()
+	api.addLocationsEndpoint()
 
 	err := api.Router.Walk(func(route *mux.Route, _ *mux.Router, _ []*mux.Route) error {
 		methods, err := route.GetMethods()
@@ -210,3 +211,9 @@ func (apiHandler *apiHandler) addPostureCheckEndpoint() {
 	apiHandler.Router.HandleFunc("/posture-checks/{postureCheckId}", postureCheckHandler.GetPostureCheck).Methods("GET", "OPTIONS")
 	apiHandler.Router.HandleFunc("/posture-checks/{postureCheckId}", postureCheckHandler.DeletePostureCheck).Methods("DELETE", "OPTIONS")
 }
+
+func (apiHandler *apiHandler) addLocationsEndpoint() {
+	locationHandler := NewLocationsHandlerHandler(apiHandler.AccountManager, apiHandler.AuthCfg)
+	apiHandler.Router.HandleFunc("/locations/countries", locationHandler.GetAllCountries).Methods("GET", "OPTIONS")
+	apiHandler.Router.HandleFunc("/locations/{country}/cities", locationHandler.GetCitiesByCountry).Methods("GET", "OPTIONS")
+}
diff --git a/management/server/http/locations_handler.go b/management/server/http/locations_handler.go
new file mode 100644
index 00000000000..dd2c1ade3b9
--- /dev/null
+++ b/management/server/http/locations_handler.go
@@ -0,0 +1,31 @@
+package http
+
+import (
+	"net/http"
+
+	"github.com/netbirdio/netbird/management/server"
+	"github.com/netbirdio/netbird/management/server/jwtclaims"
+)
+
+// LocationsHandler is a handler that returns locations.
+type LocationsHandler struct {
+	accountManager  server.AccountManager
+	claimsExtractor *jwtclaims.ClaimsExtractor
+}
+
+// NewLocationsHandlerHandler creates a new Location handler
+func NewLocationsHandlerHandler(accountManager server.AccountManager, authCfg AuthCfg) *LocationsHandler {
+	return &LocationsHandler{
+		accountManager: accountManager,
+		claimsExtractor: jwtclaims.NewClaimsExtractor(
+			jwtclaims.WithAudience(authCfg.Audience),
+			jwtclaims.WithUserIDClaim(authCfg.UserIDClaim),
+		),
+	}
+}
+
+func (l *LocationsHandler) GetAllCountries(w http.ResponseWriter, r *http.Request) {
+}
+
+func (l *LocationsHandler) GetCitiesByCountry(w http.ResponseWriter, r *http.Request) {
+}

From ec1fc8142df01b9d4a84531a304b0dee236b67ef Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Thu, 1 Feb 2024 15:07:35 +0300
Subject: [PATCH 02/19] Add sqlite3 check and database generation in geolite
 script

---
 infrastructure_files/download-geolite2.sh | 110 +++++++++++++++-------
 1 file changed, 75 insertions(+), 35 deletions(-)

diff --git a/infrastructure_files/download-geolite2.sh b/infrastructure_files/download-geolite2.sh
index 2bd1c09045e..e09ae0084f8 100755
--- a/infrastructure_files/download-geolite2.sh
+++ b/infrastructure_files/download-geolite2.sh
@@ -22,41 +22,81 @@ then
     exit 1
 fi
 
-DATABASE_URL="https://download.maxmind.com/geoip/databases/GeoLite2-City/download?suffix=tar.gz"
-SIGNATURE_URL="https://download.maxmind.com/geoip/databases/GeoLite2-City/download?suffix=tar.gz.sha256"
-
-# Download the database and signature files
-echo "Downloading database file..."
-DATABASE_FILE=$(curl -s -u "$MM_ACCOUNT_ID":"$MM_LICENSE_KEY" -L -O -J "$DATABASE_URL" -w "%{filename_effective}")
-echo "Downloading signature file..."
-SIGNATURE_FILE=$(curl -s -u "$MM_ACCOUNT_ID":"$MM_LICENSE_KEY" -L -O -J "$SIGNATURE_URL" -w "%{filename_effective}")
-
-# Verify the signature
-echo "Verifying signature..."
-if sha256sum -c --status "$SIGNATURE_FILE"; then
-    echo "Signature is valid."
-else
-    echo "Signature is invalid. Aborting."
+if ! command -v sqlite3 &> /dev/null
+then
+    echo "sqlite3 is not installed or not in PATH, please install with your package manager. e.g. sudo apt install sqlite3" > /dev/stderr
     exit 1
 fi
 
-# Unpack the database file
-EXTRACTION_DIR=$(basename "$DATABASE_FILE" .tar.gz)
-echo "Unpacking $DATABASE_FILE..."
-mkdir -p "$EXTRACTION_DIR"
-tar -xzvf "$DATABASE_FILE"
-
-# Create a SHA256 signature file
-MMDB_FILE="GeoLite2-City.mmdb"
-cd "$EXTRACTION_DIR"
-sha256sum "$MMDB_FILE" > "$MMDB_FILE.sha256"
-echo "SHA256 signature created for $MMDB_FILE."
-cd -
-
-# Remove downloaded files
-rm "$DATABASE_FILE" "$SIGNATURE_FILE"
-
-# Done. Print next steps
-echo "Process completed successfully."
-echo "Now you can place $EXTRACTION_DIR/$MMDB_FILE to 'datadir' of management service."
-echo -e "Example:\n\tdocker compose cp $EXTRACTION_DIR/$MMDB_FILE management:/var/lib/netbird/"
+download_geolite_mmdb() {
+  DATABASE_URL="https://download.maxmind.com/geoip/databases/GeoLite2-City/download?suffix=tar.gz"
+  SIGNATURE_URL="https://download.maxmind.com/geoip/databases/GeoLite2-City/download?suffix=tar.gz.sha256"
+
+  # Download the database and signature files
+  echo "Downloading database file..."
+  DATABASE_FILE=$(curl -s -u "$MM_ACCOUNT_ID":"$MM_LICENSE_KEY" -L -O -J "$DATABASE_URL" -w "%{filename_effective}")
+  echo "Downloading signature file..."
+  SIGNATURE_FILE=$(curl -s -u "$MM_ACCOUNT_ID":"$MM_LICENSE_KEY" -L -O -J "$SIGNATURE_URL" -w "%{filename_effective}")
+
+  # Verify the signature
+  echo "Verifying signature..."
+  if sha256sum -c --status "$SIGNATURE_FILE"; then
+      echo "Signature is valid."
+  else
+      echo "Signature is invalid. Aborting."
+      exit 1
+  fi
+
+  # Unpack the database file
+  EXTRACTION_DIR=$(basename "$DATABASE_FILE" .tar.gz)
+  echo "Unpacking $DATABASE_FILE..."
+  mkdir -p "$EXTRACTION_DIR"
+  tar -xzvf "$DATABASE_FILE"
+
+  # Create a SHA256 signature file
+  MMDB_FILE="GeoLite2-City.mmdb"
+  cd "$EXTRACTION_DIR"
+  sha256sum "$MMDB_FILE" > "$MMDB_FILE.sha256"
+  echo "SHA256 signature created for $MMDB_FILE."
+  cd -
+
+  # Remove downloaded files
+  rm "$DATABASE_FILE" "$SIGNATURE_FILE"
+
+  # Done. Print next steps
+  echo "Process completed successfully."
+  echo "Now you can place $EXTRACTION_DIR/$MMDB_FILE to 'datadir' of management service."
+  echo -e "Example:\n\tdocker compose cp $EXTRACTION_DIR/$MMDB_FILE management:/var/lib/netbird/"
+}
+
+
+download_geolite_csv_and_create_sqlite_db() {
+  DATABASE_URL="https://download.maxmind.com/geoip/databases/GeoLite2-City-CSV/download?suffix=zip"
+
+  # Download the database file
+  echo "Downloading database file..."
+  DATABASE_FILE=$(curl -s -u "$MM_ACCOUNT_ID":"$MM_LICENSE_KEY" -L -O -J "$DATABASE_URL" -w "%{filename_effective}")
+
+  unzip "$DATABASE_FILE"
+
+  EXTRACTION_DIR=$(basename "$DATABASE_FILE" .zip)
+  DB_NAME="geonames.db"
+
+# Create SQLite database and import data from CSV
+sqlite3 "$DB_NAME" <<EOF
+.mode csv
+.import "$EXTRACTION_DIR/GeoLite2-City-Locations-en.csv" geonames
+EOF
+
+
+  # Remove downloaded and extracted files
+  rm -r -r "$EXTRACTION_DIR"
+  rm  "$DATABASE_FILE"
+
+  echo "SQLite database '$DB_NAME' created successfully."
+  echo "Now you can place $DB_NAME to 'datadir' of management service."
+  echo -e "Example:\n\tdocker compose cp $DB_NAME management:/var/lib/netbird/"
+}
+
+download_geolite_mmdb
+download_geolite_csv_and_create_sqlite_db
\ No newline at end of file

From 52f02b3fe46deeec5a9c040c12c61755a419f01e Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Thu, 1 Feb 2024 16:53:56 +0300
Subject: [PATCH 03/19] Add SQLite storage for geolocation data

---
 management/server/geolocation/manager.go | 40 +++++++++++++++
 management/server/geolocation/store.go   | 63 ++++++++++++++++++++++++
 2 files changed, 103 insertions(+)
 create mode 100644 management/server/geolocation/manager.go
 create mode 100644 management/server/geolocation/store.go

diff --git a/management/server/geolocation/manager.go b/management/server/geolocation/manager.go
new file mode 100644
index 00000000000..15d6c5a291b
--- /dev/null
+++ b/management/server/geolocation/manager.go
@@ -0,0 +1,40 @@
+package geolocation
+
+type Manager struct {
+	Store *SqliteStore
+}
+
+func Newmanager(store *SqliteStore) *Manager {
+	return &Manager{
+		Store: store,
+	}
+}
+
+func (m *Manager) GetAllCountries() ([]string, error) {
+	countries, err := m.Store.GetAllCountries()
+	if err != nil {
+		return nil, err
+	}
+
+	var validCountries []string
+	for _, country := range countries {
+		if country != "" {
+			validCountries = append(validCountries, country)
+		}
+	}
+	return validCountries, nil
+}
+
+func (m *Manager) GetCitiesByCountry(countryISOCode string) ([]string, error) {
+	cities, err := m.Store.GetCitiesByCountry(countryISOCode)
+	if err != nil {
+		return nil, err
+	}
+	var validCities []string
+	for _, country := range cities {
+		if country != "" {
+			validCities = append(validCities, country)
+		}
+	}
+	return validCities, nil
+}
diff --git a/management/server/geolocation/store.go b/management/server/geolocation/store.go
new file mode 100644
index 00000000000..8cc1def3171
--- /dev/null
+++ b/management/server/geolocation/store.go
@@ -0,0 +1,63 @@
+package geolocation
+
+import (
+	"path/filepath"
+	"runtime"
+
+	"gorm.io/driver/sqlite"
+	"gorm.io/gorm"
+	"gorm.io/gorm/logger"
+)
+
+// SqliteStore represents an account storage backed by a Sqlite DB persisted to disk
+type SqliteStore struct {
+	db        *gorm.DB
+	storeFile string
+}
+
+func NewSqliteStore(dataDir string) (*SqliteStore, error) {
+	storeStr := "geonames.db?cache=shared"
+	if runtime.GOOS == "windows" {
+		// Vo avoid `The process cannot access the file because it is being used by another process` on Windows
+		storeStr = "geonames.db"
+	}
+
+	file := filepath.Join(dataDir, storeStr)
+	db, err := gorm.Open(sqlite.Open(file), &gorm.Config{
+		Logger:      logger.Default.LogMode(logger.Silent),
+		PrepareStmt: true,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	sql, err := db.DB()
+	if err != nil {
+		return nil, err
+	}
+	conns := runtime.NumCPU()
+	sql.SetMaxOpenConns(conns) // TODO: make it configurable
+
+	return &SqliteStore{db: db, storeFile: file}, nil
+}
+
+func (s *SqliteStore) GetAllCountries() ([]string, error) {
+	var countries []string
+	result := s.db.Table("geonames").Distinct("country_iso_code").Pluck("country_iso_code", &countries)
+	if result.Error != nil {
+		return nil, result.Error
+	}
+	return countries, nil
+}
+
+func (s *SqliteStore) GetCitiesByCountry(countryISOCode string) ([]string, error) {
+	var cities []string
+	result := s.db.Table("geonames").
+		Where("country_iso_code = ?", countryISOCode).
+		Distinct("city_name").
+		Pluck("city_name", &cities)
+	if result.Error != nil {
+		return nil, result.Error
+	}
+	return cities, nil
+}

From dc3f2b5bd36d31bf9f05ea5ce7051b17d23dc185 Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Thu, 1 Feb 2024 21:33:15 +0300
Subject: [PATCH 04/19] Refactor file existence check into a separate function

---
 management/server/geolocation/geolocation.go | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/management/server/geolocation/geolocation.go b/management/server/geolocation/geolocation.go
index 4461bc1865f..8154e276b82 100644
--- a/management/server/geolocation/geolocation.go
+++ b/management/server/geolocation/geolocation.go
@@ -71,11 +71,8 @@ func NewGeolocation(datadir string) (*Geolocation, error) {
 }
 
 func openDB(mmdbPath string) (*maxminddb.Reader, error) {
-	_, err := os.Stat(mmdbPath)
-
-	if os.IsNotExist(err) {
-		return nil, fmt.Errorf("%v does not exist", mmdbPath)
-	} else if err != nil {
+	_, err := fileExists(mmdbPath)
+	if err != nil {
 		return nil, err
 	}
 
@@ -187,3 +184,14 @@ func (gl *Geolocation) reload(newSha256sum []byte) error {
 
 	return nil
 }
+
+func fileExists(filePath string) (bool, error) {
+	_, err := os.Stat(filePath)
+	if err == nil {
+		return true, nil
+	}
+	if os.IsNotExist(err) {
+		return false, fmt.Errorf("%v does not exist", filePath)
+	}
+	return false, err
+}

From 3d4c1fed530eac599df5da8d0e7d258ea42cf5d8 Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Thu, 1 Feb 2024 21:38:42 +0300
Subject: [PATCH 05/19] Integrate geolocation services into management
 application

---
 management/cmd/management.go                | 11 +++-
 management/server/geolocation/manager.go    | 28 +++++-----
 management/server/geolocation/store.go      |  6 ++-
 management/server/http/api/openapi.yml      | 59 +++++++++++++++++++++
 management/server/http/handler.go           | 26 +++++----
 management/server/http/locations_handler.go | 55 ++++++++++++++++++-
 6 files changed, 159 insertions(+), 26 deletions(-)

diff --git a/management/cmd/management.go b/management/cmd/management.go
index dcafe8d5ef9..4c44b930aba 100644
--- a/management/cmd/management.go
+++ b/management/cmd/management.go
@@ -31,6 +31,7 @@ import (
 
 	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/realip"
 	"github.com/netbirdio/management-integrations/integrations"
+
 	"github.com/netbirdio/netbird/encryption"
 	mgmtProto "github.com/netbirdio/netbird/management/proto"
 	"github.com/netbirdio/netbird/management/server"
@@ -162,6 +163,14 @@ var (
 				}
 			}
 
+			geolocationStore, err := geolocation.NewSqliteStore(config.Datadir)
+			if err != nil {
+				log.Warnf("could not  geo location database, we proceed without location endpoints")
+			} else {
+				log.Infof("geo location database has been initialized from %s", config.Datadir)
+			}
+			geolocationManager := geolocation.NewManager(geolocationStore)
+
 			geo, err := geolocation.NewGeolocation(config.Datadir)
 			if err != nil {
 				log.Warnf("could not initialize geo location service, we proceed without geo support")
@@ -230,7 +239,7 @@ var (
 				UserIDClaim:  config.HttpConfig.AuthUserIDClaim,
 				KeysLocation: config.HttpConfig.AuthKeysLocation,
 			}
-			httpAPIHandler, err := httpapi.APIHandler(accountManager, *jwtValidator, appMetrics, httpAPIAuthCfg)
+			httpAPIHandler, err := httpapi.APIHandler(accountManager, geolocationManager, *jwtValidator, appMetrics, httpAPIAuthCfg)
 			if err != nil {
 				return fmt.Errorf("failed creating HTTP API handler: %v", err)
 			}
diff --git a/management/server/geolocation/manager.go b/management/server/geolocation/manager.go
index 15d6c5a291b..0312ff48d16 100644
--- a/management/server/geolocation/manager.go
+++ b/management/server/geolocation/manager.go
@@ -4,37 +4,41 @@ type Manager struct {
 	Store *SqliteStore
 }
 
-func Newmanager(store *SqliteStore) *Manager {
+func NewManager(store *SqliteStore) *Manager {
+	if store == nil {
+		return nil
+	}
 	return &Manager{
 		Store: store,
 	}
 }
 
 func (m *Manager) GetAllCountries() ([]string, error) {
-	countries, err := m.Store.GetAllCountries()
+	allCountries, err := m.Store.GetAllCountries()
 	if err != nil {
 		return nil, err
 	}
 
-	var validCountries []string
-	for _, country := range countries {
+	countries := make([]string, 0)
+	for _, country := range allCountries {
 		if country != "" {
-			validCountries = append(validCountries, country)
+			countries = append(countries, country)
 		}
 	}
-	return validCountries, nil
+	return countries, nil
 }
 
 func (m *Manager) GetCitiesByCountry(countryISOCode string) ([]string, error) {
-	cities, err := m.Store.GetCitiesByCountry(countryISOCode)
+	allCities, err := m.Store.GetCitiesByCountry(countryISOCode)
 	if err != nil {
 		return nil, err
 	}
-	var validCities []string
-	for _, country := range cities {
-		if country != "" {
-			validCities = append(validCities, country)
+
+	cities := make([]string, 0)
+	for _, city := range allCities {
+		if city != "" {
+			cities = append(cities, city)
 		}
 	}
-	return validCities, nil
+	return cities, nil
 }
diff --git a/management/server/geolocation/store.go b/management/server/geolocation/store.go
index 8cc1def3171..3578590e1c1 100644
--- a/management/server/geolocation/store.go
+++ b/management/server/geolocation/store.go
@@ -18,10 +18,14 @@ type SqliteStore struct {
 func NewSqliteStore(dataDir string) (*SqliteStore, error) {
 	storeStr := "geonames.db?cache=shared"
 	if runtime.GOOS == "windows" {
-		// Vo avoid `The process cannot access the file because it is being used by another process` on Windows
 		storeStr = "geonames.db"
 	}
 
+	_, err := fileExists(filepath.Join(dataDir, "geonames.db"))
+	if err != nil {
+		return nil, err
+	}
+
 	file := filepath.Join(dataDir, storeStr)
 	db, err := gorm.Open(sqlite.Open(file), &gorm.Config{
 		Logger:      logger.Default.LogMode(logger.Silent),
diff --git a/management/server/http/api/openapi.yml b/management/server/http/api/openapi.yml
index e35d12ce106..a01e2bd9dc3 100644
--- a/management/server/http/api/openapi.yml
+++ b/management/server/http/api/openapi.yml
@@ -2733,3 +2733,62 @@ paths:
           "$ref": "#/components/responses/forbidden"
         '500':
           "$ref": "#/components/responses/internal_error"
+  /api/locations/countries:
+    get:
+      summary: List all country codes
+      description: Get list of all country in 2-letter ISO 3166-1 alpha-2 codes
+      tags: [ "Geo Locations" ]
+      security:
+        - BearerAuth: [ ]
+        - TokenAuth: [ ]
+      responses:
+        '200':
+          description: List of country codes
+          content:
+            application/json:
+              schema:
+                type: array
+                items:
+                  type: string
+                  example: "DE"
+        '400':
+          "$ref": "#/components/responses/bad_request"
+        '401':
+          "$ref": "#/components/responses/requires_authentication"
+        '403':
+          "$ref": "#/components/responses/forbidden"
+        '500':
+          "$ref": "#/components/responses/internal_error"
+  /api/locations/countries/{country}/cities:
+    get:
+      summary: List all city names by country
+      description: Get a list of all English city names for a given country code
+      tags: [ "Geo Locations" ]
+      security:
+        - BearerAuth: [ ]
+        - TokenAuth: [ ]
+      parameters:
+        - in: path
+          name: country
+          required: true
+          schema:
+            type: string
+          description: The 2-letter ISO 3166-1 alpha-2 country code
+      responses:
+        '200':
+          description: List of city names
+          content:
+            application/json:
+              schema:
+                type: array
+                items:
+                  type: string
+                  example: "Berlin"
+        '400':
+          "$ref": "#/components/responses/bad_request"
+        '401':
+          "$ref": "#/components/responses/requires_authentication"
+        '403':
+          "$ref": "#/components/responses/forbidden"
+        '500':
+          "$ref": "#/components/responses/internal_error"
\ No newline at end of file
diff --git a/management/server/http/handler.go b/management/server/http/handler.go
index 3f64867cf45..b95a9ff89cc 100644
--- a/management/server/http/handler.go
+++ b/management/server/http/handler.go
@@ -9,6 +9,7 @@ import (
 	"github.com/netbirdio/management-integrations/integrations"
 
 	s "github.com/netbirdio/netbird/management/server"
+	"github.com/netbirdio/netbird/management/server/geolocation"
 	"github.com/netbirdio/netbird/management/server/http/middleware"
 	"github.com/netbirdio/netbird/management/server/jwtclaims"
 	"github.com/netbirdio/netbird/management/server/telemetry"
@@ -23,9 +24,10 @@ type AuthCfg struct {
 }
 
 type apiHandler struct {
-	Router         *mux.Router
-	AccountManager s.AccountManager
-	AuthCfg        AuthCfg
+	Router          *mux.Router
+	AccountManager  s.AccountManager
+	LocationManager *geolocation.Manager
+	AuthCfg         AuthCfg
 }
 
 // EmptyObject is an empty struct used to return empty JSON object
@@ -33,7 +35,7 @@ type emptyObject struct {
 }
 
 // APIHandler creates the Management service HTTP API handler registering all the available endpoints.
-func APIHandler(accountManager s.AccountManager, jwtValidator jwtclaims.JWTValidator, appMetrics telemetry.AppMetrics, authCfg AuthCfg) (http.Handler, error) {
+func APIHandler(accountManager s.AccountManager, LocationManager *geolocation.Manager, jwtValidator jwtclaims.JWTValidator, appMetrics telemetry.AppMetrics, authCfg AuthCfg) (http.Handler, error) {
 	claimsExtractor := jwtclaims.NewClaimsExtractor(
 		jwtclaims.WithAudience(authCfg.Audience),
 		jwtclaims.WithUserIDClaim(authCfg.UserIDClaim),
@@ -63,9 +65,10 @@ func APIHandler(accountManager s.AccountManager, jwtValidator jwtclaims.JWTValid
 	router.Use(metricsMiddleware.Handler, corsMiddleware.Handler, authMiddleware.Handler, acMiddleware.Handler)
 
 	api := apiHandler{
-		Router:         router,
-		AccountManager: accountManager,
-		AuthCfg:        authCfg,
+		Router:          router,
+		AccountManager:  accountManager,
+		LocationManager: LocationManager,
+		AuthCfg:         authCfg,
 	}
 
 	integrations.RegisterHandlers(api.Router, accountManager, claimsExtractor)
@@ -213,7 +216,10 @@ func (apiHandler *apiHandler) addPostureCheckEndpoint() {
 }
 
 func (apiHandler *apiHandler) addLocationsEndpoint() {
-	locationHandler := NewLocationsHandlerHandler(apiHandler.AccountManager, apiHandler.AuthCfg)
-	apiHandler.Router.HandleFunc("/locations/countries", locationHandler.GetAllCountries).Methods("GET", "OPTIONS")
-	apiHandler.Router.HandleFunc("/locations/{country}/cities", locationHandler.GetCitiesByCountry).Methods("GET", "OPTIONS")
+	// enable location endpoints if location manager is enabled
+	if apiHandler.LocationManager != nil {
+		locationHandler := NewLocationsHandlerHandler(apiHandler.AccountManager, apiHandler.LocationManager, apiHandler.AuthCfg)
+		apiHandler.Router.HandleFunc("/locations/countries", locationHandler.GetAllCountries).Methods("GET", "OPTIONS")
+		apiHandler.Router.HandleFunc("/locations/countries/{country}/cities", locationHandler.GetCitiesByCountry).Methods("GET", "OPTIONS")
+	}
 }
diff --git a/management/server/http/locations_handler.go b/management/server/http/locations_handler.go
index dd2c1ade3b9..fb5367e913a 100644
--- a/management/server/http/locations_handler.go
+++ b/management/server/http/locations_handler.go
@@ -3,20 +3,27 @@ package http
 import (
 	"net/http"
 
+	"github.com/gorilla/mux"
+
 	"github.com/netbirdio/netbird/management/server"
+	"github.com/netbirdio/netbird/management/server/geolocation"
+	"github.com/netbirdio/netbird/management/server/http/util"
 	"github.com/netbirdio/netbird/management/server/jwtclaims"
+	"github.com/netbirdio/netbird/management/server/status"
 )
 
 // LocationsHandler is a handler that returns locations.
 type LocationsHandler struct {
 	accountManager  server.AccountManager
+	locationManager *geolocation.Manager
 	claimsExtractor *jwtclaims.ClaimsExtractor
 }
 
 // NewLocationsHandlerHandler creates a new Location handler
-func NewLocationsHandlerHandler(accountManager server.AccountManager, authCfg AuthCfg) *LocationsHandler {
+func NewLocationsHandlerHandler(accountManager server.AccountManager, locationManager *geolocation.Manager, authCfg AuthCfg) *LocationsHandler {
 	return &LocationsHandler{
-		accountManager: accountManager,
+		accountManager:  accountManager,
+		locationManager: locationManager,
 		claimsExtractor: jwtclaims.NewClaimsExtractor(
 			jwtclaims.WithAudience(authCfg.Audience),
 			jwtclaims.WithUserIDClaim(authCfg.UserIDClaim),
@@ -24,8 +31,52 @@ func NewLocationsHandlerHandler(accountManager server.AccountManager, authCfg Au
 	}
 }
 
+// GetAllCountries retrieves a list of all countries
 func (l *LocationsHandler) GetAllCountries(w http.ResponseWriter, r *http.Request) {
+	if err := l.authenticateUser(r); err != nil {
+		util.WriteError(err, w)
+		return
+	}
+
+	countries, err := l.locationManager.GetAllCountries()
+	if err != nil {
+		util.WriteError(err, w)
+		return
+	}
+	util.WriteJSONObject(w, countries)
 }
 
+// GetCitiesByCountry retrieves a list of cities based on the given country code
 func (l *LocationsHandler) GetCitiesByCountry(w http.ResponseWriter, r *http.Request) {
+	if err := l.authenticateUser(r); err != nil {
+		util.WriteError(err, w)
+		return
+	}
+
+	vars := mux.Vars(r)
+	countryCode := vars["country"]
+	if !countryCodeRegex.MatchString(countryCode) {
+		util.WriteError(status.Errorf(status.InvalidArgument, "invalid country code"), w)
+		return
+	}
+
+	cities, err := l.locationManager.GetCitiesByCountry(countryCode)
+	if err != nil {
+		util.WriteError(err, w)
+		return
+	}
+	util.WriteJSONObject(w, cities)
+}
+
+func (l *LocationsHandler) authenticateUser(r *http.Request) error {
+	claims := l.claimsExtractor.FromRequestContext(r)
+	_, user, err := l.accountManager.GetAccountFromToken(claims)
+	if err != nil {
+		return err
+	}
+
+	if !user.HasAdminPower() {
+		return status.Errorf(status.PermissionDenied, "user is not allowed to perform this action")
+	}
+	return nil
 }

From c9ef72650cbca0a69a61a271be853abf1731667a Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Fri, 2 Feb 2024 11:06:08 +0300
Subject: [PATCH 06/19] Refactoring

---
 management/server/geolocation/manager.go |  3 +++
 management/server/geolocation/store.go   | 11 ++++++-----
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/management/server/geolocation/manager.go b/management/server/geolocation/manager.go
index 0312ff48d16..b78aa6fbc46 100644
--- a/management/server/geolocation/manager.go
+++ b/management/server/geolocation/manager.go
@@ -4,6 +4,7 @@ type Manager struct {
 	Store *SqliteStore
 }
 
+// NewManager creates a new Manager instance with the given SqliteStore.
 func NewManager(store *SqliteStore) *Manager {
 	if store == nil {
 		return nil
@@ -13,6 +14,7 @@ func NewManager(store *SqliteStore) *Manager {
 	}
 }
 
+// GetAllCountries retrieves a list of all countries.
 func (m *Manager) GetAllCountries() ([]string, error) {
 	allCountries, err := m.Store.GetAllCountries()
 	if err != nil {
@@ -28,6 +30,7 @@ func (m *Manager) GetAllCountries() ([]string, error) {
 	return countries, nil
 }
 
+// GetCitiesByCountry retrieves a list of cities in a specific country based on the country's ISO code.
 func (m *Manager) GetCitiesByCountry(countryISOCode string) ([]string, error) {
 	allCities, err := m.Store.GetCitiesByCountry(countryISOCode)
 	if err != nil {
diff --git a/management/server/geolocation/store.go b/management/server/geolocation/store.go
index 3578590e1c1..643081dbd43 100644
--- a/management/server/geolocation/store.go
+++ b/management/server/geolocation/store.go
@@ -9,10 +9,9 @@ import (
 	"gorm.io/gorm/logger"
 )
 
-// SqliteStore represents an account storage backed by a Sqlite DB persisted to disk
+// SqliteStore represents a location storage backed by a Sqlite DB.
 type SqliteStore struct {
-	db        *gorm.DB
-	storeFile string
+	db *gorm.DB
 }
 
 func NewSqliteStore(dataDir string) (*SqliteStore, error) {
@@ -40,11 +39,12 @@ func NewSqliteStore(dataDir string) (*SqliteStore, error) {
 		return nil, err
 	}
 	conns := runtime.NumCPU()
-	sql.SetMaxOpenConns(conns) // TODO: make it configurable
+	sql.SetMaxOpenConns(conns)
 
-	return &SqliteStore{db: db, storeFile: file}, nil
+	return &SqliteStore{db: db}, nil
 }
 
+// GetAllCountries returns a list of all countries in the store.
 func (s *SqliteStore) GetAllCountries() ([]string, error) {
 	var countries []string
 	result := s.db.Table("geonames").Distinct("country_iso_code").Pluck("country_iso_code", &countries)
@@ -54,6 +54,7 @@ func (s *SqliteStore) GetAllCountries() ([]string, error) {
 	return countries, nil
 }
 
+// GetCitiesByCountry retrieves a list of cities from the store based on the given country ISO code.
 func (s *SqliteStore) GetCitiesByCountry(countryISOCode string) ([]string, error) {
 	var cities []string
 	result := s.db.Table("geonames").

From dc97dae100a4c1e0efdb7236ccb49ab938bd79eb Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Fri, 2 Feb 2024 17:56:43 +0300
Subject: [PATCH 07/19] Refactor city retrieval to include Geonames ID

---
 management/server/geolocation/manager.go    | 20 +++++++-------------
 management/server/geolocation/store.go      | 14 ++++++++------
 management/server/http/api/openapi.yml      | 20 ++++++++++++++++----
 management/server/http/api/types.gen.go     |  9 +++++++++
 management/server/http/locations_handler.go | 15 ++++++++++++++-
 5 files changed, 54 insertions(+), 24 deletions(-)

diff --git a/management/server/geolocation/manager.go b/management/server/geolocation/manager.go
index b78aa6fbc46..a7d4ee70150 100644
--- a/management/server/geolocation/manager.go
+++ b/management/server/geolocation/manager.go
@@ -1,5 +1,10 @@
 package geolocation
 
+type City struct {
+	GeoNameID int `gorm:"column:geoname_id"`
+	CityName  string
+}
+
 type Manager struct {
 	Store *SqliteStore
 }
@@ -31,17 +36,6 @@ func (m *Manager) GetAllCountries() ([]string, error) {
 }
 
 // GetCitiesByCountry retrieves a list of cities in a specific country based on the country's ISO code.
-func (m *Manager) GetCitiesByCountry(countryISOCode string) ([]string, error) {
-	allCities, err := m.Store.GetCitiesByCountry(countryISOCode)
-	if err != nil {
-		return nil, err
-	}
-
-	cities := make([]string, 0)
-	for _, city := range allCities {
-		if city != "" {
-			cities = append(cities, city)
-		}
-	}
-	return cities, nil
+func (m *Manager) GetCitiesByCountry(countryISOCode string) ([]City, error) {
+	return m.Store.GetCitiesByCountry(countryISOCode)
 }
diff --git a/management/server/geolocation/store.go b/management/server/geolocation/store.go
index 643081dbd43..af8a19fb164 100644
--- a/management/server/geolocation/store.go
+++ b/management/server/geolocation/store.go
@@ -15,9 +15,9 @@ type SqliteStore struct {
 }
 
 func NewSqliteStore(dataDir string) (*SqliteStore, error) {
-	storeStr := "geonames.db?cache=shared"
+	storeStr := "geonames.db?cache=shared&mode=ro"
 	if runtime.GOOS == "windows" {
-		storeStr = "geonames.db"
+		storeStr = "geonames.db?&mode=ro"
 	}
 
 	_, err := fileExists(filepath.Join(dataDir, "geonames.db"))
@@ -55,14 +55,16 @@ func (s *SqliteStore) GetAllCountries() ([]string, error) {
 }
 
 // GetCitiesByCountry retrieves a list of cities from the store based on the given country ISO code.
-func (s *SqliteStore) GetCitiesByCountry(countryISOCode string) ([]string, error) {
-	var cities []string
+func (s *SqliteStore) GetCitiesByCountry(countryISOCode string) ([]City, error) {
+	var cities []City
 	result := s.db.Table("geonames").
+		Select("geoname_id", "city_name").
 		Where("country_iso_code = ?", countryISOCode).
-		Distinct("city_name").
-		Pluck("city_name", &cities)
+		Group("city_name").
+		Scan(&cities)
 	if result.Error != nil {
 		return nil, result.Error
 	}
+
 	return cities, nil
 }
diff --git a/management/server/http/api/openapi.yml b/management/server/http/api/openapi.yml
index a01e2bd9dc3..ea599b6c984 100644
--- a/management/server/http/api/openapi.yml
+++ b/management/server/http/api/openapi.yml
@@ -918,6 +918,21 @@ components:
       required:
         - country_code
         - city_name
+    City:
+      description: Describe city geographical location information
+      type: object
+      properties:
+        geoname_id:
+          description: Integer ID of the record in GeoNames database
+          type: integer
+          example: 2950158
+        city_name:
+          description: Commonly used English name of the city
+          type: string
+          example: "Berlin"
+      required:
+        - geoname_id
+        - city_name
     PostureCheckUpdate:
       type: object
       properties:
@@ -2780,10 +2795,7 @@ paths:
           content:
             application/json:
               schema:
-                type: array
-                items:
-                  type: string
-                  example: "Berlin"
+                $ref: '#/components/schemas/City'
         '400':
           "$ref": "#/components/responses/bad_request"
         '401':
diff --git a/management/server/http/api/types.gen.go b/management/server/http/api/types.gen.go
index 34d4ef3ed9e..d6dc3c9698d 100644
--- a/management/server/http/api/types.gen.go
+++ b/management/server/http/api/types.gen.go
@@ -194,6 +194,15 @@ type Checks struct {
 	OsVersionCheck *OSVersionCheck `json:"os_version_check,omitempty"`
 }
 
+// City Describe city geographical location information
+type City struct {
+	// CityName Commonly used English name of the city
+	CityName string `json:"city_name"`
+
+	// GeonameId Integer ID of the record in GeoNames database
+	GeonameId int `json:"geoname_id"`
+}
+
 // DNSSettings defines model for DNSSettings.
 type DNSSettings struct {
 	// DisabledManagementGroups Groups whose DNS management is disabled
diff --git a/management/server/http/locations_handler.go b/management/server/http/locations_handler.go
index fb5367e913a..acc271f6e8d 100644
--- a/management/server/http/locations_handler.go
+++ b/management/server/http/locations_handler.go
@@ -7,6 +7,7 @@ import (
 
 	"github.com/netbirdio/netbird/management/server"
 	"github.com/netbirdio/netbird/management/server/geolocation"
+	"github.com/netbirdio/netbird/management/server/http/api"
 	"github.com/netbirdio/netbird/management/server/http/util"
 	"github.com/netbirdio/netbird/management/server/jwtclaims"
 	"github.com/netbirdio/netbird/management/server/status"
@@ -60,11 +61,16 @@ func (l *LocationsHandler) GetCitiesByCountry(w http.ResponseWriter, r *http.Req
 		return
 	}
 
-	cities, err := l.locationManager.GetCitiesByCountry(countryCode)
+	allCities, err := l.locationManager.GetCitiesByCountry(countryCode)
 	if err != nil {
 		util.WriteError(err, w)
 		return
 	}
+
+	cities := make([]api.City, 0, len(allCities))
+	for _, city := range allCities {
+		cities = append(cities, toCityResponse(city))
+	}
 	util.WriteJSONObject(w, cities)
 }
 
@@ -80,3 +86,10 @@ func (l *LocationsHandler) authenticateUser(r *http.Request) error {
 	}
 	return nil
 }
+
+func toCityResponse(city geolocation.City) api.City {
+	return api.City{
+		CityName:  city.CityName,
+		GeonameId: city.GeoNameID,
+	}
+}

From c7479f9b4ea78a2c4a7eae51bdc30a9bdccfad92 Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Mon, 5 Feb 2024 13:23:17 +0300
Subject: [PATCH 08/19] Add signature verification for GeoLite2 database
 download

---
 infrastructure_files/download-geolite2.sh | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/infrastructure_files/download-geolite2.sh b/infrastructure_files/download-geolite2.sh
index e09ae0084f8..7cc5528ae29 100755
--- a/infrastructure_files/download-geolite2.sh
+++ b/infrastructure_files/download-geolite2.sh
@@ -72,16 +72,30 @@ download_geolite_mmdb() {
 
 download_geolite_csv_and_create_sqlite_db() {
   DATABASE_URL="https://download.maxmind.com/geoip/databases/GeoLite2-City-CSV/download?suffix=zip"
+  SIGNATURE_URL="https://download.maxmind.com/geoip/databases/GeoLite2-City-CSV/download?suffix=zip.sha256"
+
 
   # Download the database file
   echo "Downloading database file..."
   DATABASE_FILE=$(curl -s -u "$MM_ACCOUNT_ID":"$MM_LICENSE_KEY" -L -O -J "$DATABASE_URL" -w "%{filename_effective}")
+  echo "Downloading signature file..."
+  SIGNATURE_FILE=$(curl -s -u "$MM_ACCOUNT_ID":"$MM_LICENSE_KEY" -L -O -J "$SIGNATURE_URL" -w "%{filename_effective}")
 
-  unzip "$DATABASE_FILE"
+  # Verify the signature
+  echo "Verifying signature..."
+  if sha256sum -c --status "$SIGNATURE_FILE"; then
+      echo "Signature is valid."
+  else
+      echo "Signature is invalid. Aborting."
+      exit 1
+  fi
 
+  # Unpack the database file
   EXTRACTION_DIR=$(basename "$DATABASE_FILE" .zip)
   DB_NAME="geonames.db"
 
+  unzip "$DATABASE_FILE"
+
 # Create SQLite database and import data from CSV
 sqlite3 "$DB_NAME" <<EOF
 .mode csv
@@ -91,7 +105,7 @@ EOF
 
   # Remove downloaded and extracted files
   rm -r -r "$EXTRACTION_DIR"
-  rm  "$DATABASE_FILE"
+  rm  "$DATABASE_FILE" "$SIGNATURE_FILE"
 
   echo "SQLite database '$DB_NAME' created successfully."
   echo "Now you can place $DB_NAME to 'datadir' of management service."

From 8eb158c599423cf9dfc5b9229a06a3f4567bdfc6 Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Mon, 5 Feb 2024 16:08:53 +0300
Subject: [PATCH 09/19] Change to in-memory database for geolocation store

---
 management/server/geolocation/store.go | 46 +++++++++++++++++++++-----
 1 file changed, 38 insertions(+), 8 deletions(-)

diff --git a/management/server/geolocation/store.go b/management/server/geolocation/store.go
index af8a19fb164..7cac3324b77 100644
--- a/management/server/geolocation/store.go
+++ b/management/server/geolocation/store.go
@@ -1,6 +1,7 @@
 package geolocation
 
 import (
+	"fmt"
 	"path/filepath"
 	"runtime"
 
@@ -15,18 +16,18 @@ type SqliteStore struct {
 }
 
 func NewSqliteStore(dataDir string) (*SqliteStore, error) {
-	storeStr := "geonames.db?cache=shared&mode=ro"
-	if runtime.GOOS == "windows" {
-		storeStr = "geonames.db?&mode=ro"
-	}
-
-	_, err := fileExists(filepath.Join(dataDir, "geonames.db"))
+	file := filepath.Join(dataDir, "geonames.db")
+	_, err := fileExists(file)
 	if err != nil {
 		return nil, err
 	}
 
-	file := filepath.Join(dataDir, storeStr)
-	db, err := gorm.Open(sqlite.Open(file), &gorm.Config{
+	storeStr := ":memory:?cache=shared&mode=ro"
+	if runtime.GOOS == "windows" {
+		storeStr = ":memory:?&mode=ro"
+	}
+
+	db, err := gorm.Open(sqlite.Open(storeStr), &gorm.Config{
 		Logger:      logger.Default.LogMode(logger.Silent),
 		PrepareStmt: true,
 	})
@@ -34,6 +35,10 @@ func NewSqliteStore(dataDir string) (*SqliteStore, error) {
 		return nil, err
 	}
 
+	if err := setupInMemoryDBFromFile(db, file); err != nil {
+		return nil, err
+	}
+
 	sql, err := db.DB()
 	if err != nil {
 		return nil, err
@@ -68,3 +73,28 @@ func (s *SqliteStore) GetCitiesByCountry(countryISOCode string) ([]City, error)
 
 	return cities, nil
 }
+
+// setupInMemoryDBFromFile prepares an in-memory DB by attaching a file database and,
+// copies the data from the attached database to the in-memory database.
+func setupInMemoryDBFromFile(db *gorm.DB, source string) error {
+	// Attach the on-disk database to the in-memory database
+	attachStmt := fmt.Sprintf("ATTACH DATABASE '%s' AS source;", source)
+	if err := db.Exec(attachStmt).Error; err != nil {
+		return err
+	}
+
+	err := db.Exec(`
+		CREATE TABLE geonames AS SELECT * FROM source.geonames;
+	`).Error
+	if err != nil {
+		return err
+	}
+
+	// Detach the on-disk database from the in-memory database
+	err = db.Exec("DETACH DATABASE source;").Error
+	if err != nil {
+		return err
+	}
+
+	return nil
+}

From 06b080fe19bc8c8f348facfe71fa278f1540428b Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Mon, 5 Feb 2024 16:47:44 +0300
Subject: [PATCH 10/19] Merge manager to geolocation

---
 management/cmd/management.go                  | 10 +-
 management/server/geolocation/geolocation.go  | 91 +++++++++++++------
 management/server/geolocation/manager.go      | 41 ---------
 ...ons_handler.go => geolocations_handler.go} | 28 +++---
 management/server/http/handler.go             | 22 ++---
 5 files changed, 88 insertions(+), 104 deletions(-)
 delete mode 100644 management/server/geolocation/manager.go
 rename management/server/http/{locations_handler.go => geolocations_handler.go} (71%)

diff --git a/management/cmd/management.go b/management/cmd/management.go
index 4c44b930aba..0ce94db7da9 100644
--- a/management/cmd/management.go
+++ b/management/cmd/management.go
@@ -163,14 +163,6 @@ var (
 				}
 			}
 
-			geolocationStore, err := geolocation.NewSqliteStore(config.Datadir)
-			if err != nil {
-				log.Warnf("could not  geo location database, we proceed without location endpoints")
-			} else {
-				log.Infof("geo location database has been initialized from %s", config.Datadir)
-			}
-			geolocationManager := geolocation.NewManager(geolocationStore)
-
 			geo, err := geolocation.NewGeolocation(config.Datadir)
 			if err != nil {
 				log.Warnf("could not initialize geo location service, we proceed without geo support")
@@ -239,7 +231,7 @@ var (
 				UserIDClaim:  config.HttpConfig.AuthUserIDClaim,
 				KeysLocation: config.HttpConfig.AuthKeysLocation,
 			}
-			httpAPIHandler, err := httpapi.APIHandler(accountManager, geolocationManager, *jwtValidator, appMetrics, httpAPIAuthCfg)
+			httpAPIHandler, err := httpapi.APIHandler(accountManager, geo, *jwtValidator, appMetrics, httpAPIAuthCfg)
 			if err != nil {
 				return fmt.Errorf("failed creating HTTP API handler: %v", err)
 			}
diff --git a/management/server/geolocation/geolocation.go b/management/server/geolocation/geolocation.go
index 8154e276b82..39b2f1e667c 100644
--- a/management/server/geolocation/geolocation.go
+++ b/management/server/geolocation/geolocation.go
@@ -22,6 +22,7 @@ type Geolocation struct {
 	mux                 *sync.RWMutex
 	sha256sum           []byte
 	db                  *maxminddb.Reader
+	locationDB          *SqliteStore
 	stopCh              chan struct{}
 	reloadCheckInterval time.Duration
 }
@@ -43,6 +44,11 @@ type Record struct {
 	} `maxminddb:"country"`
 }
 
+type City struct {
+	GeoNameID int `gorm:"column:geoname_id"`
+	CityName  string
+}
+
 func NewGeolocation(datadir string) (*Geolocation, error) {
 	mmdbPath := path.Join(datadir, mmdbFileName)
 
@@ -56,11 +62,17 @@ func NewGeolocation(datadir string) (*Geolocation, error) {
 		return nil, err
 	}
 
+	locationDB, err := NewSqliteStore(datadir)
+	if err != nil {
+		return nil, err
+	}
+
 	geo := &Geolocation{
 		mmdbPath:            mmdbPath,
 		mux:                 &sync.RWMutex{},
 		sha256sum:           sha256sum,
 		db:                  db,
+		locationDB:          locationDB,
 		reloadCheckInterval: 60 * time.Second, // TODO: make configurable
 		stopCh:              make(chan struct{}),
 	}
@@ -70,35 +82,6 @@ func NewGeolocation(datadir string) (*Geolocation, error) {
 	return geo, nil
 }
 
-func openDB(mmdbPath string) (*maxminddb.Reader, error) {
-	_, err := fileExists(mmdbPath)
-	if err != nil {
-		return nil, err
-	}
-
-	db, err := maxminddb.Open(mmdbPath)
-	if err != nil {
-		return nil, fmt.Errorf("%v could not be opened: %w", mmdbPath, err)
-	}
-
-	return db, nil
-}
-
-func getSha256sum(mmdbPath string) ([]byte, error) {
-	f, err := os.Open(mmdbPath)
-	if err != nil {
-		return nil, err
-	}
-	defer f.Close()
-
-	h := sha256.New()
-	if _, err := io.Copy(h, f); err != nil {
-		return nil, err
-	}
-
-	return h.Sum(nil), nil
-}
-
 func (gl *Geolocation) Lookup(ip string) (*Record, error) {
 	gl.mux.RLock()
 	defer gl.mux.RUnlock()
@@ -117,6 +100,27 @@ func (gl *Geolocation) Lookup(ip string) (*Record, error) {
 	return &record, nil
 }
 
+// GetAllCountries retrieves a list of all countries.
+func (gl *Geolocation) GetAllCountries() ([]string, error) {
+	allCountries, err := gl.locationDB.GetAllCountries()
+	if err != nil {
+		return nil, err
+	}
+
+	countries := make([]string, 0)
+	for _, country := range allCountries {
+		if country != "" {
+			countries = append(countries, country)
+		}
+	}
+	return countries, nil
+}
+
+// GetCitiesByCountry retrieves a list of cities in a specific country based on the country's ISO code.
+func (gl *Geolocation) GetCitiesByCountry(countryISOCode string) ([]City, error) {
+	return gl.locationDB.GetCitiesByCountry(countryISOCode)
+}
+
 func (gl *Geolocation) Stop() error {
 	close(gl.stopCh)
 	if gl.db != nil {
@@ -185,6 +189,35 @@ func (gl *Geolocation) reload(newSha256sum []byte) error {
 	return nil
 }
 
+func openDB(mmdbPath string) (*maxminddb.Reader, error) {
+	_, err := fileExists(mmdbPath)
+	if err != nil {
+		return nil, err
+	}
+
+	db, err := maxminddb.Open(mmdbPath)
+	if err != nil {
+		return nil, fmt.Errorf("%v could not be opened: %w", mmdbPath, err)
+	}
+
+	return db, nil
+}
+
+func getSha256sum(mmdbPath string) ([]byte, error) {
+	f, err := os.Open(mmdbPath)
+	if err != nil {
+		return nil, err
+	}
+	defer f.Close()
+
+	h := sha256.New()
+	if _, err := io.Copy(h, f); err != nil {
+		return nil, err
+	}
+
+	return h.Sum(nil), nil
+}
+
 func fileExists(filePath string) (bool, error) {
 	_, err := os.Stat(filePath)
 	if err == nil {
diff --git a/management/server/geolocation/manager.go b/management/server/geolocation/manager.go
deleted file mode 100644
index a7d4ee70150..00000000000
--- a/management/server/geolocation/manager.go
+++ /dev/null
@@ -1,41 +0,0 @@
-package geolocation
-
-type City struct {
-	GeoNameID int `gorm:"column:geoname_id"`
-	CityName  string
-}
-
-type Manager struct {
-	Store *SqliteStore
-}
-
-// NewManager creates a new Manager instance with the given SqliteStore.
-func NewManager(store *SqliteStore) *Manager {
-	if store == nil {
-		return nil
-	}
-	return &Manager{
-		Store: store,
-	}
-}
-
-// GetAllCountries retrieves a list of all countries.
-func (m *Manager) GetAllCountries() ([]string, error) {
-	allCountries, err := m.Store.GetAllCountries()
-	if err != nil {
-		return nil, err
-	}
-
-	countries := make([]string, 0)
-	for _, country := range allCountries {
-		if country != "" {
-			countries = append(countries, country)
-		}
-	}
-	return countries, nil
-}
-
-// GetCitiesByCountry retrieves a list of cities in a specific country based on the country's ISO code.
-func (m *Manager) GetCitiesByCountry(countryISOCode string) ([]City, error) {
-	return m.Store.GetCitiesByCountry(countryISOCode)
-}
diff --git a/management/server/http/locations_handler.go b/management/server/http/geolocations_handler.go
similarity index 71%
rename from management/server/http/locations_handler.go
rename to management/server/http/geolocations_handler.go
index acc271f6e8d..b45d9e586bf 100644
--- a/management/server/http/locations_handler.go
+++ b/management/server/http/geolocations_handler.go
@@ -13,18 +13,18 @@ import (
 	"github.com/netbirdio/netbird/management/server/status"
 )
 
-// LocationsHandler is a handler that returns locations.
-type LocationsHandler struct {
-	accountManager  server.AccountManager
-	locationManager *geolocation.Manager
-	claimsExtractor *jwtclaims.ClaimsExtractor
+// GeolocationsHandler is a handler that returns locations.
+type GeolocationsHandler struct {
+	accountManager     server.AccountManager
+	geolocationManager *geolocation.Geolocation
+	claimsExtractor    *jwtclaims.ClaimsExtractor
 }
 
 // NewLocationsHandlerHandler creates a new Location handler
-func NewLocationsHandlerHandler(accountManager server.AccountManager, locationManager *geolocation.Manager, authCfg AuthCfg) *LocationsHandler {
-	return &LocationsHandler{
-		accountManager:  accountManager,
-		locationManager: locationManager,
+func NewLocationsHandlerHandler(accountManager server.AccountManager, geolocationManager *geolocation.Geolocation, authCfg AuthCfg) *GeolocationsHandler {
+	return &GeolocationsHandler{
+		accountManager:     accountManager,
+		geolocationManager: geolocationManager,
 		claimsExtractor: jwtclaims.NewClaimsExtractor(
 			jwtclaims.WithAudience(authCfg.Audience),
 			jwtclaims.WithUserIDClaim(authCfg.UserIDClaim),
@@ -33,13 +33,13 @@ func NewLocationsHandlerHandler(accountManager server.AccountManager, locationMa
 }
 
 // GetAllCountries retrieves a list of all countries
-func (l *LocationsHandler) GetAllCountries(w http.ResponseWriter, r *http.Request) {
+func (l *GeolocationsHandler) GetAllCountries(w http.ResponseWriter, r *http.Request) {
 	if err := l.authenticateUser(r); err != nil {
 		util.WriteError(err, w)
 		return
 	}
 
-	countries, err := l.locationManager.GetAllCountries()
+	countries, err := l.geolocationManager.GetAllCountries()
 	if err != nil {
 		util.WriteError(err, w)
 		return
@@ -48,7 +48,7 @@ func (l *LocationsHandler) GetAllCountries(w http.ResponseWriter, r *http.Reques
 }
 
 // GetCitiesByCountry retrieves a list of cities based on the given country code
-func (l *LocationsHandler) GetCitiesByCountry(w http.ResponseWriter, r *http.Request) {
+func (l *GeolocationsHandler) GetCitiesByCountry(w http.ResponseWriter, r *http.Request) {
 	if err := l.authenticateUser(r); err != nil {
 		util.WriteError(err, w)
 		return
@@ -61,7 +61,7 @@ func (l *LocationsHandler) GetCitiesByCountry(w http.ResponseWriter, r *http.Req
 		return
 	}
 
-	allCities, err := l.locationManager.GetCitiesByCountry(countryCode)
+	allCities, err := l.geolocationManager.GetCitiesByCountry(countryCode)
 	if err != nil {
 		util.WriteError(err, w)
 		return
@@ -74,7 +74,7 @@ func (l *LocationsHandler) GetCitiesByCountry(w http.ResponseWriter, r *http.Req
 	util.WriteJSONObject(w, cities)
 }
 
-func (l *LocationsHandler) authenticateUser(r *http.Request) error {
+func (l *GeolocationsHandler) authenticateUser(r *http.Request) error {
 	claims := l.claimsExtractor.FromRequestContext(r)
 	_, user, err := l.accountManager.GetAccountFromToken(claims)
 	if err != nil {
diff --git a/management/server/http/handler.go b/management/server/http/handler.go
index b95a9ff89cc..fd0a4eb21c2 100644
--- a/management/server/http/handler.go
+++ b/management/server/http/handler.go
@@ -24,10 +24,10 @@ type AuthCfg struct {
 }
 
 type apiHandler struct {
-	Router          *mux.Router
-	AccountManager  s.AccountManager
-	LocationManager *geolocation.Manager
-	AuthCfg         AuthCfg
+	Router             *mux.Router
+	AccountManager     s.AccountManager
+	geolocationManager *geolocation.Geolocation
+	AuthCfg            AuthCfg
 }
 
 // EmptyObject is an empty struct used to return empty JSON object
@@ -35,7 +35,7 @@ type emptyObject struct {
 }
 
 // APIHandler creates the Management service HTTP API handler registering all the available endpoints.
-func APIHandler(accountManager s.AccountManager, LocationManager *geolocation.Manager, jwtValidator jwtclaims.JWTValidator, appMetrics telemetry.AppMetrics, authCfg AuthCfg) (http.Handler, error) {
+func APIHandler(accountManager s.AccountManager, LocationManager *geolocation.Geolocation, jwtValidator jwtclaims.JWTValidator, appMetrics telemetry.AppMetrics, authCfg AuthCfg) (http.Handler, error) {
 	claimsExtractor := jwtclaims.NewClaimsExtractor(
 		jwtclaims.WithAudience(authCfg.Audience),
 		jwtclaims.WithUserIDClaim(authCfg.UserIDClaim),
@@ -65,10 +65,10 @@ func APIHandler(accountManager s.AccountManager, LocationManager *geolocation.Ma
 	router.Use(metricsMiddleware.Handler, corsMiddleware.Handler, authMiddleware.Handler, acMiddleware.Handler)
 
 	api := apiHandler{
-		Router:          router,
-		AccountManager:  accountManager,
-		LocationManager: LocationManager,
-		AuthCfg:         authCfg,
+		Router:             router,
+		AccountManager:     accountManager,
+		geolocationManager: LocationManager,
+		AuthCfg:            authCfg,
 	}
 
 	integrations.RegisterHandlers(api.Router, accountManager, claimsExtractor)
@@ -217,8 +217,8 @@ func (apiHandler *apiHandler) addPostureCheckEndpoint() {
 
 func (apiHandler *apiHandler) addLocationsEndpoint() {
 	// enable location endpoints if location manager is enabled
-	if apiHandler.LocationManager != nil {
-		locationHandler := NewLocationsHandlerHandler(apiHandler.AccountManager, apiHandler.LocationManager, apiHandler.AuthCfg)
+	if apiHandler.geolocationManager != nil {
+		locationHandler := NewLocationsHandlerHandler(apiHandler.AccountManager, apiHandler.geolocationManager, apiHandler.AuthCfg)
 		apiHandler.Router.HandleFunc("/locations/countries", locationHandler.GetAllCountries).Methods("GET", "OPTIONS")
 		apiHandler.Router.HandleFunc("/locations/countries/{country}/cities", locationHandler.GetCitiesByCountry).Methods("GET", "OPTIONS")
 	}

From 3c4f45f75956eaf8ff1f51e32dd85a70316a4fe5 Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Mon, 5 Feb 2024 20:34:31 +0300
Subject: [PATCH 11/19] Update GetAllCountries to return Country name and iso
 code

---
 management/server/geolocation/geolocation.go  | 24 +++++++++++++++----
 management/server/geolocation/store.go        |  9 ++++---
 management/server/http/api/openapi.yml        | 18 ++++++++++++--
 management/server/http/api/types.gen.go       |  9 +++++++
 .../server/http/geolocations_handler.go       | 14 ++++++++++-
 5 files changed, 64 insertions(+), 10 deletions(-)

diff --git a/management/server/geolocation/geolocation.go b/management/server/geolocation/geolocation.go
index 39b2f1e667c..1550f058905 100644
--- a/management/server/geolocation/geolocation.go
+++ b/management/server/geolocation/geolocation.go
@@ -49,6 +49,11 @@ type City struct {
 	CityName  string
 }
 
+type Country struct {
+	CountryISOCode string `gorm:"column:country_iso_code"`
+	CountryName    string
+}
+
 func NewGeolocation(datadir string) (*Geolocation, error) {
 	mmdbPath := path.Join(datadir, mmdbFileName)
 
@@ -101,15 +106,15 @@ func (gl *Geolocation) Lookup(ip string) (*Record, error) {
 }
 
 // GetAllCountries retrieves a list of all countries.
-func (gl *Geolocation) GetAllCountries() ([]string, error) {
+func (gl *Geolocation) GetAllCountries() ([]Country, error) {
 	allCountries, err := gl.locationDB.GetAllCountries()
 	if err != nil {
 		return nil, err
 	}
 
-	countries := make([]string, 0)
+	countries := make([]Country, 0)
 	for _, country := range allCountries {
-		if country != "" {
+		if country.CountryName != "" {
 			countries = append(countries, country)
 		}
 	}
@@ -118,7 +123,18 @@ func (gl *Geolocation) GetAllCountries() ([]string, error) {
 
 // GetCitiesByCountry retrieves a list of cities in a specific country based on the country's ISO code.
 func (gl *Geolocation) GetCitiesByCountry(countryISOCode string) ([]City, error) {
-	return gl.locationDB.GetCitiesByCountry(countryISOCode)
+	allCities, err := gl.locationDB.GetCitiesByCountry(countryISOCode)
+	if err != nil {
+		return nil, err
+	}
+
+	cities := make([]City, 0)
+	for _, city := range allCities {
+		if city.CityName != "" {
+			cities = append(cities, city)
+		}
+	}
+	return cities, nil
 }
 
 func (gl *Geolocation) Stop() error {
diff --git a/management/server/geolocation/store.go b/management/server/geolocation/store.go
index 7cac3324b77..3ea934d8f45 100644
--- a/management/server/geolocation/store.go
+++ b/management/server/geolocation/store.go
@@ -50,9 +50,12 @@ func NewSqliteStore(dataDir string) (*SqliteStore, error) {
 }
 
 // GetAllCountries returns a list of all countries in the store.
-func (s *SqliteStore) GetAllCountries() ([]string, error) {
-	var countries []string
-	result := s.db.Table("geonames").Distinct("country_iso_code").Pluck("country_iso_code", &countries)
+func (s *SqliteStore) GetAllCountries() ([]Country, error) {
+	var countries []Country
+	result := s.db.Table("geonames").
+		Select("country_iso_code", "country_name").
+		Group("country_name").
+		Scan(&countries)
 	if result.Error != nil {
 		return nil, result.Error
 	}
diff --git a/management/server/http/api/openapi.yml b/management/server/http/api/openapi.yml
index ea599b6c984..7a16da59b3b 100644
--- a/management/server/http/api/openapi.yml
+++ b/management/server/http/api/openapi.yml
@@ -918,6 +918,21 @@ components:
       required:
         - country_code
         - city_name
+    Country:
+      description: Describe country geographical location information
+      type: object
+      properties:
+        country_name:
+          description: Commonly used English name of the country
+          type: string
+          example: "Germany"
+        country_code:
+          description: 2-letter ISO 3166-1 alpha-2 code that represents the country
+          type: string
+          example: "DE"
+      required:
+        - country_name
+        - country_code
     City:
       description: Describe city geographical location information
       type: object
@@ -2787,8 +2802,7 @@ paths:
           name: country
           required: true
           schema:
-            type: string
-          description: The 2-letter ISO 3166-1 alpha-2 country code
+            $ref: '#/components/schemas/Country'
       responses:
         '200':
           description: List of city names
diff --git a/management/server/http/api/types.gen.go b/management/server/http/api/types.gen.go
index d6dc3c9698d..27b91b71f69 100644
--- a/management/server/http/api/types.gen.go
+++ b/management/server/http/api/types.gen.go
@@ -203,6 +203,15 @@ type City struct {
 	GeonameId int `json:"geoname_id"`
 }
 
+// Country Describe country geographical location information
+type Country struct {
+	// CountryCode 2-letter ISO 3166-1 alpha-2 code that represents the country
+	CountryCode string `json:"country_code"`
+
+	// CountryName Commonly used English name of the country
+	CountryName string `json:"country_name"`
+}
+
 // DNSSettings defines model for DNSSettings.
 type DNSSettings struct {
 	// DisabledManagementGroups Groups whose DNS management is disabled
diff --git a/management/server/http/geolocations_handler.go b/management/server/http/geolocations_handler.go
index b45d9e586bf..45d14460ff3 100644
--- a/management/server/http/geolocations_handler.go
+++ b/management/server/http/geolocations_handler.go
@@ -39,11 +39,16 @@ func (l *GeolocationsHandler) GetAllCountries(w http.ResponseWriter, r *http.Req
 		return
 	}
 
-	countries, err := l.geolocationManager.GetAllCountries()
+	allCountries, err := l.geolocationManager.GetAllCountries()
 	if err != nil {
 		util.WriteError(err, w)
 		return
 	}
+
+	countries := make([]api.Country, 0, len(allCountries))
+	for _, country := range allCountries {
+		countries = append(countries, toCountryResponse(country))
+	}
 	util.WriteJSONObject(w, countries)
 }
 
@@ -87,6 +92,13 @@ func (l *GeolocationsHandler) authenticateUser(r *http.Request) error {
 	return nil
 }
 
+func toCountryResponse(country geolocation.Country) api.Country {
+	return api.Country{
+		CountryName: country.CountryName,
+		CountryCode: country.CountryISOCode,
+	}
+}
+
 func toCityResponse(city geolocation.City) api.City {
 	return api.City{
 		CityName:  city.CityName,

From b59c426fb195bceeeedff670f0f939ba282b15e5 Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Mon, 5 Feb 2024 20:57:10 +0300
Subject: [PATCH 12/19] fix tests

---
 management/server/geolocation/geolocation_test.go | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/management/server/geolocation/geolocation_test.go b/management/server/geolocation/geolocation_test.go
index 8d4c71599be..8fb8d0e2406 100644
--- a/management/server/geolocation/geolocation_test.go
+++ b/management/server/geolocation/geolocation_test.go
@@ -3,10 +3,12 @@ package geolocation
 import (
 	"os"
 	"path"
+	"sync"
 	"testing"
 
-	"github.com/netbirdio/netbird/util"
 	"github.com/stretchr/testify/assert"
+
+	"github.com/netbirdio/netbird/util"
 )
 
 // from https://github.com/maxmind/MaxMind-DB/blob/main/test-data/GeoLite2-City-Test.mmdb
@@ -24,8 +26,14 @@ func TestGeoLite_Lookup(t *testing.T) {
 		}
 	}()
 
-	geo, err := NewGeolocation(tempDir)
+	db, err := openDB(mmdbPath)
 	assert.NoError(t, err)
+
+	geo := &Geolocation{
+		mux:    &sync.RWMutex{},
+		db:     db,
+		stopCh: make(chan struct{}),
+	}
 	assert.NotNil(t, geo)
 	defer func() {
 		err = geo.Stop()

From b9bdfa90c7eaf5e0568c151deebcfbadf055b92e Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Mon, 5 Feb 2024 20:58:00 +0300
Subject: [PATCH 13/19] Add reload to SqliteStore

---
 management/server/geolocation/geolocation.go |  2 +-
 management/server/geolocation/store.go       | 78 ++++++++++++++------
 2 files changed, 57 insertions(+), 23 deletions(-)

diff --git a/management/server/geolocation/geolocation.go b/management/server/geolocation/geolocation.go
index 1550f058905..86f217a2ffd 100644
--- a/management/server/geolocation/geolocation.go
+++ b/management/server/geolocation/geolocation.go
@@ -202,7 +202,7 @@ func (gl *Geolocation) reload(newSha256sum []byte) error {
 
 	log.Infof("Successfully reloaded '%s'", gl.mmdbPath)
 
-	return nil
+	return gl.locationDB.reload()
 }
 
 func openDB(mmdbPath string) (*maxminddb.Reader, error) {
diff --git a/management/server/geolocation/store.go b/management/server/geolocation/store.go
index 3ea934d8f45..5ab4f110b5f 100644
--- a/management/server/geolocation/store.go
+++ b/management/server/geolocation/store.go
@@ -4,7 +4,9 @@ import (
 	"fmt"
 	"path/filepath"
 	"runtime"
+	"sync"
 
+	log "github.com/sirupsen/logrus"
 	"gorm.io/driver/sqlite"
 	"gorm.io/gorm"
 	"gorm.io/gorm/logger"
@@ -12,7 +14,9 @@ import (
 
 // SqliteStore represents a location storage backed by a Sqlite DB.
 type SqliteStore struct {
-	db *gorm.DB
+	db       *gorm.DB
+	filePath string
+	mux      *sync.RWMutex
 }
 
 func NewSqliteStore(dataDir string) (*SqliteStore, error) {
@@ -22,31 +26,16 @@ func NewSqliteStore(dataDir string) (*SqliteStore, error) {
 		return nil, err
 	}
 
-	storeStr := ":memory:?cache=shared&mode=ro"
-	if runtime.GOOS == "windows" {
-		storeStr = ":memory:?&mode=ro"
-	}
-
-	db, err := gorm.Open(sqlite.Open(storeStr), &gorm.Config{
-		Logger:      logger.Default.LogMode(logger.Silent),
-		PrepareStmt: true,
-	})
+	db, err := connectDB(file)
 	if err != nil {
 		return nil, err
 	}
 
-	if err := setupInMemoryDBFromFile(db, file); err != nil {
-		return nil, err
-	}
-
-	sql, err := db.DB()
-	if err != nil {
-		return nil, err
-	}
-	conns := runtime.NumCPU()
-	sql.SetMaxOpenConns(conns)
-
-	return &SqliteStore{db: db}, nil
+	return &SqliteStore{
+		db:       db,
+		filePath: file,
+		mux:      &sync.RWMutex{},
+	}, nil
 }
 
 // GetAllCountries returns a list of all countries in the store.
@@ -77,6 +66,51 @@ func (s *SqliteStore) GetCitiesByCountry(countryISOCode string) ([]City, error)
 	return cities, nil
 }
 
+func (s *SqliteStore) reload() error {
+	s.mux.Lock()
+	defer s.mux.Unlock()
+
+	log.Infof("Reloading '%s'", s.filePath)
+
+	db, err := connectDB(s.filePath)
+	if err != nil {
+		return err
+	}
+	s.db = db
+
+	log.Infof("Successfully reloaded '%s'", s.filePath)
+	return nil
+}
+
+// connectDB connects to an SQLite database and prepares it by setting up an in-memory database.
+func connectDB(source string) (*gorm.DB, error) {
+	storeStr := ":memory:?cache=shared&mode=ro"
+	if runtime.GOOS == "windows" {
+		storeStr = ":memory:?&mode=ro"
+	}
+
+	db, err := gorm.Open(sqlite.Open(storeStr), &gorm.Config{
+		Logger:      logger.Default.LogMode(logger.Silent),
+		PrepareStmt: true,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	if err := setupInMemoryDBFromFile(db, source); err != nil {
+		return nil, err
+	}
+
+	sql, err := db.DB()
+	if err != nil {
+		return nil, err
+	}
+	conns := runtime.NumCPU()
+	sql.SetMaxOpenConns(conns)
+
+	return db, nil
+}
+
 // setupInMemoryDBFromFile prepares an in-memory DB by attaching a file database and,
 // copies the data from the attached database to the in-memory database.
 func setupInMemoryDBFromFile(db *gorm.DB, source string) error {

From 1bd46448a01d0053e2776181b0b5ac743fe3d964 Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Mon, 5 Feb 2024 21:12:14 +0300
Subject: [PATCH 14/19] Add geoname indexes

---
 management/server/geolocation/store.go | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/management/server/geolocation/store.go b/management/server/geolocation/store.go
index 5ab4f110b5f..01a70b25149 100644
--- a/management/server/geolocation/store.go
+++ b/management/server/geolocation/store.go
@@ -133,5 +133,16 @@ func setupInMemoryDBFromFile(db *gorm.DB, source string) error {
 		return err
 	}
 
+	// index geoname_id and country_iso_code field
+	err = db.Exec("CREATE INDEX idx_geonames_country_iso_code ON geonames(country_iso_code);").Error
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	err = db.Exec("CREATE INDEX idx_geonames_geoname_id ON geonames(geoname_id);").Error
+	if err != nil {
+		log.Fatal(err)
+	}
+
 	return nil
 }

From bb88dfa7207c80f30d82b6b2a437a88067dcbd92 Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Wed, 7 Feb 2024 15:36:39 +0300
Subject: [PATCH 15/19] move db file check to connectDB

---
 management/server/geolocation/store.go | 24 ++++++++++++++----------
 1 file changed, 14 insertions(+), 10 deletions(-)

diff --git a/management/server/geolocation/store.go b/management/server/geolocation/store.go
index 01a70b25149..724292e6995 100644
--- a/management/server/geolocation/store.go
+++ b/management/server/geolocation/store.go
@@ -12,6 +12,10 @@ import (
 	"gorm.io/gorm/logger"
 )
 
+const (
+	geoSqliteDBFile = "geonames.db"
+)
+
 // SqliteStore represents a location storage backed by a Sqlite DB.
 type SqliteStore struct {
 	db       *gorm.DB
@@ -20,20 +24,14 @@ type SqliteStore struct {
 }
 
 func NewSqliteStore(dataDir string) (*SqliteStore, error) {
-	file := filepath.Join(dataDir, "geonames.db")
-	_, err := fileExists(file)
-	if err != nil {
-		return nil, err
-	}
-
-	db, err := connectDB(file)
+	db, err := connectDB(dataDir)
 	if err != nil {
 		return nil, err
 	}
 
 	return &SqliteStore{
 		db:       db,
-		filePath: file,
+		filePath: filepath.Join(dataDir, geoSqliteDBFile),
 		mux:      &sync.RWMutex{},
 	}, nil
 }
@@ -83,7 +81,13 @@ func (s *SqliteStore) reload() error {
 }
 
 // connectDB connects to an SQLite database and prepares it by setting up an in-memory database.
-func connectDB(source string) (*gorm.DB, error) {
+func connectDB(dataDir string) (*gorm.DB, error) {
+	file := filepath.Join(dataDir, geoSqliteDBFile)
+	_, err := fileExists(file)
+	if err != nil {
+		return nil, err
+	}
+
 	storeStr := ":memory:?cache=shared&mode=ro"
 	if runtime.GOOS == "windows" {
 		storeStr = ":memory:?&mode=ro"
@@ -97,7 +101,7 @@ func connectDB(source string) (*gorm.DB, error) {
 		return nil, err
 	}
 
-	if err := setupInMemoryDBFromFile(db, source); err != nil {
+	if err := setupInMemoryDBFromFile(db, file); err != nil {
 		return nil, err
 	}
 

From d61408e56427eec0f9d5d95c334e83dfae0eed4f Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Wed, 7 Feb 2024 16:20:18 +0300
Subject: [PATCH 16/19] Add concurrency safety to SQL queries and database
 reloading

The commit adds mutex locks to the GetAllCountries and GetCitiesByCountry functions to ensure thread-safety during database queries. Additionally, it introduces a mechanism to safely close the old database connection before a new connection is established upon reloading, which improves the reliability of database operations. Lastly, it moves the checking of database file existence to the connectDB function.
---
 management/server/geolocation/store.go | 29 ++++++++++++++++++++++++--
 1 file changed, 27 insertions(+), 2 deletions(-)

diff --git a/management/server/geolocation/store.go b/management/server/geolocation/store.go
index 724292e6995..14dcd2fee14 100644
--- a/management/server/geolocation/store.go
+++ b/management/server/geolocation/store.go
@@ -5,6 +5,7 @@ import (
 	"path/filepath"
 	"runtime"
 	"sync"
+	"time"
 
 	log "github.com/sirupsen/logrus"
 	"gorm.io/driver/sqlite"
@@ -38,6 +39,9 @@ func NewSqliteStore(dataDir string) (*SqliteStore, error) {
 
 // GetAllCountries returns a list of all countries in the store.
 func (s *SqliteStore) GetAllCountries() ([]Country, error) {
+	s.mux.Lock()
+	defer s.mux.Unlock()
+
 	var countries []Country
 	result := s.db.Table("geonames").
 		Select("country_iso_code", "country_name").
@@ -51,6 +55,9 @@ func (s *SqliteStore) GetAllCountries() ([]Country, error) {
 
 // GetCitiesByCountry retrieves a list of cities from the store based on the given country ISO code.
 func (s *SqliteStore) GetCitiesByCountry(countryISOCode string) ([]City, error) {
+	s.mux.Lock()
+	defer s.mux.Unlock()
+
 	var cities []City
 	result := s.db.Table("geonames").
 		Select("geoname_id", "city_name").
@@ -70,18 +77,36 @@ func (s *SqliteStore) reload() error {
 
 	log.Infof("Reloading '%s'", s.filePath)
 
-	db, err := connectDB(s.filePath)
+	newDb, err := connectDB(s.filePath)
 	if err != nil {
 		return err
 	}
-	s.db = db
+
+	_ = s.close()
+	s.db = newDb
 
 	log.Infof("Successfully reloaded '%s'", s.filePath)
 	return nil
 }
 
+// close closes the database connection.
+// It retrieves the underlying *sql.DB object from the *gorm.DB object
+// and calls the Close() method on it.
+func (s *SqliteStore) close() error {
+	sqlDB, err := s.db.DB()
+	if err != nil {
+		return err
+	}
+	return sqlDB.Close()
+}
+
 // connectDB connects to an SQLite database and prepares it by setting up an in-memory database.
 func connectDB(dataDir string) (*gorm.DB, error) {
+	start := time.Now()
+	defer func() {
+		log.Debugf("took %v to setup geoname db", time.Since(start))
+	}()
+
 	file := filepath.Join(dataDir, geoSqliteDBFile)
 	_, err := fileExists(file)
 	if err != nil {

From 1e4677478adbaa9cd10f2d91b72a91f2d35b344c Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Wed, 7 Feb 2024 16:50:02 +0300
Subject: [PATCH 17/19] Add sha256 sum check to geolocation store before reload

---
 management/server/geolocation/geolocation.go |  6 ++-
 management/server/geolocation/store.go       | 56 +++++++++++++-------
 2 files changed, 43 insertions(+), 19 deletions(-)

diff --git a/management/server/geolocation/geolocation.go b/management/server/geolocation/geolocation.go
index 25400aba833..40d8205ca48 100644
--- a/management/server/geolocation/geolocation.go
+++ b/management/server/geolocation/geolocation.go
@@ -178,6 +178,10 @@ func (gl *Geolocation) reloader() {
 		case <-gl.stopCh:
 			return
 		case <-time.After(gl.reloadCheckInterval):
+			if err := gl.locationDB.reload(); err != nil {
+				log.Errorf("reload failed: %s", err)
+			}
+
 			newSha256sum1, err := getSha256sum(gl.mmdbPath)
 			if err != nil {
 				log.Errorf("failed to calculate sha256 sum for '%s': %s", gl.mmdbPath, err)
@@ -229,7 +233,7 @@ func (gl *Geolocation) reload(newSha256sum []byte) error {
 
 	log.Infof("Successfully reloaded '%s'", gl.mmdbPath)
 
-	return gl.locationDB.reload()
+	return nil
 }
 
 func fileExists(filePath string) (bool, error) {
diff --git a/management/server/geolocation/store.go b/management/server/geolocation/store.go
index 14dcd2fee14..acd1190a516 100644
--- a/management/server/geolocation/store.go
+++ b/management/server/geolocation/store.go
@@ -1,6 +1,7 @@
 package geolocation
 
 import (
+	"bytes"
 	"fmt"
 	"path/filepath"
 	"runtime"
@@ -19,21 +20,30 @@ const (
 
 // SqliteStore represents a location storage backed by a Sqlite DB.
 type SqliteStore struct {
-	db       *gorm.DB
-	filePath string
-	mux      *sync.RWMutex
+	db        *gorm.DB
+	filePath  string
+	mux       *sync.RWMutex
+	sha256sum []byte
 }
 
 func NewSqliteStore(dataDir string) (*SqliteStore, error) {
-	db, err := connectDB(dataDir)
+	file := filepath.Join(dataDir, geoSqliteDBFile)
+
+	db, err := connectDB(file)
+	if err != nil {
+		return nil, err
+	}
+
+	sha256sum, err := getSha256sum(file)
 	if err != nil {
 		return nil, err
 	}
 
 	return &SqliteStore{
-		db:       db,
-		filePath: filepath.Join(dataDir, geoSqliteDBFile),
-		mux:      &sync.RWMutex{},
+		db:        db,
+		filePath:  file,
+		mux:       &sync.RWMutex{},
+		sha256sum: sha256sum,
 	}, nil
 }
 
@@ -71,21 +81,32 @@ func (s *SqliteStore) GetCitiesByCountry(countryISOCode string) ([]City, error)
 	return cities, nil
 }
 
+// reload attempts to reload the SqliteStore's database if the database file has changed.
 func (s *SqliteStore) reload() error {
 	s.mux.Lock()
 	defer s.mux.Unlock()
 
-	log.Infof("Reloading '%s'", s.filePath)
-
-	newDb, err := connectDB(s.filePath)
+	sha256sum, err := getSha256sum(s.filePath)
 	if err != nil {
-		return err
+		log.Errorf("failed to calculate sha256 sum for '%s': %s", s.filePath, err)
 	}
 
-	_ = s.close()
-	s.db = newDb
+	if !bytes.Equal(s.sha256sum, sha256sum) {
+		log.Infof("Reloading '%s'", s.filePath)
+
+		newDb, err := connectDB(s.filePath)
+		if err != nil {
+			return err
+		}
+
+		_ = s.close()
+		s.db = newDb
+
+		log.Infof("Successfully reloaded '%s'", s.filePath)
+	} else {
+		log.Debugf("No changes in '%s', no need to reload", s.filePath)
+	}
 
-	log.Infof("Successfully reloaded '%s'", s.filePath)
 	return nil
 }
 
@@ -101,14 +122,13 @@ func (s *SqliteStore) close() error {
 }
 
 // connectDB connects to an SQLite database and prepares it by setting up an in-memory database.
-func connectDB(dataDir string) (*gorm.DB, error) {
+func connectDB(filePath string) (*gorm.DB, error) {
 	start := time.Now()
 	defer func() {
 		log.Debugf("took %v to setup geoname db", time.Since(start))
 	}()
 
-	file := filepath.Join(dataDir, geoSqliteDBFile)
-	_, err := fileExists(file)
+	_, err := fileExists(filePath)
 	if err != nil {
 		return nil, err
 	}
@@ -126,7 +146,7 @@ func connectDB(dataDir string) (*gorm.DB, error) {
 		return nil, err
 	}
 
-	if err := setupInMemoryDBFromFile(db, file); err != nil {
+	if err := setupInMemoryDBFromFile(db, filePath); err != nil {
 		return nil, err
 	}
 

From a96632254293114ba13fc4993595f55f9bbad1a0 Mon Sep 17 00:00:00 2001
From: Yury Gargay <yury.gargay@gmail.com>
Date: Wed, 7 Feb 2024 16:14:06 +0100
Subject: [PATCH 18/19] Use read lock

---
 management/server/geolocation/store.go | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/management/server/geolocation/store.go b/management/server/geolocation/store.go
index acd1190a516..5f9e7753662 100644
--- a/management/server/geolocation/store.go
+++ b/management/server/geolocation/store.go
@@ -49,8 +49,8 @@ func NewSqliteStore(dataDir string) (*SqliteStore, error) {
 
 // GetAllCountries returns a list of all countries in the store.
 func (s *SqliteStore) GetAllCountries() ([]Country, error) {
-	s.mux.Lock()
-	defer s.mux.Unlock()
+	s.mux.RLock()
+	defer s.mux.RUnlock()
 
 	var countries []Country
 	result := s.db.Table("geonames").
@@ -65,8 +65,8 @@ func (s *SqliteStore) GetAllCountries() ([]Country, error) {
 
 // GetCitiesByCountry retrieves a list of cities from the store based on the given country ISO code.
 func (s *SqliteStore) GetCitiesByCountry(countryISOCode string) ([]City, error) {
-	s.mux.Lock()
-	defer s.mux.Unlock()
+	s.mux.RLock()
+	defer s.mux.RUnlock()
 
 	var cities []City
 	result := s.db.Table("geonames").

From 20635abdee82b9fcc1a1d4986beb59727cf5b53f Mon Sep 17 00:00:00 2001
From: Yury Gargay <yury.gargay@gmail.com>
Date: Wed, 7 Feb 2024 16:23:24 +0100
Subject: [PATCH 19/19] Check SHA256 twice when reload geonames db

---
 management/server/geolocation/geolocation.go |  4 ++--
 management/server/geolocation/store.go       | 15 +++++++++++++--
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/management/server/geolocation/geolocation.go b/management/server/geolocation/geolocation.go
index 40d8205ca48..69e45232ace 100644
--- a/management/server/geolocation/geolocation.go
+++ b/management/server/geolocation/geolocation.go
@@ -179,7 +179,7 @@ func (gl *Geolocation) reloader() {
 			return
 		case <-time.After(gl.reloadCheckInterval):
 			if err := gl.locationDB.reload(); err != nil {
-				log.Errorf("reload failed: %s", err)
+				log.Errorf("geonames db reload failed: %s", err)
 			}
 
 			newSha256sum1, err := getSha256sum(gl.mmdbPath)
@@ -202,7 +202,7 @@ func (gl *Geolocation) reloader() {
 				}
 				err = gl.reload(newSha256sum2)
 				if err != nil {
-					log.Errorf("reload failed: %s", err)
+					log.Errorf("mmdb reload failed: %s", err)
 				}
 			} else {
 				log.Debugf("No changes in '%s', no need to reload. Next check is in %.0f seconds.",
diff --git a/management/server/geolocation/store.go b/management/server/geolocation/store.go
index 5f9e7753662..93f8eacdd5c 100644
--- a/management/server/geolocation/store.go
+++ b/management/server/geolocation/store.go
@@ -86,12 +86,23 @@ func (s *SqliteStore) reload() error {
 	s.mux.Lock()
 	defer s.mux.Unlock()
 
-	sha256sum, err := getSha256sum(s.filePath)
+	newSha256sum1, err := getSha256sum(s.filePath)
 	if err != nil {
 		log.Errorf("failed to calculate sha256 sum for '%s': %s", s.filePath, err)
 	}
 
-	if !bytes.Equal(s.sha256sum, sha256sum) {
+	if !bytes.Equal(s.sha256sum, newSha256sum1) {
+		// we check sum twice just to avoid possible case when we reload during update of the file
+		// considering the frequency of file update (few times a week) checking sum twice should be enough
+		time.Sleep(50 * time.Millisecond)
+		newSha256sum2, err := getSha256sum(s.filePath)
+		if err != nil {
+			return fmt.Errorf("failed to calculate sha256 sum for '%s': %s", s.filePath, err)
+		}
+		if !bytes.Equal(newSha256sum1, newSha256sum2) {
+			return fmt.Errorf("sha256 sum changed during reloading of '%s'", s.filePath)
+		}
+
 		log.Infof("Reloading '%s'", s.filePath)
 
 		newDb, err := connectDB(s.filePath)