diff --git a/common/geodata/cache.go b/common/geodata/cache.go index 4e0cbfa53..d4bc2f91a 100644 --- a/common/geodata/cache.go +++ b/common/geodata/cache.go @@ -32,13 +32,13 @@ func (g GeoIPCache) Set(key string, value *router.GeoIP) { } func (g GeoIPCache) Unmarshal(filename, code string) (*router.GeoIP, error) { - filename = platform.GetAssetLocation(filename) - idx := strings.ToUpper(filename + "|" + code) + asset := platform.GetAssetLocation(filename) + idx := strings.ToUpper(asset + "|" + code) if g.Has(idx) { return g.Get(idx), nil } - geoipBytes, err := Decode(filename, code) + geoipBytes, err := Decode(asset, code) switch err { case nil: var geoip router.GeoIP @@ -48,10 +48,13 @@ func (g GeoIPCache) Unmarshal(filename, code string) (*router.GeoIP, error) { g.Set(idx, &geoip) return &geoip, nil + case errCodeNotFound: + return nil, newError(code, " not found in ", filename) + case errFailedToReadBytes, errFailedToReadExpectedLenBytes, errInvalidGeodataFile, errInvalidGeodataVarintLength: newError("failed to decode geodata file: ", filename, ". Fallback to the original ReadFile method.").AtWarning().WriteToLog() - geoipBytes, err = ioutil.ReadFile(filename) + geoipBytes, err = ioutil.ReadFile(asset) if err != nil { return nil, err } @@ -72,7 +75,7 @@ func (g GeoIPCache) Unmarshal(filename, code string) (*router.GeoIP, error) { return nil, err } - return nil, nil + return nil, newError(code, " not found in ", filename) } type GeoSiteCache map[string]*router.GeoSite @@ -96,13 +99,13 @@ func (g GeoSiteCache) Set(key string, value *router.GeoSite) { } func (g GeoSiteCache) Unmarshal(filename, code string) (*router.GeoSite, error) { - filename = platform.GetAssetLocation(filename) - idx := strings.ToUpper(filename + "|" + code) + asset := platform.GetAssetLocation(filename) + idx := strings.ToUpper(asset + "|" + code) if g.Has(idx) { return g.Get(idx), nil } - geositeBytes, err := Decode(filename, code) + geositeBytes, err := Decode(asset, code) switch err { case nil: var geosite router.GeoSite @@ -112,10 +115,13 @@ func (g GeoSiteCache) Unmarshal(filename, code string) (*router.GeoSite, error) g.Set(idx, &geosite) return &geosite, nil + case errCodeNotFound: + return nil, newError(code, " not found in ", filename) + case errFailedToReadBytes, errFailedToReadExpectedLenBytes, errInvalidGeodataFile, errInvalidGeodataVarintLength: newError("failed to decode geodata file: ", filename, ". Fallback to the original ReadFile method.").AtWarning().WriteToLog() - geositeBytes, err = ioutil.ReadFile(filename) + geositeBytes, err = ioutil.ReadFile(asset) if err != nil { return nil, err } @@ -136,5 +142,5 @@ func (g GeoSiteCache) Unmarshal(filename, code string) (*router.GeoSite, error) return nil, err } - return nil, nil + return nil, newError(code, " not found in ", filename) } diff --git a/common/geodata/decode.go b/common/geodata/decode.go index 7ee645a57..7a8ab9b58 100644 --- a/common/geodata/decode.go +++ b/common/geodata/decode.go @@ -10,6 +10,7 @@ package geodata import ( + "io" "os" "runtime" "strings" @@ -26,6 +27,7 @@ var ( errFailedToReadExpectedLenBytes = errors.New("failed to read expected length of bytes") errInvalidGeodataFile = errors.New("invalid geodata file") errInvalidGeodataVarintLength = errors.New("invalid geodata varint length") + errCodeNotFound = errors.New("code not found") ) func emitBytes(f *os.File, code string) ([]byte, error) { @@ -41,6 +43,9 @@ Loop: for { container := make([]byte, advancedN) bytesRead, err := f.Read(container) + if err == io.EOF { + return nil, errCodeNotFound + } if err != nil { return nil, errFailedToReadBytes }