reader.go

package geoip2

import (
	"bytes"
	"errors"
	"net"
	"strconv"
)

var ErrNotFound = errors.New("not found")

type reader struct {
	metadata          *Metadata
	buffer            []byte
	decoderBuffer     []byte
	nodeBuffer        []byte
	ipV4Start         uint
	ipV4StartBitDepth uint
	nodeOffsetMult    uint
}

func (r *reader) getOffsetWithPrefix(ip net.IP) (uint, uint, error) {
	pointer, prefix, err := r.lookupPointer(ip)
	if err != nil {
		return 0, 0, err
	}
	offset := pointer - uint(r.metadata.NodeCount) - uint(dataSectionSeparatorSize)
	if offset >= uint(len(r.buffer)) {
		return 0, 0, errors.New("the MaxMind DB search tree is corrupt: " + strconv.Itoa(int(pointer)))
	}
	return offset, prefix, nil
}

func (r *reader) getOffset(ip net.IP) (uint, error) {
	offset, _, err := r.getOffsetWithPrefix(ip)
	if err != nil {
		return 0, err
	}
	return offset, nil
}

func (r *reader) lookupPointer(ip net.IP) (uint, uint, error) {
	if ip == nil {
		return 0, 0, errors.New("IP cannot be nil")
	}
	ipV4 := ip.To4()
	if ipV4 != nil {
		ip = ipV4
	}
	if len(ip) == 16 && r.metadata.IPVersion == 4 {
		return 0, 0, errors.New("cannot look up an IPv6 address in an IPv4-only database")
	}
	bitCount := uint(len(ip)) * 8
	node := uint(0)
	if bitCount == 32 {
		node = r.ipV4Start
	}
	nodeCount := uint(r.metadata.NodeCount)
	i := uint(0)
	for ; i < bitCount && node < nodeCount; i++ {
		bit := 1 & (ip[i>>3] >> (7 - (i % 8)))
		offset := node * r.nodeOffsetMult
		if bit == 0 {
			node = r.readLeft(offset)
		} else {
			node = r.readRight(offset)
		}
	}
	if node == nodeCount {
		return 0, 0, ErrNotFound
	} else if node > nodeCount {
		return node, i, nil
	}
	return 0, 0, errors.New("invalid node in search tree")
}

func (r *reader) readLeft(nodeNumber uint) uint {
	switch r.metadata.RecordSize {
	case 28:
		return ((uint(r.nodeBuffer[nodeNumber+3]) & 0xF0) << 20) | (uint(r.nodeBuffer[nodeNumber]) << 16) | (uint(r.nodeBuffer[nodeNumber+1]) << 8) | uint(r.nodeBuffer[nodeNumber+2])
	case 24:
		return (uint(r.nodeBuffer[nodeNumber]) << 16) | (uint(r.nodeBuffer[nodeNumber+1]) << 8) | uint(r.nodeBuffer[nodeNumber+2])
	default: // case 32:
		return (uint(r.nodeBuffer[nodeNumber]) << 24) | (uint(r.nodeBuffer[nodeNumber+1]) << 16) | (uint(r.nodeBuffer[nodeNumber+2]) << 8) | uint(r.nodeBuffer[nodeNumber+3])
	}
}

func (r *reader) readRight(nodeNumber uint) uint {
	switch r.metadata.RecordSize {
	case 28:
		return ((uint(r.nodeBuffer[nodeNumber+3]) & 0x0F) << 24) | (uint(r.nodeBuffer[nodeNumber+4]) << 16) | (uint(r.nodeBuffer[nodeNumber+5]) << 8) | uint(r.nodeBuffer[nodeNumber+6])
	case 24:
		return (uint(r.nodeBuffer[nodeNumber+3]) << 16) | (uint(r.nodeBuffer[nodeNumber+4]) << 8) | uint(r.nodeBuffer[nodeNumber+5])
	default: // case 32:
		return (uint(r.nodeBuffer[nodeNumber+4]) << 24) | (uint(r.nodeBuffer[nodeNumber+5]) << 16) | (uint(r.nodeBuffer[nodeNumber+6]) << 8) | uint(r.nodeBuffer[nodeNumber+7])
	}
}

func newReader(buffer []byte) (*reader, error) {
	if len(buffer) == 0 {
		return nil, errors.New("buffer is empty")
	}

	metadataStart := bytes.LastIndex(buffer, metadataStartMarker)
	metadata, err := readMetadata(buffer[metadataStart+len(metadataStartMarker):])
	if err != nil {
		return nil, err
	}
	nodeOffsetMult := uint(metadata.RecordSize) / 4
	searchTreeSize := uint(metadata.NodeCount) * nodeOffsetMult
	dataSectionStart := searchTreeSize + dataSectionSeparatorSize
	if dataSectionStart > uint(metadataStart) {
		return nil, errors.New("the MaxMind DB contains invalid metadata")
	}
	reader := &reader{
		metadata:       metadata,
		buffer:         buffer,
		decoderBuffer:  buffer[searchTreeSize+dataSectionSeparatorSize : metadataStart],
		nodeBuffer:     buffer[:searchTreeSize],
		nodeOffsetMult: nodeOffsetMult,
	}
	if metadata.IPVersion == 6 {
		node := uint(0)
		i := uint(0)
		for ; i < 96 && node < uint(metadata.NodeCount); i++ {
			node = reader.readLeft(node * nodeOffsetMult)
		}
		reader.ipV4Start = node
		reader.ipV4StartBitDepth = i
	}
	return reader, nil
}