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  }