From 2071d002fc5c3b5ff7a3fca6a5c99f5517196853 Mon Sep 17 00:00:00 2001
From: Giovanni <561184+wargio@users.noreply.github.com>
Date: Fri, 14 Feb 2025 04:01:43 +0800
Subject: [PATCH] feat: add ability to set custom Geolite DB URL

---
 backend/internal/common/env_config.go       |  7 +++++--
 backend/internal/service/geolite_service.go | 20 ++++++++++++++------
 backend/internal/service/test_service.go    |  2 +-
 backend/internal/service/user_service.go    |  2 +-
 4 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/backend/internal/common/env_config.go b/backend/internal/common/env_config.go
index b8a2025..390a19d 100644
--- a/backend/internal/common/env_config.go
+++ b/backend/internal/common/env_config.go
@@ -10,8 +10,9 @@ import (
 type DbProvider string
 
 const (
-	DbProviderSqlite   DbProvider = "sqlite"
-	DbProviderPostgres DbProvider = "postgres"
+	DbProviderSqlite      DbProvider = "sqlite"
+	DbProviderPostgres    DbProvider = "postgres"
+	MaxMindGeoLiteCityUrl string     = "https://download.maxmind.com/app/geoip_download?edition_id=GeoLite2-City&license_key=%s&suffix=tar.gz"
 )
 
 type EnvConfigSchema struct {
@@ -25,6 +26,7 @@ type EnvConfigSchema struct {
 	Host                     string     `env:"HOST"`
 	MaxMindLicenseKey        string     `env:"MAXMIND_LICENSE_KEY"`
 	GeoLiteDBPath            string     `env:"GEOLITE_DB_PATH"`
+	GeoLiteDBUrl             string     `env:"GEOLITE_DB_URL"`
 	UiConfigDisabled         bool       `env:"PUBLIC_UI_CONFIG_DISABLED"`
 }
 
@@ -39,6 +41,7 @@ var EnvConfig = &EnvConfigSchema{
 	Host:                     "0.0.0.0",
 	MaxMindLicenseKey:        "",
 	GeoLiteDBPath:            "data/GeoLite2-City.mmdb",
+	GeoLiteDBUrl:             MaxMindGeoLiteCityUrl,
 	UiConfigDisabled:         false,
 }
 
diff --git a/backend/internal/service/geolite_service.go b/backend/internal/service/geolite_service.go
index 66554e9..6c280d2 100644
--- a/backend/internal/service/geolite_service.go
+++ b/backend/internal/service/geolite_service.go
@@ -21,7 +21,8 @@ import (
 )
 
 type GeoLiteService struct {
-	mutex sync.Mutex
+	disableUpdater bool
+	mutex          sync.Mutex
 }
 
 var localhostIPNets = []*net.IPNet{
@@ -43,6 +44,12 @@ var tailscaleIPNets = []*net.IPNet{
 func NewGeoLiteService() *GeoLiteService {
 	service := &GeoLiteService{}
 
+	if common.EnvConfig.MaxMindLicenseKey == "" && common.EnvConfig.GeoLiteDBUrl == common.MaxMindGeoLiteCityUrl {
+		// Warn the user, and disable the updater.
+		log.Println("MAXMIND_LICENSE_KEY environment variable is empty. The GeoLite2 City database won't be updated.")
+		service.disableUpdater = true
+	}
+
 	go func() {
 		if err := service.updateDatabase(); err != nil {
 			log.Printf("Failed to update GeoLite2 City database: %v\n", err)
@@ -104,18 +111,19 @@ func (s *GeoLiteService) GetLocationByIP(ipAddress string) (country, city string
 
 // UpdateDatabase checks the age of the database and updates it if it's older than 14 days.
 func (s *GeoLiteService) updateDatabase() error {
+	if s.disableUpdater {
+		// Avoid updating the GeoLite2 City database.
+		return nil
+	}
+
 	if s.isDatabaseUpToDate() {
 		log.Println("GeoLite2 City database is up-to-date.")
 		return nil
 	}
 
 	log.Println("Updating GeoLite2 City database...")
+	downloadUrl := fmt.Sprintf(common.EnvConfig.GeoLiteDBUrl, common.EnvConfig.MaxMindLicenseKey)
 
-	// Download and extract the database
-	downloadUrl := fmt.Sprintf(
-		"https://download.maxmind.com/app/geoip_download?edition_id=GeoLite2-City&license_key=%s&suffix=tar.gz",
-		common.EnvConfig.MaxMindLicenseKey,
-	)
 	// Download the database tar.gz file
 	resp, err := http.Get(downloadUrl)
 	if err != nil {
diff --git a/backend/internal/service/test_service.go b/backend/internal/service/test_service.go
index 146b838..f6f4816 100644
--- a/backend/internal/service/test_service.go
+++ b/backend/internal/service/test_service.go
@@ -11,8 +11,8 @@ import (
 	"time"
 
 	"github.com/fxamacker/cbor/v2"
-	"github.com/pocket-id/pocket-id/backend/resources"
 	datatype "github.com/pocket-id/pocket-id/backend/internal/model/types"
+	"github.com/pocket-id/pocket-id/backend/resources"
 
 	"github.com/go-webauthn/webauthn/protocol"
 	"github.com/pocket-id/pocket-id/backend/internal/common"
diff --git a/backend/internal/service/user_service.go b/backend/internal/service/user_service.go
index fcdae16..060b370 100644
--- a/backend/internal/service/user_service.go
+++ b/backend/internal/service/user_service.go
@@ -11,9 +11,9 @@ import (
 	"github.com/pocket-id/pocket-id/backend/internal/common"
 	"github.com/pocket-id/pocket-id/backend/internal/dto"
 	"github.com/pocket-id/pocket-id/backend/internal/model"
+	datatype "github.com/pocket-id/pocket-id/backend/internal/model/types"
 	"github.com/pocket-id/pocket-id/backend/internal/utils"
 	"github.com/pocket-id/pocket-id/backend/internal/utils/email"
-	datatype "github.com/pocket-id/pocket-id/backend/internal/model/types"
 	"gorm.io/gorm"
 )