Add network field support to ASN database queries

reader.go

@@ -19,28 +19,36 @@ type reader struct {
 	nodeOffsetMult    uint
 }
 
-func (r *reader) getOffset(ip net.IP) (uint, error) {
-	pointer, err := r.lookupPointer(ip)
+func (r *reader) getOffsetWithPrefix(ip net.IP) (uint, uint, error) {
+	pointer, prefix, err := r.lookupPointer(ip)
 	if err != nil {
-		return 0, err
+		return 0, 0, err
 	}
 	offset := pointer - uint(r.metadata.NodeCount) - uint(dataSectionSeparatorSize)
 	if offset >= uint(len(r.buffer)) {
-		return 0, errors.New("the MaxMind DB search tree is corrupt: " + strconv.Itoa(int(pointer)))
+		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, error) {
+func (r *reader) lookupPointer(ip net.IP) (uint, uint, error) {
 	if ip == nil {
-		return 0, errors.New("IP cannot be 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, errors.New("cannot look up an IPv6 address in an IPv4-only database")
+		return 0, 0, errors.New("cannot look up an IPv6 address in an IPv4-only database")
 	}
 	bitCount := uint(len(ip)) * 8
 	node := uint(0)
@@ -59,11 +67,11 @@ func (r *reader) lookupPointer(ip net.IP) (uint, error) {
 		}
 	}
 	if node == nodeCount {
-		return 0, ErrNotFound
+		return 0, 0, ErrNotFound
 	} else if node > nodeCount {
-		return node, nil
+		return node, i, nil
 	}
-	return 0, errors.New("invalid node in search tree")
+	return 0, 0, errors.New("invalid node in search tree")
 }
 
 func (r *reader) readLeft(nodeNumber uint) uint {

reader_asn.go

@@ -2,6 +2,7 @@ package geoip2
 
 import (
 	"errors"
+	"fmt"
 	"io/ioutil"
 	"net"
 	"strconv"
@@ -12,7 +13,7 @@ type ASNReader struct {
 }
 
 func (r *ASNReader) Lookup(ip net.IP) (*ASN, error) {
-	offset, err := r.getOffset(ip)
+	offset, prefix, err := r.getOffsetWithPrefix(ip)
 	if err != nil {
 		return nil, err
 	}
@@ -21,6 +22,8 @@ func (r *ASNReader) Lookup(ip net.IP) (*ASN, error) {
 		return nil, err
 	}
 	result := &ASN{}
+	_, network, err := net.ParseCIDR(fmt.Sprintf("%s/%s", ip.String(), strconv.Itoa(int(prefix))))
+	result.Network = network.String()
 	switch dataType {
 	case dataTypeMap:
 		_, err = readASNMap(result, r.decoderBuffer, size, offset)

reader_test.go

@@ -396,6 +396,9 @@ func TestASN(t *testing.T) {
 	if record.AutonomousSystemOrganization != "Telstra Pty Ltd" {
 		t.Fatal()
 	}
+	if record.Network != "1.128.0.0/11" {
+		t.Fatal()
+	}
 
 	record, err = reader.Lookup(net.ParseIP("2600:6000::"))
 	if err != nil {
@@ -407,6 +410,10 @@ func TestASN(t *testing.T) {
 	if record.AutonomousSystemOrganization != "Merit Network Inc." {
 		t.Fatal()
 	}
+	if record.Network != "2600:6000::/20" {
+		t.Fatal()
+	}
+
 }
 
 func TestDBIPCity(t *testing.T) {
@@ -487,4 +494,7 @@ func TestDBIPASN(t *testing.T) {
 	if record.AutonomousSystemOrganization != "Comcast Cable Communications, LLC" {
 		t.Fatal()
 	}
+	if record.Network != "66.30.0.0/15" {
+		t.Fatal()
+	}
 }

types.go

@@ -123,6 +123,7 @@ type AnonymousIP struct {
 type ASN struct {
 	AutonomousSystemNumber       uint32
 	AutonomousSystemOrganization string
+	Network                      string
 }
 
 type Domain struct {