Merge pull request #51 from safing/feature/new-geoipdb

Update to new geoip database

intel/geoip/database.go

@@ -14,13 +14,13 @@ import (
 )
 
 var (
-	dbCityFile *updater.File
-	dbASNFile  *updater.File
-	dbFileLock sync.Mutex
+	geoDBv4File *updater.File
+	geoDBv6File *updater.File
+	dbFileLock  sync.Mutex
 
-	dbCity *maxminddb.Reader
-	dbASN  *maxminddb.Reader
-	dbLock sync.Mutex
+	geoDBv4Reader *maxminddb.Reader
+	geoDBv6Reader *maxminddb.Reader
+	dbLock        sync.Mutex
 
 	dbInUse    = abool.NewBool(false) // only activate if used for first time
 	dbDoReload = abool.NewBool(true)  // if database should be reloaded
@@ -51,7 +51,11 @@ func doReload() error {
 	// reload if needed
 	if dbDoReload.SetToIf(true, false) {
 		closeDBs()
-		return openDBs()
+		if err := openDBs(); err != nil {
+			// try again the next time
+			dbDoReload.SetTo(true)
+			return err
+		}
 	}
 
 	return nil
@@ -59,23 +63,33 @@ func doReload() error {
 
 func openDBs() error {
 	var err error
-	dbCityFile, err = updates.GetFile("intel/geoip/geoip-city.mmdb")
+
+	geoDBv4File, err = updates.GetFile("intel/geoip/geoipv4.mmdb.gz")
 	if err != nil {
-		return fmt.Errorf("could not get GeoIP City database file: %s", err)
+		return fmt.Errorf("could not get GeoIP v4 database file: %s", err)
 	}
-	dbCity, err = maxminddb.Open(dbCityFile.Path())
+	unpackedV4, err := geoDBv4File.Unpack(".gz", updater.UnpackGZIP)
+	if err != nil {
+		return err
+	}
+	geoDBv4Reader, err = maxminddb.Open(unpackedV4)
 	if err != nil {
 		return err
 	}
 
-	dbASNFile, err = updates.GetFile("intel/geoip/geoip-asn.mmdb")
+	geoDBv6File, err = updates.GetFile("intel/geoip/geoipv6.mmdb.gz")
 	if err != nil {
-		return fmt.Errorf("could not get GeoIP ASN database file: %s", err)
+		return fmt.Errorf("could not get GeoIP v6 database file: %s", err)
 	}
-	dbASN, err = maxminddb.Open(dbASNFile.Path())
+	unpackedV6, err := geoDBv6File.Unpack(".gz", updater.UnpackGZIP)
 	if err != nil {
 		return err
 	}
+	geoDBv6Reader, err = maxminddb.Open(unpackedV6)
+	if err != nil {
+		return err
+	}
+
 	return nil
 }
 
@@ -85,19 +99,19 @@ func handleError(err error) {
 }
 
 func closeDBs() {
-	if dbCity != nil {
-		err := dbCity.Close()
+	if geoDBv4Reader != nil {
+		err := geoDBv4Reader.Close()
 		if err != nil {
 			log.Warningf("network/geoip: failed to close database: %s", err)
 		}
 	}
-	dbCity = nil
+	geoDBv4Reader = nil
 
-	if dbASN != nil {
-		err := dbASN.Close()
+	if geoDBv6Reader != nil {
+		err := geoDBv6Reader.Close()
 		if err != nil {
 			log.Warningf("network/geoip: failed to close database: %s", err)
 		}
 	}
-	dbASN = nil
+	geoDBv6Reader = nil
 }

intel/geoip/lookup.go

@@ -2,8 +2,18 @@ package geoip
 
 import (
 	"net"
+
+	"github.com/oschwald/maxminddb-golang"
+	"github.com/safing/portbase/log"
 )
 
+func getReader(ip net.IP) *maxminddb.Reader {
+	if v4 := ip.To4(); v4 != nil {
+		return geoDBv4Reader
+	}
+	return geoDBv6Reader
+}
+
 // GetLocation returns Location data of an IP address
 func GetLocation(ip net.IP) (record *Location, err error) {
 	dbLock.Lock()
@@ -14,13 +24,12 @@ func GetLocation(ip net.IP) (record *Location, err error) {
 		return nil, err
 	}
 
+	db := getReader(ip)
+
 	record = &Location{}
 
 	// fetch
-	err = dbCity.Lookup(ip, record)
-	if err == nil {
-		err = dbASN.Lookup(ip, record)
-	}
+	err = db.Lookup(ip, record)
 
 	// retry
 	if err != nil {
@@ -30,17 +39,17 @@ func GetLocation(ip net.IP) (record *Location, err error) {
 		if err != nil {
 			return nil, err
 		}
+		db = getReader(ip)
 
 		// refetch
-		err = dbCity.Lookup(ip, record)
-		if err == nil {
-			err = dbASN.Lookup(ip, record)
-		}
+		err = db.Lookup(ip, record)
 	}
 
 	if err != nil {
 		return nil, err
 	}
 
+	log.Tracef("geoip: record: %+v", record)
+
 	return record, nil
 }

intel/geoip/module.go

@@ -27,10 +27,10 @@ func prep() error {
 func upgradeDatabases(_ context.Context, _ interface{}) error {
 	dbFileLock.Lock()
 	reload := false
-	if dbCityFile != nil && dbCityFile.UpgradeAvailable() {
+	if geoDBv4File != nil && geoDBv4File.UpgradeAvailable() {
 		reload = true
 	}
-	if dbASNFile != nil && dbASNFile.UpgradeAvailable() {
+	if geoDBv6File != nil && geoDBv6File.UpgradeAvailable() {
 		reload = true
 	}
 	dbFileLock.Unlock()

profile/endpoints/endpoints_test.go

@@ -322,7 +322,7 @@ func TestEndpointMatching(t *testing.T) {
 	}
 
 	testEndpointMatch(t, ep, (&intel.Entity{
-		IP: net.ParseIP("1.1.1.1"),
+		IP: net.ParseIP("1.1.1.2"),
 	}).Init(), Permitted)
 	testEndpointMatch(t, ep, (&intel.Entity{
 		IP: net.ParseIP("8.8.8.8"),