net/netip: add missing encoding.BinaryUnmarshaler to AddrPort and Prefix

The Addr type got an encoding.BinaryUnmarshaler implementation, but not
AddrPort and Prefix. This commit adds the missing implementation of that
interface to these types. It also adds two round trip tests that follow
the template of the existing one for Addr.

Updates #49298.

Change-Id: Iac633aed8aac579960815bb64d06ff3181214841
Reviewed-on: https://go-review.googlesource.com/c/go/+/360875
Trust: Jason A. Donenfeld <[email protected]>
Run-TryBot: Jason A. Donenfeld <[email protected]>
TryBot-Result: Go Bot <[email protected]>
Reviewed-by: Brad Fitzpatrick <[email protected]>

Author: zx2c4

src/net/netip/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)
+}

src/net/netip/netip.go

@@ -1170,6 +1170,34 @@ 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, err := p.Addr().MarshalBinary()
+	if err != nil {
+		return nil, err
+	}
+	b = append(b, 0, 0)
+	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.
@@ -1401,6 +1429,32 @@ 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, err := p.Addr().MarshalBinary()
+	if err != nil {
+		return nil, err
+	}
+	return append(b, uint8(p.Bits())), 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() {

src/net/netip/netip_test.go

@@ -25,6 +25,7 @@ type uint128 = Uint128
 var (
 	mustPrefix = MustParsePrefix
 	mustIP     = MustParseAddr
+	mustIPPort = MustParseAddrPort
 )
 
 func TestParseAddr(t *testing.T) {
@@ -332,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)
 		}
 	}
 }