dhcpv6: use uio buffer in DHCPv6 message parsing

4 files changed, 32 insertions, 45 deletions

diff --git a/dhcpv6/dhcpv6.go b/dhcpv6/dhcpv6.go
index 70c2bff..25fdcd0 100644
--- a/dhcpv6/dhcpv6.go
+++ b/dhcpv6/dhcpv6.go
@@ -7,6 +7,7 @@ import (
 	"strings"
 
 	"github.com/insomniacslk/dhcp/iana"
+	"github.com/u-root/u-root/pkg/uio"
 )
 
 type DHCPv6 interface {
@@ -29,36 +30,19 @@ type DHCPv6 interface {
 type Modifier func(d DHCPv6) DHCPv6
 
 func FromBytes(data []byte) (DHCPv6, error) {
-	var (
-		isRelay     = false
-		headerSize  int
-		messageType = MessageType(data[0])
-	)
+	buf := uio.NewBigEndianBuffer(data)
+	messageType := MessageType(buf.Read8())
+
 	if messageType == MessageTypeRelayForward || messageType == MessageTypeRelayReply {
-		isRelay = true
-	}
-	if isRelay {
-		headerSize = RelayHeaderSize
-	} else {
-		headerSize = MessageHeaderSize
-	}
-	if len(data) < headerSize {
-		return nil, fmt.Errorf("Invalid header size: shorter than %v bytes", headerSize)
-	}
-	if isRelay {
-		var (
-			linkAddr, peerAddr []byte
-		)
 		d := DHCPv6Relay{
 			messageType: messageType,
-			hopCount:    uint8(data[1]),
+			hopCount:    buf.Read8(),
 		}
-		linkAddr = append(linkAddr, data[2:18]...)
-		d.linkAddr = linkAddr
-		peerAddr = append(peerAddr, data[18:34]...)
-		d.peerAddr = peerAddr
+		d.linkAddr = net.IP(buf.CopyN(net.IPv6len))
+		d.peerAddr = net.IP(buf.CopyN(net.IPv6len))
+
 		// TODO fail if no OptRelayMessage is present
-		if err := d.options.FromBytes(data[34:]); err != nil {
+		if err := d.options.FromBytes(buf.Data()); err != nil {
 			return nil, err
 		}
 		return &d, nil
@@ -66,8 +50,8 @@ func FromBytes(data []byte) (DHCPv6, error) {
 		d := DHCPv6Message{
 			messageType: messageType,
 		}
-		copy(d.transactionID[:], data[1:4])
-		if err := d.options.FromBytes(data[4:]); err != nil {
+		buf.ReadBytes(d.transactionID[:])
+		if err := d.options.FromBytes(buf.Data()); err != nil {
 			return nil, err
 		}
 		return &d, nil
diff --git a/dhcpv6/dhcpv6message.go b/dhcpv6/dhcpv6message.go
index 85d93a3..078f72a 100644
--- a/dhcpv6/dhcpv6message.go
+++ b/dhcpv6/dhcpv6message.go
@@ -9,6 +9,7 @@ import (
 	"time"
 
 	"github.com/insomniacslk/dhcp/iana"
+	"github.com/u-root/u-root/pkg/uio"
 )
 
 const MessageHeaderSize = 4
@@ -307,16 +308,16 @@ func (d *DHCPv6Message) Summary() string {
 	return ret
 }
 
-// Convert a DHCPv6Message structure into its binary representation, suitable for being
-// sent over the network
+// ToBytes returns the serialized version of this message as defined by RFC
+// 3315, Section 5.
 func (d *DHCPv6Message) ToBytes() []byte {
-	var ret []byte
-	ret = append(ret, byte(d.messageType))
-	ret = append(ret, d.transactionID[:]...) // discard the first byte
+	buf := uio.NewBigEndianBuffer(nil)
+	buf.Write8(uint8(d.messageType))
+	buf.WriteBytes(d.transactionID[:])
 	for _, opt := range d.options {
-		ret = append(ret, opt.ToBytes()...)
+		buf.WriteBytes(opt.ToBytes())
 	}
-	return ret
+	return buf.Data()
 }
 
 func (d *DHCPv6Message) Length() int {
diff --git a/dhcpv6/dhcpv6relay.go b/dhcpv6/dhcpv6relay.go
index b51c139..e2024f4 100644
--- a/dhcpv6/dhcpv6relay.go
+++ b/dhcpv6/dhcpv6relay.go
@@ -4,6 +4,8 @@ import (
 	"errors"
 	"fmt"
 	"net"
+
+	"github.com/u-root/u-root/pkg/uio"
 )
 
 const RelayHeaderSize = 34
@@ -49,17 +51,18 @@ func (r *DHCPv6Relay) Summary() string {
 	return ret
 }
 
+// ToBytes returns the serialized version of this relay message as defined by
+// RFC 3315, Section 6.
 func (r *DHCPv6Relay) ToBytes() []byte {
-	ret := make([]byte, RelayHeaderSize)
-	ret[0] = byte(r.messageType)
-	ret[1] = byte(r.hopCount)
-	copy(ret[2:18], r.linkAddr)
-	copy(ret[18:34], r.peerAddr)
+	buf := uio.NewBigEndianBuffer(make([]byte, 0, RelayHeaderSize))
+	buf.Write8(byte(r.messageType))
+	buf.Write8(byte(r.hopCount))
+	buf.WriteBytes(r.linkAddr.To16())
+	buf.WriteBytes(r.peerAddr.To16())
 	for _, opt := range r.options {
-		ret = append(ret, opt.ToBytes()...)
+		buf.WriteBytes(opt.ToBytes())
 	}
-
-	return ret
+	return buf.Data()
 }
 
 func (r *DHCPv6Relay) SetMessageType(messageType MessageType) {
diff --git a/dhcpv6/option_prefixdelegation.go b/dhcpv6/option_prefixdelegation.go
index da5dff6..5875162 100644
--- a/dhcpv6/option_prefixdelegation.go
+++ b/dhcpv6/option_prefixdelegation.go
@@ -1,13 +1,12 @@
 package dhcpv6
 
-// This module defines the OptIAForPrefixDelegation structure.
-// https://www.ietf.org/rfc/rfc3633.txt
-
 import (
 	"encoding/binary"
 	"fmt"
 )
 
+// OptIAForPrefixDelegation implements the identity association for prefix
+// delegation option defined by RFC 3633, Section 9.
 type OptIAForPrefixDelegation struct {
 	IaId    [4]byte
 	T1      uint32