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 <Jason@zx2c4.com>
Run-TryBot: Jason A. Donenfeld <Jason@zx2c4.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>

diff --git a/src/net/netip/leaf_alts.go b/src/net/netip/leaf_alts.go
index c51f7dfa54ee433bf84acfe73e8634a9b2093d3b..70513abfd920cec4780e4f55818b268e146a50ef 100644 (file)
--- a/src/net/netip/leaf_alts.go
+++ b/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)
+}

diff --git a/src/net/netip/netip.go b/src/net/netip/netip.go
index 9e08be94fcb2c6de46be7f2c5d0f466c36c3cbd1..90672e045d0cad3916fa5d8ec3b64b9d1308e7a8 100644 (file)
--- a/src/net/netip/netip.go
+++ b/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() {

diff --git a/src/net/netip/netip_test.go b/src/net/netip/netip_test.go
index c39b1ec201efeaa9752b3c546b1d28a031093944..63af853cb3e103b25be4bd60ba348910d8f2bc03 100644 (file)
--- a/src/net/netip/netip_test.go
+++ b/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)
                }
        }
 }