Merge pull request strukturag#40 from strukturag/geoip-from-file

Support loading a GeoIP database from a local file.

.github/workflows/test.yml

@@ -8,6 +8,8 @@ on:
 
 jobs:
   go:
+    env:
+      MAXMIND_GEOLITE2_LICENSE: ${{ secrets.MAXMIND_GEOLITE2_LICENSE }}
     strategy:
       matrix:
         go-version:

server.conf.in

@@ -138,11 +138,17 @@ url =
 [geoip]
 # License key to use when downloading the MaxMind GeoIP database. You can
 # register an account at "https://www.maxmind.com/en/geolite2/signup" for
-# free. See "https://dev.maxmind.com/geoip/geoip2/geolite2/"" for further
+# free. See "https://dev.maxmind.com/geoip/geoip2/geolite2/" for further
 # information.
 # Leave empty to disable GeoIP lookups.
 #license =
 
+# Optional URL to download a MaxMind GeoIP database from. Will be generated if
+# "license" is provided above. Can be a "file://" url if a local file should
+# be used. Please note that the database must provide a country field when
+# looking up IP addresses.
+#url =
+
 [stats]
 # Comma-separated list of IP addresses that are allowed to access the stats
 # endpoint. Leave empty (or commented) to only allow access from "127.0.0.1".

src/signaling/geoip.go

@@ -31,6 +31,7 @@ import (
 	"net"
 	"net/http"
 	"net/url"
+	"os"
 	"strings"
 	"sync"
 	"time"
@@ -56,20 +57,35 @@ func GetGeoIpDownloadUrl(license string) string {
 
 type GeoLookup struct {
 	url    string
+	isFile bool
 	client http.Client
 	mu     sync.Mutex
 
-	lastModified string
-	reader       *maxminddb.Reader
+	lastModifiedHeader string
+	lastModifiedTime   time.Time
+
+	reader *maxminddb.Reader
 }
 
-func NewGeoLookup(url string) (*GeoLookup, error) {
+func NewGeoLookupFromUrl(url string) (*GeoLookup, error) {
 	geoip := &GeoLookup{
 		url: url,
 	}
 	return geoip, nil
 }
 
+func NewGeoLookupFromFile(filename string) (*GeoLookup, error) {
+	geoip := &GeoLookup{
+		url:    filename,
+		isFile: true,
+	}
+	if err := geoip.Update(); err != nil {
+		geoip.Close()
+		return nil, err
+	}
+	return geoip, nil
+}
+
 func (g *GeoLookup) Close() {
 	g.mu.Lock()
 	if g.reader != nil {
@@ -80,12 +96,52 @@ func (g *GeoLookup) Close() {
 }
 
 func (g *GeoLookup) Update() error {
+	if g.isFile {
+		return g.updateFile()
+	} else {
+		return g.updateUrl()
+	}
+}
+
+func (g *GeoLookup) updateFile() error {
+	info, err := os.Stat(g.url)
+	if err != nil {
+		return err
+	}
+
+	if info.ModTime().Equal(g.lastModifiedTime) {
+		return nil
+	}
+
+	reader, err := maxminddb.Open(g.url)
+	if err != nil {
+		return err
+	}
+
+	if err := reader.Verify(); err != nil {
+		return err
+	}
+
+	metadata := reader.Metadata
+	log.Printf("Using %s GeoIP database from %s (built on %s)", metadata.DatabaseType, g.url, time.Unix(int64(metadata.BuildEpoch), 0).UTC())
+
+	g.mu.Lock()
+	if g.reader != nil {
+		g.reader.Close()
+	}
+	g.reader = reader
+	g.lastModifiedTime = info.ModTime()
+	g.mu.Unlock()
+	return nil
+}
+
+func (g *GeoLookup) updateUrl() error {
 	request, err := http.NewRequest("GET", g.url, nil)
 	if err != nil {
 		return err
 	}
-	if g.lastModified != "" {
-		request.Header.Add("If-Modified-Since", g.lastModified)
+	if g.lastModifiedHeader != "" {
+		request.Header.Add("If-Modified-Since", g.lastModifiedHeader)
 	}
 	response, err := g.client.Do(request)
 	if err != nil {
@@ -150,7 +206,7 @@ func (g *GeoLookup) Update() error {
 		g.reader.Close()
 	}
 	g.reader = reader
-	g.lastModified = response.Header.Get("Last-Modified")
+	g.lastModifiedHeader = response.Header.Get("Last-Modified")
 	g.mu.Unlock()
 	return nil
 }

src/signaling/geoip_test.go

@@ -22,32 +22,24 @@
 package signaling
 
 import (
+	"archive/tar"
+	"compress/gzip"
+	"io"
+	"io/ioutil"
 	"net"
+	"net/http"
 	"os"
+	"strings"
 	"testing"
 )
 
-func TestGeoLookup(t *testing.T) {
-	license := os.Getenv("MAXMIND_GEOLITE2_LICENSE")
-	if license == "" {
-		t.Skip("No MaxMind GeoLite2 license was set in MAXMIND_GEOLITE2_LICENSE environment variable.")
-	}
-
+func testGeoLookupReader(t *testing.T, reader *GeoLookup) {
 	tests := map[string]string{
 		// Example from maxminddb-golang code.
 		"81.2.69.142": "GB",
 		// Local addresses don't have a country assigned.
 		"127.0.0.1": "",
 	}
-	reader, err := NewGeoLookup(GetGeoIpDownloadUrl(license))
-	if err != nil {
-		t.Fatal(err)
-	}
-	defer reader.Close()
-
-	if err := reader.Update(); err != nil {
-		t.Fatal(err)
-	}
 
 	for ip, expected := range tests {
 		country, err := reader.LookupCountry(net.ParseIP(ip))
@@ -62,13 +54,32 @@ func TestGeoLookup(t *testing.T) {
 	}
 }
 
+func TestGeoLookup(t *testing.T) {
+	license := os.Getenv("MAXMIND_GEOLITE2_LICENSE")
+	if license == "" {
+		t.Skip("No MaxMind GeoLite2 license was set in MAXMIND_GEOLITE2_LICENSE environment variable.")
+	}
+
+	reader, err := NewGeoLookupFromUrl(GetGeoIpDownloadUrl(license))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer reader.Close()
+
+	if err := reader.Update(); err != nil {
+		t.Fatal(err)
+	}
+
+	testGeoLookupReader(t, reader)
+}
+
 func TestGeoLookupCaching(t *testing.T) {
 	license := os.Getenv("MAXMIND_GEOLITE2_LICENSE")
 	if license == "" {
 		t.Skip("No MaxMind GeoLite2 license was set in MAXMIND_GEOLITE2_LICENSE environment variable.")
 	}
 
-	reader, err := NewGeoLookup(GetGeoIpDownloadUrl(license))
+	reader, err := NewGeoLookupFromUrl(GetGeoIpDownloadUrl(license))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -110,9 +121,74 @@ func TestGeoLookupContinent(t *testing.T) {
 }
 
 func TestGeoLookupCloseEmpty(t *testing.T) {
-	reader, err := NewGeoLookup("ignore-url")
+	reader, err := NewGeoLookupFromUrl("ignore-url")
 	if err != nil {
 		t.Fatal(err)
 	}
 	reader.Close()
 }
+
+func TestGeoLookupFromFile(t *testing.T) {
+	license := os.Getenv("MAXMIND_GEOLITE2_LICENSE")
+	if license == "" {
+		t.Skip("No MaxMind GeoLite2 license was set in MAXMIND_GEOLITE2_LICENSE environment variable.")
+	}
+
+	url := GetGeoIpDownloadUrl(license)
+	resp, err := http.Get(url)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer resp.Body.Close()
+
+	body := resp.Body
+	if strings.HasSuffix(url, ".gz") {
+		body, err = gzip.NewReader(body)
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	tmpfile, err := ioutil.TempFile("", "geoipdb")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.Remove(tmpfile.Name())
+
+	tarfile := tar.NewReader(body)
+	foundDatabase := false
+	for {
+		header, err := tarfile.Next()
+		if err == io.EOF {
+			break
+		} else if err != nil {
+			t.Fatal(err)
+		}
+
+		if !strings.HasSuffix(header.Name, ".mmdb") {
+			continue
+		}
+
+		if _, err := io.Copy(tmpfile, tarfile); err != nil {
+			tmpfile.Close()
+			t.Fatal(err)
+		}
+		if err := tmpfile.Close(); err != nil {
+			t.Fatal(err)
+		}
+		foundDatabase = true
+		break
+	}
+
+	if !foundDatabase {
+		t.Fatal("Did not find MaxMind database in tarball")
+	}
+
+	reader, err := NewGeoLookupFromFile(tmpfile.Name())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer reader.Close()
+
+	testGeoLookupReader(t, reader)
+}

src/signaling/hub.go

@@ -209,8 +209,14 @@ func NewHub(config *goconf.ConfigFile, nats NatsClient, r *mux.Router, version s
 
 	var geoip *GeoLookup
 	if geoipUrl != "" {
-		log.Printf("Downloading GeoIP database from %s", geoipUrl)
-		geoip, err = NewGeoLookup(geoipUrl)
+		if strings.HasPrefix(geoipUrl, "file://") {
+			geoipUrl = geoipUrl[7:]
+			log.Printf("Using GeoIP database from %s", geoipUrl)
+			geoip, err = NewGeoLookupFromFile(geoipUrl)
+		} else {
+			log.Printf("Downloading GeoIP database from %s", geoipUrl)
+			geoip, err = NewGeoLookupFromUrl(geoipUrl)
+		}
 		if err != nil {
 			return nil, err
 		}