From 12eb59f6118301c58c19a11318a2bdd0129d1802 Mon Sep 17 00:00:00 2001
From: Tim Cooper <tim.cooper@layeh.com>
Date: Tue, 12 Oct 2021 15:21:37 -0500
Subject: [PATCH] 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.
---
 core/chat/server.go |  5 ++-
 geoip/geoip.go      | 87 +++++++++++++++++++++++----------------------
 2 files changed, 48 insertions(+), 44 deletions(-)

diff --git a/core/chat/server.go b/core/chat/server.go
index 76ee44059..37c7d9fd9 100644
--- a/core/chat/server.go
+++ b/core/chat/server.go
@@ -38,6 +38,8 @@ type Server struct {
 
 	// unregister requests from clients.
 	unregister chan uint // the ChatClient id
+
+	geoipClient *geoip.Client
 }
 
 // NewChat will return a new instance of the chat server.
@@ -51,6 +53,7 @@ func NewChat() *Server {
 		inbound:                  make(chan chatClientEvent),
 		unregister:               make(chan uint),
 		maxSocketConnectionLimit: maximumConcurrentConnectionLimit,
+		geoipClient:              geoip.NewClient(),
 	}
 
 	return server
@@ -117,7 +120,7 @@ func (s *Server) Addclient(conn *websocket.Conn, user *user.User, accessToken st
 
 	// Asynchronously, optionally, fetch GeoIP data.
 	go func(client *Client) {
-		client.Geo = geoip.GetGeoFromIP(ipAddress)
+		client.Geo = s.geoipClient.GetGeoFromIP(ipAddress)
 	}(client)
 
 	return client
diff --git a/geoip/geoip.go b/geoip/geoip.go
index 0f02d2611..575af7b39 100644
--- a/geoip/geoip.go
+++ b/geoip/geoip.go
@@ -6,14 +6,27 @@ package geoip
 
 import (
 	"net"
+	"sync"
+	"sync/atomic"
 
 	"github.com/oschwald/geoip2-golang"
 	log "github.com/sirupsen/logrus"
 )
 
-var _geoIPCache = map[string]GeoDetails{}
-var _enabled = true // Try to use GeoIP support it by default.
-var geoIPDatabasePath = "data/GeoLite2-City.mmdb"
+const 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.
 type GeoDetails struct {
@@ -24,9 +37,9 @@ type GeoDetails struct {
 
 // GetGeoFromIP returns geo details associated with an IP address if we
 // have previously fetched it.
-func GetGeoFromIP(ip string) *GeoDetails {
-	if cachedGeoDetails, ok := _geoIPCache[ip]; ok {
-		return &cachedGeoDetails
+func (c *Client) GetGeoFromIP(ip string) *GeoDetails {
+	if cachedGeoDetails, ok := c.cache.Load(ip); ok {
+		return cachedGeoDetails.(*GeoDetails)
 	}
 
 	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.
-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 !_enabled {
+	if atomic.LoadInt32(&c.enabled) == 0 {
 		return nil
 	}
 
-	// Don't re-fetch if we already have it.
-	if geoDetails, ok := _geoIPCache[ip]; ok {
-		return &geoDetails
-	}
-
 	db, err := geoip2.Open(geoIPDatabasePath)
 	if err != nil {
-		log.Traceln("GeoIP support is disabled. visit http://owncast.online/docs/geoip to learn how to enable.", err)
-		_enabled = false
+		log.Traceln("GeoIP support is disabled. visit https://owncast.online/docs/geoip to learn how to enable.", err)
+		atomic.StoreInt32(&c.enabled, 0)
 		return nil
 	}
-
 	defer db.Close()
 
+	var response *GeoDetails
 	ipObject := net.ParseIP(ip)
 
 	record, err := db.City(ipObject)
-	if err != nil {
-		log.Warnln(err)
-		return nil
-	}
+	if err == nil {
+		// If no country is available then exit
+		// 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
-	if record.Country.IsoCode == "" {
-		return nil
-	}
-
-	// 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
+			response = &GeoDetails{
+				CountryCode: record.Country.IsoCode,
+				RegionName:  regionName,
+				TimeZone:    record.Location.TimeZone,
+			}
 		}
+	} else {
+		log.Warnln(err)
 	}
 
-	response := GeoDetails{
-		CountryCode: record.Country.IsoCode,
-		RegionName:  regionName,
-		TimeZone:    record.Location.TimeZone,
-	}
+	c.cache.Store(ip, response)
 
-	_geoIPCache[ip] = response
-
-	return &response
+	return response
 }