Add flags.json to bindata and use map instead of slices

.github/workflows/go.yml

@@ -30,7 +30,7 @@ jobs:
         if: startsWith(github.ref, 'refs/tags/')
         run: |
           go get -u github.com/go-bindata/go-bindata/...
-          go-bindata -o internal/bindata/geoip/geoip.go -pkg bingeoip  assets/GeoLite2-City.mmdb
+          go-bindata -o internal/bindata/geoip/geoip.go -pkg bingeoip  assets/GeoLite2-City.mmdb assets/flags.json
           go-bindata -o internal/bindata/html/html.go -pkg binhtml  assets/html/
 
       - name: Build

internal/bindata/geoip/geoip.go

@@ -1,6 +1,7 @@
 // Code generated for package bingeoip by go-bindata DO NOT EDIT. (@generated)
 // sources:
 // assets/GeoLite2-City.mmdb
+// assets/flags.json
 package bingeoip
 
 import (
@@ -43,6 +44,24 @@ func assetsGeolite2CityMmdb() (*asset, error) {
 	return a, err
 }
 
+// assetsFlagsJson reads file data from disk. It returns an error on failure.
+func assetsFlagsJson() (*asset, error) {
+	path := "assets/flags.json"
+	name := "assets/flags.json"
+	bytes, err := bindataRead(path, name)
+	if err != nil {
+		return nil, err
+	}
+
+	fi, err := os.Stat(path)
+	if err != nil {
+		err = fmt.Errorf("Error reading asset info %s at %s: %v", name, path, err)
+	}
+
+	a := &asset{bytes: bytes, info: fi}
+	return a, err
+}
+
 // Asset loads and returns the asset for the given name.
 // It returns an error if the asset could not be found or
 // could not be loaded.
@@ -96,6 +115,7 @@ func AssetNames() []string {
 // _bindata is a table, holding each asset generator, mapped to its name.
 var _bindata = map[string]func() (*asset, error){
 	"assets/GeoLite2-City.mmdb": assetsGeolite2CityMmdb,
+	"assets/flags.json":         assetsFlagsJson,
 }
 
 // AssetDir returns the file names below a certain
@@ -141,6 +161,7 @@ type bintree struct {
 var _bintree = &bintree{nil, map[string]*bintree{
 	"assets": &bintree{nil, map[string]*bintree{
 		"GeoLite2-City.mmdb": &bintree{assetsGeolite2CityMmdb, map[string]*bintree{}},
+		"flags.json":         &bintree{assetsFlagsJson, map[string]*bintree{}},
 	}},
 }}
 

pkg/proxy/geoip.go

@@ -19,45 +19,61 @@ func InitGeoIpDB() {
 	if err != nil {
 		panic(err)
 	}
-	geoIp = NewGeoIP("assets/GeoLite2-City.mmdb")
+	err = bingeoip.RestoreAsset("", "assets/flags.json")
+	if err != nil {
+		panic(err)
+	}
+	geoIp = NewGeoIP("assets/GeoLite2-City.mmdb", "assets/flags.json")
 }
 
 // GeoIP2
 type GeoIP struct {
-	db    *geoip2.Reader
-	flags cclist
+	db       *geoip2.Reader
+	emojiMap map[string]string
 }
 
-type CountryCode struct {
-	Code     string `json:"code"`
-	Emoji    string `json:"emoji"`
-	Unicode  string `json:"unicode`
-	Name     string `json:"name"`
-	Title    string `json:"title"`
-	Dialcode string `json:"dialCode`
+type CountryEmoji struct {
+	Code  string `json:"code"`
+	Emoji string `json:"emoji"`
 }
 
 // new geoip from db file
-func NewGeoIP(filePath string) (geoip GeoIP) {
-	var countrycodes cclist
+func NewGeoIP(geodb, flags string) (geoip GeoIP) {
 	// 判断文件是否存在
-	_, err := os.Stat(filePath)
+	_, err := os.Stat(geodb)
 	if err != nil && os.IsNotExist(err) {
-		log.Println("文件不存在，请自行下载 Geoip2 City库，并保存在", filePath)
+		log.Println("文件不存在，请自行下载 Geoip2 City库，并保存在", geodb)
 		os.Exit(1)
 	} else {
-		db, err := geoip2.Open(filePath)
+		db, err := geoip2.Open(geodb)
 		if err != nil {
 			log.Fatal(err)
 		}
-		data, err := ioutil.ReadFile("assets/flags.json")
+		geoip.db = db
+	}
+
+	_, err = os.Stat(flags)
+	if err != nil && os.IsNotExist(err) {
+		log.Println("flags 文件不存在，请自行下载 flags.json，并保存在", flags)
+		os.Exit(1)
+	} else {
+		data, err := ioutil.ReadFile(flags)
 		if err != nil {
 			log.Fatal(err)
 			return
 		}
+		var countryEmojiList = make([]CountryEmoji, 0)
+		err = json.Unmarshal(data, &countryEmojiList)
+		if err != nil {
+			log.Fatalln(err.Error())
+			return
+		}
 
-		json.Unmarshal(data, &countrycodes)
-		geoip = GeoIP{db: db, flags: countrycodes}
+		emojiMap := make(map[string]string)
+		for _, i := range countryEmojiList {
+			emojiMap[i.Code] = i.Emoji
+		}
+		geoip.emojiMap = emojiMap
 	}
 	return
 }
@@ -68,59 +84,22 @@ func (g GeoIP) Find(ipORdomain string) (ip, country string, err error) {
 	if err != nil {
 		return "", "", err
 	}
-	ipData := net.ParseIP(ips[0].String())
-	record, err := g.db.City(ipData)
+	ip = ips[0].String()
+
+	var record *geoip2.City
+	record, err = g.db.City(ips[0])
 	if err != nil {
-		return "", "", err
+		return
 	}
-
-	countrycode := record.Country.IsoCode
-	if countrycode == "" {
-		// handle unknown country code
-		return ips[0].String(), fmt.Sprintf("🏁 ZZ"), nil
+	countryIsoCode := record.Country.IsoCode
+	if countryIsoCode == "" {
+		country = fmt.Sprintf("🏁 ZZ")
 	}
-	countryflag := g.getFlag(countrycode)
-	country = fmt.Sprintf("%v %v", countryflag, countrycode)
-
-	return ips[0].String(), country, nil
-}
-
-// getFlag method take country code as input, return its corresponding country/region flag
-func (g GeoIP) getFlag(countrycode string) string {
-	result := find(g.flags, countrycode)
-	return result.Emoji
-}
-
-type cclist []CountryCode
-
-func (c cclist) in(code string, left bool) bool {
-	length := len(c)
-	if left {
-		return c[length-1].Code >= code
-	}
-	return c[0].Code <= code
-}
-
-// find will find corresponding country flag emoji from flags for a given country code
-func find(list cclist, target string) CountryCode {
-	var result CountryCode
-
-	length := len(list)
-	if length == 1 {
-		if list[0].Code == target {
-			return list[0]
-		}
-		return CountryCode{Emoji: "🏁"}
-	}
-	split := length / 2
-	left, right := list[:split], list[split:]
-
-	if left.in(target, true) {
-		result = find(left, target)
-	}
-	if right.in(target, false) {
-		result = find(right, target)
-	}
-
-	return result
+	emoji, found := g.emojiMap[countryIsoCode]
+	if found {
+		country = fmt.Sprintf("%v %v", emoji, countryIsoCode)
+	} else {
+		country = fmt.Sprintf("🏁 ZZ")
+	}
+	return
 }