github.com/infraboard/keyauth@v0.8.1/apps/ip2region/reader/reader.go (about) 1 package reader 2 3 import ( 4 "errors" 5 "fmt" 6 "io" 7 "io/ioutil" 8 "strconv" 9 "strings" 10 11 "github.com/infraboard/keyauth/apps/ip2region" 12 ) 13 14 const ( 15 INDEX_BLOCK_LENGTH = 12 16 TOTAL_HEADER_LENGTH = 8192 17 ) 18 19 // New todo 20 func New(db io.Reader) (*IPReader, error) { 21 dbBinStr, err := ioutil.ReadAll(db) 22 if err != nil { 23 return nil, fmt.Errorf("load db file to memory error") 24 } 25 26 return &IPReader{ 27 dbBinStr: dbBinStr, 28 firstIndexPtr: getLong(dbBinStr, 0), 29 lastIndexPtr: getLong(dbBinStr, 4), 30 }, nil 31 } 32 33 // IPReader tood 34 type IPReader struct { 35 // super block index info 36 firstIndexPtr int64 37 lastIndexPtr int64 38 dbBinStr []byte 39 } 40 41 // TotalBlocks todo 42 func (r *IPReader) TotalBlocks() int64 { 43 return (r.lastIndexPtr-r.firstIndexPtr)/INDEX_BLOCK_LENGTH + 1 44 } 45 46 // MemorySearch todo 47 func (r *IPReader) MemorySearch(ipStr string) (*ip2region.IPInfo, error) { 48 ipInfo := ip2region.IPInfo{} 49 50 ip, err := ip2long(ipStr) 51 if err != nil { 52 return nil, err 53 } 54 55 h := r.TotalBlocks() 56 var dataPtr, l int64 57 for l <= h { 58 m := (l + h) >> 1 59 p := r.firstIndexPtr + m*INDEX_BLOCK_LENGTH 60 sip := getLong(r.dbBinStr, p) 61 if ip < sip { 62 h = m - 1 63 } else { 64 eip := getLong(r.dbBinStr, p+4) 65 if ip > eip { 66 l = m + 1 67 } else { 68 dataPtr = getLong(r.dbBinStr, p+8) 69 break 70 } 71 } 72 } 73 if dataPtr == 0 { 74 return &ipInfo, errors.New("not found") 75 } 76 77 dataLen := ((dataPtr >> 24) & 0xFF) 78 dataPtr = (dataPtr & 0x00FFFFFF) 79 ipInfo = getIPInfo(getLong(r.dbBinStr, dataPtr), r.dbBinStr[(dataPtr)+4:dataPtr+dataLen]) 80 return &ipInfo, nil 81 } 82 83 func getLong(b []byte, offset int64) int64 { 84 85 val := (int64(b[offset]) | 86 int64(b[offset+1])<<8 | 87 int64(b[offset+2])<<16 | 88 int64(b[offset+3])<<24) 89 90 return val 91 92 } 93 94 func ip2long(IPStr string) (int64, error) { 95 if IPStr == "" { 96 return 0, fmt.Errorf("ip ins \"\"") 97 } 98 99 bits := strings.Split(IPStr, ".") 100 if len(bits) != 4 { 101 return 0, fmt.Errorf("ip [%s] format error", IPStr) 102 } 103 104 var sum int64 105 for i, n := range bits { 106 bit, _ := strconv.ParseInt(n, 10, 64) 107 sum += bit << uint(24-8*i) 108 } 109 110 return sum, nil 111 } 112 113 func getIPInfo(cityID int64, line []byte) ip2region.IPInfo { 114 lineSlice := strings.Split(string(line), "|") 115 ipInfo := ip2region.IPInfo{} 116 length := len(lineSlice) 117 ipInfo.CityID = cityID 118 if length < 5 { 119 for i := 0; i <= 5-length; i++ { 120 lineSlice = append(lineSlice, "") 121 } 122 } 123 124 ipInfo.Country = lineSlice[0] 125 ipInfo.Region = lineSlice[1] 126 ipInfo.Province = lineSlice[2] 127 ipInfo.City = lineSlice[3] 128 ipInfo.ISP = lineSlice[4] 129 return ipInfo 130 }