Implement uio.Unmarshaler for Message/RelayMessage/Options and use them

hugelgupf · hugelgupf · commit 50801df7ed2a · 2023-03-05T14:45:58.000-08:00
Moves towards avoiding allocating a new buf.NewBigEndianBuffer in every
single FromBytes / ToBytes function.

Signed-off-by: Chris Koch &lt;chrisko@google.com&gt;
diff --git a/dhcpv6/dhcpv6message.go b/dhcpv6/dhcpv6message.go
@@ -364,21 +364,21 @@ type Message struct {
 // FromBytes parses a DHCPv6 message from a byte stream.
 func (m *Message) FromBytes(data []byte) error {
 	buf := uio.NewBigEndianBuffer(data)
+	m.Unmarshal(buf)
+	return buf.FinError()
+}
+
+// Unmarshal parses a DHCPv6 message from buf.
+func (m *Message) Unmarshal(buf *uio.Lexer) {
 	messageType := MessageType(buf.Read8())
 
 	if messageType == MessageTypeRelayForward || messageType == MessageTypeRelayReply {
-		return fmt.Errorf("wrong message type")
+		buf.SetError(fmt.Errorf("wrong message type"))
 	}
 
 	m.MessageType = messageType
 	buf.ReadBytes(m.TransactionID[:])
-	if buf.Error() != nil {
-		return fmt.Errorf("failed to parse DHCPv6 header: %w", buf.Error())
-	}
-	if err := m.Options.FromBytes(buf.Data()); err != nil {
-		return err
-	}
-	return nil
+	m.Options.Unmarshal(buf)
 }
 
 var randomRead = rand.Read
diff --git a/dhcpv6/dhcpv6relay.go b/dhcpv6/dhcpv6relay.go
@@ -82,25 +82,25 @@ type RelayMessage struct {
 // FromBytes parses a relay message from a byte stream.
 func (r *RelayMessage) FromBytes(data []byte) error {
 	buf := uio.NewBigEndianBuffer(data)
+	r.Unmarshal(buf)
+	return buf.FinError()
+}
+
+// Unmarshal parses a relay message from a buf.
+func (r *RelayMessage) Unmarshal(buf *uio.Lexer) {
 	messageType := MessageType(buf.Read8())
 
 	if messageType != MessageTypeRelayForward && messageType != MessageTypeRelayReply {
-		return fmt.Errorf("wrong message type")
+		buf.SetError(fmt.Errorf("wrong message type"))
 	}
 
 	r.MessageType = messageType
 	r.HopCount = buf.Read8()
 	r.LinkAddr = net.IP(buf.CopyN(net.IPv6len))
 	r.PeerAddr = net.IP(buf.CopyN(net.IPv6len))
 
-	if buf.Error() != nil {
-		return fmt.Errorf("error parsing RelayMessage header: %w", buf.Error())
-	}
 	// TODO: fail if no OptRelayMessage is present.
-	if err := r.Options.FromBytes(buf.Data()); err != nil {
-		return err
-	}
-	return nil
+	r.Options.Unmarshal(buf)
 }
 
 func write16(b *uio.Lexer, ip net.IP) {
diff --git a/dhcpv6/option_iaaddress.go b/dhcpv6/option_iaaddress.go
@@ -80,9 +80,6 @@ func (op *OptIAAddress) FromBytes(data []byte) error {
 	t2.Unmarshal(buf)
 	op.PreferredLifetime = t1.Duration
 	op.ValidLifetime = t2.Duration
-
-	if err := op.Options.FromBytes(buf.ReadAll()); err != nil {
-		return err
-	}
+	op.Options.Unmarshal(buf)
 	return buf.FinError()
 }
diff --git a/dhcpv6/option_iaaddress_test.go b/dhcpv6/option_iaaddress_test.go
@@ -90,9 +90,9 @@ func TestIAAddressParseAndGetter(t *testing.T) {
 			},
 		},
 		{
-			buf:  []byte{0, 3, 0, 1, 0},
+			buf:  []byte{0, 5, 0, 1, 0},
 			want: nil,
-			err:  uio.ErrUnreadBytes,
+			err:  uio.ErrBufferTooShort,
 		},
 		{
 			buf: []byte{
diff --git a/dhcpv6/option_iapd.go b/dhcpv6/option_iapd.go
@@ -93,9 +93,6 @@ func (op *OptIAPD) FromBytes(data []byte) error {
 	t2.Unmarshal(buf)
 	op.T1 = t1.Duration
 	op.T2 = t2.Duration
-
-	if err := op.Options.FromBytes(buf.ReadAll()); err != nil {
-		return err
-	}
+	op.Options.Unmarshal(buf)
 	return buf.FinError()
 }
diff --git a/dhcpv6/option_iapd_test.go b/dhcpv6/option_iapd_test.go
@@ -111,7 +111,7 @@ func TestIAPDParseAndGetter(t *testing.T) {
 		{
 			buf:  []byte{0, 25, 0, 1, 0},
 			want: nil,
-			err:  uio.ErrUnreadBytes,
+			err:  uio.ErrBufferTooShort,
 		},
 		{
 			buf: []byte{
diff --git a/dhcpv6/option_iaprefix.go b/dhcpv6/option_iaprefix.go
@@ -96,8 +96,6 @@ func (op *OptIAPrefix) FromBytes(data []byte) error {
 			IP:   ip,
 		}
 	}
-	if err := op.Options.FromBytes(buf.ReadAll()); err != nil {
-		return err
-	}
+	op.Options.Unmarshal(buf)
 	return buf.FinError()
 }
diff --git a/dhcpv6/option_iaprefix_test.go b/dhcpv6/option_iaprefix_test.go
@@ -96,9 +96,9 @@ func TestIAPrefixParseAndGetter(t *testing.T) {
 			},
 		},
 		{
-			buf:  []byte{0, 3, 0, 1, 0},
+			buf:  []byte{0, 26, 0, 1, 0},
 			want: nil,
-			err:  uio.ErrUnreadBytes,
+			err:  uio.ErrBufferTooShort,
 		},
 		{
 			buf: []byte{
diff --git a/dhcpv6/option_nontemporaryaddress.go b/dhcpv6/option_nontemporaryaddress.go
@@ -113,9 +113,6 @@ func (op *OptIANA) FromBytes(data []byte) error {
 	t2.Unmarshal(buf)
 	op.T1 = t1.Duration
 	op.T2 = t2.Duration
-
-	if err := op.Options.FromBytes(buf.ReadAll()); err != nil {
-		return err
-	}
+	op.Options.Unmarshal(buf)
 	return buf.FinError()
 }
diff --git a/dhcpv6/option_nontemporaryaddress_test.go b/dhcpv6/option_nontemporaryaddress_test.go
@@ -97,7 +97,7 @@ func TestIANAParseAndGetter(t *testing.T) {
 		{
 			buf:  []byte{0, 3, 0, 1, 0},
 			want: nil,
-			err:  uio.ErrUnreadBytes,
+			err:  uio.ErrBufferTooShort,
 		},
 		{
 			buf: []byte{
diff --git a/dhcpv6/option_temporaryaddress.go b/dhcpv6/option_temporaryaddress.go
@@ -43,9 +43,6 @@ func (op *OptIATA) LongString(indentSpace int) string {
 func (op *OptIATA) FromBytes(data []byte) error {
 	buf := uio.NewBigEndianBuffer(data)
 	buf.ReadBytes(op.IaId[:])
-
-	if err := op.Options.FromBytes(buf.ReadAll()); err != nil {
-		return err
-	}
+	op.Options.Unmarshal(buf)
 	return buf.FinError()
 }
diff --git a/dhcpv6/option_temporaryaddress_test.go b/dhcpv6/option_temporaryaddress_test.go
@@ -85,7 +85,7 @@ func TestIATAParseAndGetter(t *testing.T) {
 		{
 			buf:  []byte{0, 4, 0, 1, 0},
 			want: nil,
-			err:  uio.ErrUnreadBytes,
+			err:  uio.ErrBufferTooShort,
 		},
 		{
 			buf: []byte{
@@ -94,7 +94,7 @@ func TestIATAParseAndGetter(t *testing.T) {
 				1, 0, 0, // IAID too short
 			},
 			want: nil,
-			err:  uio.ErrUnreadBytes,
+			err:  uio.ErrBufferTooShort,
 		},
 		{
 			buf: []byte{
diff --git a/dhcpv6/option_vendor_opts.go b/dhcpv6/option_vendor_opts.go
@@ -43,9 +43,7 @@ func (op *OptVendorOpts) LongString(indent int) string {
 func (op *OptVendorOpts) FromBytes(data []byte) error {
 	buf := uio.NewBigEndianBuffer(data)
 	op.EnterpriseNumber = buf.Read32()
-	if err := op.VendorOpts.FromBytesWithParser(buf.ReadAll(), vendParseOption); err != nil {
-		return err
-	}
+	op.VendorOpts.UnmarshalWithParser(buf, vendParseOption)
 	return buf.FinError()
 }
 
diff --git a/dhcpv6/option_vendor_opts_test.go b/dhcpv6/option_vendor_opts_test.go
@@ -54,7 +54,7 @@ func TestVendorOptsParseAndGetter(t *testing.T) {
 		{
 			buf:  []byte{0, 17, 0, 1, 0},
 			want: nil,
-			err:  uio.ErrUnreadBytes,
+			err:  uio.ErrBufferTooShort,
 		},
 		{
 			buf:  []byte{0, 17, 0},
diff --git a/dhcpv6/options.go b/dhcpv6/options.go
@@ -34,6 +34,11 @@ func (og *OptionGeneric) String() string {
 	return fmt.Sprintf("%s: %v", og.OptionCode, og.OptionData)
 }
 
+// Unmarshal copies all data from buf into OptionData.
+func (og *OptionGeneric) Unmarshal(buf *uio.Lexer) {
+	og.OptionData = buf.ReadAll()
+}
+
 // FromBytes resets OptionData to p.
 func (og *OptionGeneric) FromBytes(p []byte) error {
 	og.OptionData = append([]byte(nil), p...)
@@ -217,21 +222,34 @@ func (o *Options) FromBytes(data []byte) error {
 	return o.FromBytesWithParser(data, ParseOption)
 }
 
+// Unmarshal reads data into o and returns an error if the options are not a
+// valid serialized representation of DHCPv6 options per RFC 3315.
+func (o *Options) Unmarshal(buf *uio.Lexer) {
+	o.UnmarshalWithParser(buf, ParseOption)
+}
+
 // OptionParser is a function signature for option parsing
 type OptionParser func(code OptionCode, data []byte) (Option, error)
 
 // FromBytesWithParser parses Options from byte sequences using the parsing
 // function that is passed in as a paremeter
 func (o *Options) FromBytesWithParser(data []byte, parser OptionParser) error {
+	buf := uio.NewBigEndianBuffer(data)
+	o.UnmarshalWithParser(buf, parser)
+	return buf.FinError()
+}
+
+// UnmarshalWithParser parses Options from byte sequences using the parsing
+// function that is passed in as a paremeter
+func (o *Options) UnmarshalWithParser(buf *uio.Lexer, parser OptionParser) {
 	if *o == nil {
 		*o = make(Options, 0, 10)
 	}
-	if len(data) == 0 {
+	if len(buf.Data()) == 0 {
 		// no options, no party
-		return nil
+		return
 	}
 
-	buf := uio.NewBigEndianBuffer(data)
 	for buf.Has(4) {
 		code := OptionCode(buf.Read16())
 		length := int(buf.Read16())
@@ -242,9 +260,9 @@ func (o *Options) FromBytesWithParser(data []byte, parser OptionParser) error {
 
 		opt, err := parser(code, optData)
 		if err != nil {
-			return err
+			buf.SetError(err)
+			return
 		}
 		*o = append(*o, opt)
 	}
-	return buf.FinError()
 }

Original file line number	Diff line number	Diff line change
`@@ -90,9 +90,9 @@ func TestIAAddressParseAndGetter(t *testing.T) {`
`90`	`90`	`},`
`91`	`91`	`},`
`92`	`92`	`{`
`93`		`- buf: []byte{0, 3, 0, 1, 0},`
	`93`	`+ buf: []byte{0, 5, 0, 1, 0},`
`94`	`94`	`want: nil,`
`95`		`- err: uio.ErrUnreadBytes,`
	`95`	`+ err: uio.ErrBufferTooShort,`
`96`	`96`	`},`
`97`	`97`	`{`
`98`	`98`	`buf: []byte{`
Original file line number	Diff line number	Diff line change
`@@ -111,7 +111,7 @@ func TestIAPDParseAndGetter(t *testing.T) {`
`111`	`111`	`{`
`112`	`112`	`buf: []byte{0, 25, 0, 1, 0},`
`113`	`113`	`want: nil,`
`114`		`- err: uio.ErrUnreadBytes,`
	`114`	`+ err: uio.ErrBufferTooShort,`
`115`	`115`	`},`
`116`	`116`	`{`
`117`	`117`	`buf: []byte{`
Original file line number	Diff line number	Diff line change
`@@ -96,8 +96,6 @@ func (op *OptIAPrefix) FromBytes(data []byte) error {`
`96`	`96`	`IP: ip,`
`97`	`97`	`}`
`98`	`98`	`}`
`99`		`- if err := op.Options.FromBytes(buf.ReadAll()); err != nil {`
`100`		`- return err`
`101`		`- }`
	`99`	`+ op.Options.Unmarshal(buf)`
`102`	`100`	`return buf.FinError()`
`103`	`101`	`}`
Original file line number	Diff line number	Diff line change
`@@ -96,9 +96,9 @@ func TestIAPrefixParseAndGetter(t *testing.T) {`
`96`	`96`	`},`
`97`	`97`	`},`
`98`	`98`	`{`
`99`		`- buf: []byte{0, 3, 0, 1, 0},`
	`99`	`+ buf: []byte{0, 26, 0, 1, 0},`
`100`	`100`	`want: nil,`
`101`		`- err: uio.ErrUnreadBytes,`
	`101`	`+ err: uio.ErrBufferTooShort,`
`102`	`102`	`},`
`103`	`103`	`{`
`104`	`104`	`buf: []byte{`
Original file line number	Diff line number	Diff line change
`@@ -97,7 +97,7 @@ func TestIANAParseAndGetter(t *testing.T) {`
`97`	`97`	`{`
`98`	`98`	`buf: []byte{0, 3, 0, 1, 0},`
`99`	`99`	`want: nil,`
`100`		`- err: uio.ErrUnreadBytes,`
	`100`	`+ err: uio.ErrBufferTooShort,`
`101`	`101`	`},`
`102`	`102`	`{`
`103`	`103`	`buf: []byte{`