geoip2: watch and update

http_test.go

@@ -15,7 +15,6 @@ import (
 )
 
 func TestHTTP(t *testing.T) {
-
  geoprovider, err := geoip2.New(geoip2.FindDB())
  if err == nil {
  targeting.Setup(geoprovider)
@@ -43,5 +42,4 @@ func TestHTTP(t *testing.T) {
  t.Log("/version didn't start with 'GeoDNS '")
  t.Fail()
  }
-
 }

targeting/geoip2/geoip2.go

@@ -2,45 +2,35 @@ package geoip2
 
 import (
  "fmt"
+ "io/fs"
  "log"
  "net"
  "os"
  "path/filepath"
  "strings"
  "sync"
+ "time"
 
  "github.com/abh/geodns/countries"
  "github.com/abh/geodns/targeting/geo"
- geoip2 "github.com/oschwald/geoip2-golang"
+ gdb "github.com/oschwald/geoip2-golang"
 )
 
-type geoType uint8
-
-const (
- countryDB = iota
- cityDB
- asnDB
-)
-
-var dbFiles map[geoType][]string
-
 // GeoIP2 contains the geoip implementation of the GeoDNS geo
 // targeting interface
 type GeoIP2 struct {
- dir string
-
- country *geoip2.Reader
- city *geoip2.Reader
- asn *geoip2.Reader
- mu sync.RWMutex
+ dir string
+ country geodb
+ city geodb
+ asn geodb
 }
 
-func init() {
- dbFiles = map[geoType][]string{
-  countryDB: []string{"GeoIP2-Country.mmdb", "GeoLite2-Country.mmdb"},
-  asnDB: []string{"GeoIP2-ASN.mmdb", "GeoLite2-ASN.mmdb"},
-  cityDB: []string{"GeoIP2-City.mmdb", "GeoLite2-City.mmdb"},
- }
+type geodb struct {
+ active bool
+ lastModified int64 // Epoch time
+ fp  string // FilePath
+ db  *gdb.Reader // Database reader
+ l sync.RWMutex // Individual lock for separate DB access and reload -- Future?
 }
 
 // FindDB returns a guess at a directory path for GeoIP data files
@@ -63,97 +53,119 @@ func FindDB() string {
  return ""
 }
 
-func (g *GeoIP2) open(t geoType, db string) (*geoip2.Reader, error) {
-
- fileName := filepath.Join(g.dir, db)
-
- if len(db) == 0 {
- found := false
- for _, f := range dbFiles[t] {
- fileName = filepath.Join(g.dir, f)
- if _, err := os.Stat(fileName); err == nil {
-  found = true
- break
+// open will create a filehandle for the provided GeoIP2 database. If opened once before and a newer modification time is present, the function will reopen the file with its new contents
+func (g *GeoIP2) open(v *geodb, fns ...string) error {
+ var fi fs.FileInfo
+ var err error
+ if v.fp == "" {
+ // We're opening this file for the first time
+ for _, i := range fns {
+ fp := filepath.Join(g.dir, i)
+ fi, err = os.Stat(fp)
+ if err != nil {
+ continue
  }
- }
- if !found {
- return nil, fmt.Errorf("could not find '%s' in '%s'", dbFiles[t], g.dir)
+ v.fp = fp
  }
  }
-
- n, err := geoip2.Open(fileName)
+ if v.fp == "" { // Recheck for empty string in case none of the DB files are found
+ return fmt.Errorf("no files found for db")
+ }
+ if fi == nil { // We have not set fileInfo and v.fp is set
+ fi, err = os.Stat(v.fp)
+ }
  if err != nil {
- return nil, err
+ return 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
+ if v.lastModified >= fi.ModTime().UTC().Unix() { // No update to existing file
+ return nil
  }
- return n, nil
-}
-
-func (g *GeoIP2) get(t geoType, db string) (*geoip2.Reader, error) {
- g.mu.RLock()
+ // Delay the lock to here because we're only
+ v.l.Lock()
+ defer v.l.Unlock()
 
- var r *geoip2.Reader
-
- switch t {
- case countryDB:
- r = g.country
- case cityDB:
- r = g.city
- case asnDB:
- r = g.asn
+ o, e := gdb.Open(v.fp)
+ if e != nil {
+ return e
  }
+ v.db = o
+ v.active = true
+ v.lastModified = fi.ModTime().UTC().Unix()
 
- // unlock so the g.open() call below won't lock
- g.mu.RUnlock()
+ return nil
+}
 
- if r != nil {
- return r, nil
+// watchFiles spawns a goroutine to check for new files every minute, reloading if the modtime is newer than the original file's modtime
+func (g *GeoIP2) watchFiles() {
+ // Not worried about goroutines leaking because only one geoip2.New call is made in main (outside of testing)
+ ticker := time.NewTicker(1 * time.Minute)
+ for { // We forever-loop here because we only run this function in a separate goroutine
+ select {
+ case <-ticker.C:
+ // Iterate through each db, check modtime. If new, reload file
+ cityErr := g.open(&g.city, "GeoIP2-City.mmdb", "GeoLite2-City.mmdb")
+ if cityErr != nil {
+ log.Printf("Failed to update City: %v\n", cityErr)
+ }
+ countryErr := g.open(&g.country, "GeoIP2-Country.mmdb", "GeoLite2-Country.mmdb")
+ if countryErr != nil {
+ log.Printf("failed to update Country: %v\n", countryErr)
+ }
+ asnErr := g.open(&g.asn, "GeoIP2-ASN.mmdb", "GeoLite2-ASN.mmdb")
+ if asnErr != nil {
+ log.Printf("failed to update ASN: %v\n", asnErr)
+ }
+ }
  }
+}
 
- return g.open(t, db)
+func (g *GeoIP2) anyActive() bool {
+ return g.country.active || g.city.active || g.asn.active
 }
 
 // New returns a new GeoIP2 provider
-func New(dir string) (*GeoIP2, error) {
- g := &GeoIP2{
+func New(dir string) (g *GeoIP2, err error) {
+ g = &GeoIP2{
  dir: dir,
  }
- _, err := g.open(countryDB, "")
- if err != nil {
+ // This routine MUST load the database files at least once.
+ cityErr := g.open(&g.city, "GeoIP2-City.mmdb", "GeoLite2-City.mmdb")
+ if cityErr != nil {
+ log.Printf("failed to load City DB: %v\n", cityErr)
+ err = cityErr
+ }
+ countryErr := g.open(&g.country, "GeoIP2-Country.mmdb", "GeoLite2-Country.mmdb")
+ if countryErr != nil {
+ log.Printf("failed to load Country DB: %v\n", countryErr)
+ err = countryErr
+ }
+ asnErr := g.open(&g.asn, "GeoIP2-ASN.mmdb", "GeoLite2-ASN.mmdb")
+ if asnErr != nil {
+ log.Printf("failed to load ASN DB: %v\n", asnErr)
+ err = asnErr
+ }
+ if !g.anyActive() {
  return nil, err
  }
-
- return g, nil
+ go g.watchFiles() // Launch goroutine to load and monitor
+ return
 }
 
 // HasASN returns if we can do ASN lookups
 func (g *GeoIP2) HasASN() (bool, error) {
- r, err := g.get(asnDB, "")
- if r != nil && err == nil {
- return true, nil
- }
- return false, err
+ return g.asn.active, nil
 }
 
 // GetASN returns the ASN for the IP (as a "as123" string) and the netmask
 func (g *GeoIP2) GetASN(ip net.IP) (string, int, error) {
- r, err := g.get(asnDB, "")
- log.Printf("GetASN for %s, got DB? %s", ip, err)
- if err != nil {
- return "", 0, err
+ g.asn.l.RLock()
+ defer g.asn.l.RUnlock()
+
+ if !g.asn.active {
+ return "", 0, fmt.Errorf("ASN db not active")
  }
 
- c, err := r.ASN(ip)
+ c, err := g.asn.db.ASN(ip)
  if err != nil {
  return "", 0, fmt.Errorf("lookup ASN for '%s': %s", ip.String(), err)
  }
@@ -167,17 +179,20 @@ func (g *GeoIP2) GetASN(ip net.IP) (string, int, error) {
 
 // HasCountry checks if the GeoIP country database is available
 func (g *GeoIP2) HasCountry() (bool, error) {
- r, err := g.get(countryDB, "")
- if r != nil && err == nil {
- return true, nil
- }
- return false, err
+ return g.country.active, nil
 }
 
 // GetCountry returns the country, continent and netmask for the given IP
 func (g *GeoIP2) GetCountry(ip net.IP) (country, continent string, netmask int) {
- r, err := g.get(countryDB, "")
- c, err := r.Country(ip)
+ // Need a read-lock because return value of Country is a pointer, not copy of the struct/object
+ g.country.l.RLock()
+ defer g.country.l.RUnlock()
+
+ if !g.country.active {
+ return "", "", 0
+ }
+
+ c, err := g.country.db.Country(ip)
  if err != nil {
  log.Printf("Could not lookup country for '%s': %s", ip.String(), err)
  return "", "", 0
@@ -193,19 +208,22 @@ func (g *GeoIP2) GetCountry(ip net.IP) (country, continent string, netmask int)
  return country, continent, 0
 }
 
-// HasLocation returns if the city database is available to
-// return lat/lon information for an IP
+// HasLocation returns if the city database is available to return lat/lon information for an IP
 func (g *GeoIP2) HasLocation() (bool, error) {
- r, err := g.get(cityDB, "")
- if r != nil && err == nil {
- return true, nil
- }
- return false, err
+ return g.city.active, nil
 }
 
 // GetLocation returns a geo.Location object for the given IP
 func (g *GeoIP2) GetLocation(ip net.IP) (l *geo.Location, err error) {
- c, err := g.city.City(ip)
+ // Need a read-lock because return value of City is a pointer, not copy of the struct/object
+ g.city.l.RLock()
+ defer g.city.l.RUnlock()
+
+ if !g.city.active {
+ return nil, fmt.Errorf("city db not active")
+ }
+
+ c, err := g.city.db.City(ip)
  if err != nil {
  log.Printf("Could not lookup CountryRegion for '%s': %s", ip.String(), err)
  return
@@ -229,5 +247,4 @@ func (g *GeoIP2) GetLocation(ip net.IP) (l *geo.Location, err error) {
  }
 
  return
-
 }

targeting/targeting_test.go

@@ -31,9 +31,9 @@ func TestTargetParse(t *testing.T) {
  }
 
  tests := [][]string{
- []string{"@ continent country asn", "@ continent country asn"},
- []string{"asn country", "country asn"},
- []string{"continent @ country", "@ continent country"},
+ {"@ continent country asn", "@ continent country asn"},
+ {"asn country", "country asn"},
+ {"continent @ country", "@ continent country"},
  }
 
  for _, strs := range tests {
@@ -102,7 +102,8 @@ func TestGetTargets(t *testing.T) {
 
  if ok, _ := g.HasASN(); ok {
  tests = append(tests,
- test{"@ continent regiongroup country region asn ip",
+ test{
+ "@ continent regiongroup country region asn ip",
  []string{"[98.248.0.1]", "[98.248.0.0]", "as7922", "us-ca", "us-west", "us", "north-america", "@"},
  "98.248.0.1",
  },