diff --git a/.travis.yml b/.travis.yml index 3d6ff3a6..0ccc0ed8 100644 --- a/.travis.yml +++ b/.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 diff --git a/config.go b/config.go index cf8c1fae..04965f51 100644 --- a/config.go +++ b/config.go @@ -16,6 +16,9 @@ type AppConfig struct { Flags struct { HasStatHat bool } + GeoIP struct { + Directory string + } } var Config = new(AppConfig) diff --git a/countries/regiongroups.go b/countries/regiongroups.go new file mode 100644 index 00000000..1b95d5da --- /dev/null +++ b/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 "" +} diff --git a/dns/geodns.conf.sample b/dns/geodns.conf.sample index 67e29fb6..17773c9e 100644 --- a/dns/geodns.conf.sample +++ b/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 diff --git a/geoip.go b/geoip.go index 4e14d252..47e6a735 100644 --- a/geoip.go +++ b/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 + } diff --git a/serve.go b/serve.go index 428d356f..4760ad96 100644 --- a/serve.go +++ b/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,7 +59,7 @@ 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 } } } @@ -69,24 +67,10 @@ func serve(w dns.ResponseWriter, req *dns.Msg, z *Zone) { } 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()) diff --git a/targeting.go b/targeting.go index 2ad1e415..fe8e6cc3 100644 --- a/targeting.go +++ b/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 { diff --git a/targeting_test.go b/targeting_test.go index a39e3321..8a88f954 100644 --- a/targeting_test.go +++ b/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", "@"}) + +} diff --git a/zones.go b/zones.go index d730b49d..3ac85cc6 100644 --- a/zones.go +++ b/zones.go @@ -205,6 +205,13 @@ func readZoneFile(zoneName, fileName string) (zone *Zone, zerr error) { //log.Println("IP", string(Zone.Regions["0.us"].IPv4[0].ip)) + switch { + case zone.Options.Targeting >= TargetRegionGroup: + geoIP.setupGeoIPCity() + case zone.Options.Targeting >= TargetContinent: + geoIP.setupGeoIPCountry() + } + return zone, nil }