Merge pull request #207 from rod-hynes/master

Support multiple GeoIP database files

psiphon/server/config.go

@@ -90,10 +90,13 @@ type Config struct {
 	// used to determine a unique discovery strategy.
 	DiscoveryValueHMACKey string
 
-	// GeoIPDatabaseFilename is the path of the GeoIP2/GeoLite2
-	// MaxMind database file. when blank, no GeoIP lookups are
-	// performed.
-	GeoIPDatabaseFilename string
+	// GeoIPDatabaseFilenames ares paths of GeoIP2/GeoLite2
+	// MaxMind database files. When empty, no GeoIP lookups are
+	// performed. Each file is queried, in order, for the
+	// logged fields: country code, city, and ISP. Multiple
+	// file support accomodates the MaxMind distribution where
+	// ISP data in a separate file.
+	GeoIPDatabaseFilenames []string
 
 	// PsinetDatabaseFilename is the path of the Psiphon automation
 	// jsonpickle format Psiphon API data file.
@@ -462,7 +465,7 @@ func GenerateConfig(params *GenerateConfigParams) ([]byte, []byte, []byte, error
 	config := &Config{
 		LogLevel:                       "info",
 		Fail2BanFormat:                 "Authentication failure for psiphon-client from %s",
-		GeoIPDatabaseFilename:          "",
+		GeoIPDatabaseFilenames:         nil,
 		HostID:                         "example-host-id",
 		ServerIPAddress:                params.ServerIPAddress,
 		DiscoveryValueHMACKey:          discoveryValueHMACKey,

psiphon/server/geoip.go

@@ -60,53 +60,74 @@ func NewGeoIPData() GeoIPData {
 // supports hot reloading of MaxMind data while the server is
 // running.
 type GeoIPService struct {
-	psiphon.ReloadableFile
-	maxMindReader         *maxminddb.Reader
+	databases             []*geoIPDatabase
 	sessionCache          *cache.Cache
 	discoveryValueHMACKey string
 }
 
+type geoIPDatabase struct {
+	psiphon.ReloadableFile
+	maxMindReader *maxminddb.Reader
+}
+
 // NewGeoIPService initializes a new GeoIPService.
-func NewGeoIPService(filename, discoveryValueHMACKey string) (*GeoIPService, error) {
+func NewGeoIPService(
+	databaseFilenames []string,
+	discoveryValueHMACKey string) (*GeoIPService, error) {
 
 	geoIP := &GeoIPService{
-		maxMindReader:         nil,
+		databases:             make([]*geoIPDatabase, len(databaseFilenames)),
 		sessionCache:          cache.New(GEOIP_SESSION_CACHE_TTL, 1*time.Minute),
 		discoveryValueHMACKey: discoveryValueHMACKey,
 	}
 
-	geoIP.ReloadableFile = psiphon.NewReloadableFile(
-		"geoip database",
-		filename,
-		func(filename string) error {
-			maxMindReader, err := maxminddb.Open(filename)
-			if err != nil {
-				// On error, geoIP state remains the same
-				return psiphon.ContextError(err)
-			}
-			geoIP.maxMindReader = maxMindReader
-			return nil
-		})
-
-	_, err := geoIP.Reload()
-	if err != nil {
-		return nil, psiphon.ContextError(err)
+	for i, filename := range databaseFilenames {
+
+		database := &geoIPDatabase{}
+		database.ReloadableFile = psiphon.NewReloadableFile(
+			filename,
+			func(filename string) error {
+				maxMindReader, err := maxminddb.Open(filename)
+				if err != nil {
+					// On error, database state remains the same
+					return psiphon.ContextError(err)
+				}
+				if database.maxMindReader != nil {
+					database.maxMindReader.Close()
+				}
+				database.maxMindReader = maxMindReader
+				return nil
+			})
+
+		_, err := database.Reload()
+		if err != nil {
+			return nil, psiphon.ContextError(err)
+		}
+
+		geoIP.databases[i] = database
 	}
 
-	return geoIP, err
+	return geoIP, nil
+}
+
+// Reloaders gets the list of reloadable databases in use
+// by the GeoIPService. This list is used to hot reload
+// these databases.
+func (geoIP *GeoIPService) Reloaders() []psiphon.Reloader {
+	reloaders := make([]psiphon.Reloader, len(geoIP.databases))
+	for i, database := range geoIP.databases {
+		reloaders[i] = database
+	}
+	return reloaders
 }
 
 // Lookup determines a GeoIPData for a given client IP address.
 func (geoIP *GeoIPService) Lookup(ipAddress string) GeoIPData {
-	geoIP.ReloadableFile.RLock()
-	defer geoIP.ReloadableFile.RUnlock()
-
 	result := NewGeoIPData()
 
 	ip := net.ParseIP(ipAddress)
 
-	// Note: maxMindReader is nil when config.GeoIPDatabaseFilename is blank.
-	if ip == nil || geoIP.maxMindReader == nil {
+	if ip == nil || len(geoIP.databases) == 0 {
 		return result
 	}
 
@@ -120,9 +141,16 @@ func (geoIP *GeoIPService) Lookup(ipAddress string) GeoIPData {
 		ISP string `maxminddb:"isp"`
 	}
 
-	err := geoIP.maxMindReader.Lookup(ip, &geoIPFields)
-	if err != nil {
-		log.WithContextFields(LogFields{"error": err}).Warning("GeoIP lookup failed")
+	// Each database will populate geoIPFields with the values it contains. In the
+	// currnt MaxMind deployment, the City database populates Country and City and
+	// the separate ISP database populates ISP.
+	for _, database := range geoIP.databases {
+		database.ReloadableFile.RLock()
+		err := database.maxMindReader.Lookup(ip, &geoIPFields)
+		database.ReloadableFile.RUnlock()
+		if err != nil {
+			log.WithContextFields(LogFields{"error": err}).Warning("GeoIP lookup failed")
+		}
 	}
 
 	if geoIPFields.Country.ISOCode != "" {

psiphon/server/psinet/psinet.go

@@ -131,7 +131,6 @@ func NewDatabase(filename string) (*Database, error) {
 	database := &Database{}
 
 	database.ReloadableFile = psiphon.NewReloadableFile(
-		"psinet database",
 		filename,
 		func(filename string) error {
 			psinetJSON, err := ioutil.ReadFile(filename)

psiphon/server/services.go

@@ -196,7 +196,7 @@ func NewSupportServices(config *Config) (*SupportServices, error) {
 	}
 
 	geoIPService, err := NewGeoIPService(
-		config.GeoIPDatabaseFilename, config.DiscoveryValueHMACKey)
+		config.GeoIPDatabaseFilenames, config.DiscoveryValueHMACKey)
 	if err != nil {
 		return nil, psiphon.ContextError(err)
 	}
@@ -217,10 +217,9 @@ func NewSupportServices(config *Config) (*SupportServices, error) {
 // established clients.
 func (support *SupportServices) Reload() {
 
-	reloaders := []psiphon.Reloader{
-		support.TrafficRulesSet,
-		support.PsinetDatabase,
-		support.GeoIPService}
+	reloaders := append(
+		[]psiphon.Reloader{support.TrafficRulesSet, support.PsinetDatabase},
+		support.GeoIPService.Reloaders()...)
 
 	for _, reloader := range reloaders {
 

psiphon/server/trafficRules.go

@@ -131,7 +131,6 @@ func NewTrafficRulesSet(filename string) (*TrafficRulesSet, error) {
 	set := &TrafficRulesSet{}
 
 	set.ReloadableFile = psiphon.NewReloadableFile(
-		"traffic rules set",
 		filename,
 		func(filename string) error {
 			configJSON, err := ioutil.ReadFile(filename)

psiphon/utils.go

@@ -334,21 +334,19 @@ type Reloader interface {
 //
 type ReloadableFile struct {
 	sync.RWMutex
-	logDescription string
-	fileName       string
-	fileInfo       os.FileInfo
-	reloadAction   func(string) error
+	fileName     string
+	fileInfo     os.FileInfo
+	reloadAction func(string) error
 }
 
 // NewReloadableFile initializes a new ReloadableFile
 func NewReloadableFile(
-	logDescription, fileName string,
+	fileName string,
 	reloadAction func(string) error) ReloadableFile {
 
 	return ReloadableFile{
-		logDescription: logDescription,
-		fileName:       fileName,
-		reloadAction:   reloadAction,
+		fileName:     fileName,
+		reloadAction: reloadAction,
 	}
 }
 
@@ -391,5 +389,5 @@ func (reloadable *ReloadableFile) Reload() (bool, error) {
 }
 
 func (reloadable *ReloadableFile) LogDescription() string {
-	return reloadable.logDescription
+	return reloadable.fileName
 }