优化ip2region查询代码

internal/iplibrary/ip2Region.go

@@ -0,0 +1,130 @@
+// 源码改自：https://github.com/lionsoul2014/ip2region/blob/master/binding/golang/ip2region/ip2Region.go
+
+package iplibrary
+
+import (
+	"errors"
+	"io/ioutil"
+	"strconv"
+	"strings"
+)
+
+const (
+	IndexBlockLength = 12
+)
+
+var err error
+
+type IP2Region struct {
+	headerSip []int64
+	headerPtr []int64
+	headerLen int64
+
+	// super block index info
+	firstIndexPtr int64
+	lastIndexPtr  int64
+	totalBlocks   int64
+
+	dbData []byte
+}
+
+type IpInfo struct {
+	CityId   int64
+	Country  string
+	Region   string
+	Province string
+	City     string
+	ISP      string
+}
+
+func (ip IpInfo) String() string {
+	return strconv.FormatInt(ip.CityId, 10) + "|" + ip.Country + "|" + ip.Region + "|" + ip.Province + "|" + ip.City + "|" + ip.ISP
+}
+
+func getIpInfo(cityId int64, line []byte) *IpInfo {
+	lineSlice := strings.Split(string(line), "|")
+	ipInfo := &IpInfo{}
+	length := len(lineSlice)
+	ipInfo.CityId = cityId
+	if length < 5 {
+		for i := 0; i <= 5-length; i++ {
+			lineSlice = append(lineSlice, "")
+		}
+	}
+
+	ipInfo.Country = lineSlice[0]
+	ipInfo.Region = lineSlice[1]
+	ipInfo.Province = lineSlice[2]
+	ipInfo.City = lineSlice[3]
+	ipInfo.ISP = lineSlice[4]
+	return ipInfo
+}
+
+func NewIP2Region(path string) (*IP2Region, error) {
+	var region = &IP2Region{}
+	region.dbData, err = ioutil.ReadFile(path)
+
+	if err != nil {
+		return nil, err
+	}
+
+	region.firstIndexPtr = region.ipLongAtOffset(0)
+	region.lastIndexPtr = region.ipLongAtOffset(4)
+	region.totalBlocks = (region.lastIndexPtr-region.firstIndexPtr)/IndexBlockLength + 1
+	return region, nil
+}
+
+func (this *IP2Region) MemorySearch(ipStr string) (ipInfo *IpInfo, err error) {
+	ip, err := ip2long(ipStr)
+	if err != nil {
+		return nil, err
+	}
+
+	h := this.totalBlocks
+	var dataPtr, l int64
+	for l <= h {
+		m := (l + h) >> 1
+		p := this.firstIndexPtr + m*IndexBlockLength
+		sip := this.ipLongAtOffset(p)
+		if ip < sip {
+			h = m - 1
+		} else {
+			eip := this.ipLongAtOffset(p + 4)
+			if ip > eip {
+				l = m + 1
+			} else {
+				dataPtr = this.ipLongAtOffset(p + 8)
+				break
+			}
+		}
+	}
+	if dataPtr == 0 {
+		return nil, errors.New("not found")
+	}
+
+	dataLen := (dataPtr >> 24) & 0xFF
+	dataPtr = dataPtr & 0x00FFFFFF
+	return getIpInfo(this.ipLongAtOffset(dataPtr), this.dbData[(dataPtr)+4:dataPtr+dataLen]), nil
+}
+
+func (this *IP2Region) ipLongAtOffset(offset int64) int64 {
+	return int64(this.dbData[offset]) |
+		int64(this.dbData[offset+1])<<8 |
+		int64(this.dbData[offset+2])<<16 |
+		int64(this.dbData[offset+3])<<24
+}
+
+func ip2long(IpStr string) (int64, error) {
+	bits := strings.Split(IpStr, ".")
+	if len(bits) != 4 {
+		return 0, errors.New("ip format error")
+	}
+
+	var sum int64
+	for i, n := range bits {
+		bit, _ := strconv.ParseInt(n, 10, 64)
+		sum += bit << uint(24-8*i)
+	}
+
+	return sum, nil
+}

internal/iplibrary/library_ip2region.go

@@ -4,17 +4,16 @@ import (
 	"fmt"
 	"github.com/TeaOSLab/EdgeAPI/internal/errors"
 	"github.com/TeaOSLab/EdgeAPI/internal/remotelogs"
-	"github.com/lionsoul2014/ip2region/binding/golang/ip2region"
 	"net"
 	"strings"
 )
 
 type IP2RegionLibrary struct {
-	db *ip2region.Ip2Region
+	db *IP2Region
 }
 
 func (this *IP2RegionLibrary) Load(dbPath string) error {
-	db, err := ip2region.New(dbPath)
+	db, err := NewIP2Region(dbPath)
 	if err != nil {
 		return err
 	}
@@ -77,6 +76,6 @@ func (this *IP2RegionLibrary) Lookup(ip string) (*Result, error) {
 
 func (this *IP2RegionLibrary) Close() {
 	if this.db != nil {
-		this.db.Close()
+		return
 	}
 }