From 65dfadf195c810970e244c73aa7184d2af4dfb11 Mon Sep 17 00:00:00 2001
From: Heitor Danilo <heitor.danilo@ossystems.com.br>
Date: Thu, 31 Oct 2024 15:16:56 -0300
Subject: [PATCH] feat(pkg): add support for downloading GeoIP databases from a
 mirror

- Added `SHELLHUB_MAXMIND_MIRROR` environment variable to allow
  specifying an alternate mirror for downloading MaxMind databases. When
  `SHELLHUB_MAXMIND_MIRROR` is set, it takes precedence over
  `SHELLHUB_MAXMIND_LICENSE`, directing the API to use the mirror for
  database downloads.
---
 .env                          | 11 ++++++++++-
 api/server.go                 | 25 ++++++++++++++++++++-----
 docker-compose.yml            |  1 +
 pkg/geoip/geolite2/fetcher.go | 30 ++++++++++++++++++++++++++++++
 4 files changed, 61 insertions(+), 6 deletions(-)

diff --git a/.env b/.env
index 38348161893..2137d168e41 100644
--- a/.env
+++ b/.env
@@ -47,8 +47,17 @@ SHELLHUB_NETWORK=shellhub_network
 # Values: A valid domain name
 SHELLHUB_PUBLIC_URL_DOMAIN=
 
+
+# Specifies an alternative mirror URL for downloading the GeoIP databases. This
+# field takes precedence over SHELLHUB_MAXMIND_LICENSE; when both are
+# configured, SHELLHUB_MAXMIND_MIRROR will be used as the primary source for
+# database downloads. Leave both blank to disable the feature.
+SHELLHUB_MAXMIND_MIRROR=
+
 # Specifies a MaxMind license key used to authenticate requests for downloading
-# the GeoIP database directly from MaxMind. Leave it blank to disable the feature.
+# the GeoIP database directly from MaxMind. If SHELLHUB_MAXMIND_MIRROR is not
+# set, this license key will be used as the fallback method for fetching the
+# database. Leave both blank to disable the feature.
 SHELLHUB_MAXMIND_LICENSE=
 
 # The schedule for worker tasks.
diff --git a/api/server.go b/api/server.go
index 94a9de2259d..b8413012ca6 100644
--- a/api/server.go
+++ b/api/server.go
@@ -106,8 +106,14 @@ type config struct {
 	// Check [https://github.com/hibiken/asynq/wiki/Task-aggregation] for more information.
 	AsynqGroupMaxSize int `env:"ASYNQ_GROUP_MAX_SIZE,default=1000"`
 
-	// GeoipMaxmindLicense is the Maxmind license key used to authenticate requests for
-	// downloading the GeoIP database directly from MaxMind.
+	// GeoipMirror specifies an alternative mirror URL for downloading the GeoIP databases.
+	// This field takes precedence over [GeoipMaxmindLicense]; when both are configured,
+	// GeoipMirror will be used as the primary source for database downloads.
+	GeoipMirror string `env:"MAXMIND_MIRROR,default="`
+
+	// GeoipMaxmindLicense is the MaxMind license key used to authenticate requests for
+	// downloading the GeoIP database directly from MaxMind. If [GeoipMirror] is not set,
+	// this license key will be used as the fallback method for fetching the database.
 	GeoipMaxmindLicense string `env:"MAXMIND_LICENSE,default="`
 }
 
@@ -148,15 +154,24 @@ func startServer(ctx context.Context, cfg *config, store store.Store, cache stor
 
 	servicesOptions := []services.Option{}
 
-	if cfg.GeoipMaxmindLicense != "" {
+	if cfg.GeoipMirror != "" {
 		log.Info("GeoIP feature is enable")
 
-		locator, err := geolite2.NewLocator(ctx, geolite2.FetchFromLicenseKey(cfg.GeoipMaxmindLicense))
+		locator, err := geolite2.NewLocator(ctx, geolite2.FetchFromMirror()(cfg.GeoipMirror))
 		if err != nil {
 			log.WithError(err).Fatal("Failed to init GeoIP")
+		} else {
+			servicesOptions = append(servicesOptions, services.WithLocator(locator))
 		}
+	} else if cfg.GeoipMaxmindLicense != "" {
+		log.Info("GeoIP feature is enable")
 
-		servicesOptions = append(servicesOptions, services.WithLocator(locator))
+		locator, err := geolite2.NewLocator(ctx, geolite2.FetchFromLicenseKey(cfg.GeoipMaxmindLicense))
+		if err != nil {
+			log.WithError(err).Fatal("Failed to init GeoIP")
+		} else {
+			servicesOptions = append(servicesOptions, services.WithLocator(locator))
+		}
 	}
 
 	service := services.NewService(store, nil, nil, cache, apiClient, servicesOptions...)
diff --git a/docker-compose.yml b/docker-compose.yml
index b63edc1d97b..eb0618b62bf 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -34,6 +34,7 @@ services:
       - SHELLHUB_ENTERPRISE=${SHELLHUB_ENTERPRISE}
       - SHELLHUB_BILLING=${SHELLHUB_BILLING}
       - SHELLHUB_CLOUD=${SHELLHUB_CLOUD}
+      - MAXMIND_MIRROR=${SHELLHUB_MAXMIND_MIRROR}
       - MAXMIND_LICENSE=${SHELLHUB_MAXMIND_LICENSE}
       - RECORD_RETENTION=${SHELLHUB_RECORD_RETENTION}
       - TELEMETRY=${SHELLHUB_TELEMETRY:-}
diff --git a/pkg/geoip/geolite2/fetcher.go b/pkg/geoip/geolite2/fetcher.go
index b25fc470814..053bc368da5 100644
--- a/pkg/geoip/geolite2/fetcher.go
+++ b/pkg/geoip/geolite2/fetcher.go
@@ -41,3 +41,33 @@ func FetchFromLicenseKey(licenseKey string) GeoliteFetcher {
 		return nil
 	}
 }
+
+func FetchFromMirror(mirror string) GeoliteFetcher {
+	return func(ctx context.Context) error {
+		urls := []string{}
+		for _, id := range []string{dbCountryID, dbCityID} {
+			_, err := os.Stat(filepath.Join(dbPath, id+dbExtension))
+			switch {
+			case errors.Is(err, fs.ErrNotExist):
+				u, err := url.Parse(mirror)
+				if err != nil {
+					return err
+				}
+
+				u.Path = "/" + id + ".tar.gz"
+
+				urls = append(urls, u.String())
+			default:
+				return err
+			}
+		}
+
+		if len(urls) > 0 {
+			if err := fetchDBs(ctx, urls); err != nil {
+				return err
+			}
+		}
+
+		return nil
+	}
+}