0

refactor geoip (#1442)

- Introduce a new Client type to remove the global variables from the file
- Use the sync package to prevent race conditions with the cache and
  enabled flag
- Cache results for IPs, even if the result is nil

There are still data races around the client.Geo variable, but that can be
resolved in a future commit.
This commit is contained in:
Tim Cooper 2021-10-12 15:21:37 -05:00 committed by GitHub
parent 01b3489287
commit 12eb59f611
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 48 additions and 44 deletions

View File

@ -38,6 +38,8 @@ type Server struct {
// unregister requests from clients. // unregister requests from clients.
unregister chan uint // the ChatClient id unregister chan uint // the ChatClient id
geoipClient *geoip.Client
} }
// NewChat will return a new instance of the chat server. // NewChat will return a new instance of the chat server.
@ -51,6 +53,7 @@ func NewChat() *Server {
inbound: make(chan chatClientEvent), inbound: make(chan chatClientEvent),
unregister: make(chan uint), unregister: make(chan uint),
maxSocketConnectionLimit: maximumConcurrentConnectionLimit, maxSocketConnectionLimit: maximumConcurrentConnectionLimit,
geoipClient: geoip.NewClient(),
} }
return server return server
@ -117,7 +120,7 @@ func (s *Server) Addclient(conn *websocket.Conn, user *user.User, accessToken st
// Asynchronously, optionally, fetch GeoIP data. // Asynchronously, optionally, fetch GeoIP data.
go func(client *Client) { go func(client *Client) {
client.Geo = geoip.GetGeoFromIP(ipAddress) client.Geo = s.geoipClient.GetGeoFromIP(ipAddress)
}(client) }(client)
return client return client

View File

@ -6,14 +6,27 @@ package geoip
import ( import (
"net" "net"
"sync"
"sync/atomic"
"github.com/oschwald/geoip2-golang" "github.com/oschwald/geoip2-golang"
log "github.com/sirupsen/logrus" log "github.com/sirupsen/logrus"
) )
var _geoIPCache = map[string]GeoDetails{} const geoIPDatabasePath = "data/GeoLite2-City.mmdb"
var _enabled = true // Try to use GeoIP support it by default.
var geoIPDatabasePath = "data/GeoLite2-City.mmdb" // Client can look up geography information for IP addresses.
type Client struct {
cache sync.Map
enabled int32
}
// NewClient creates a new Client.
func NewClient() *Client {
return &Client{
enabled: 1, // Try to use GeoIP support by default.
}
}
// GeoDetails stores details about a location. // GeoDetails stores details about a location.
type GeoDetails struct { type GeoDetails struct {
@ -24,9 +37,9 @@ type GeoDetails struct {
// GetGeoFromIP returns geo details associated with an IP address if we // GetGeoFromIP returns geo details associated with an IP address if we
// have previously fetched it. // have previously fetched it.
func GetGeoFromIP(ip string) *GeoDetails { func (c *Client) GetGeoFromIP(ip string) *GeoDetails {
if cachedGeoDetails, ok := _geoIPCache[ip]; ok { if cachedGeoDetails, ok := c.cache.Load(ip); ok {
return &cachedGeoDetails return cachedGeoDetails.(*GeoDetails)
} }
if ip == "::1" || ip == "127.0.0.1" { if ip == "::1" || ip == "127.0.0.1" {
@ -37,62 +50,50 @@ func GetGeoFromIP(ip string) *GeoDetails {
} }
} }
return fetchGeoForIP(ip) return c.fetchGeoForIP(ip)
} }
// fetchGeoForIP makes an API call to get geo details for an IP address. // fetchGeoForIP makes an API call to get geo details for an IP address.
func fetchGeoForIP(ip string) *GeoDetails { func (c *Client) fetchGeoForIP(ip string) *GeoDetails {
// If GeoIP has been disabled then don't try to access it. // If GeoIP has been disabled then don't try to access it.
if !_enabled { if atomic.LoadInt32(&c.enabled) == 0 {
return nil return nil
} }
// Don't re-fetch if we already have it.
if geoDetails, ok := _geoIPCache[ip]; ok {
return &geoDetails
}
db, err := geoip2.Open(geoIPDatabasePath) db, err := geoip2.Open(geoIPDatabasePath)
if err != nil { if err != nil {
log.Traceln("GeoIP support is disabled. visit http://owncast.online/docs/geoip to learn how to enable.", err) log.Traceln("GeoIP support is disabled. visit https://owncast.online/docs/geoip to learn how to enable.", err)
_enabled = false atomic.StoreInt32(&c.enabled, 0)
return nil return nil
} }
defer db.Close() defer db.Close()
var response *GeoDetails
ipObject := net.ParseIP(ip) ipObject := net.ParseIP(ip)
record, err := db.City(ipObject) record, err := db.City(ipObject)
if err != nil { if err == nil {
log.Warnln(err) // If no country is available then exit
return nil // If we believe this IP to be anonymous then no reason to report it
} if record.Country.IsoCode != "" && !record.Traits.IsAnonymousProxy {
var regionName = "Unknown"
if len(record.Subdivisions) > 0 {
if region, ok := record.Subdivisions[0].Names["en"]; ok {
regionName = region
}
}
// If no country is available then exit response = &GeoDetails{
if record.Country.IsoCode == "" { CountryCode: record.Country.IsoCode,
return nil RegionName: regionName,
} TimeZone: record.Location.TimeZone,
}
// If we believe this IP to be anonymous then no reason to report it
if record.Traits.IsAnonymousProxy {
return nil
}
var regionName = "Unknown"
if len(record.Subdivisions) > 0 {
if region, ok := record.Subdivisions[0].Names["en"]; ok {
regionName = region
} }
} else {
log.Warnln(err)
} }
response := GeoDetails{ c.cache.Store(ip, response)
CountryCode: record.Country.IsoCode,
RegionName: regionName,
TimeZone: record.Location.TimeZone,
}
_geoIPCache[ip] = response return response
return &response
} }