From b55f4371e3feb2493d81a4755b92e14a89469ea0 Mon Sep 17 00:00:00 2001
From: Maxim Pogozhiy <foxdalas@gmail.com>
Date: Mon, 28 Dec 2020 15:28:16 +1000
Subject: [PATCH] Add GeoIP Local mirror support

---
 .../templates/controller-daemonset.yaml             |  3 +++
 cmd/nginx/flags.go                                  |  3 ++-
 cmd/nginx/flags_test.go                             | 13 +++++++++++++
 internal/nginx/maxmind.go                           | 12 +++++++++++-
 4 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/charts/ingress-nginx/templates/controller-daemonset.yaml b/charts/ingress-nginx/templates/controller-daemonset.yaml
index aa8693b47..27d2ed101 100644
--- a/charts/ingress-nginx/templates/controller-daemonset.yaml
+++ b/charts/ingress-nginx/templates/controller-daemonset.yaml
@@ -98,6 +98,9 @@ spec:
             - --validating-webhook-certificate={{ .Values.controller.admissionWebhooks.certificate }}
             - --validating-webhook-key={{ .Values.controller.admissionWebhooks.key }}
           {{- end }}
+          {{- if .Values.controller.maxmindMirror }}
+            - --maxmind-mirror={{ .Values.controller.maxmindMirror }}
+          {{- end}}
           {{- if .Values.controller.maxmindLicenseKey }}
             - --maxmind-license-key={{ .Values.controller.maxmindLicenseKey }}
           {{- end }}
diff --git a/cmd/nginx/flags.go b/cmd/nginx/flags.go
index 26408e61b..0bb44c483 100644
--- a/cmd/nginx/flags.go
+++ b/cmd/nginx/flags.go
@@ -173,6 +173,7 @@ Takes the form "<host>:port". If not provided, no admission controller is starte
 		statusUpdateInterval = flags.Int("status-update-interval", status.UpdateInterval, "Time interval in seconds in which the status should check if an update is required. Default is 60 seconds")
 	)
 
+	flags.StringVar(&nginx.MaxmindMirror, "maxmind-mirror", "", `Maxmind mirror url (example: http://geoip.local/databases`)
 	flags.StringVar(&nginx.MaxmindLicenseKey, "maxmind-license-key", "", `Maxmind license key to download GeoLite2 Databases.
 https://blog.maxmind.com/2019/12/18/significant-changes-to-accessing-and-using-geolite2-databases`)
 	flags.StringVar(&nginx.MaxmindEditionIDs, "maxmind-edition-ids", "GeoLite2-City,GeoLite2-ASN", `Maxmind edition ids to download GeoLite2 Databases.`)
@@ -299,7 +300,7 @@ https://blog.maxmind.com/2019/12/18/significant-changes-to-accessing-and-using-g
 		config.RootCAFile = *rootCAFile
 	}
 
-	if nginx.MaxmindLicenseKey != "" && nginx.MaxmindEditionIDs != "" {
+	if (nginx.MaxmindLicenseKey != "" || nginx.MaxmindMirror != "") && nginx.MaxmindEditionIDs != "" {
 		if err := nginx.ValidateGeoLite2DBEditions(); err != nil {
 			return false, nil, err
 		}
diff --git a/cmd/nginx/flags_test.go b/cmd/nginx/flags_test.go
index 411c7c993..4286d4a20 100644
--- a/cmd/nginx/flags_test.go
+++ b/cmd/nginx/flags_test.go
@@ -92,3 +92,16 @@ func TestMaxmindEdition(t *testing.T) {
 		t.Fatalf("Expected an error parsing flags but none returned")
 	}
 }
+
+func TestMaxmindMirror(t *testing.T) {
+	resetForTesting(func() { t.Fatal("Parsing failed") })
+
+	oldArgs := os.Args
+	defer func() { os.Args = oldArgs }()
+	os.Args = []string{"cmd", "--publish-service", "namespace/test", "--http-port", "0", "--https-port", "0", "--maxmind-mirror", "http://geoip.local", "--maxmind-license-key", "0000000", "--maxmind-edition-ids", "GeoLite2-City, TestCheck"}
+
+	_, _, err := parseFlags()
+	if err == nil {
+		t.Fatalf("Expected an error parsing flags but none returned")
+	}
+}
diff --git a/internal/nginx/maxmind.go b/internal/nginx/maxmind.go
index a303e8b2e..0f3c37016 100644
--- a/internal/nginx/maxmind.go
+++ b/internal/nginx/maxmind.go
@@ -36,6 +36,9 @@ var MaxmindEditionIDs = ""
 // MaxmindEditionFiles maxmind databases on disk
 var MaxmindEditionFiles []string
 
+// MaxmindMirror maxmind database mirror url (http://geoip.local)
+var MaxmindMirror = ""
+
 const (
 	geoIPPath   = "/etc/nginx/geoip"
 	dbExtension = ".mmdb"
@@ -68,8 +71,15 @@ func DownloadGeoLite2DB() error {
 	return nil
 }
 
+func createURL(mirror, licenseKey, dbName string) string {
+	if len(mirror) > 0 {
+		return fmt.Sprintf("%s/%s.tar.gz", mirror, dbName)
+	}
+	return fmt.Sprintf(maxmindURL, licenseKey, dbName)
+}
+
 func downloadDatabase(dbName string) error {
-	url := fmt.Sprintf(maxmindURL, MaxmindLicenseKey, dbName)
+	url := createURL(MaxmindMirror, MaxmindLicenseKey, dbName)
 	req, err := http.NewRequest(http.MethodGet, url, nil)
 	if err != nil {
 		return err