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

Moves towards avoiding allocating a new buf.NewBigEndianBuffer in every single FromBytes / ToBytes function. Signed-off-by: Chris Koch <chrisko@google.com>
insomniacslk · Mar 5, 2023 · 50801df · 50801df
1 parent d75d3f1
commit 50801df
Show file tree

Hide file tree

Showing 15 changed files with 54 additions and 52 deletions.
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()
 }