mirror of
				https://github.com/TeaOSLab/EdgeCommon.git
				synced 2025-11-04 05:00:24 +08:00 
			
		
		
		
	
		
			
				
	
	
		
			348 lines
		
	
	
		
			6.8 KiB
		
	
	
	
		
			Go
		
	
	
	
	
	
			
		
		
	
	
			348 lines
		
	
	
		
			6.8 KiB
		
	
	
	
		
			Go
		
	
	
	
	
	
// Copyright 2022 Liuxiangchao iwind.liu@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()
 | 
						|
}
 |