Browse Source

Fix: code not found in geo files (#953)

Loyalsoldier 4 years ago
parent
commit
a585ca28a8
2 changed files with 21 additions and 10 deletions
  1. 16 10
      common/geodata/cache.go
  2. 5 0
      common/geodata/decode.go

+ 16 - 10
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) {
 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) {
 	if g.Has(idx) {
 		return g.Get(idx), nil
 		return g.Get(idx), nil
 	}
 	}
 
 
-	geoipBytes, err := Decode(filename, code)
+	geoipBytes, err := Decode(asset, code)
 	switch err {
 	switch err {
 	case nil:
 	case nil:
 		var geoip router.GeoIP
 		var geoip router.GeoIP
@@ -48,10 +48,13 @@ func (g GeoIPCache) Unmarshal(filename, code string) (*router.GeoIP, error) {
 		g.Set(idx, &geoip)
 		g.Set(idx, &geoip)
 		return &geoip, nil
 		return &geoip, nil
 
 
+	case errCodeNotFound:
+		return nil, newError(code, " not found in ", filename)
+
 	case errFailedToReadBytes, errFailedToReadExpectedLenBytes,
 	case errFailedToReadBytes, errFailedToReadExpectedLenBytes,
 		errInvalidGeodataFile, errInvalidGeodataVarintLength:
 		errInvalidGeodataFile, errInvalidGeodataVarintLength:
 		newError("failed to decode geodata file: ", filename, ". Fallback to the original ReadFile method.").AtWarning().WriteToLog()
 		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 {
 		if err != nil {
 			return nil, err
 			return nil, err
 		}
 		}
@@ -72,7 +75,7 @@ func (g GeoIPCache) Unmarshal(filename, code string) (*router.GeoIP, error) {
 		return nil, err
 		return nil, err
 	}
 	}
 
 
-	return nil, nil
+	return nil, newError(code, " not found in ", filename)
 }
 }
 
 
 type GeoSiteCache map[string]*router.GeoSite
 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) {
 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) {
 	if g.Has(idx) {
 		return g.Get(idx), nil
 		return g.Get(idx), nil
 	}
 	}
 
 
-	geositeBytes, err := Decode(filename, code)
+	geositeBytes, err := Decode(asset, code)
 	switch err {
 	switch err {
 	case nil:
 	case nil:
 		var geosite router.GeoSite
 		var geosite router.GeoSite
@@ -112,10 +115,13 @@ func (g GeoSiteCache) Unmarshal(filename, code string) (*router.GeoSite, error)
 		g.Set(idx, &geosite)
 		g.Set(idx, &geosite)
 		return &geosite, nil
 		return &geosite, nil
 
 
+	case errCodeNotFound:
+		return nil, newError(code, " not found in ", filename)
+
 	case errFailedToReadBytes, errFailedToReadExpectedLenBytes,
 	case errFailedToReadBytes, errFailedToReadExpectedLenBytes,
 		errInvalidGeodataFile, errInvalidGeodataVarintLength:
 		errInvalidGeodataFile, errInvalidGeodataVarintLength:
 		newError("failed to decode geodata file: ", filename, ". Fallback to the original ReadFile method.").AtWarning().WriteToLog()
 		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 {
 		if err != nil {
 			return nil, err
 			return nil, err
 		}
 		}
@@ -136,5 +142,5 @@ func (g GeoSiteCache) Unmarshal(filename, code string) (*router.GeoSite, error)
 		return nil, err
 		return nil, err
 	}
 	}
 
 
-	return nil, nil
+	return nil, newError(code, " not found in ", filename)
 }
 }

+ 5 - 0
common/geodata/decode.go

@@ -10,6 +10,7 @@
 package geodata
 package geodata
 
 
 import (
 import (
+	"io"
 	"os"
 	"os"
 	"runtime"
 	"runtime"
 	"strings"
 	"strings"
@@ -26,6 +27,7 @@ var (
 	errFailedToReadExpectedLenBytes = errors.New("failed to read expected length of bytes")
 	errFailedToReadExpectedLenBytes = errors.New("failed to read expected length of bytes")
 	errInvalidGeodataFile           = errors.New("invalid geodata file")
 	errInvalidGeodataFile           = errors.New("invalid geodata file")
 	errInvalidGeodataVarintLength   = errors.New("invalid geodata varint length")
 	errInvalidGeodataVarintLength   = errors.New("invalid geodata varint length")
+	errCodeNotFound                 = errors.New("code not found")
 )
 )
 
 
 func emitBytes(f *os.File, code string) ([]byte, error) {
 func emitBytes(f *os.File, code string) ([]byte, error) {
@@ -41,6 +43,9 @@ Loop:
 	for {
 	for {
 		container := make([]byte, advancedN)
 		container := make([]byte, advancedN)
 		bytesRead, err := f.Read(container)
 		bytesRead, err := f.Read(container)
+		if err == io.EOF {
+			return nil, errCodeNotFound
+		}
 		if err != nil {
 		if err != nil {
 			return nil, errFailedToReadBytes
 			return nil, errFailedToReadBytes
 		}
 		}