diff --git a/src/main.go b/src/main.go index 670dbbd..e114930 100644 --- a/src/main.go +++ b/src/main.go @@ -46,7 +46,6 @@ import ( "imuslab.com/zoraxy/mod/utils" ) - /* SIGTERM handler, do shutdown sequences before closing */ func SetupCloseHandler() { c := make(chan os.Signal, 2) diff --git a/src/mod/geodb/geodb.go b/src/mod/geodb/geodb.go index 6432b40..314a75e 100644 --- a/src/mod/geodb/geodb.go +++ b/src/mod/geodb/geodb.go @@ -3,6 +3,7 @@ package geodb import ( _ "embed" "net/http" + "sync" "time" "imuslab.com/zoraxy/mod/database" @@ -21,10 +22,10 @@ type Store struct { geotrie *trie geotrieIpv6 *trie sysdb *database.Database - slowLookupCacheIpv4 map[string]string //Cache for slow lookup - slowLookupCacheIpv6 map[string]string //Cache for slow lookup - cacheClearTicker *time.Ticker //Ticker for clearing cache - cacheClearTickerStopChan chan bool //Stop channel for cache clear ticker + slowLookupCacheIpv4 sync.Map //Cache for slow lookup, ip -> cc + slowLookupCacheIpv6 sync.Map //Cache for slow lookup ipv6, ip -> cc + cacheClearTicker *time.Ticker //Ticker for clearing cache + cacheClearTickerStopChan chan bool //Stop channel for cache clear ticker option *StoreOptions } @@ -61,7 +62,7 @@ func NewGeoDb(sysdb *database.Database, option *StoreOptions) (*Store, error) { } if option.SlowLookupCacheClearInterval == 0 { - option.SlowLookupCacheClearInterval = 15 * time.Minute + option.SlowLookupCacheClearInterval = 30 * time.Minute } //Create a new store @@ -71,8 +72,8 @@ func NewGeoDb(sysdb *database.Database, option *StoreOptions) (*Store, error) { geodbIpv6: parsedGeoDataIpv6, geotrieIpv6: ipv6Trie, sysdb: sysdb, - slowLookupCacheIpv4: make(map[string]string), - slowLookupCacheIpv6: make(map[string]string), + slowLookupCacheIpv4: sync.Map{}, + slowLookupCacheIpv6: sync.Map{}, cacheClearTicker: time.NewTicker(option.SlowLookupCacheClearInterval), cacheClearTickerStopChan: make(chan bool), option: option, @@ -86,8 +87,8 @@ func NewGeoDb(sysdb *database.Database, option *StoreOptions) (*Store, error) { case <-store.cacheClearTickerStopChan: return case <-thisGeoDBStore.cacheClearTicker.C: - thisGeoDBStore.slowLookupCacheIpv4 = make(map[string]string) - thisGeoDBStore.slowLookupCacheIpv6 = make(map[string]string) + thisGeoDBStore.slowLookupCacheIpv4 = sync.Map{} + thisGeoDBStore.slowLookupCacheIpv6 = sync.Map{} } } }(thisGeoDBStore) diff --git a/src/mod/geodb/geodb_test.go b/src/mod/geodb/geodb_test.go index 4e3a784..d778e8a 100644 --- a/src/mod/geodb/geodb_test.go +++ b/src/mod/geodb/geodb_test.go @@ -42,7 +42,7 @@ func TestTrieConstruct(t *testing.T) { func TestResolveCountryCodeFromIP(t *testing.T) { // Create a new store store, err := geodb.NewGeoDb(nil, &geodb.StoreOptions{ - false, + true, true, 0, }) @@ -84,4 +84,24 @@ func TestResolveCountryCodeFromIP(t *testing.T) { if info.CountryIsoCode != expected { t.Errorf("expected country code %s, but got %s for IP %s", expected, info.CountryIsoCode, ip) } + + // Test for issue #401 + // Create 100 concurrent goroutines to resolve country code for random IP addresses in the test cases above + for i := 0; i < 100; i++ { + go func() { + for _, testcase := range knownIpCountryMap { + ip := testcase[0] + expected := testcase[1] + info, err := store.ResolveCountryCodeFromIP(ip) + if err != nil { + t.Errorf("error resolving country code for IP %s: %v", ip, err) + return + } + if info.CountryIsoCode != expected { + t.Errorf("expected country code %s, but got %s for IP %s", expected, info.CountryIsoCode, ip) + } + } + }() + } + } diff --git a/src/mod/geodb/slowSearch.go b/src/mod/geodb/slowSearch.go index 8ebc977..8c79a4f 100644 --- a/src/mod/geodb/slowSearch.go +++ b/src/mod/geodb/slowSearch.go @@ -58,7 +58,8 @@ func (s *Store) slowSearchIpv4(ipAddr string) string { } //Check if already in cache - if cc, ok := s.slowLookupCacheIpv4[ipAddr]; ok { + cc := s.GetSlowSearchCachedIpv4(ipAddr) + if cc != "" { return cc } @@ -70,7 +71,7 @@ func (s *Store) slowSearchIpv4(ipAddr string) string { inRange, _ := isIPv4InRange(startIp, endIp, ipAddr) if inRange { //Add to cache - s.slowLookupCacheIpv4[ipAddr] = cc + s.slowLookupCacheIpv4.Store(ipAddr, cc) return cc } } @@ -83,7 +84,8 @@ func (s *Store) slowSearchIpv6(ipAddr string) string { } //Check if already in cache - if cc, ok := s.slowLookupCacheIpv6[ipAddr]; ok { + cc := s.GetSlowSearchCachedIpv6(ipAddr) + if cc != "" { return cc } @@ -95,9 +97,27 @@ func (s *Store) slowSearchIpv6(ipAddr string) string { inRange, _ := isIPv6InRange(startIp, endIp, ipAddr) if inRange { //Add to cache - s.slowLookupCacheIpv6[ipAddr] = cc + s.slowLookupCacheIpv6.Store(ipAddr, cc) return cc } } return "" } + +// GetSlowSearchCachedIpv4 return the country code for the given ipv4 address, return empty string if not found +func (s *Store) GetSlowSearchCachedIpv4(ipAddr string) string { + cc, ok := s.slowLookupCacheIpv4.Load(ipAddr) + if ok { + return cc.(string) + } + return "" +} + +// GetSlowSearchCachedIpv6 return the country code for the given ipv6 address, return empty string if not found +func (s *Store) GetSlowSearchCachedIpv6(ipAddr string) string { + cc, ok := s.slowLookupCacheIpv6.Load(ipAddr) + if ok { + return cc.(string) + } + return "" +}