Repository: inetaf/netaddr
Pull request #12 (2020-04-13 07:26:38)

netaddr: add IPPrefix type

crawshaw: These commits can all be squashed into the original before merging, just experimenting to see if leaving them unsquashed on my end improves the github PR experience for reviewers.

Patch:
diff --git a/netaddr.go b/netaddr.go
index e772308..e076c61 100644
--- a/netaddr.go
+++ b/netaddr.go
@@ -11,7 +11,9 @@ import (
 	"bytes"
 	"errors"
 	"fmt"
+	"math"
 	"net"
+	"strconv"
 	"strings"
 )
 
@@ -323,3 +325,98 @@ func (p IPPort) TCPAddr() *net.TCPAddr {
 		Zone: zone,
 	}
 }
+
+// IPPrefix is an IP address prefix representing an IP network.
+//
+// The first Bits of IP are specified, the remaining bits match any address.
+// The range of Bits is [0,32] for IPv4 or [0,128] for IPv6.
+type IPPrefix struct {
+	IP   IP
+	Bits uint8
+}
+
+// ParseIPPrefix parses s as an IP address prefix.
+// The string can be in the form "192.168.1.0/24" or "2001::db8::/32",
+// the CIDR notation defined in RFC 4632 and RFC 4291.
+func ParseIPPrefix(s string) (IPPrefix, error) {
+	i := strings.IndexByte(s, '/')
+	if i < 0 {
+		return IPPrefix{}, fmt.Errorf("netaddr.ParseIPPrefix(%q): no '/'", s)
+	}
+	ip, err := ParseIP(s[:i])
+	if err != nil {
+		return IPPrefix{}, fmt.Errorf("netaddr.ParseIPPrefix(%q): %w", s, err)
+	}
+	s = s[i+1:]
+	bits, err := strconv.Atoi(s)
+	if err != nil {
+		return IPPrefix{}, fmt.Errorf("netaddr.ParseIPPrefix(%q): bad prefix: %w", s, err)
+	}
+	maxBits := 32
+	if ip.Is6() {
+		maxBits = 128
+	}
+	if bits < 0 || bits > maxBits {
+		return IPPrefix{}, fmt.Errorf("netaddr.ParseIPPrefix(%q): prefix length out of range", s)
+	}
+	return IPPrefix{
+		IP:   ip,
+		Bits: uint8(bits),
+	}, nil
+}
+
+// IPNet returns the net.IPNet representation of an IPPrefix.
+// The returned value is always non-nil.
+// Any zone identifier is dropped in the conversion.
+func (p IPPrefix) IPNet() *net.IPNet {
+	bits := 128
+	if p.IP.Is4() {
+		bits = 32
+	}
+	stdIP, _ := p.IP.ipZone()
+	return &net.IPNet{
+		IP:   stdIP,
+		Mask: net.CIDRMask(int(p.Bits), bits),
+	}
+}
+
+// Contains reports whether the network p includes addr.
+//
+// An IPv4 address will not match an IPv6 prefix.
+// A 4-in-6 IP will not match an IPv4 prefix.
+func (p IPPrefix) Contains(addr IP) bool {
+	var nn, ip []byte // these do not escape and so do not allocate
+	if p.IP.is4() {
+		if !addr.is4() {
+			return false
+		}
+		a1 := p.IP.ipImpl.(v4Addr)
+		a2 := addr.ipImpl.(v4Addr)
+		nn, ip = a1[:], a2[:]
+	} else {
+		if addr.is4() {
+			return false
+		}
+		a1 := p.IP.ipImpl.(v6Addr)
+		a2 := addr.ipImpl.(v6Addr)
+		nn, ip = a1[:], a2[:]
+	}
+	bits := p.Bits
+	for i := 0; bits > 0 && i < len(nn); i++ {
+		m := uint8(math.MaxUint8)
+		if bits < 8 {
+			zeros := 8 - bits
+			m = m >> zeros << zeros
+		}
+		if nn[i]&m != ip[i]&m {
+			return false
+		}
+		bits -= 8
+	}
+	return true
+}
+
+// Strings returns the CIDR notation of p: "<ip>/<bits>".
+func (p IPPrefix) String() string {
+	return fmt.Sprintf("%s/%d", p.IP, p.Bits)
+}


Test patch:
diff --git a/netaddr_test.go b/netaddr_test.go
index 868ef5b..a7b459c 100644
--- a/netaddr_test.go
+++ b/netaddr_test.go
@@ -10,6 +10,7 @@ import (
 	"net"
 	"reflect"
 	"sort"
+	"strings"
 	"testing"
 )
 
@@ -272,6 +273,144 @@ func TestIs4In6(t *testing.T) {
 	}
 }
 
+func TestIPPrefix(t *testing.T) {
+	tests := []struct {
+		prefix      string
+		ip          IP
+		bits        uint8
+		contains    []IP
+		notContains []IP
+	}{
+		{
+			prefix:      "192.168.0.0/24",
+			ip:          mustIP("192.168.0.0"),
+			bits:        24,
+			contains:    mustIPs("192.168.0.1", "192.168.0.55"),
+			notContains: mustIPs("192.168.1.1", "1.1.1.1"),
+		},
+		{
+			prefix:      "192.168.1.1/32",
+			ip:          mustIP("192.168.1.1"),
+			bits:        32,
+			contains:    mustIPs("192.168.1.1"),
+			notContains: mustIPs("192.168.1.2"),
+		},
+		{
+			prefix:      "2001:db8::/96",
+			ip:          mustIP("2001:db8::"),
+			bits:        96,
+			contains:    mustIPs("2001:db8::aaaa:bbbb", "2001:db8::1"),
+			notContains: mustIPs("2001:db8::1:aaaa:bbbb", "2001:db9::"),
+		},
+		{
+			prefix:   "0.0.0.0/0",
+			ip:       mustIP("0.0.0.0"),
+			bits:     0,
+			contains: mustIPs("192.168.0.1", "1.1.1.1"),
+		},
+		{
+			prefix:   "::/0",
+			ip:       mustIP("::"),
+			bits:     0,
+			contains: mustIPs("::1", "2001:db8::1"),
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.prefix, func(t *testing.T) {
+			prefix, err := ParseIPPrefix(test.prefix)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if prefix.IP != test.ip {
+				t.Errorf("IP=%s, want %s", prefix.IP, test.ip)
+			}
+			if prefix.Bits != test.bits {
+				t.Errorf("bits=%d, want %d", prefix.Bits, test.bits)
+			}
+			for _, ip := range test.contains {
+				if !prefix.Contains(ip) {
+					t.Errorf("does not contain %s", ip)
+				}
+			}
+			for _, ip := range test.notContains {
+				if prefix.Contains(ip) {
+					t.Errorf("contains %s", ip)
+				}
+			}
+			if got := prefix.String(); got != test.prefix {
+				t.Errorf("prefix.String()=%q, want %q", got, test.prefix)
+			}
+		})
+	}
+}
+
+func TestParseIPPrefixAllocs(t *testing.T) {
+	tests := []struct {
+		ip    string
+		slash string
+	}{
+		{"192.168.1.0", "/24"},
+		{"aaaa:bbbb:cccc::", "/24"},
+	}
+	for _, test := range tests {
+		prefix := test.ip + test.slash
+		t.Run(prefix, func(t *testing.T) {
+			ipAllocs := int(testing.AllocsPerRun(5, func() {
+				ParseIP(test.ip)
+			}))
+			prefixAllocs := int(testing.AllocsPerRun(5, func() {
+				ParseIPPrefix(prefix)
+			}))
+			if got := prefixAllocs - ipAllocs; got != 0 {
+				t.Errorf("allocs=%d, want 0", got)
+			}
+		})
+	}
+}
+
+func TestParseIPPrefixError(t *testing.T) {
+	tests := []struct {
+		prefix string
+		errstr string
+	}{
+		{
+			prefix: "192.168.0.0",
+			errstr: "no '/'",
+		},
+		{
+			prefix: "1.257.1.1/24",
+			errstr: "lookup 1.257.1.1: no such host", // TODO: improve ParseIP error
+		},
+		{
+			prefix: "1.1.1.0/q",
+			errstr: "bad prefix",
+		},
+		{
+			prefix: "1.1.1.0/-1",
+			errstr: "out of range",
+		},
+		{
+			prefix: "1.1.1.0/33",
+			errstr: "out of range",
+		},
+		{
+			prefix: "2001::/129",
+			errstr: "out of range",
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.prefix, func(t *testing.T) {
+			_, err := ParseIPPrefix(test.prefix)
+			if err == nil {
+				t.Fatal("no error")
+			}
+			if got := err.Error(); !strings.Contains(got, test.errstr) {
+				t.Errorf("error is missing substring %q: %s", test.errstr, got)
+			}
+		})
+	}
+}
+
 func mustIP(s string) IP {
 	ip, err := ParseIP(s)
 	if err != nil {
@@ -281,6 +420,14 @@ func mustIP(s string) IP {
 	return ip
 }
 
+func mustIPs(strs ...string) []IP {
+	var res []IP
+	for _, s := range strs {
+		res = append(res, mustIP(s))
+	}
+	return res
+}
+
 func BenchmarkStdIPv4(b *testing.B) {
 	b.ReportAllocs()
 	ips := []net.IP{}
@@ -357,3 +504,15 @@ func BenchmarkIPv6(b *testing.B) {
 		}
 	}
 }
+
+func BenchmarkIPv4Contains(b *testing.B) {
+	b.ReportAllocs()
+	prefix := IPPrefix{
+		IP:   IPv4(192, 168, 1, 0),
+		Bits: 24,
+	}
+	ip := IPv4(192, 168, 1, 1)
+	for i := 0; i < b.N; i++ {
+		prefix.Contains(ip)
+	}
+}