Refactor GeoIP lookups to support IPv6 better

geo/db.go

@@ -0,0 +1,80 @@
+package geo
+
+import (
+	"log"
+	"net"
+	"strings"
+	"time"
+
+	"github.com/abh/geoip"
+
+	"github.com/abh/geodns/countries"
+)
+
+type geodb struct {
+	db       *geoip.GeoIP
+	loaded   bool
+	lastLoad time.Time
+}
+
+type geodbs struct {
+	isv6    bool
+	country geodb
+	city    geodb
+	asn     geodb
+}
+
+func (g *geodbs) GetCountry(ip net.IP) (country, continent string, netmask int) {
+	if !g.country.loaded {
+		return "", "", 0
+	}
+
+	country, netmask = g.country.db.GetCountry(ip.String())
+	if len(country) > 0 {
+		country = strings.ToLower(country)
+		continent = countries.CountryContinent[country]
+	}
+	return
+}
+
+func (g *geodbs) GetCountryRegion(ip net.IP) (country, continent, regionGroup, region string, netmask int) {
+	if !g.city.loaded {
+		log.Println("No city database available")
+		country, continent, netmask = g.GetCountry(ip)
+		return
+	}
+
+	record := g.city.db.GetRecord(ip.String())
+	if record == nil {
+		return
+	}
+
+	country = record.CountryCode
+	region = record.Region
+	if len(country) > 0 {
+		country = strings.ToLower(country)
+		continent = countries.CountryContinent[country]
+
+		if len(region) > 0 {
+			region = country + "-" + strings.ToLower(region)
+			regionGroup = countries.CountryRegionGroup(country, region)
+		}
+
+	}
+	return
+}
+
+func (g *geodbs) GetASN(ip net.IP) (asn string, netmask int) {
+	if !g.asn.loaded {
+		log.Println("No asn database available")
+		return
+	}
+	name, netmask := g.asn.db.GetName(ip.String())
+	if len(name) > 0 {
+		index := strings.Index(name, " ")
+		if index > 0 {
+			asn = strings.ToLower(name[:index])
+		}
+	}
+	return
+}

geo/geo_test.go

@@ -0,0 +1,78 @@
+package geo
+
+import (
+	"net"
+	"testing"
+)
+
+func TestGeoIPv4(t *testing.T) {
+	ip := net.ParseIP("207.171.1.1")
+
+	geoIP := New()
+	geoIP.(*GeoIP).DisableV6 = true
+
+	if geoIP.SetupCity() == nil {
+		_, _, regionGroup, _, _ := geoIP.GetCountryRegion(ip)
+		if regionGroup != "us-west" {
+			t.Errorf("Expected regionGroup '%s', got '%s'", "us-west", regionGroup)
+		}
+	}
+	if geoIP.SetupCountry() == nil {
+		country, continent, _ := geoIP.GetCountry(ip)
+		if country != "us" {
+			t.Errorf("Expected country '%s', got '%s'", "us", country)
+		}
+		if continent != "north-america" {
+			t.Errorf("Expected continent '%s', got '%s'", "north-america", continent)
+		}
+	}
+	if geoIP.SetupASN() == nil {
+		asn, _ := geoIP.GetASN(ip)
+		if asn != "as7012" {
+			t.Errorf("Expected ASN '%s', got '%s'", "as7012", asn)
+		}
+	}
+
+	asn, _ := geoIP.GetASN(ip)
+	if asn != "as7012" {
+		t.Errorf("Expected ASN '%s', got '%s'", "as7012", asn)
+	}
+}
+
+func TestGeoIPv6(t *testing.T) {
+	// los angeles
+	ip := net.ParseIP("2607:f238:3::1")
+
+	// lux
+	// 2001:888:2156::3:18:64
+
+	geoIP := New()
+
+	if geoIP.SetupCity() == nil {
+		_, _, regionGroup, _, _ := geoIP.GetCountryRegion(ip)
+		if regionGroup != "us-west" {
+			t.Errorf("Expected regionGroup '%s', got '%s'", "us-west", regionGroup)
+		}
+	}
+	if geoIP.SetupCountry() == nil {
+		country, continent, _ := geoIP.GetCountry(ip)
+		if country != "us" {
+			t.Errorf("Expected country '%s', got '%s'", "us", country)
+		}
+		if continent != "north-america" {
+			t.Errorf("Expected continent '%s', got '%s'", "north-america", continent)
+		}
+	}
+	if geoIP.SetupASN() == nil {
+		asn, _ := geoIP.GetASN(ip)
+		if asn != "as7012" {
+			t.Errorf("Expected ASN '%s', got '%s'", "as7012", asn)
+		}
+	}
+
+	asn, _ := geoIP.GetASN(ip)
+	if asn != "as7012" {
+		t.Errorf("Expected ASN '%s', got '%s'", "as7012", asn)
+	}
+
+}

geo/geoip.go

@@ -0,0 +1,150 @@
+package geo
+
+import (
+	"log"
+	"net"
+	"time"
+
+	"github.com/abh/geoip"
+)
+
+type Geotargeter interface {
+	SetDirectory(string)
+
+	GetCountry(net.IP) (country, continent string, netmask int)
+	GetCountryRegion(ip net.IP) (country, continent, regionGroup, region string, netmask int)
+	GetASN(ip net.IP) (asn string, netmask int)
+
+	SetupCity() error
+	SetupCountry() error
+	SetupASN() error
+}
+
+type GeoIP struct {
+	DisableV6 bool
+
+	v4 geodbs
+	v6 geodbs
+}
+
+func New() Geotargeter {
+	g := new(GeoIP)
+	return g
+}
+
+func (g *GeoIP) GetCountry(ip net.IP) (country, continent string, netmask int) {
+	db := g.ipdb(ip)
+	if db == nil {
+		return
+	}
+	return db.GetCountry(ip)
+}
+
+func (g *GeoIP) GetCountryRegion(ip net.IP) (country, continent, regionGroup, region string, netmask int) {
+	db := g.ipdb(ip)
+	if db == nil {
+		return
+	}
+	return db.GetCountryRegion(ip)
+}
+
+func (g *GeoIP) GetASN(ip net.IP) (asn string, netmask int) {
+	db := g.ipdb(ip)
+	if db == nil {
+		return
+	}
+	return db.GetASN(ip)
+}
+
+func (g *GeoIP) ipdb(ip net.IP) *geodbs {
+	switch isv6(ip) {
+	case false:
+		return &g.v4
+	case true:
+		return &g.v6
+	}
+	panic("impossible")
+}
+
+func isv6(ip net.IP) bool {
+	ip4 := ip.To4()
+	rv := ip4 == nil
+	return rv
+}
+
+func (g *GeoIP) SetDirectory(dir string) {
+	if len(dir) > 0 {
+		geoip.SetCustomDirectory(dir)
+	}
+}
+
+func loadDB(db *geodb, geotype int) error {
+	if db.loaded {
+		return nil
+	}
+
+	gi, err := geoip.OpenType(geotype)
+	if gi == nil || err != nil {
+		log.Printf("Could not open GeoIP database (%d): %s\n", geotype, err)
+		return err
+	}
+	db.lastLoad = time.Now()
+	db.loaded = true
+	db.db = gi
+
+	return nil
+}
+
+func (g *GeoIP) SetupCity() error {
+	var err, err6 error
+
+	if g.v4.city.loaded == false {
+		err = loadDB(&g.v4.city, geoip.GEOIP_CITY_EDITION_REV1)
+	}
+
+	if g.v6.city.loaded == false && g.DisableV6 == false {
+		err6 = loadDB(&g.v6.city, geoip.GEOIP_CITY_EDITION_REV1_V6)
+	}
+
+	if err != nil {
+		return err
+	}
+
+	return err6
+}
+
+func (g *GeoIP) SetupCountry() error {
+	var err, err6 error
+
+	if g.v4.country.loaded == false {
+		err = loadDB(&g.v4.country, geoip.GEOIP_COUNTRY_EDITION)
+	}
+
+	if g.v6.country.loaded == false && g.DisableV6 == false {
+		err6 = loadDB(&g.v6.country, geoip.GEOIP_COUNTRY_EDITION_V6)
+	}
+
+	if err != nil {
+		return err
+	}
+
+	return err6
+}
+
+func (g *GeoIP) SetupASN() error {
+	var err, err6 error
+
+	if g.v4.asn.loaded == false {
+		err = loadDB(&g.v4.asn, geoip.GEOIP_ASNUM_EDITION)
+	}
+
+	if g.v6.asn.loaded == false && g.DisableV6 == false {
+		err6 = loadDB(&g.v6.asn, geoip.GEOIP_ASNUM_EDITION_V6)
+	}
+
+	if err != nil {
+		return err
+	}
+
+	return err6
+}

geoip.go

@@ -1,141 +0,0 @@
-package main
-
-import (
-	"github.com/abh/geodns/countries"
-	"github.com/abh/geoip"
-	"log"
-	"net"
-	"strings"
-	"time"
-)
-
-type GeoIP struct {
-	country         *geoip.GeoIP
-	hasCountry      bool
-	countryLastLoad time.Time
-
-	city         *geoip.GeoIP
-	cityLastLoad time.Time
-	hasCity      bool
-
-	asn         *geoip.GeoIP
-	hasAsn      bool
-	asnLastLoad time.Time
-}
-
-var geoIP = new(GeoIP)
-
-func (g *GeoIP) GetCountry(ip net.IP) (country, continent string, netmask int) {
-	if g.country == nil {
-		return "", "", 0
-	}
-
-	country, netmask = geoIP.country.GetCountry(ip.String())
-	if len(country) > 0 {
-		country = strings.ToLower(country)
-		continent = countries.CountryContinent[country]
-	}
-	return
-}
-
-func (g *GeoIP) GetCountryRegion(ip net.IP) (country, continent, regionGroup, region string, netmask int) {
-	if g.city == nil {
-		log.Println("No city database available")
-		country, continent, netmask = g.GetCountry(ip)
-		return
-	}
-
-	record := geoIP.city.GetRecord(ip.String())
-	if record == nil {
-		return
-	}
-
-	country = record.CountryCode
-	region = record.Region
-	if len(country) > 0 {
-		country = strings.ToLower(country)
-		continent = countries.CountryContinent[country]
-
-		if len(region) > 0 {
-			region = country + "-" + strings.ToLower(region)
-			regionGroup = countries.CountryRegionGroup(country, region)
-		}
-
-	}
-	return
-}
-
-func (g *GeoIP) GetASN(ip net.IP) (asn string, netmask int) {
-	if g.asn == nil {
-		log.Println("No asn database available")
-		return
-	}
-	name, netmask := g.asn.GetName(ip.String())
-	if len(name) > 0 {
-		index := strings.Index(name, " ")
-		if index > 0 {
-			asn = strings.ToLower(name[:index])
-		}
-	}
-	return
-}
-
-func (g *GeoIP) setDirectory() {
-	if len(Config.GeoIP.Directory) > 0 {
-		geoip.SetCustomDirectory(Config.GeoIP.Directory)
-	}
-}
-
-func (g *GeoIP) setupGeoIPCountry() {
-	if g.country != nil {
-		return
-	}
-
-	g.setDirectory()
-
-	gi, err := geoip.OpenType(geoip.GEOIP_COUNTRY_EDITION)
-	if gi == nil || err != nil {
-		log.Printf("Could not open country GeoIP database: %s\n", err)
-		return
-	}
-	g.countryLastLoad = time.Now()
-	g.hasCity = true
-	g.country = gi
-
-}
-
-func (g *GeoIP) setupGeoIPCity() {
-	if g.city != nil {
-		return
-	}
-
-	g.setDirectory()
-
-	gi, err := geoip.OpenType(geoip.GEOIP_CITY_EDITION_REV1)
-	if gi == nil || err != nil {
-		log.Printf("Could not open city GeoIP database: %s\n", err)
-		return
-	}
-	g.cityLastLoad = time.Now()
-	g.hasCity = true
-	g.city = gi
-
-}
-
-func (g *GeoIP) setupGeoIPASN() {
-	if g.asn != nil {
-		return
-	}
-
-	g.setDirectory()
-
-	gi, err := geoip.OpenType(geoip.GEOIP_ASNUM_EDITION)
-	if gi == nil || err != nil {
-		log.Printf("Could not open ASN GeoIP database: %s\n", err)
-		return
-	}
-	g.asnLastLoad = time.Now()
-	g.hasAsn = true
-	g.asn = gi
-
-}

targeting_test.go

@@ -2,6 +2,8 @@ package main
 
 import (
 	"net"
+
+	"github.com/abh/geodns/geo"
 	. "gopkg.in/check.v1"
 )
 
@@ -38,18 +40,21 @@ func (s *TargetingSuite) TestTargetParse(c *C) {
 func (s *TargetingSuite) TestGetTargets(c *C) {
 	ip := net.ParseIP("207.171.1.1")
 
-	geoIP.setupGeoIPCity()
-	geoIP.setupGeoIPCountry()
-	geoIP.setupGeoIPASN()
-
-	tgt, _ := parseTargets("@ continent country")
-	targets, _ := tgt.GetTargets(ip)
-	c.Check(targets, DeepEquals, []string{"us", "north-america", "@"})
+	if geoIP == nil {
+		geoIP = geo.New()
+	}
 
-	if geoIP.city == nil {
+	err := geoIP.SetupCity()
+	if err != nil {
 		c.Log("City GeoIP database requred for these tests")
 		return
 	}
+	geoIP.SetupCountry()
+	geoIP.SetupASN()
+
+	tgt, _ := parseTargets("@ continent country")
+	targets, _ := tgt.GetTargets(ip)
+	c.Check(targets, DeepEquals, []string{"us", "north-america", "@"})
 
 	tgt, _ = parseTargets("@ continent country region ")
 	targets, _ = tgt.GetTargets(ip)

zones.go

@@ -16,8 +16,15 @@ import (
 
 	"github.com/abh/dns"
 	"github.com/abh/errorutil"
+	"github.com/abh/geodns/geo"
 )
 
+var geoIP geo.Geotargeter
+
+func init() {
+	geoIP = geo.New()
+}
+
 // Zones maps domain names to zone data
 type Zones map[string]*Zone
 
@@ -205,14 +212,18 @@ func readZoneFile(zoneName, fileName string) (zone *Zone, zerr error) {
 
 	//log.Println("IP", string(Zone.Regions["0.us"].IPv4[0].ip))
 
+	if len(Config.GeoIP.Directory) > 0 {
+		geoIP.SetDirectory(Config.GeoIP.Directory)
+	}
+
 	switch {
 	case zone.Options.Targeting >= TargetRegionGroup:
-		geoIP.setupGeoIPCity()
+		geoIP.SetupCity()
 	case zone.Options.Targeting >= TargetContinent:
-		geoIP.setupGeoIPCountry()
+		geoIP.SetupCountry()
 	}
 	if zone.Options.Targeting&TargetASN > 0 {
-		geoIP.setupGeoIPASN()
+		geoIP.SetupASN()
 	}
 
 	return zone, nil