Add vendor specific information option

4 files changed, 176 insertions, 96 deletions

diff --git a/dhcpv4/bsdp/boot_image_test.go b/dhcpv4/bsdp/boot_image_test.go
index ebf02c3..86a2e9f 100644
--- a/dhcpv4/bsdp/boot_image_test.go
+++ b/dhcpv4/bsdp/boot_image_test.go
@@ -8,9 +8,6 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-/*
- * BootImageID
- */
 func TestBootImageIDToBytes(t *testing.T) {
 	b := BootImageID{
 		IsInstall: true,
@@ -54,6 +51,11 @@ func TestBootImageIDFromBytesFail(t *testing.T) {
 	require.Error(t, err)
 }
 
+func TestBootImageIDString(t *testing.T) {
+	b := BootImageID{IsInstall: false, ImageType: BootImageTypeMacOSX, Index: 1001}
+	require.Equal(t, "[1001] uninstallable macOS image", b.String())
+}
+
 /*
  * BootImage
  */
@@ -128,3 +130,15 @@ func TestBootImageFromBytesShortBootImage(t *testing.T) {
 	require.Nil(t, b)
 	require.Error(t, err)
 }
+
+func TestBootImageString(t *testing.T) {
+	b := BootImage{
+		ID: BootImageID{
+			IsInstall: false,
+			ImageType: BootImageTypeMacOSX,
+			Index:     0x1010,
+		},
+		Name: "bsdp-21",
+	}
+	require.Equal(t, "bsdp-21 [4112] uninstallable macOS image", b.String())
+}
diff --git a/dhcpv4/bsdp/bsdp.go b/dhcpv4/bsdp/bsdp.go
index 1c2b100..f04e255 100644
--- a/dhcpv4/bsdp/bsdp.go
+++ b/dhcpv4/bsdp/bsdp.go
@@ -9,7 +9,6 @@ package bsdp
 import (
 	"errors"
 	"fmt"
-	"log"
 	"net"
 	"syscall"
 
@@ -33,71 +32,15 @@ func makeVendorClassIdentifier() (string, error) {
 	return fmt.Sprintf("AAPLBSDPC/i386/%s", hwModel), nil
 }
 
-// ParseBootImagesFromOption parses data from the BSDPOptionBootImageList
-// option and returns a list of BootImages.
-func ParseBootImagesFromOption(data []byte) ([]BootImage, error) {
-	// Should at least have the # bytes of boot images.
-	if len(data) < 4 {
-		return nil, fmt.Errorf("invalid length boot image list")
-	}
-
-	var (
-		readByteCount = 0
-		start         = data
-		bootImages    []BootImage
-	)
-	for {
-		bootImage, err := BootImageFromBytes(start)
-		if err != nil {
-			return nil, err
-		}
-		bootImages = append(bootImages, *bootImage)
-		// Read BootImageID + name length + name
-		readByteCount += 4 + 1 + len(bootImage.Name)
-		if readByteCount+1 >= len(data) {
-			break
-		}
-		start = start[readByteCount:]
-	}
-
-	return bootImages, nil
-}
-
-// ParseVendorOptionsFromOptions extracts the sub-options list of the vendor-
-// specific options from the larger DHCP options list.
-// TODO: Implement options.GetOneOption for dhcpv4.
-func ParseVendorOptionsFromOptions(options []dhcpv4.Option) []dhcpv4.Option {
-	var (
-		vendorOpts []dhcpv4.Option
-		err        error
-	)
-	for _, opt := range options {
-		if opt.Code() == dhcpv4.OptionVendorSpecificInformation {
-			vendorOpts, err = dhcpv4.OptionsFromBytesWithoutMagicCookie(opt.(*dhcpv4.OptionGeneric).Data)
-			if err != nil {
-				log.Println("Warning: could not parse vendor options in DHCP options")
-				return []dhcpv4.Option{}
-			}
-			break
-		}
-	}
-	return vendorOpts
-}
-
 // ParseBootImageListFromAck parses the list of boot images presented in the
 // ACK[LIST] packet and returns them as a list of BootImages.
 func ParseBootImageListFromAck(ack dhcpv4.DHCPv4) ([]BootImage, error) {
 	var bootImages []BootImage
-	for _, opt := range ParseVendorOptionsFromOptions(ack.Options()) {
-		if opt.Code() == OptionBootImageList {
-			images, err := ParseBootImagesFromOption(opt.(*dhcpv4.OptionGeneric).Data)
-			if err != nil {
-				return nil, err
-			}
-			bootImages = append(bootImages, images...)
-		}
+	o := ack.GetOption(dhcpv4.OptionVendorSpecificInformation)
+	vendorOpt := o.(*OptVendorSpecificInformation)
+	for _, o := range vendorOpt.GetOptions(OptionBootImageList) {
+		bootImages = append(bootImages, o.(*OptBootImageList).BootImages...)
 	}
-
 	return bootImages, nil
 }
 
@@ -123,21 +66,13 @@ func NewInformListForInterface(iface string, replyPort uint16) (*dhcpv4.DHCPv4,
 		&OptMessageType{MessageTypeList},
 		&OptVersion{Version1_1},
 	}
-
 	if needsReplyPort(replyPort) {
 		vendorOpts = append(vendorOpts, &OptReplyPort{replyPort})
 	}
-	var vendorOptsBytes []byte
-	for _, opt := range vendorOpts {
-		vendorOptsBytes = append(vendorOptsBytes, opt.ToBytes()...)
-	}
-	d.AddOption(&dhcpv4.OptionGeneric{
-		OptionCode: dhcpv4.OptionVendorSpecificInformation,
-		Data:       vendorOptsBytes,
-	})
+	d.AddOption(&OptVendorSpecificInformation{vendorOpts})
 
 	d.AddOption(&dhcpv4.OptParameterRequestList{
-		RequestedOpts: []dhcpv4.OptionCode{
+		[]dhcpv4.OptionCode{
 			dhcpv4.OptionVendorSpecificInformation,
 			dhcpv4.OptionClassIdentifier,
 		},
@@ -194,7 +129,7 @@ func InformSelectForAck(ack dhcpv4.DHCPv4, replyPort uint16, selectedImage BootI
 	if serverIP.To4() == nil {
 		return nil, fmt.Errorf("could not parse server identifier from ACK")
 	}
-	vendorOpts = append(vendorOpts, &dhcpv4.OptServerIdentifier{ServerID: serverIP})
+	vendorOpts = append(vendorOpts, &dhcpv4.OptServerIdentifier{serverIP})
 
 	// Validate replyPort if requested.
 	if needsReplyPort(replyPort) {
@@ -216,14 +151,7 @@ func InformSelectForAck(ack dhcpv4.DHCPv4, replyPort uint16, selectedImage BootI
 		},
 	})
 	d.AddOption(&dhcpv4.OptMessageType{dhcpv4.MessageTypeInform})
-	var vendorOptsBytes []byte
-	for _, opt := range vendorOpts {
-		vendorOptsBytes = append(vendorOptsBytes, opt.ToBytes()...)
-	}
-	d.AddOption(&dhcpv4.OptionGeneric{
-		OptionCode: dhcpv4.OptionVendorSpecificInformation,
-		Data:       vendorOptsBytes,
-	})
+	d.AddOption(&OptVendorSpecificInformation{vendorOpts})
 	d.AddOption(&dhcpv4.OptionGeneric{OptionCode: dhcpv4.OptionEnd})
 	return d, nil
 }
diff --git a/dhcpv4/bsdp/option_vendor_specific_information.go b/dhcpv4/bsdp/option_vendor_specific_information.go
new file mode 100644
index 0000000..b6e27ac
--- /dev/null
+++ b/dhcpv4/bsdp/option_vendor_specific_information.go
@@ -0,0 +1,138 @@
+// +build darwin
+
+package bsdp
+
+import (
+	"errors"
+	"fmt"
+
+	"github.com/insomniacslk/dhcp/dhcpv4"
+)
+
+// OptVendorSpecificInformation encapsulates the BSDP-specific options used for
+// the protocol.
+type OptVendorSpecificInformation struct {
+	Options []dhcpv4.Option
+}
+
+// parseOption is similar to dhcpv4.ParseOption, except that it switches based
+// on the BSDP specific options.
+func parseOption(data []byte) (dhcpv4.Option, error) {
+	if len(data) == 0 {
+		return nil, dhcpv4.ErrZeroLengthByteStream
+	}
+	var (
+		opt dhcpv4.Option
+		err error
+	)
+	switch dhcpv4.OptionCode(data[0]) {
+	case OptionMessageType:
+		opt, err = ParseOptMessageType(data)
+	case OptionVersion:
+		opt, err = ParseOptVersion(data)
+	case OptionReplyPort:
+		opt, err = ParseOptReplyPort(data)
+	case OptionSelectedBootImageID:
+		opt, err = ParseOptSelectedBootImageID(data)
+	default:
+		opt, err = dhcpv4.ParseOptionGeneric(data)
+	}
+	if err != nil {
+		return nil, err
+	}
+	return opt, nil
+}
+
+// ParseOptVendorSpecificInformation constructs an OptVendorSpecificInformation struct from a sequence of
+// bytes and returns it, or an error.
+func ParseOptVendorSpecificInformation(data []byte) (*OptVendorSpecificInformation, error) {
+	// Should at least have code + length
+	if len(data) < 2 {
+		return nil, dhcpv4.ErrShortByteStream
+	}
+	code := dhcpv4.OptionCode(data[0])
+	if code != dhcpv4.OptionVendorSpecificInformation {
+		return nil, fmt.Errorf("expected option %v, got %v instead", dhcpv4.OptionVendorSpecificInformation, code)
+	}
+	length := int(data[1])
+	if len(data) < length+2 {
+		return nil, fmt.Errorf("expected length 2, got %d instead", length)
+	}
+
+	options := make([]dhcpv4.Option, 0, 10)
+	idx := 0
+	for {
+		if idx == len(data) {
+			break
+		}
+		// This should never happen.
+		if idx > len(data) {
+			return nil, errors.New("read past the end of options")
+		}
+		opt, err := parseOption(data[idx:])
+		if err != nil {
+			return nil, err
+		}
+		options = append(options, opt)
+
+		// Account for code + length bytes
+		idx += 2 + opt.Length()
+	}
+
+	return &OptVendorSpecificInformation{options}, nil
+}
+
+// Code returns the option code.
+func (o *OptVendorSpecificInformation) Code() dhcpv4.OptionCode {
+	return dhcpv4.OptionVendorSpecificInformation
+}
+
+// ToBytes returns a serialized stream of bytes for this option.
+func (o *OptVendorSpecificInformation) ToBytes() []byte {
+	bs := []byte{byte(o.Code()), byte(o.Length())}
+
+	// Append data section
+	for _, opt := range o.Options {
+		bs = append(bs, opt.ToBytes()...)
+	}
+	return bs
+}
+
+// String returns a human-readable string for this option.
+func (o *OptVendorSpecificInformation) String() string {
+	s := "Vendor Specific Information ->\n"
+	for _, opt := range o.Options {
+		s += "  " + opt.String() + "\n"
+	}
+	return s
+}
+
+// Length returns the length of the data portion of this option. Take into
+// account code + data length bytes for each sub option.
+func (o *OptVendorSpecificInformation) Length() int {
+	var length int
+	for _, opt := range o.Options {
+		length += 2 + opt.Length()
+	}
+	return length
+}
+
+// GetOptions returns all suboptions that match the given OptionCode code.
+func (o *OptVendorSpecificInformation) GetOptions(code dhcpv4.OptionCode) []dhcpv4.Option {
+	var opts []dhcpv4.Option
+	for _, opt := range o.Options {
+		if o.Code() == code {
+			opts = append(opts, opt)
+		}
+	}
+	return opts
+}
+
+// GetOption returns the first suboption that matches the OptionCode code.
+func (o *OptVendorSpecificInformation) GetOption(code dhcpv4.OptionCode) dhcpv4.Option {
+	opts := o.GetOptions(code)
+	if len(opts) == 0 {
+		return nil
+	}
+	return opts[0]
+}
diff --git a/dhcpv4/bsdp/types.go b/dhcpv4/bsdp/types.go
index d04d6ef..691c903 100644
--- a/dhcpv4/bsdp/types.go
+++ b/dhcpv4/bsdp/types.go
@@ -25,18 +25,18 @@ const (
 // OptionCodeToString maps BSDP OptionCodes to human-readable strings
 // describing what they are.
 var OptionCodeToString = map[dhcpv4.OptionCode]string{
-	OptionMessageType:                   " Message Type",
-	OptionVersion:                       " Version",
-	OptionServerIdentifier:              " Server Identifier",
-	OptionServerPriority:                " Server Priority",
-	OptionReplyPort:                     " Reply Port",
+	OptionMessageType:                   "Message Type",
+	OptionVersion:                       "Version",
+	OptionServerIdentifier:              "Server Identifier",
+	OptionServerPriority:                "Server Priority",
+	OptionReplyPort:                     "Reply Port",
 	OptionBootImageListPath:             "", // Not used
-	OptionDefaultBootImageID:            " Default Boot Image ID",
-	OptionSelectedBootImageID:           " Selected Boot Image ID",
-	OptionBootImageList:                 " Boot Image List",
-	OptionNetboot1_0Firmware:            " Netboot 1.0 Firmware",
-	OptionBootImageAttributesFilterList: " Boot Image Attributes Filter List",
-	OptionShadowMountPath:               " Shadow Mount Path",
-	OptionShadowFilePath:                " Shadow File Path",
-	OptionMachineName:                   " Machine Name",
+	OptionDefaultBootImageID:            "Default Boot Image ID",
+	OptionSelectedBootImageID:           "Selected Boot Image ID",
+	OptionBootImageList:                 "Boot Image List",
+	OptionNetboot1_0Firmware:            "Netboot 1.0 Firmware",
+	OptionBootImageAttributesFilterList: "Boot Image Attributes Filter List",
+	OptionShadowMountPath:               "Shadow Mount Path",
+	OptionShadowFilePath:                "Shadow File Path",
+	OptionMachineName:                   "Machine Name",
 }