优化查找IP地址信息的逻辑

go.mod

@@ -23,8 +23,7 @@ require (
 	github.com/smartwalle/alipay/v3 v3.1.7
 	github.com/spf13/viper v1.10.1
 	github.com/ugorji/go v1.2.7 // indirect
-	github.com/yinheli/mahonia v0.0.0-20131226213531-0eef680515cc // indirect
-	github.com/yinheli/qqwry v0.0.0-20160229183603-f50680010f4a
+	github.com/yinheli/mahonia v0.0.0-20131226213531-0eef680515cc
 	golang.org/x/crypto v0.0.0-20220307211146-efcb8507fb70 // indirect
 	golang.org/x/image v0.0.0-20220413100746-70e8d0d3baa9 // indirect
 	golang.org/x/net v0.0.0-20220127200216-cd36cc0744dd // indirect

go.sum

@@ -675,8 +675,6 @@ github.com/xlab/treeprint v0.0.0-20180616005107-d6fb6747feb6/go.mod h1:ce1O1j6Ut
 github.com/xordataexchange/crypt v0.0.3-0.20170626215501-b2862e3d0a77/go.mod h1:aYKd//L2LvnjZzWKhF00oedf4jCCReLcmhLdhm1A27Q=
 github.com/yinheli/mahonia v0.0.0-20131226213531-0eef680515cc h1:7VHQaaNwHymWbj8lAcXMYX1qopebSBHwYC3ceXLWONU=
 github.com/yinheli/mahonia v0.0.0-20131226213531-0eef680515cc/go.mod h1:Pcc297eVCbkDBBVq8FbnI+qDUeIMrHy4Bo7nveAuCAs=
-github.com/yinheli/qqwry v0.0.0-20160229183603-f50680010f4a h1:VUPXGL4N1B5xqomxI4vZn0momgleh0hcH0PE/oIOsAI=
-github.com/yinheli/qqwry v0.0.0-20160229183603-f50680010f4a/go.mod h1:Zva9ErVtC2arl+9xtHwiXujgejX1S3VpOYExADJ9kio=
 github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.32/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=

pkg/util/ip.go

@@ -1,12 +1,8 @@
 package util
 
-import (
-	"github.com/yinheli/qqwry"
-)
+import "github.com/rocboss/paopao-ce/pkg/util/iploc"
 
 func GetIPLoc(ip string) string {
-	q := qqwry.NewQQwry("qqwry.dat")
-	q.Find(ip)
-
-	return q.Country
+	country, _ := iploc.Find(ip)
+	return country
 }

pkg/util/iploc/iploc.go

@@ -0,0 +1,128 @@
+package iploc
+
+import (
+	_ "embed"
+	"encoding/binary"
+	"net"
+
+	"github.com/yinheli/mahonia"
+)
+
+// Note: This file is a modified version of https://github.com/yinheli/qqwry.
+
+//go:embed qqwry.dat
+var qqwry []byte
+
+const (
+	cIndexLen      = 7
+	cRedirectMode1 = 0x01
+	cRedirectMode2 = 0x02
+)
+
+// Find get country and city base ip
+func Find(ip string) (string, string) {
+	offset := searchIndex(binary.BigEndian.Uint32(net.ParseIP(ip).To4()))
+	if offset <= 0 {
+		return "", ""
+	}
+	var country, area []byte
+	mode := readMode(offset + 4)
+	if mode == cRedirectMode1 {
+		countryOffset := readUInt24(offset + 5)
+		mode = readMode(countryOffset)
+		if mode == cRedirectMode2 {
+			c := readUInt24(countryOffset + 1)
+			country = readString(c)
+			countryOffset += 4
+		} else {
+			country = readString(countryOffset)
+			countryOffset += uint32(len(country) + 1)
+		}
+		area = readArea(countryOffset)
+	} else if mode == cRedirectMode2 {
+		countryOffset := readUInt24(offset + 5)
+		country = readString(countryOffset)
+		area = readArea(offset + 8)
+	} else {
+		country = readString(offset + 4)
+		area = readArea(offset + uint32(5+len(country)))
+	}
+	enc := mahonia.NewDecoder("gbk")
+	Country := enc.ConvertString(string(country))
+	City := enc.ConvertString(string(area))
+	return Country, City
+}
+
+func readMode(offset uint32) byte {
+	return qqwry[offset]
+}
+
+func readArea(offset uint32) []byte {
+	mode := readMode(offset)
+	if mode == cRedirectMode1 || mode == cRedirectMode2 {
+		areaOffset := readUInt24(offset + 1)
+		if areaOffset == 0 {
+			return []byte("")
+		} else {
+			return readString(areaOffset)
+		}
+	} else {
+		return readString(offset)
+	}
+}
+
+func readString(offset uint32) []byte {
+	data := make([]byte, 0, 30)
+	for {
+		if qqwry[offset] == 0 {
+			break
+		}
+		data = append(data, qqwry[offset])
+		offset++
+	}
+	return data
+}
+
+func searchIndex(ip uint32) uint32 {
+	header := qqwry[:8]
+	start := binary.LittleEndian.Uint32(header[:4])
+	end := binary.LittleEndian.Uint32(header[4:])
+	for {
+		mid := getMiddleOffset(start, end)
+		buf := qqwry[mid : mid+cIndexLen]
+		ipaddr := binary.LittleEndian.Uint32(buf[:4])
+		if end-start == cIndexLen {
+			offset := byte3ToUInt32(buf[4:])
+			// TODO: 这里可能有bug，需要优化
+			if ip < binary.LittleEndian.Uint32(qqwry[end:end+4]) {
+				return offset
+			} else {
+				return 0
+			}
+		}
+		// 找到的比较大，向前移
+		if ipaddr > ip {
+			end = mid
+		} else if ipaddr < ip { // 找到的比较小，向后移
+			start = mid
+		} else {
+			return byte3ToUInt32(buf[4:])
+		}
+	}
+}
+
+func readUInt24(offset uint32) uint32 {
+	return byte3ToUInt32(qqwry[offset : offset+3])
+}
+
+func getMiddleOffset(start uint32, end uint32) uint32 {
+	records := ((end - start) / cIndexLen) >> 1
+	return start + records*cIndexLen
+}
+
+func byte3ToUInt32(data []byte) uint32 {
+	i := uint32(data[0]) & 0xff
+	i |= (uint32(data[1]) << 8) & 0xff00
+	i |= (uint32(data[2]) << 16) & 0xff0000
+	return i
+}

pkg/util/iploc/iploc_test.go

@@ -0,0 +1,25 @@
+package iploc
+
+import (
+	"testing"
+)
+
+func TestFind(t *testing.T) {
+	for _, data := range []struct {
+		ip      string
+		country string
+		city    string
+	}{
+		{ip: "127.0.0.1", country: "本机地址", city: " CZ88.NET"},
+		{ip: "180.89.94.9", country: "北京市", city: "鹏博士宽带"},
+	} {
+		country, city := Find(data.ip)
+		t.Logf("ip:%v, country:%v, city:%v", data.ip, country, city)
+		if country != data.country {
+			t.Errorf("find ip:%s expect country: %s got: %s", data.ip, data.country, country)
+		}
+		if city != data.city {
+			t.Errorf("find ip:%s expect city: %s got: %s", data.ip, data.city, city)
+		}
+	}
+}