Refactor geoip infrastructure

geodns.go

@@ -34,6 +34,8 @@ import (
 	"github.com/abh/geodns/monitor"
 	"github.com/abh/geodns/querylog"
 	"github.com/abh/geodns/server"
+	"github.com/abh/geodns/targeting"
+	"github.com/abh/geodns/targeting/geoip2"
 	"github.com/abh/geodns/zones"
 	"github.com/pborman/uuid"
 )
@@ -200,6 +202,16 @@ func main() {
 		log.Println("StatHat integration has been removed in favor of more generic metrics")
 	}
 
+	if len(Config.GeoIPDirectory()) > 0 {
+		geoProvider, err := geoip2.New(Config.GeoIPDirectory())
+		if err != nil {
+			log.Printf("Configuring geo provider: %s", err)
+		}
+		if geoProvider != nil {
+			targeting.Setup(geoProvider)
+		}
+	}
+
 	mon := monitor.NewMonitor(serverInfo)
 	go mon.Run()
 

targeting/geo/geo.go

@@ -0,0 +1,43 @@
+package geo
+
+import (
+	"math"
+	"net"
+
+	"github.com/golang/geo/s2"
+)
+
+type Provider interface {
+	HasCountry() (bool, error)
+	GetCountry(ip net.IP) (country, continent string, netmask int)
+	HasASN() (bool, error)
+	GetASN(net.IP) (asn string, netmask int, err error)
+	HasLocation() (bool, error)
+	GetLocation(ip net.IP) (location *Location, err error)
+}
+
+const MAX_DISTANCE = 360
+
+type Location struct {
+	Country     string
+	Continent   string
+	RegionGroup string
+	Region      string
+	Latitude    float64
+	Longitude   float64
+	Netmask     int
+}
+
+func (l *Location) MaxDistance() float64 {
+	return MAX_DISTANCE
+}
+
+func (l *Location) Distance(to *Location) float64 {
+	if to == nil {
+		return MAX_DISTANCE
+	}
+	ll1 := s2.LatLngFromDegrees(l.Latitude, l.Longitude)
+	ll2 := s2.LatLngFromDegrees(to.Latitude, to.Longitude)
+	angle := ll1.Distance(ll2)
+	return math.Abs(angle.Degrees())
+}

targeting/geoip.go

@@ -1,168 +0,0 @@
-package targeting
-
-import (
-	"log"
-	"math"
-	"net"
-	"strings"
-	"time"
-
-	"github.com/abh/geodns/countries"
-
-	"github.com/abh/geoip"
-	"github.com/golang/geo/s2"
-)
-
-type GeoIPData 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
-}
-
-const MAX_DISTANCE = 360
-
-type Location struct {
-	Latitude  float64
-	Longitude float64
-}
-
-func (l *Location) MaxDistance() float64 {
-	return MAX_DISTANCE
-}
-
-func (l *Location) Distance(to *Location) float64 {
-	if to == nil {
-		return MAX_DISTANCE
-	}
-	ll1 := s2.LatLngFromDegrees(l.Latitude, l.Longitude)
-	ll2 := s2.LatLngFromDegrees(to.Latitude, to.Longitude)
-	angle := ll1.Distance(ll2)
-	return math.Abs(angle.Degrees())
-}
-
-var geoIP = &GeoIPData{}
-
-func GeoIP() *GeoIPData {
-	// mutex this and allow it to reload as needed?
-	return geoIP
-}
-
-func (g *GeoIPData) GetCountry(ip net.IP) (country, continent string, netmask int) {
-	if g.country == nil {
-		return "", "", 0
-	}
-
-	country, netmask = g.country.GetCountry(ip.String())
-	if len(country) > 0 {
-		country = strings.ToLower(country)
-		continent = countries.CountryContinent[country]
-	}
-	return
-}
-
-func (g *GeoIPData) GetCountryRegion(ip net.IP) (country, continent, regionGroup, region string, netmask int, location *Location) {
-	if g.city == nil {
-		log.Println("No city database available")
-		country, continent, netmask = g.GetCountry(ip)
-		return
-	}
-
-	record := g.city.GetRecord(ip.String())
-	if record == nil {
-		return
-	}
-
-	location = &Location{float64(record.Latitude), float64(record.Longitude)}
-
-	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 *GeoIPData) 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 *GeoIPData) SetDirectory(directory string) {
-	// directory := Config.GeoIPDataDirectory()
-	if len(directory) > 0 {
-		geoip.SetCustomDirectory(directory)
-	}
-}
-
-func (g *GeoIPData) SetupGeoIPCountry() {
-	if g.country != nil {
-		return
-	}
-
-	gi, err := geoip.OpenType(geoip.GEOIP_COUNTRY_EDITION)
-	if gi == nil || err != nil {
-		log.Printf("Could not open country GeoIPData database: %s\n", err)
-		return
-	}
-	g.countryLastLoad = time.Now()
-	g.hasCity = true
-	g.country = gi
-
-}
-
-func (g *GeoIPData) SetupGeoIPCity() {
-	if g.city != nil {
-		return
-	}
-
-	gi, err := geoip.OpenType(geoip.GEOIP_CITY_EDITION_REV1)
-	if gi == nil || err != nil {
-		log.Printf("Could not open city GeoIPData database: %s\n", err)
-		return
-	}
-	g.cityLastLoad = time.Now()
-	g.hasCity = true
-	g.city = gi
-
-}
-
-func (g *GeoIPData) SetupGeoIPASN() {
-	if g.asn != nil {
-		return
-	}
-
-	gi, err := geoip.OpenType(geoip.GEOIP_ASNUM_EDITION)
-	if gi == nil || err != nil {
-		log.Printf("Could not open ASN GeoIPData database: %s\n", err)
-		return
-	}
-	g.asnLastLoad = time.Now()
-	g.hasAsn = true
-	g.asn = gi
-
-}

targeting/geoip2/geoip2.go

@@ -0,0 +1,171 @@
+package geoip2
+
+import (
+	"fmt"
+	"log"
+	"net"
+	"path/filepath"
+	"strings"
+	"sync"
+
+	"github.com/abh/geodns/countries"
+	"github.com/abh/geodns/targeting/geo"
+	geoip2 "github.com/oschwald/geoip2-golang"
+)
+
+type geoType uint8
+
+const (
+	countryDB = iota
+	cityDB
+	asnDB
+)
+
+type g2 struct {
+	dir string
+
+	country *geoip2.Reader
+	city    *geoip2.Reader
+	asn     *geoip2.Reader
+	mu      sync.RWMutex
+}
+
+func (g *g2) open(t geoType, db string) (*geoip2.Reader, error) {
+	f := filepath.Join(g.dir, db)
+	n, err := geoip2.Open(f)
+	if err != nil {
+		return nil, err
+	}
+	g.mu.Lock()
+	defer g.mu.Unlock()
+
+	switch t {
+	case countryDB:
+		g.country = n
+	case cityDB:
+		g.city = n
+	case asnDB:
+		g.asn = n
+	}
+	return n, nil
+}
+
+func (g *g2) get(t geoType, db string) (*geoip2.Reader, error) {
+	g.mu.RLock()
+
+	var r *geoip2.Reader
+
+	switch t {
+	case countryDB:
+		r = g.country
+	case cityDB:
+		r = g.city
+	case asnDB:
+		r = g.asn
+	}
+
+	// unlock so the g.open() call below won't lock
+	g.mu.RUnlock()
+
+	if r != nil {
+		return r, nil
+	}
+
+	return g.open(t, db)
+}
+
+func New(dir string) (*g2, error) {
+	g := &g2{
+		dir: dir,
+	}
+	_, err := g.open(countryDB, "GeoIP2-Country.mmdb")
+	if err != nil {
+		return nil, err
+	}
+
+	return g, nil
+}
+
+func (g *g2) HasASN() (bool, error) {
+	r, err := g.get(asnDB, "GeoIP2-ASN.mmdb")
+	if r != nil && err == nil {
+		return true, nil
+	}
+	return false, err
+}
+
+func (g *g2) GetASN(ip net.IP) (string, int, error) {
+	r, err := g.get(asnDB, "GeoIP2-ASN.mmdb")
+	if err != nil {
+		return "", 0, err
+	}
+
+	c, err := r.ISP(ip)
+	if err != nil {
+		return "", 0, fmt.Errorf("lookup ASN for '%s': %s", ip.String(), err)
+	}
+	asn := c.AutonomousSystemNumber
+	return fmt.Sprintf("as%d", asn), 0, nil
+}
+
+func (g *g2) HasCountry() (bool, error) {
+	r, err := g.get(countryDB, "GeoIP2-Country.mmdb")
+	if r != nil && err == nil {
+		return true, nil
+	}
+	return false, err
+}
+
+func (g *g2) GetCountry(ip net.IP) (country, continent string, netmask int) {
+	r, err := g.get(countryDB, "GeoIP2.mmdb")
+	c, err := r.Country(ip)
+	if err != nil {
+		log.Printf("Could not lookup country for '%s': %s", ip.String(), err)
+		return "", "", 0
+	}
+
+	country = c.Country.IsoCode
+
+	if len(country) > 0 {
+		country = strings.ToLower(country)
+		continent = countries.CountryContinent[country]
+	}
+
+	return country, continent, 0
+}
+
+func (g *g2) HasLocation() (bool, error) {
+	r, err := g.get(cityDB, "GeoIP2-City.mmdb")
+	if r != nil && err == nil {
+		return true, nil
+	}
+	return false, err
+}
+
+func (g *g2) GetLocation(ip net.IP) (l *geo.Location, err error) {
+	c, err := g.city.City(ip)
+	if err != nil {
+		log.Printf("Could not lookup CountryRegion for '%s': %s", ip.String(), err)
+		return
+	}
+
+	l = &geo.Location{
+		Latitude:  float64(c.Location.Latitude),
+		Longitude: float64(c.Location.Longitude),
+		Country:   strings.ToLower(c.Country.IsoCode),
+	}
+
+	if len(c.Subdivisions) > 0 {
+		l.Region = strings.ToLower(c.Subdivisions[0].IsoCode)
+	}
+	if len(l.Country) > 0 {
+		l.Continent = countries.CountryContinent[l.Country]
+		if len(l.Region) > 0 {
+			l.Region = l.Country + "-" + l.Region
+			l.RegionGroup = countries.CountryRegionGroup(l.Country, l.Region)
+		}
+	}
+
+	return
+
+}

targeting/targeting.go

@@ -2,8 +2,11 @@ package targeting
 
 import (
 	"fmt"
+	"log"
 	"net"
 	"strings"
+
+	"github.com/abh/geodns/targeting/geo"
 )
 
 type TargetOptions int
@@ -24,24 +27,20 @@ func init() {
 	cidr48Mask = net.CIDRMask(48, 128)
 }
 
-func (t TargetOptions) GetTargets(ip net.IP, hasClosest bool) ([]string, int, *Location) {
+var g geo.Provider
 
-	targets := make([]string, 0)
+func Setup(gn geo.Provider) error {
+	g = gn
+	return nil
+}
 
-	var country, continent, region, regionGroup, asn string
-	var netmask int
-	var location *Location
+func Geo() geo.Provider {
+	return g
+}
 
-	g := GeoIP()
+func (t TargetOptions) GetTargets(ip net.IP, hasClosest bool) ([]string, int, *geo.Location) {
 
-	if t&TargetASN > 0 {
-		asn, netmask = g.GetASN(ip)
-	}
-	if t&TargetRegion > 0 || t&TargetRegionGroup > 0 || hasClosest {
-		country, continent, regionGroup, region, netmask, location = geoIP.GetCountryRegion(ip)
-	} else if t&TargetCountry > 0 || t&TargetContinent > 0 {
-		country, continent, netmask = g.GetCountry(ip)
-	}
+	targets := make([]string, 0)
 
 	if t&TargetIP > 0 {
 		ipStr := ip.String()
@@ -59,8 +58,35 @@ func (t TargetOptions) GetTargets(ip net.IP, hasClosest bool) ([]string, int, *L
 		}
 	}
 
-	if t&TargetASN > 0 && len(asn) > 0 {
-		targets = append(targets, asn)
+	if t&TargetASN > 0 {
+		asn, _, err := g.GetASN(ip)
+		if err != nil {
+			log.Printf("GetASN error: %s", err)
+		}
+		if len(asn) > 0 {
+			targets = append(targets, asn)
+		}
+	}
+
+	var country, continent, region, regionGroup string
+	var netmask int
+	var location *geo.Location
+
+	if t&TargetRegion > 0 || t&TargetRegionGroup > 0 || hasClosest {
+		var err error
+		location, err = g.GetLocation(ip)
+		if location == nil || err != nil {
+			return targets, 0, nil
+		}
+		log.Printf("Location for '%s' (%s): %+v", ip, err, location)
+		country = location.Country
+		continent = location.Continent
+		region = location.Region
+		regionGroup = location.RegionGroup
+		// continent, regionGroup, region, netmask,
+
+	} else if t&TargetCountry > 0 || t&TargetContinent > 0 {
+		country, continent, netmask = g.GetCountry(ip)
 	}
 
 	if t&TargetRegion > 0 && len(region) > 0 {

targeting/targeting_test.go

@@ -4,6 +4,8 @@ import (
 	"net"
 	"reflect"
 	"testing"
+
+	"github.com/abh/geodns/targeting/geoip2"
 )
 
 func TestTargetString(t *testing.T) {
@@ -50,28 +52,36 @@ func TestTargetParse(t *testing.T) {
 func TestGetTargets(t *testing.T) {
 	ip := net.ParseIP("207.171.1.1")
 
-	GeoIP().SetDirectory("../db")
+	g, err := geoip2.New("../db")
+	if err != nil {
+		t.Fatalf("opening geoip2: %s", err)
+	}
+	Setup(g)
 
-	GeoIP().SetupGeoIPCity()
-	GeoIP().SetupGeoIPCountry()
-	GeoIP().SetupGeoIPASN()
+	// GeoIP().SetDirectory("../db")
+	// GeoIP().SetupGeoIPCity()
+	// GeoIP().SetupGeoIPCountry()
+	// GeoIP().SetupGeoIPASN()
 
 	tgt, _ := ParseTargets("@ continent country")
 	targets, _, _ := tgt.GetTargets(ip, false)
-	if !reflect.DeepEqual(targets, []string{"us", "north-america", "@"}) {
-		t.Fatalf("Unexpected parse results of targets")
+	expect := []string{"us", "north-america", "@"}
+	if !reflect.DeepEqual(targets, expect) {
+		t.Fatalf("Unexpected parse results of targets, got '%s', expected '%s'", targets, expect)
 	}
 
-	if geoIP.city == nil {
+	if !g.HasLocation() {
 		t.Log("City GeoIP database requred for these tests")
 		return
 	}
 
-	tests := []struct {
+	type test struct {
 		Str     string
 		Targets []string
 		IP      string
-	}{
+	}
+
+	tests := []test{
 		{
 			"@ continent country region ",
 			[]string{"us-ca", "us", "north-america", "@"},
@@ -82,11 +92,6 @@ func TestGetTargets(t *testing.T) {
 			[]string{"us-ca", "us-west", "us", "north-america", "@"},
 			"",
 		},
-		{
-			"@ continent regiongroup country region asn ip",
-			[]string{"[207.171.1.1]", "[207.171.1.0]", "as7012", "us-ca", "us-west", "us", "north-america", "@"},
-			"",
-		},
 		{
 			"ip",
 			[]string{"[2607:f238:2::ff:4]", "[2607:f238:2::]"},
@@ -94,6 +99,14 @@ func TestGetTargets(t *testing.T) {
 		},
 	}
 
+	if g.HasASN() {
+		tests = append(tests,
+			test{"@ continent regiongroup country region asn ip",
+				[]string{"[207.171.1.1]", "[207.171.1.0]", "as7012", "us-ca", "us-west", "us", "north-america", "@"},
+				"",
+			})
+	}
+
 	for _, test := range tests {
 		if len(test.IP) > 0 {
 			ip = net.ParseIP(test.IP)
@@ -103,7 +116,7 @@ func TestGetTargets(t *testing.T) {
 		targets, _, _ = tgt.GetTargets(ip, false)
 
 		if !reflect.DeepEqual(targets, test.Targets) {
-			t.Logf("For targets '%s' expected '%s', got '%s'", test.Str, test.Targets, targets)
+			t.Logf("For IP '%s' targets '%s' expected '%s', got '%s'", ip, test.Str, test.Targets, targets)
 			t.Fail()
 		}
 

zones/muxmanager.go

@@ -60,7 +60,7 @@ func (mm *MuxManager) Run() {
 		if err != nil {
 			log.Printf("error reading zones: %s", err)
 		}
-		time.Sleep(5 * time.Second)
+		time.Sleep(2 * time.Second)
 	}
 }
 

zones/picker.go

@@ -4,7 +4,7 @@ import (
 	"math/rand"
 
 	"github.com/abh/geodns/health"
-	"github.com/abh/geodns/targeting"
+	"github.com/abh/geodns/targeting/geo"
 
 	"github.com/miekg/dns"
 )
@@ -23,7 +23,7 @@ func (zone *Zone) filterHealth(servers Records) (Records, int) {
 	return tmpServers, sum
 }
 
-func (zone *Zone) Picker(label *Label, qtype uint16, max int, location *targeting.Location) Records {
+func (zone *Zone) Picker(label *Label, qtype uint16, max int, location *geo.Location) Records {
 
 	if qtype == dns.TypeANY {
 		var result Records

zones/reader.go

@@ -116,14 +116,30 @@ func (zone *Zone) ReadZoneFile(fileName string) (zerr error) {
 
 	//log.Println("IP", string(Zone.Regions["0.us"].IPv4[0].ip))
 
+	if zone.Options.Targeting == 0 && !zone.HasClosest {
+		// no targeting requested
+		return nil
+	}
+
+	if targeting.Geo() == nil {
+		log.Printf("'%s': No geo provider configured", zone.Origin)
+		return nil
+	}
+
 	switch {
 	case zone.Options.Targeting >= targeting.TargetRegionGroup || zone.HasClosest:
-		targeting.GeoIP().SetupGeoIPCity()
+		if ok, err := targeting.Geo().HasLocation(); !ok {
+			log.Printf("Zone '%s' requested location/city targeting but geo provider isn't available: %s", zone.Origin, err)
+		}
 	case zone.Options.Targeting >= targeting.TargetContinent:
-		targeting.GeoIP().SetupGeoIPCountry()
+		if ok, err := targeting.Geo().HasCountry(); !ok {
+			log.Printf("Zone '%s' requested country targeting but geo provider isn't available: %s", zone.Origin, err)
+		}
 	}
 	if zone.Options.Targeting&targeting.TargetASN > 0 {
-		targeting.GeoIP().SetupGeoIPASN()
+		if ok, err := targeting.Geo().HasASN(); !ok {
+			log.Printf("Zone '%s' requested ASN targeting but geo provider isn't available: %s", zone.Origin, err)
+		}
 	}
 
 	if zone.HasClosest {

zones/zone.go

@@ -10,6 +10,7 @@ import (
 	"github.com/abh/geodns/applog"
 	"github.com/abh/geodns/health"
 	"github.com/abh/geodns/targeting"
+	"github.com/abh/geodns/targeting/geo"
 
 	"github.com/miekg/dns"
 	"github.com/rcrowley/go-metrics"
@@ -36,7 +37,7 @@ type ZoneLogging struct {
 type Record struct {
 	RR     dns.RR
 	Weight int
-	Loc    *targeting.Location
+	Loc    *geo.Location
 	Test   string
 }
 
@@ -279,24 +280,25 @@ func (z *Zone) FindLabels(s string, targets []string, qts []uint16) []LabelMatch
 // Find the locations of all the A records within a zone. If we were being really clever
 // here we could use LOC records too. But for the time being we'll just use GeoIP
 func (z *Zone) SetLocations() {
-	qtypes := []uint16{dns.TypeA}
-	for _, label := range z.Labels {
-		if label.Closest {
-			for _, qtype := range qtypes {
-				if label.Records[qtype] != nil && len(label.Records[qtype]) > 0 {
-					for i := range label.Records[qtype] {
-						label.Records[qtype][i].Loc = nil
-						rr := label.Records[qtype][i].RR
-						if a, ok := rr.(*dns.A); ok {
-							ip := a.A
-							_, _, _, _, _, location := targeting.GeoIP().GetCountryRegion(ip)
-							label.Records[qtype][i].Loc = location
-						}
-					}
-				}
-			}
-		}
-	}
+	log.Println("todo: SetLocations()")
+	// qtypes := []uint16{dns.TypeA}
+	// for _, label := range z.Labels {
+	// 	if label.Closest {
+	// 		for _, qtype := range qtypes {
+	// 			if label.Records[qtype] != nil && len(label.Records[qtype]) > 0 {
+	// 				for i := range label.Records[qtype] {
+	// 					label.Records[qtype][i].Loc = nil
+	// 					rr := label.Records[qtype][i].RR
+	// 					if a, ok := rr.(*dns.A); ok {
+	// 						ip := a.A
+	// 						_, _, _, _, _, location := targeting.GeoIP().GetCountryRegion(ip)
+	// 						label.Records[qtype][i].Loc = location
+	// 					}
+	// 				}
+	// 			}
+	// 		}
+	// 	}
+	// }
 }
 
 func (z *Zone) addHealthReference(l *Label, data interface{}) {