go-common/vendor/github.com/ipipdotnet/ipdb-go/reader.go

251 lines
4.9 KiB
Go
Raw Normal View History

2019-04-22 10:49:16 +00:00
package ipdb
import (
"os"
"encoding/binary"
"errors"
"encoding/json"
"io/ioutil"
"net"
"strings"
"reflect"
"unsafe"
"time"
)
const IPv4 = 0x01
const IPv6 = 0x02
var (
ErrFileSize = errors.New("IP Database file size error.")
ErrMetaData = errors.New("IP Database metadata error.")
ErrReadFull = errors.New("IP Database ReadFull error.")
ErrDatabaseError = errors.New("database error")
ErrIPFormat = errors.New("Query IP Format error.")
ErrNoSupportLanguage = errors.New("language not support")
ErrNoSupportIPv4 = errors.New("IPv4 not support")
ErrNoSupportIPv6 = errors.New("IPv6 not support")
ErrDataNotExists = errors.New("data is not exists")
)
type MetaData struct {
Build int64 `json:"build"`
IPVersion uint16 `json:"ip_version"`
Languages map[string]int `json:"languages"`
NodeCount int `json:"node_count"`
TotalSize int `json:"total_size"`
Fields []string `json:"fields"`
}
type reader struct {
fileSize int
nodeCount int
v4offset int
meta MetaData
data []byte
refType map[string]string
}
func newReader(name string, obj interface{}) (*reader, error) {
var err error
var fileInfo os.FileInfo
fileInfo, err = os.Stat(name)
if err != nil {
return nil, err
}
fileSize := int(fileInfo.Size())
body, err := ioutil.ReadFile(name)
if err != nil {
return nil, ErrReadFull
}
var meta MetaData
metaLength := int(binary.BigEndian.Uint32(body[0:4]))
if err := json.Unmarshal(body[4:4+metaLength], &meta); err != nil {
return nil, err
}
if len(meta.Languages) == 0 || len(meta.Fields) == 0 {
return nil, ErrMetaData
}
if fileSize != (4+metaLength+meta.TotalSize) {
return nil, ErrFileSize
}
var dm map[string]string
if obj != nil {
t := reflect.TypeOf(obj).Elem()
dm = make(map[string]string, t.NumField())
for i := 0; i < t.NumField(); i++ {
k := t.Field(i).Tag.Get("json")
dm[k] = t.Field(i).Name
}
}
db := &reader{
fileSize: fileSize,
nodeCount: meta.NodeCount,
meta:meta,
refType: dm,
data: body[4+metaLength:],
}
if db.v4offset == 0 {
node := 0
for i := 0; i < 96 && node < db.nodeCount; i++ {
if i >= 80 {
node = db.readNode(node, 1)
} else {
node = db.readNode(node, 0)
}
}
db.v4offset = node
}
return db, nil
}
func (db *reader) Find(addr, language string) ([]string, error) {
return db.find1(addr, language)
}
func (db *reader) FindMap(addr, language string) (map[string]string, error) {
data, err := db.find1(addr, language)
if err != nil {
return nil, err
}
info := make(map[string]string, len(db.meta.Fields))
for k, v := range data {
info[db.meta.Fields[k]] = v
}
return info, nil
}
func (db *reader) find0(addr string) ([]byte, error) {
var err error
var node int
ipv := net.ParseIP(addr)
if ip := ipv.To4(); ip != nil {
if !db.IsIPv4Support() {
return nil, ErrNoSupportIPv4
}
node, err = db.search(ip, 32)
} else if ip := ipv.To16(); ip != nil {
if !db.IsIPv6Support() {
return nil, ErrNoSupportIPv6
}
node, err = db.search(ip, 128)
} else {
return nil, ErrIPFormat
}
if err != nil || node < 0 {
return nil, err
}
body, err := db.resolve(node)
if err != nil {
return nil, err
}
return body, nil
}
func (db *reader) find1(addr, language string) ([]string, error) {
off, ok := db.meta.Languages[language]
if !ok {
return nil, ErrNoSupportLanguage
}
body, err := db.find0(addr)
if err != nil {
return nil, err
}
str := (*string)(unsafe.Pointer(&body))
tmp := strings.Split(*str, "\t")
if (off + len(db.meta.Fields)) > len(tmp) {
return nil, ErrDatabaseError
}
return tmp[off:off+len(db.meta.Fields)], nil
}
func (db *reader) search(ip net.IP, bitCount int) (int, error) {
var node int
if bitCount == 32 {
node = db.v4offset
} else {
node = 0;
}
for i := 0; i < bitCount; i++ {
if node > db.nodeCount {
break
}
node = db.readNode(node, ((0xFF & int(ip[i >> 3])) >> uint(7 - (i % 8))) & 1)
}
if node > db.nodeCount {
return node, nil
}
return -1, ErrDataNotExists
}
func (db *reader) readNode(node, index int) int {
off := node * 8 + index * 4
return int(binary.BigEndian.Uint32(db.data[off:off+4]))
}
func (db *reader) resolve(node int) ([]byte, error) {
resolved := node - db.nodeCount + db.nodeCount * 8
if resolved >= db.fileSize {
return nil, ErrDatabaseError
}
size := int(binary.BigEndian.Uint16(db.data[resolved:resolved+2]))
if (resolved+2+size) > len(db.data) {
return nil, ErrDatabaseError
}
bytes := db.data[resolved+2:resolved+2+size]
return bytes, nil
}
func (db *reader) IsIPv4Support() bool {
return (int(db.meta.IPVersion) & IPv4) == IPv4
}
func (db *reader) IsIPv6Support() bool {
return (int(db.meta.IPVersion) & IPv6) == IPv6
}
func (db *reader) Build() time.Time {
return time.Unix(db.meta.Build, 0).In(time.UTC)
}
func (db *reader) Languages() []string {
ls := make([]string, 0, len(db.meta.Languages))
for k := range db.meta.Languages {
ls = append(ls, k)
}
return ls
}