Support targeting by region/state with GeoIPCity

See also issue #39

.travis.yml

@@ -2,6 +2,8 @@ language: go
 before_install:
   - sudo apt-get install libgeoip-dev bzr
 install:
+  - mkdir -p $TRAVIS_BUILD_DIR/db
+  - curl -s http://geodns.bitnames.com/geoip/GeoLiteCity.dat.gz  | gzip -cd > $TRAVIS_BUILD_DIR/db/GeoIPCity.dat
   - go get github.com/miekg/dns
   - go get github.com/abh/geoip
   - go get launchpad.net/gocheck

config.go

@@ -16,6 +16,9 @@ type AppConfig struct {
 	Flags struct {
 		HasStatHat bool
 	}
+	GeoIP struct {
+		Directory string
+	}
 }
 
 var Config = new(AppConfig)

countries/regiongroups.go

@@ -0,0 +1,74 @@
+package countries
+
+import (
+	"log"
+)
+
+func CountryRegionGroup(country, region string) string {
+
+	if country != "us" {
+		return ""
+	}
+
+	regions := map[string]string{
+		"us-ak": "us-west",
+		"us-az": "us-west",
+		"us-ca": "us-west",
+		"us-co": "us-west",
+		"us-hi": "us-west",
+		"us-id": "us-west",
+		"us-mt": "us-west",
+		"us-nm": "us-west",
+		"us-nv": "us-west",
+		"us-or": "us-west",
+		"us-ut": "us-west",
+		"us-wa": "us-west",
+		"us-wy": "us-west",
+
+		"us-ar": "us-central",
+		"us-ia": "us-central",
+		"us-in": "us-central",
+		"us-ks": "us-central",
+		"us-la": "us-central",
+		"us-mn": "us-central",
+		"us-mo": "us-central",
+		"us-nd": "us-central",
+		"us-ne": "us-central",
+		"us-ok": "us-central",
+		"us-sd": "us-central",
+		"us-tx": "us-central",
+		"us-wi": "us-central",
+
+		"us-al": "us-east",
+		"us-ct": "us-east",
+		"us-dc": "us-east",
+		"us-de": "us-east",
+		"us-fl": "us-east",
+		"us-ga": "us-east",
+		"us-ky": "us-east",
+		"us-ma": "us-east",
+		"us-md": "us-east",
+		"us-me": "us-east",
+		"us-mi": "us-east",
+		"us-ms": "us-east",
+		"us-nc": "us-east",
+		"us-nh": "us-east",
+		"us-nj": "us-east",
+		"us-ny": "us-east",
+		"us-oh": "us-east",
+		"us-pa": "us-east",
+		"us-ri": "us-east",
+		"us-sc": "us-east",
+		"us-tn": "us-east",
+		"us-va": "us-east",
+		"us-vt": "us-east",
+		"us-wv": "us-east",
+	}
+
+	if group, ok := regions[region]; ok {
+		return group
+	}
+
+	log.Printf("Did not find a region group for '%s'/'%s'", country, region)
+	return ""
+}

dns/geodns.conf.sample

@@ -3,6 +3,10 @@
 ; It is recommended to distribute the configuration file globally
 ; with your .json zone files.
 
+[geoip]
+;; Directory containing the GeoIP .dat database files
+;directory=/usr/local/share/GeoIP/
+
 [stathat]
 ;; Add an API key to send query counts and other metrics to stathat
 ;apikey=abc123

geoip.go

@@ -1,16 +1,101 @@
 package main
 
 import (
+	"github.com/abh/geodns/countries"
 	"github.com/abh/geoip"
 	"log"
+	"net"
+	"strings"
+	"time"
 )
 
-func setupGeoIP() *geoip.GeoIP {
+type GeoIP struct {
+	country         *geoip.GeoIP
+	hasCountry      bool
+	countryLastLoad time.Time
 
-	gi, err := geoip.Open()
+	city         *geoip.GeoIP
+	cityLastLoad time.Time
+	hasCity      bool
+}
+
+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())
+
+	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) 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 GeoIP database: %s\n", err)
-		return nil
+		log.Printf("Could not open country GeoIP database: %s\n", err)
+		return
 	}
-	return gi
+	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.countryLastLoad = time.Now()
+	g.hasCity = true
+	g.city = gi
+
 }

serve.go

@@ -4,7 +4,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"github.com/abh/dns"
-	"github.com/abh/geodns/countries"
 	"log"
 	"net"
 	"os"
@@ -19,8 +18,6 @@ func getQuestionName(z *Zone, req *dns.Msg) string {
 	return strings.ToLower(strings.Join(ql, "."))
 }
 
-var geoIP = setupGeoIP()
-
 func serve(w dns.ResponseWriter, req *dns.Msg, z *Zone) {
 
 	qtype := req.Question[0].Qtype
@@ -41,9 +38,10 @@ func serve(w dns.ResponseWriter, req *dns.Msg, z *Zone) {
 	z.Metrics.LabelStats.Add(label)
 
 	realIp, _, _ := net.SplitHostPort(w.RemoteAddr().String())
+
 	z.Metrics.ClientStats.Add(realIp)
 
-	var ip string // EDNS or real IP
+	var ip net.IP // EDNS or real IP
 	var edns *dns.EDNS0_SUBNET
 	var opt_rr *dns.OPT
 
@@ -61,32 +59,18 @@ func serve(w dns.ResponseWriter, req *dns.Msg, z *Zone) {
 					logPrintln("Got edns", e.Address, e.Family, e.SourceNetmask, e.SourceScope)
 					if e.Address != nil {
 						edns = e
-						ip = e.Address.String()
+						ip = e.Address
 					}
 				}
 			}
 		}
 	}
 
 	if len(ip) == 0 { // no edns subnet
-		ip = realIp
+		ip = net.ParseIP(realIp)
 	}
 
-	var targets []string
-	var country string
-	var netmask int
-	if geoIP != nil {
-		country, netmask = geoIP.GetCountry(ip)
-		country = strings.ToLower(country)
-		if len(country) > 0 {
-			targets = append(targets, country)
-			continent := countries.CountryContinent[country]
-			if len(continent) > 0 {
-				targets = append(targets, continent)
-			}
-		}
-		targets = append(targets, "@")
-	}
+	targets, netmask := z.Options.Targeting.GetTargets(ip)
 
 	m := new(dns.Msg)
 	m.SetReply(req)
@@ -131,13 +115,17 @@ func serve(w dns.ResponseWriter, req *dns.Msg, z *Zone) {
 				h := dns.RR_Header{Ttl: 1, Class: dns.ClassINET, Rrtype: dns.TypeTXT}
 				h.Name = label + "." + z.Origin + "."
 
+				txt := []string{
+					w.RemoteAddr().String(),
+					ip.String(),
+				}
+
+				targets, netmask := z.Options.Targeting.GetTargets(ip)
+				txt = append(txt, strings.Join(targets, " "))
+				txt = append(txt, fmt.Sprintf("/%d", netmask))
+
 				m.Answer = []dns.RR{&dns.TXT{Hdr: h,
-					Txt: []string{
-						w.RemoteAddr().String(),
-						ip,
-						string(country),
-						string(countries.CountryContinent[country]),
-					},
+					Txt: txt,
 				}}
 			} else {
 				m.Ns = append(m.Ns, z.SoaRR())

targeting.go

@@ -2,8 +2,8 @@ package main
 
 import (
 	"fmt"
+	"net"
 	"strings"
-	// "github.com/abh/geodns/countries"
 )
 
 type TargetOptions int
@@ -16,6 +16,43 @@ const (
 	TargetRegion
 )
 
+func (t TargetOptions) GetTargets(ip net.IP) ([]string, int) {
+
+	targets := make([]string, 0)
+
+	var country, continent string
+	var netmask int
+
+	switch {
+	case t >= TargetRegionGroup:
+		var region, regionGroup string
+		country, continent, regionGroup, region, netmask = geoIP.GetCountryRegion(ip)
+		if t&TargetRegion > 0 && len(region) > 0 {
+			targets = append(targets, region)
+		}
+		if t&TargetRegionGroup > 0 && len(regionGroup) > 0 {
+			targets = append(targets, regionGroup)
+		}
+
+	case t >= TargetContinent:
+		country, continent, netmask = geoIP.GetCountry(ip)
+	}
+
+	if len(country) > 0 {
+		if t&TargetCountry > 0 {
+			targets = append(targets, country)
+		}
+		if t&TargetContinent > 0 && len(continent) > 0 {
+			targets = append(targets, continent)
+		}
+	}
+
+	if t&TargetGlobal > 0 {
+		targets = append(targets, "@")
+	}
+	return targets, netmask
+}
+
 func (t TargetOptions) String() string {
 	targets := make([]string, 0)
 	if t&TargetGlobal > 0 {

targeting_test.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	. "launchpad.net/gocheck"
+	"net"
 )
 
 type TargetingSuite struct {
@@ -10,6 +11,7 @@ type TargetingSuite struct {
 var _ = Suite(&TargetingSuite{})
 
 func (s *TargetingSuite) SetUpSuite(c *C) {
+	Config.GeoIP.Directory = "db"
 }
 
 func (s *TargetingSuite) TestTargetString(c *C) {
@@ -32,3 +34,28 @@ func (s *TargetingSuite) TestTargetParse(c *C) {
 	str = tgt.String()
 	c.Check(str, Equals, "@ continent country")
 }
+func (s *TargetingSuite) TestGetTargets(c *C) {
+
+	ip := net.ParseIP("207.171.7.51")
+
+	geoIP.setupGeoIPCity()
+	geoIP.setupGeoIPCountry()
+
+	tgt, _ := parseTargets("@ continent country")
+	targets, _ := tgt.GetTargets(ip)
+	c.Check(targets, DeepEquals, []string{"us", "north-america", "@"})
+
+	if geoIP.city == nil {
+		c.Log("City GeoIP database requred for these tests")
+		return
+	}
+
+	tgt, _ = parseTargets("@ continent country region ")
+	targets, _ = tgt.GetTargets(ip)
+	c.Check(targets, DeepEquals, []string{"us-ca", "us", "north-america", "@"})
+
+	tgt, _ = parseTargets("@ continent regiongroup country region ")
+	targets, _ = tgt.GetTargets(ip)
+	c.Check(targets, DeepEquals, []string{"us-ca", "us-west", "us", "north-america", "@"})
+
+}