Import from e83a204HEADmaster

Signed-off-by: Jason A. Donenfeld <Jason@zx2c4.com>

3 files changed, 254 insertions, 69 deletions

diff --git a/leaf_alts.go b/leaf_alts.go
index c51f7df..70513ab 100644
--- a/leaf_alts.go
+++ b/leaf_alts.go
@@ -41,3 +41,14 @@ func bePutUint32(b []byte, v uint32) {
 	b[2] = byte(v >> 8)
 	b[3] = byte(v)
 }
+
+func leUint16(b []byte) uint16 {
+	_ = b[1] // bounds check hint to compiler; see golang.org/issue/14808
+	return uint16(b[0]) | uint16(b[1])<<8
+}
+
+func lePutUint16(b []byte, v uint16) {
+	_ = b[1] // early bounds check to guarantee safety of writes below
+	b[0] = byte(v)
+	b[1] = byte(v >> 8)
+}
diff --git a/netip.go b/netip.go
index 678197f..82cd4cd 100644
--- a/netip.go
+++ b/netip.go
@@ -2,9 +2,9 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
-// Package netip defines a IP address type that's a small value type.
+// Package netip defines an IP address type that's a small value type.
 // Building on that Addr type, the package also defines AddrPort (an
-// IP address and a port), and Prefix (a IP address and a bit length
+// IP address and a port), and Prefix (an IP address and a bit length
 // prefix).
 //
 // Compared to the net.IP type, this package's Addr type takes less
@@ -32,6 +32,9 @@ import (
 //
 // Unlike net.IP or net.IPAddr, Addr is a comparable value
 // type (it supports == and can be a map key) and is immutable.
+//
+// The zero Addr is not a valid IP address.
+// Addr{} is distinct from both 0.0.0.0 and ::.
 type Addr struct {
 	// addr is the hi and lo bits of an IPv6 address. If z==z4,
 	// hi and lo contain the IPv4-mapped IPv6 address.
@@ -75,6 +78,9 @@ func IPv6LinkLocalAllNodes() Addr { return AddrFrom16([16]byte{0: 0xff, 1: 0x02,
 // IPv6Unspecified returns the IPv6 unspecified address "::".
 func IPv6Unspecified() Addr { return Addr{z: z6noz} }
 
+// IPv4Unspecified returns the IPv4 unspecified address "0.0.0.0".
+func IPv4Unspecified() Addr { return AddrFrom4([4]byte{}) }
+
 // AddrFrom4 returns the address of the IPv4 address given by the bytes in addr.
 func AddrFrom4(addr [4]byte) Addr {
 	return Addr{
@@ -155,9 +161,14 @@ func (err parseAddrError) Error() string {
 func parseIPv4(s string) (ip Addr, err error) {
 	var fields [4]uint8
 	var val, pos int
+	var digLen int // number of digits in current octet
 	for i := 0; i < len(s); i++ {
 		if s[i] >= '0' && s[i] <= '9' {
+			if digLen == 1 && val == 0 {
+				return Addr{}, parseAddrError{in: s, msg: "IPv4 field has octet with leading zero"}
+			}
 			val = val*10 + int(s[i]) - '0'
+			digLen++
 			if val > 255 {
 				return Addr{}, parseAddrError{in: s, msg: "IPv4 field has value >255"}
 			}
@@ -175,6 +186,7 @@ func parseIPv4(s string) (ip Addr, err error) {
 			fields[pos] = uint8(val)
 			pos++
 			val = 0
+			digLen = 0
 		} else {
 			return Addr{}, parseAddrError{in: s, msg: "unexpected character", at: s[i:]}
 		}
@@ -442,31 +454,6 @@ func (ip Addr) Less(ip2 Addr) bool { return ip.Compare(ip2) == -1 }
 
 func (ip Addr) lessOrEq(ip2 Addr) bool { return ip.Compare(ip2) <= 0 }
 
-// ipZone returns the standard library net.IP from ip, as well
-// as the zone.
-// The optional reuse IP provides memory to reuse.
-func (ip Addr) ipZone(reuse []byte) (stdIP []byte, zone string) {
-	base := reuse[:0]
-	switch {
-	case ip.z == z0:
-		return nil, ""
-	case ip.Is4():
-		a4 := ip.As4()
-		return append(base, a4[:]...), ""
-	default:
-		a16 := ip.As16()
-		return append(base, a16[:]...), ip.Zone()
-	}
-}
-
-// IPAddrParts returns the net.IPAddr representation of an Addr.
-//
-// The slice will be nil if ip is the zero Addr.
-// The zone is the empty string if there is no zone.
-func (ip Addr) IPAddrParts() (slice []byte, zone string) {
-	return ip.ipZone(nil)
-}
-
 // Is4 reports whether ip is an IPv4 address.
 //
 // It returns false for IP4-mapped IPv6 addresses. See IP.Unmap.
@@ -614,7 +601,7 @@ func (ip Addr) IsGlobalUnicast() bool {
 
 	// Match package net's IsGlobalUnicast logic. Notably private IPv4 addresses
 	// and ULA IPv6 addresses are still considered "global unicast".
-	if ip.Is4() && (ip == AddrFrom4([4]byte{}) || ip == AddrFrom4([4]byte{255, 255, 255, 255})) {
+	if ip.Is4() && (ip == IPv4Unspecified() || ip == AddrFrom4([4]byte{255, 255, 255, 255})) {
 		return false
 	}
 
@@ -652,7 +639,7 @@ func (ip Addr) IsPrivate() bool {
 //
 // Note that the zero Addr is not an unspecified address.
 func (ip Addr) IsUnspecified() bool {
-	return ip == AddrFrom4([4]byte{}) || ip == IPv6Unspecified()
+	return ip == IPv4Unspecified() || ip == IPv6Unspecified()
 }
 
 // Prefix keeps only the top b bits of IP, producing a Prefix
@@ -692,21 +679,19 @@ const (
 // IPv6 addresses with zones are returned without their zone (use the
 // Zone method to get it).
 // The ip zero value returns all zeroes.
-func (ip Addr) As16() [16]byte {
-	var ret [16]byte
-	bePutUint64(ret[:8], ip.addr.hi)
-	bePutUint64(ret[8:], ip.addr.lo)
-	return ret
+func (ip Addr) As16() (a16 [16]byte) {
+	bePutUint64(a16[:8], ip.addr.hi)
+	bePutUint64(a16[8:], ip.addr.lo)
+	return a16
 }
 
 // As4 returns an IPv4 or IPv4-in-IPv6 address in its 4-byte representation.
 // If ip is the zero Addr or an IPv6 address, As4 panics.
 // Note that 0.0.0.0 is not the zero Addr.
-func (ip Addr) As4() [4]byte {
+func (ip Addr) As4() (a4 [4]byte) {
 	if ip.z == z4 || ip.Is4In6() {
-		var ret [4]byte
-		bePutUint32(ret[:], uint32(ip.addr.lo))
-		return ret
+		bePutUint32(a4[:], uint32(ip.addr.lo))
+		return a4
 	}
 	if ip.z == z0 {
 		panic("As4 called on IP zero value")
@@ -714,6 +699,23 @@ func (ip Addr) As4() [4]byte {
 	panic("As4 called on IPv6 address")
 }
 
+// AsSlice returns an IPv4 or IPv6 address in its respective 4-byte or 16-byte representation.
+func (ip Addr) AsSlice() []byte {
+	switch ip.z {
+	case z0:
+		return nil
+	case z4:
+		var ret [4]byte
+		bePutUint32(ret[:], uint32(ip.addr.lo))
+		return ret[:]
+	default:
+		var ret [16]byte
+		bePutUint64(ret[:8], ip.addr.hi)
+		bePutUint64(ret[8:], ip.addr.lo)
+		return ret[:]
+	}
+}
+
 // Next returns the address following ip.
 // If there is none, it returns the zero Addr.
 func (ip Addr) Next() Addr {
@@ -966,25 +968,30 @@ func (ip *Addr) UnmarshalText(text []byte) error {
 	return err
 }
 
-// MarshalBinary implements the encoding.BinaryMarshaler interface.
-// It returns a zero-length slice for the zero Addr,
-// the 4-byte form for an IPv4 address,
-// and the 16-byte form with zone appended for an IPv6 address.
-func (ip Addr) MarshalBinary() ([]byte, error) {
+func (ip Addr) marshalBinaryWithTrailingBytes(trailingBytes int) []byte {
+	var b []byte
 	switch ip.z {
 	case z0:
-		return nil, nil
+		b = make([]byte, trailingBytes)
 	case z4:
-		b := ip.As4()
-		return b[:], nil
+		b = make([]byte, 4+trailingBytes)
+		bePutUint32(b, uint32(ip.addr.lo))
 	default:
-		b16 := ip.As16()
-		b := b16[:]
-		if z := ip.Zone(); z != "" {
-			b = append(b, []byte(z)...)
-		}
-		return b, nil
+		z := ip.Zone()
+		b = make([]byte, 16+len(z)+trailingBytes)
+		bePutUint64(b[:8], ip.addr.hi)
+		bePutUint64(b[8:], ip.addr.lo)
+		copy(b[16:], z)
 	}
+	return b
+}
+
+// MarshalBinary implements the encoding.BinaryMarshaler interface.
+// It returns a zero-length slice for the zero Addr,
+// the 4-byte form for an IPv4 address,
+// and the 16-byte form with zone appended for an IPv6 address.
+func (ip Addr) MarshalBinary() ([]byte, error) {
+	return ip.marshalBinaryWithTrailingBytes(0), nil
 }
 
 // UnmarshalBinary implements the encoding.BinaryUnmarshaler interface.
@@ -1174,6 +1181,30 @@ func (p *AddrPort) UnmarshalText(text []byte) error {
 	return err
 }
 
+// MarshalBinary implements the encoding.BinaryMarshaler interface.
+// It returns Addr.MarshalBinary with an additional two bytes appended
+// containing the port in little-endian.
+func (p AddrPort) MarshalBinary() ([]byte, error) {
+	b := p.Addr().marshalBinaryWithTrailingBytes(2)
+	lePutUint16(b[len(b)-2:], p.Port())
+	return b, nil
+}
+
+// UnmarshalBinary implements the encoding.BinaryUnmarshaler interface.
+// It expects data in the form generated by MarshalBinary.
+func (p *AddrPort) UnmarshalBinary(b []byte) error {
+	if len(b) < 2 {
+		return errors.New("unexpected slice size")
+	}
+	var addr Addr
+	err := addr.UnmarshalBinary(b[:len(b)-2])
+	if err != nil {
+		return err
+	}
+	*p = AddrPortFrom(addr, leUint16(b[len(b)-2:]))
+	return nil
+}
+
 // Prefix is an IP address prefix (CIDR) representing an IP network.
 //
 // The first Bits() of Addr() are specified. The remaining bits match any address.
@@ -1184,7 +1215,7 @@ type Prefix struct {
 	// bits is logically a uint8 (storing [0,128]) but also
 	// encodes an "invalid" bit, currently represented by the
 	// invalidPrefixBits sentinel value. It could be packed into
-	// the uint8 more with more comlicated expressions in the
+	// the uint8 more with more complicated expressions in the
 	// accessors, but the extra byte (in padding anyway) doesn't
 	// hurt and simplifies code below.
 	bits int16
@@ -1195,7 +1226,7 @@ type Prefix struct {
 // public API.
 const invalidPrefixBits = -1
 
-// PrefixFrom returns an Prefix with the provided IP address and bit
+// PrefixFrom returns a Prefix with the provided IP address and bit
 // prefix length.
 //
 // It does not allocate. Unlike Addr.Prefix, PrefixFrom does not mask
@@ -1222,7 +1253,7 @@ func (p Prefix) Addr() Addr { return p.ip }
 // It reports -1 if invalid.
 func (p Prefix) Bits() int { return int(p.bits) }
 
-// IsValid reports whether whether p.Bits() has a valid range for p.IP().
+// IsValid reports whether p.Bits() has a valid range for p.IP().
 // If p.Addr() is the zero Addr, IsValid returns false.
 // Note that if p is the zero Prefix, then p.IsValid() == false.
 func (p Prefix) IsValid() bool { return !p.ip.isZero() && p.bits >= 0 && int(p.bits) <= p.ip.BitLen() }
@@ -1405,6 +1436,30 @@ func (p *Prefix) UnmarshalText(text []byte) error {
 	return err
 }
 
+// MarshalBinary implements the encoding.BinaryMarshaler interface.
+// It returns Addr.MarshalBinary with an additional byte appended
+// containing the prefix bits.
+func (p Prefix) MarshalBinary() ([]byte, error) {
+	b := p.Addr().withoutZone().marshalBinaryWithTrailingBytes(1)
+	b[len(b)-1] = uint8(p.Bits())
+	return b, nil
+}
+
+// UnmarshalBinary implements the encoding.BinaryUnmarshaler interface.
+// It expects data in the form generated by MarshalBinary.
+func (p *Prefix) UnmarshalBinary(b []byte) error {
+	if len(b) < 1 {
+		return errors.New("unexpected slice size")
+	}
+	var addr Addr
+	err := addr.UnmarshalBinary(b[:len(b)-1])
+	if err != nil {
+		return err
+	}
+	*p = PrefixFrom(addr, int(b[len(b)-1]))
+	return nil
+}
+
 // String returns the CIDR notation of p: "<ip>/<bits>".
 func (p Prefix) String() string {
 	if !p.IsValid() {
diff --git a/netip_test.go b/netip_test.go
index 9191a0b..f113659 100644
--- a/netip_test.go
+++ b/netip_test.go
@@ -25,13 +25,15 @@ type uint128 = Uint128
 var (
 	mustPrefix = MustParsePrefix
 	mustIP     = MustParseAddr
+	mustIPPort = MustParseAddrPort
 )
 
 func TestParseAddr(t *testing.T) {
 	var validIPs = []struct {
-		in  string
-		ip  Addr   // output of ParseAddr()
-		str string // output of String(). If "", use in.
+		in      string
+		ip      Addr   // output of ParseAddr()
+		str     string // output of String(). If "", use in.
+		wantErr string
 	}{
 		// Basic zero IPv4 address.
 		{
@@ -45,15 +47,18 @@ func TestParseAddr(t *testing.T) {
 		},
 		// IPv4 address in windows-style "print all the digits" form.
 		{
-			in:  "010.000.015.001",
-			ip:  MkAddr(Mk128(0, 0xffff0a000f01), Z4),
-			str: "10.0.15.1",
+			in:      "010.000.015.001",
+			wantErr: `ParseAddr("010.000.015.001"): IPv4 field has octet with leading zero`,
 		},
 		// IPv4 address with a silly amount of leading zeros.
 		{
-			in:  "000001.00000002.00000003.000000004",
-			ip:  MkAddr(Mk128(0, 0xffff01020304), Z4),
-			str: "1.2.3.4",
+			in:      "000001.00000002.00000003.000000004",
+			wantErr: `ParseAddr("000001.00000002.00000003.000000004"): IPv4 field has octet with leading zero`,
+		},
+		// 4-in-6 with octet with leading zero
+		{
+			in:      "::ffff:1.2.03.4",
+			wantErr: `ParseAddr("::ffff:1.2.03.4"): ParseAddr("1.2.03.4"): IPv4 field has octet with leading zero (at "1.2.03.4")`,
 		},
 		// Basic zero IPv6 address.
 		{
@@ -121,10 +126,16 @@ func TestParseAddr(t *testing.T) {
 		t.Run(test.in, func(t *testing.T) {
 			got, err := ParseAddr(test.in)
 			if err != nil {
+				if err.Error() == test.wantErr {
+					return
+				}
 				t.Fatal(err)
 			}
+			if test.wantErr != "" {
+				t.Fatalf("wanted error %q; got none", test.wantErr)
+			}
 			if got != test.ip {
-				t.Errorf("ParseAddr(%q) got %#v, want %#v", test.in, got, test.ip)
+				t.Errorf("got %#v, want %#v", got, test.ip)
 			}
 
 			// Check that ParseAddr is a pure function.
@@ -322,10 +333,91 @@ func TestAddrMarshalUnmarshalBinary(t *testing.T) {
 	}
 
 	// Cannot unmarshal from unexpected IP length.
-	for _, l := range []int{3, 5} {
+	for _, n := range []int{3, 5} {
 		var ip2 Addr
-		if err := ip2.UnmarshalBinary(bytes.Repeat([]byte{1}, l)); err == nil {
-			t.Fatalf("unmarshaled from unexpected IP length %d", l)
+		if err := ip2.UnmarshalBinary(bytes.Repeat([]byte{1}, n)); err == nil {
+			t.Fatalf("unmarshaled from unexpected IP length %d", n)
+		}
+	}
+}
+
+func TestAddrPortMarshalUnmarshalBinary(t *testing.T) {
+	tests := []struct {
+		ipport   string
+		wantSize int
+	}{
+		{"1.2.3.4:51820", 4 + 2},
+		{"[fd7a:115c:a1e0:ab12:4843:cd96:626b:430b]:80", 16 + 2},
+		{"[::ffff:c000:0280]:65535", 16 + 2},
+		{"[::ffff:c000:0280%eth0]:1", 20 + 2},
+	}
+	for _, tc := range tests {
+		var ipport AddrPort
+		if len(tc.ipport) > 0 {
+			ipport = mustIPPort(tc.ipport)
+		}
+		b, err := ipport.MarshalBinary()
+		if err != nil {
+			t.Fatal(err)
+		}
+		if len(b) != tc.wantSize {
+			t.Fatalf("%q encoded to size %d; want %d", tc.ipport, len(b), tc.wantSize)
+		}
+		var ipport2 AddrPort
+		if err := ipport2.UnmarshalBinary(b); err != nil {
+			t.Fatal(err)
+		}
+		if ipport != ipport2 {
+			t.Fatalf("got %v; want %v", ipport2, ipport)
+		}
+	}
+
+	// Cannot unmarshal from unexpected lengths.
+	for _, n := range []int{3, 7} {
+		var ipport2 AddrPort
+		if err := ipport2.UnmarshalBinary(bytes.Repeat([]byte{1}, n)); err == nil {
+			t.Fatalf("unmarshaled from unexpected length %d", n)
+		}
+	}
+}
+
+func TestPrefixMarshalUnmarshalBinary(t *testing.T) {
+	type testCase struct {
+		prefix   Prefix
+		wantSize int
+	}
+	tests := []testCase{
+		{mustPrefix("1.2.3.4/24"), 4 + 1},
+		{mustPrefix("fd7a:115c:a1e0:ab12:4843:cd96:626b:430b/118"), 16 + 1},
+		{mustPrefix("::ffff:c000:0280/96"), 16 + 1},
+		{mustPrefix("::ffff:c000:0280%eth0/37"), 16 + 1}, // Zone should be stripped
+	}
+	tests = append(tests,
+		testCase{PrefixFrom(tests[0].prefix.Addr(), 33), tests[0].wantSize},
+		testCase{PrefixFrom(tests[1].prefix.Addr(), 129), tests[1].wantSize})
+	for _, tc := range tests {
+		prefix := tc.prefix
+		b, err := prefix.MarshalBinary()
+		if err != nil {
+			t.Fatal(err)
+		}
+		if len(b) != tc.wantSize {
+			t.Fatalf("%q encoded to size %d; want %d", tc.prefix, len(b), tc.wantSize)
+		}
+		var prefix2 Prefix
+		if err := prefix2.UnmarshalBinary(b); err != nil {
+			t.Fatal(err)
+		}
+		if prefix != prefix2 {
+			t.Fatalf("got %v; want %v", prefix2, prefix)
+		}
+	}
+
+	// Cannot unmarshal from unexpected lengths.
+	for _, n := range []int{3, 6} {
+		var prefix2 Prefix
+		if err := prefix2.UnmarshalBinary(bytes.Repeat([]byte{1}, n)); err == nil {
+			t.Fatalf("unmarshaled from unexpected length %d", n)
 		}
 	}
 }
@@ -963,7 +1055,7 @@ func TestIs4In6(t *testing.T) {
 		{mustIP("::ffff:192.0.2.128"), true, mustIP("192.0.2.128")},
 		{mustIP("::ffff:192.0.2.128%eth0"), true, mustIP("192.0.2.128")},
 		{mustIP("::fffe:c000:0280"), false, mustIP("::fffe:c000:0280")},
-		{mustIP("::ffff:127.001.002.003"), true, mustIP("127.1.2.3")},
+		{mustIP("::ffff:127.1.2.3"), true, mustIP("127.1.2.3")},
 		{mustIP("::ffff:7f01:0203"), true, mustIP("127.1.2.3")},
 		{mustIP("0:0:0:0:0000:ffff:127.1.2.3"), true, mustIP("127.1.2.3")},
 		{mustIP("0:0:0:0:000000:ffff:127.1.2.3"), true, mustIP("127.1.2.3")},
@@ -1796,3 +1888,30 @@ func TestInvalidAddrPortString(t *testing.T) {
 		}
 	}
 }
+
+func TestAsSlice(t *testing.T) {
+	tests := []struct {
+		in   Addr
+		want []byte
+	}{
+		{in: Addr{}, want: nil},
+		{in: mustIP("1.2.3.4"), want: []byte{1, 2, 3, 4}},
+		{in: mustIP("ffff::1"), want: []byte{0xff, 0xff, 15: 1}},
+	}
+
+	for _, test := range tests {
+		got := test.in.AsSlice()
+		if !bytes.Equal(got, test.want) {
+			t.Errorf("%v.AsSlice() = %v want %v", test.in, got, test.want)
+		}
+	}
+}
+
+var sink16 [16]byte
+
+func BenchmarkAs16(b *testing.B) {
+	addr := MustParseAddr("1::10")
+	for i := 0; i < b.N; i++ {
+		sink16 = addr.As16()
+	}
+}