dhcpv4: nicer (un-)marshaling of DHCP messages.

4 files changed, 89 insertions, 118 deletions

diff --git a/dhcpv4/dhcpv4.go b/dhcpv4/dhcpv4.go
index 4660dba..a713927 100644
--- a/dhcpv4/dhcpv4.go
+++ b/dhcpv4/dhcpv4.go
@@ -2,7 +2,6 @@ package dhcpv4
 
 import (
 	"crypto/rand"
-	"encoding/binary"
 	"errors"
 	"fmt"
 	"log"
@@ -10,13 +9,20 @@ import (
 	"strings"
 
 	"github.com/insomniacslk/dhcp/iana"
+	"github.com/u-root/u-root/pkg/uio"
 )
 
-// HeaderSize is the DHCPv4 header size in bytes.
-const HeaderSize = 236
+const (
+	// minPacketLen is the minimum DHCP header length.
+	minPacketLen = 236
 
-// MaxMessageSize is the maximum size in bytes that a DHCPv4 packet can hold.
-const MaxMessageSize = 576
+	// MaxMessageSize is the maximum size in bytes that a DHCPv4 packet can hold.
+	MaxMessageSize = 576
+)
+
+// magicCookie is the magic 4-byte value at the beginning of the list of options
+// in a DHCPv4 packet.
+var magicCookie = [4]byte{99, 130, 83, 99}
 
 // DHCPv4 represents a DHCPv4 packet header and options. See the New* functions
 // to build DHCPv4 packets.
@@ -117,11 +123,8 @@ func New() (*DHCPv4, error) {
 	copy(d.clientHwAddr[:], []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0})
 	copy(d.serverHostName[:], []byte{})
 	copy(d.bootFileName[:], []byte{})
-	options, err := OptionsFromBytes(MagicCookie)
-	if err != nil {
-		return nil, err
-	}
-	d.options = options
+
+	d.options = make([]Option, 0, 10)
 	// the End option has to be added explicitly
 	d.AddOption(&OptionGeneric{OptionCode: OptionEnd})
 	return &d, nil
@@ -268,32 +271,45 @@ func NewReplyFromRequest(request *DHCPv4, modifiers ...Modifier) (*DHCPv4, error
 
 // FromBytes encodes the DHCPv4 packet into a sequence of bytes, and returns an
 // error if the packet is not valid.
-func FromBytes(data []byte) (*DHCPv4, error) {
-	if len(data) < HeaderSize {
-		return nil, fmt.Errorf("Invalid DHCPv4 header: shorter than %v bytes", HeaderSize)
+func FromBytes(q []byte) (*DHCPv4, error) {
+	var p DHCPv4
+	buf := uio.NewBigEndianBuffer(q)
+
+	p.opcode = OpcodeType(buf.Read8())
+	p.hwType = iana.HwTypeType(buf.Read8())
+	p.hwAddrLen = buf.Read8()
+	p.hopCount = buf.Read8()
+
+	buf.ReadBytes(p.transactionID[:])
+
+	p.numSeconds = buf.Read16()
+	p.flags = buf.Read16()
+
+	p.clientIPAddr = net.IP(buf.CopyN(net.IPv4len))
+	p.yourIPAddr = net.IP(buf.CopyN(net.IPv4len))
+	p.serverIPAddr = net.IP(buf.CopyN(net.IPv4len))
+	p.gatewayIPAddr = net.IP(buf.CopyN(net.IPv4len))
+
+	buf.ReadBytes(p.clientHwAddr[:])
+	buf.ReadBytes(p.serverHostName[:])
+	buf.ReadBytes(p.bootFileName[:])
+
+	var cookie [4]byte
+	buf.ReadBytes(cookie[:])
+
+	if err := buf.Error(); err != nil {
+		return nil, err
 	}
-	d := DHCPv4{
-		opcode:        OpcodeType(data[0]),
-		hwType:        iana.HwTypeType(data[1]),
-		hwAddrLen:     data[2],
-		hopCount:      data[3],
-		numSeconds:    binary.BigEndian.Uint16(data[8:10]),
-		flags:         binary.BigEndian.Uint16(data[10:12]),
-		clientIPAddr:  net.IP(data[12:16]),
-		yourIPAddr:    net.IP(data[16:20]),
-		serverIPAddr:  net.IP(data[20:24]),
-		gatewayIPAddr: net.IP(data[24:28]),
-	}
-	copy(d.transactionID[:], data[4:8])
-	copy(d.clientHwAddr[:], data[28:44])
-	copy(d.serverHostName[:], data[44:108])
-	copy(d.bootFileName[:], data[108:236])
-	options, err := OptionsFromBytes(data[236:])
+	if cookie != magicCookie {
+		return nil, fmt.Errorf("malformed DHCP packet: got magic cookie %v, want %v", cookie[:], magicCookie[:])
+	}
+
+	opts, err := OptionsFromBytes(buf.Data())
 	if err != nil {
 		return nil, err
 	}
-	d.options = options
-	return &d, nil
+	p.options = opts
+	return &p, nil
 }
 
 // Opcode returns the OpcodeType for the packet,
@@ -722,36 +738,46 @@ func (d *DHCPv4) IsOptionRequested(requested OptionCode) bool {
 	return false
 }
 
+// In case somebody forgets to set an IP, just write 0s as default values.
+func writeIP(b *uio.Lexer, ip net.IP) {
+	var zeros [net.IPv4len]byte
+	if ip == nil {
+		b.WriteBytes(zeros[:])
+	} else {
+		b.WriteBytes(ip[:net.IPv4len])
+	}
+}
+
 // ToBytes encodes a DHCPv4 structure into a sequence of bytes in its wire
 // format.
 func (d *DHCPv4) ToBytes() []byte {
-	// This won't check if the End option is present, you've been warned
-	var ret []byte
-	u16 := make([]byte, 2)
-
-	ret = append(ret, byte(d.opcode))
-	ret = append(ret, byte(d.hwType))
-	ret = append(ret, byte(d.hwAddrLen))
-	ret = append(ret, byte(d.hopCount))
-	ret = append(ret, d.transactionID[:]...)
-	binary.BigEndian.PutUint16(u16, d.numSeconds)
-	ret = append(ret, u16...)
-	binary.BigEndian.PutUint16(u16, d.flags)
-	ret = append(ret, u16...)
-	ret = append(ret, d.clientIPAddr.To4()...)
-	ret = append(ret, d.yourIPAddr.To4()...)
-	ret = append(ret, d.serverIPAddr.To4()...)
-	ret = append(ret, d.gatewayIPAddr.To4()...)
-	ret = append(ret, d.clientHwAddr[:16]...)
-	ret = append(ret, d.serverHostName[:64]...)
-	ret = append(ret, d.bootFileName[:128]...)
-
-	d.ValidateOptions() // print warnings about broken options, if any
-	ret = append(ret, MagicCookie...)
+	buf := uio.NewBigEndianBuffer(make([]byte, 0, minPacketLen))
+	buf.Write8(uint8(d.opcode))
+	buf.Write8(uint8(d.hwType))
+
+	// HwAddrLen
+	buf.Write8(d.hwAddrLen)
+	buf.Write8(d.hopCount)
+	buf.WriteBytes(d.transactionID[:])
+	buf.Write16(d.numSeconds)
+	buf.Write16(d.flags)
+
+	writeIP(buf, d.clientIPAddr[:])
+	writeIP(buf, d.yourIPAddr[:])
+	writeIP(buf, d.serverIPAddr[:])
+	writeIP(buf, d.gatewayIPAddr[:])
+
+	buf.WriteBytes(d.clientHwAddr[:])
+	buf.WriteBytes(d.serverHostName[:])
+	buf.WriteBytes(d.bootFileName[:])
+
+	// The magic cookie.
+	buf.WriteBytes(magicCookie[:])
+
 	for _, opt := range d.options {
-		ret = append(ret, opt.ToBytes()...)
+		buf.WriteBytes(opt.ToBytes())
 	}
-	return ret
+	return buf.Data()
 }
 
 // OptionGetter is a interface that knows how to retrieve an option from a
diff --git a/dhcpv4/dhcpv4_test.go b/dhcpv4/dhcpv4_test.go
index 2e5b8b0..7a208ce 100644
--- a/dhcpv4/dhcpv4_test.go
+++ b/dhcpv4/dhcpv4_test.go
@@ -310,7 +310,7 @@ func TestNewToBytes(t *testing.T) {
 		expected = append(expected, 0)
 	}
 	// Magic Cookie
-	expected = append(expected, MagicCookie...)
+	expected = append(expected, magicCookie[:]...)
 	// End
 	expected = append(expected, 0xff)
 
diff --git a/dhcpv4/options.go b/dhcpv4/options.go
index dc4a724..a5a51be 100644
--- a/dhcpv4/options.go
+++ b/dhcpv4/options.go
@@ -1,9 +1,7 @@
 package dhcpv4
 
 import (
-	"bytes"
 	"errors"
-	"fmt"
 )
 
 // ErrShortByteStream is an error that is thrown any time a short byte stream is
@@ -14,10 +12,6 @@ var ErrShortByteStream = errors.New("short byte stream")
 // byte stream is encountered.
 var ErrZeroLengthByteStream = errors.New("zero-length byte stream")
 
-// MagicCookie is the magic 4-byte value at the beginning of the list of options
-// in a DHCPv4 packet.
-var MagicCookie = []byte{99, 130, 83, 99}
-
 // OptionCode is a single byte representing the code for a given Option.
 type OptionCode byte
 
@@ -94,26 +88,12 @@ func ParseOption(data []byte) (Option, error) {
 }
 
 // OptionsFromBytes parses a sequence of bytes until the end and builds a list
-// of options from it. The sequence must contain the Magic Cookie. Returns an
-// error if any invalid option or length is found.
+// of options from it.
+//
+// The sequence should not contain the DHCP magic cookie.
+//
+// Returns an error if any invalid option or length is found.
 func OptionsFromBytes(data []byte) ([]Option, error) {
-	if len(data) < len(MagicCookie) {
-		return nil, errors.New("invalid options: shorter than 4 bytes")
-	}
-	if !bytes.Equal(data[:len(MagicCookie)], MagicCookie) {
-		return nil, fmt.Errorf("invalid magic cookie: %v", data[:len(MagicCookie)])
-	}
-	opts, err := OptionsFromBytesWithoutMagicCookie(data[len(MagicCookie):])
-	if err != nil {
-		return nil, err
-	}
-	return opts, nil
-}
-
-// OptionsFromBytesWithoutMagicCookie parses a sequence of bytes until the end
-// and builds a list of options from it. The sequence should not contain the
-// DHCP magic cookie. Returns an error if any invalid option or length is found.
-func OptionsFromBytesWithoutMagicCookie(data []byte) ([]Option, error) {
 	return OptionsFromBytesWithParser(data, ParseOption)
 }
 
diff --git a/dhcpv4/options_test.go b/dhcpv4/options_test.go
index bf869f2..2557af5 100644
--- a/dhcpv4/options_test.go
+++ b/dhcpv4/options_test.go
@@ -189,38 +189,3 @@ func TestParseOptionShortOption(t *testing.T) {
 	_, err := ParseOption(option)
 	require.Error(t, err, "should get error from short options")
 }
-
-func TestOptionsFromBytes(t *testing.T) {
-	options := []byte{
-		99, 130, 83, 99, // Magic Cookie
-		5, 4, 192, 168, 1, 1, // DNS
-		255,     // end
-		0, 0, 0, //padding
-	}
-	opts, err := OptionsFromBytes(options)
-	require.NoError(t, err)
-	require.Equal(t, 2, len(opts))
-	require.Equal(t, opts[0].(*OptionGeneric), &OptionGeneric{OptionCode: OptionNameServer, Data: []byte{192, 168, 1, 1}})
-	require.Equal(t, opts[1].(*OptionGeneric), &OptionGeneric{OptionCode: OptionEnd})
-}
-
-func TestOptionsFromBytesZeroLength(t *testing.T) {
-	options := []byte{}
-	_, err := OptionsFromBytes(options)
-	require.Error(t, err)
-}
-
-func TestOptionsFromBytesBadMagicCookie(t *testing.T) {
-	options := []byte{1, 2, 3, 4}
-	_, err := OptionsFromBytes(options)
-	require.Error(t, err)
-}
-
-func TestOptionsFromBytesShortOption(t *testing.T) {
-	options := []byte{
-		99, 130, 83, 99, // Magic Cookie
-		5, 4, 192, 168, // DNS
-	}
-	_, err := OptionsFromBytes(options)
-	require.Error(t, err)
-}