Support loading a GeoIP database from a local file.

This commit is contained in:
Joachim Bauch 2020-08-13 14:07:26 +02:00
parent 2734197e8d
commit 5e3164b5a4
Failed to extract signature
4 changed files with 170 additions and 26 deletions

View file

@ -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".

View file

@ -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
}

View file

@ -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)
}

View file

@ -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
}