Files
EdgeCommon/pkg/iplibrary/reader.go
2024-05-17 18:28:59 +08:00

348 lines
6.8 KiB
Go

// Copyright 2022 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
package iplibrary
import (
"bytes"
"encoding/binary"
"encoding/json"
"errors"
"github.com/TeaOSLab/EdgeCommon/pkg/configutils"
"io"
"math/big"
"net"
"runtime"
"sort"
"strconv"
"strings"
)
// Reader IP库Reader
type Reader struct {
meta *Meta
regionMap map[string]*ipRegion // 缓存重复的区域用来节约内存
ipV4Items []ipv4ItemV1
ipV6Items []ipv6ItemV1
lastIPFrom uint64
lastCountryId uint16
lastProvinceId uint16
lastCityId uint32
lastTownId uint32
lastProviderId uint16
}
// NewReaderV1 创建新Reader对象
func NewReaderV1(reader io.Reader) (*Reader, error) {
var libReader = &Reader{
regionMap: map[string]*ipRegion{},
}
if runtime.NumCPU() >= 4 /** CPU数量较多的通常有着大内存 **/ {
libReader.ipV4Items = make([]ipv4ItemV1, 0, 6_000_000)
} else {
libReader.ipV4Items = make([]ipv4ItemV1, 0, 600_000)
}
err := libReader.load(reader)
if err != nil {
return nil, err
}
return libReader, nil
}
// 从Reader中加载数据
func (this *Reader) load(reader io.Reader) error {
var buf = make([]byte, 1024)
var metaLine []byte
var metaLineFound = false
var dataBuf = []byte{}
for {
n, err := reader.Read(buf)
if n > 0 {
var data = buf[:n]
dataBuf = append(dataBuf, data...)
if metaLineFound {
left, err := this.parse(dataBuf)
if err != nil {
return err
}
dataBuf = left
} else {
var index = bytes.IndexByte(dataBuf, '\n')
if index > 0 {
metaLine = dataBuf[:index]
dataBuf = dataBuf[index+1:]
metaLineFound = true
var meta = &Meta{}
err = json.Unmarshal(metaLine, &meta)
if err != nil {
return err
}
meta.Init()
this.meta = meta
left, err := this.parse(dataBuf)
if err != nil {
return err
}
dataBuf = left
}
}
}
if err != nil {
if err != io.EOF {
return err
}
break
}
}
sort.Slice(this.ipV4Items, func(i, j int) bool {
var from0 = this.ipV4Items[i].IPFrom
var to0 = this.ipV4Items[i].IPTo
var from1 = this.ipV4Items[j].IPFrom
var to1 = this.ipV4Items[j].IPTo
if from0 == from1 {
return to0 < to1
}
return from0 < from1
})
sort.Slice(this.ipV6Items, func(i, j int) bool {
var from0 = this.ipV6Items[i].IPFrom
var to0 = this.ipV6Items[i].IPTo
var from1 = this.ipV6Items[j].IPFrom
var to1 = this.ipV6Items[j].IPTo
if from0 == from1 {
return to0 < to1
}
return from0 < from1
})
// 清理内存
this.regionMap = nil
return nil
}
func (this *Reader) Lookup(ip net.IP) *QueryResult {
if ip == nil {
return &QueryResult{}
}
var ipLong = this.ip2long(ip)
var isV4 = configutils.IsIPv4(ip)
var resultItem any
if isV4 {
sort.Search(len(this.ipV4Items), func(i int) bool {
var item = this.ipV4Items[i]
if item.IPFrom <= uint32(ipLong) {
if item.IPTo >= uint32(ipLong) {
resultItem = item
return false
}
return false
}
return true
})
} else {
sort.Search(len(this.ipV6Items), func(i int) bool {
var item = this.ipV6Items[i]
if item.IPFrom <= ipLong {
if item.IPTo >= ipLong {
resultItem = item
return false
}
return false
}
return true
})
}
return &QueryResult{
item: resultItem,
meta: this.meta,
}
}
func (this *Reader) Meta() *Meta {
return this.meta
}
func (this *Reader) IPv4Items() []ipv4ItemV1 {
return this.ipV4Items
}
func (this *Reader) IPv6Items() []ipv6ItemV1 {
return this.ipV6Items
}
func (this *Reader) Destroy() {
this.meta = nil
this.regionMap = nil
this.ipV4Items = nil
this.ipV6Items = nil
}
// 分析数据
func (this *Reader) parse(data []byte) (left []byte, err error) {
if len(data) == 0 {
return
}
for {
var index = bytes.IndexByte(data, '\n')
if index >= 0 {
var line = data[:index]
err = this.parseLine(line)
if err != nil {
return nil, err
}
data = data[index+1:]
} else {
left = data
break
}
}
return
}
// 单行分析
func (this *Reader) parseLine(line []byte) error {
const maxPieces = 8
var pieces = strings.Split(string(line), "|")
var countPieces = len(pieces)
if countPieces < maxPieces { // 补足一行
for i := 0; i < maxPieces-countPieces; i++ {
pieces = append(pieces, "")
}
} else if countPieces > maxPieces {
return errors.New("invalid ip definition '" + string(line) + "'")
}
var version = pieces[0]
if len(version) == 0 {
version = "4"
}
if version != "4" && version != "6" {
return errors.New("invalid ip version '" + string(line) + "'")
}
// ip range
var ipFrom uint64
var ipTo uint64
if strings.HasPrefix(pieces[1], "+") {
ipFrom = this.lastIPFrom + this.decodeUint64(pieces[1][1:])
} else {
ipFrom = this.decodeUint64(pieces[1])
}
if len(pieces[2]) == 0 {
ipTo = ipFrom
} else {
ipTo = this.decodeUint64(pieces[2]) + ipFrom
}
this.lastIPFrom = ipFrom
// country
var countryId uint16
if pieces[3] == "+" {
countryId = this.lastCountryId
} else {
countryId = uint16(this.decodeUint64(pieces[3]))
}
this.lastCountryId = countryId
var provinceId uint16
if pieces[4] == "+" {
provinceId = this.lastProvinceId
} else {
provinceId = uint16(this.decodeUint64(pieces[4]))
}
this.lastProvinceId = provinceId
// city
var cityId uint32
if pieces[5] == "+" {
cityId = this.lastCityId
} else {
cityId = uint32(this.decodeUint64(pieces[5]))
}
this.lastCityId = cityId
// town
var townId uint32
if pieces[6] == "+" {
townId = this.lastTownId
} else {
townId = uint32(this.decodeUint64(pieces[6]))
}
this.lastTownId = townId
// provider
var providerId uint16
if pieces[7] == "+" {
providerId = this.lastProviderId
} else {
providerId = uint16(this.decodeUint64(pieces[7]))
}
this.lastProviderId = providerId
var hash = HashRegion(countryId, provinceId, cityId, townId, providerId)
region, ok := this.regionMap[hash]
if !ok {
region = &ipRegion{
CountryId: countryId,
ProvinceId: provinceId,
CityId: cityId,
TownId: townId,
ProviderId: providerId,
}
this.regionMap[hash] = region
}
if version == "4" {
this.ipV4Items = append(this.ipV4Items, ipv4ItemV1{
IPFrom: uint32(ipFrom),
IPTo: uint32(ipTo),
Region: region,
})
} else {
this.ipV6Items = append(this.ipV6Items, ipv6ItemV1{
IPFrom: ipFrom,
IPTo: ipTo,
Region: region,
})
}
return nil
}
func (this *Reader) decodeUint64(s string) uint64 {
if this.meta != nil && this.meta.Version == Version2 {
i, _ := strconv.ParseUint(s, 32, 64)
return i
}
i, _ := strconv.ParseUint(s, 10, 64)
return i
}
func (this *Reader) ip2long(netIP net.IP) uint64 {
if len(netIP) == 0 {
return 0
}
var b4 = netIP.To4()
if b4 != nil {
return uint64(binary.BigEndian.Uint32(b4.To4()))
}
var i = big.NewInt(0)
i.SetBytes(netIP.To16())
return i.Uint64()
}