EdgeCommon/pkg/iplibrary/reader.go

// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .

package iplibrary

import (
	"bytes"
	"encoding/json"
	"errors"
	"github.com/TeaOSLab/EdgeCommon/pkg/configutils"
	"github.com/iwind/TeaGo/types"
	"io"
	"net"
	"sort"
	"strings"
)

type Reader struct {
	meta *Meta

	regionMap map[string]*ipRegion

	ipV4Items []*ipItem
	ipV6Items []*ipItem
}

func NewReader(reader io.Reader) (*Reader, error) {
	var libReader = &Reader{
		regionMap: map[string]*ipRegion{},
	}
	err := libReader.load(reader)
	if err != nil {
		return nil, err
	}
	return libReader, nil
}

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
	})

	return nil
}

func (this *Reader) Lookup(ip net.IP) *QueryResult {
	if ip == nil {
		return &QueryResult{}
	}

	var ipLong = configutils.IP2Long(ip)
	var isV4 = configutils.IsIPv4(ip)
	var resultItem *ipItem
	if isV4 {
		sort.Search(len(this.ipV4Items), func(i int) bool {
			var item = this.ipV4Items[i]
			if item.IPFrom <= ipLong {
				if item.IPTo >= 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() []*ipItem {
	return this.ipV4Items
}

func (this *Reader) IPv6Items() []*ipItem {
	return this.ipV6Items
}

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]
			var pieces = strings.Split(string(line), "|")
			if len(pieces) != 8 {
				return nil, errors.New("invalid ip definition '" + string(line) + "'")
			}

			var version = pieces[0]
			if len(version) == 0 {
				version = "4"
			}

			if version != "4" && version != "6" {
				return nil, errors.New("invalid ip version '" + string(line) + "'")
			}

			var ipFrom uint64
			var ipTo uint64
			if len(pieces[2]) == 0 {
				pieces[2] = pieces[1]
				ipFrom = types.Uint64(pieces[1])
				ipTo = types.Uint64(pieces[2])
			} else {
				ipFrom = types.Uint64(pieces[1])
				ipTo = types.Uint64(pieces[2]) + ipFrom
			}

			var countryId = types.Uint32(pieces[3])
			var provinceId = types.Uint32(pieces[4])
			var cityId = types.Uint32(pieces[5])
			var townId = types.Uint32(pieces[6])
			var providerId = types.Uint32(pieces[7])
			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, &ipItem{
					IPFrom: ipFrom,
					IPTo:   ipTo,
					Region: region,
				})
			} else {
				this.ipV6Items = append(this.ipV6Items, &ipItem{
					IPFrom: ipFrom,
					IPTo:   ipTo,
					Region: region,
				})
			}

			data = data[index+1:]
		} else {
			left = data
			break
		}
	}
	return
}