dhcpv6: added modifiers (#41)

dhcpv6: added modifiers

Added support for packet modifiers, i.e. functions that can arbitrarily
manipulate a DHCPv6 packet. These modifiers are used by NewMessage,
NewSolicitForInterface, NewRequestForAdvertise, and can be used by other
packet creation functions.
A bunch of sample modifiers have been added under modifiers.go , too.
With the introduction of modifiers I also removed some options that
should not necessarily be in a standard DHCPv6 message.

5 files changed, 92 insertions, 30 deletions

diff --git a/dhcpv6/client.go b/dhcpv6/client.go
index 3e79a0a..6edf798 100644
--- a/dhcpv6/client.go
+++ b/dhcpv6/client.go
@@ -37,8 +37,10 @@ func NewClient() *Client {
 }
 
 // Exchange executes a 4-way DHCPv6 request (SOLICIT, ADVERTISE, REQUEST,
-// REPLY). If the SOLICIT packet is nil, defaults are used.
-func (c *Client) Exchange(ifname string, solicit DHCPv6) ([]DHCPv6, error) {
+// REPLY). If the SOLICIT packet is nil, defaults are used. The modifiers will
+// be applied to the Request packet. A common use is to make sure that the
+// Request packet has the right options, see modifiers.go
+func (c *Client) Exchange(ifname string, solicit DHCPv6, modifiers ...Modifier) ([]DHCPv6, error) {
 	conversation := make([]DHCPv6, 0)
 	var err error
 
@@ -56,6 +58,9 @@ func (c *Client) Exchange(ifname string, solicit DHCPv6) ([]DHCPv6, error) {
 	if request != nil {
 		conversation = append(conversation, request)
 	}
+	for _, mod := range modifiers {
+		request = mod(request)
+	}
 	if err != nil {
 		return conversation, err
 	}
@@ -121,26 +126,28 @@ func (c *Client) sendReceive(ifname string, packet DHCPv6, expectedType MessageT
 		return nil, err
 	}
 
-	// wait for an ADVERTISE response
-	buf := make([]byte, maxUDPReceivedPacketSize)
+	// wait for a reply
 	oobdata := []byte{} // ignoring oob data
 	conn.SetReadDeadline(time.Now().Add(c.ReadTimeout))
 	var (
 		adv       DHCPv6
 		isMessage bool
 	)
+	defer conn.Close()
 	msg, ok := packet.(*DHCPv6Message)
 	if ok {
 		isMessage = true
 	}
 	for {
+		buf := make([]byte, maxUDPReceivedPacketSize)
 		n, _, _, _, err := conn.ReadMsgUDP(buf, oobdata)
 		if err != nil {
 			return nil, err
 		}
 		adv, err = FromBytes(buf[:n])
 		if err != nil {
-			return nil, err
+			// skip non-DHCP packets
+			continue
 		}
 		if recvMsg, ok := adv.(*DHCPv6Message); ok && isMessage {
 			// if a regular message, check the transaction ID first
diff --git a/dhcpv6/dhcpv6.go b/dhcpv6/dhcpv6.go
index 0532690..eb98106 100644
--- a/dhcpv6/dhcpv6.go
+++ b/dhcpv6/dhcpv6.go
@@ -19,6 +19,10 @@ type DHCPv6 interface {
 	AddOption(Option)
 }
 
+// Modifier defines the signature for functions that can modify DHCPv6
+// structures. This is used to simplify packet manipulation
+type Modifier func(d DHCPv6) DHCPv6
+
 func FromBytes(data []byte) (DHCPv6, error) {
 	var (
 		isRelay     = false
@@ -73,7 +77,8 @@ func FromBytes(data []byte) (DHCPv6, error) {
 	}
 }
 
-func NewMessage() (*DHCPv6Message, error) {
+// NewMessage creates a new DHCPv6 message with default options
+func NewMessage() (DHCPv6, error) {
 	tid, err := GenerateTransactionID()
 	if err != nil {
 		return nil, err
@@ -106,9 +111,9 @@ func getOption(options []Option, code OptionCode) Option {
 	return opts[0]
 }
 
-// Decapsulate extracts the content of a relay message. It does not recurse if
-// there are nested relay messages. Returns the original packet if is not not a
-// relay message
+// DecapsulateRelay extracts the content of a relay message. It does not recurse
+// if there are nested relay messages. Returns the original packet if is not not
+// a relay message
 func DecapsulateRelay(l DHCPv6) (DHCPv6, error) {
 	if !l.IsRelay() {
 		return l, nil
@@ -124,8 +129,8 @@ func DecapsulateRelay(l DHCPv6) (DHCPv6, error) {
 	return relayOpt.RelayMessage(), nil
 }
 
-// Encapsulate creates a DHCPv6Relay message containing the passed DHCPv6 object.
-// struct as payload. The passed message type must be  either RELAY_FORW or
+// EncapsulateRelay creates a DHCPv6Relay message containing the passed DHCPv6
+// message as payload. The passed message type must be  either RELAY_FORW or
 // RELAY_REPL
 func EncapsulateRelay(d DHCPv6, mType MessageType, linkAddr, peerAddr net.IP) (DHCPv6, error) {
 	if mType != RELAY_FORW && mType != RELAY_REPL {
diff --git a/dhcpv6/dhcpv6_test.go b/dhcpv6/dhcpv6_test.go
index 265b17f..17a8896 100644
--- a/dhcpv6/dhcpv6_test.go
+++ b/dhcpv6/dhcpv6_test.go
@@ -31,9 +31,9 @@ func TestNewMessage(t *testing.T) {
 	d, err := NewMessage()
 	require.NoError(t, err)
 	require.NotNil(t, d)
-	require.Equal(t, SOLICIT, d.messageType)
-	require.NotEqual(t, 0, d.transactionID)
-	require.Empty(t, d.options)
+	require.Equal(t, SOLICIT, d.Type())
+	require.NotEqual(t, 0, d.(*DHCPv6Message).transactionID)
+	require.Empty(t, d.(*DHCPv6Message).options)
 }
 
 func TestSettersAndGetters(t *testing.T) {
diff --git a/dhcpv6/dhcpv6message.go b/dhcpv6/dhcpv6message.go
index b14daa6..0cf2360 100644
--- a/dhcpv6/dhcpv6message.go
+++ b/dhcpv6/dhcpv6message.go
@@ -57,21 +57,21 @@ func GenerateTransactionID() (*uint32, error) {
 	return tid, nil
 }
 
-// Return a time integer suitable for DUID-LLT, i.e. the current time counted in
-// seconds since January 1st, 2000, midnight UTC, modulo 2^32
+// GetTime returns a time integer suitable for DUID-LLT, i.e. the current time counted
+// in seconds since January 1st, 2000, midnight UTC, modulo 2^32
 func GetTime() uint32 {
 	now := time.Since(time.Date(2000, time.January, 1, 0, 0, 0, 0, time.UTC))
 	return uint32((now.Nanoseconds() / 1000000000) % 0xffffffff)
 }
 
-// Create a new SOLICIT message with DUID-LLT, using the given network
-// interface's hardware address and current time
-func NewSolicitForInterface(ifname string) (*DHCPv6Message, error) {
+// NewSolicitForInterface creates a new SOLICIT message with DUID-LLT, using the
+// given network interface's hardware address and current time
+func NewSolicitForInterface(ifname string, modifiers ...Modifier) (DHCPv6, error) {
 	d, err := NewMessage()
 	if err != nil {
 		return nil, err
 	}
-	d.SetMessage(SOLICIT)
+	d.(*DHCPv6Message).SetMessage(SOLICIT)
 	iface, err := net.InterfaceByName(ifname)
 	if err != nil {
 		return nil, err
@@ -98,10 +98,17 @@ func NewSolicitForInterface(ifname string) (*DHCPv6Message, error) {
 	iaNa.SetT1(0xe10)
 	iaNa.SetT2(0x1518)
 	d.AddOption(&iaNa)
+
+	// apply modifiers
+	for _, mod := range modifiers {
+		d = mod(d)
+	}
 	return d, nil
 }
 
-func NewRequestFromAdvertise(advertise DHCPv6) (DHCPv6, error) {
+// NewRequestFromAdvertise creates a new REQUEST packet based on an ADVERTISE
+// packet options.
+func NewRequestFromAdvertise(advertise DHCPv6, modifiers ...Modifier) (DHCPv6, error) {
 	if advertise == nil {
 		return nil, fmt.Errorf("ADVERTISE cannot be nil")
 	}
@@ -139,23 +146,23 @@ func NewRequestFromAdvertise(advertise DHCPv6) (DHCPv6, error) {
 	// add OptRequestedOption
 	oro := OptRequestedOption{}
 	oro.SetRequestedOptions([]OptionCode{
-		OPT_BOOTFILE_URL,
-		OPT_BOOTFILE_PARAM,
+		DNS_RECURSIVE_NAME_SERVER,
+		DOMAIN_SEARCH_LIST,
 	})
 	req.AddOption(&oro)
-	// add OPTION_NII
-	nii := OptNetworkInterfaceId{}
-	nii.SetType(1)
-	nii.SetMajor(3) // UNDI - Universal Network Device Interface
-	nii.SetMinor(2) // UNDI rev. 3.2 - second generation EFI runtime driver support, see rfc457
-	req.AddOption(&nii)
 	// add OPTION_VENDOR_CLASS, only if present in the original request
 	// TODO implement OptionVendorClass
 	vClass := adv.GetOneOption(OPTION_VENDOR_CLASS)
 	if vClass != nil {
 		req.AddOption(vClass)
 	}
-	return &req, nil
+
+	// apply modifiers
+	d := DHCPv6(&req)
+	for _, mod := range modifiers {
+		d = mod(d)
+	}
+	return d, nil
 }
 
 func (d *DHCPv6Message) Type() MessageType {
diff --git a/dhcpv6/modifiers.go b/dhcpv6/modifiers.go
new file mode 100644
index 0000000..2ad027b
--- /dev/null
+++ b/dhcpv6/modifiers.go
@@ -0,0 +1,43 @@
+package dhcpv6
+
+import (
+	"log"
+)
+
+// WithNetboot adds bootfile URL and bootfile param options to a DHCPv6 packet.
+func WithNetboot(d DHCPv6) DHCPv6 {
+	msg, ok := d.(*DHCPv6Message)
+	if !ok {
+		log.Printf("WithNetboot: not a DHCPv6Message")
+		return d
+	}
+	// add OPT_BOOTFILE_URL and OPT_BOOTFILE_PARAM
+	opt := msg.GetOneOption(OPTION_ORO)
+	if opt == nil {
+		opt = &OptRequestedOption{}
+	}
+	// TODO only add options if they are not there already
+	oro := opt.(*OptRequestedOption)
+	oro.AddRequestedOption(OPT_BOOTFILE_URL)
+	oro.AddRequestedOption(OPT_BOOTFILE_PARAM)
+	msg.UpdateOption(oro)
+	return d
+}
+
+// WithUserClass adds a user class option to the packet
+func WithUserClass(uc string) Modifier {
+	return func(d DHCPv6) DHCPv6 {
+		ouc := OptUserClass{UserClass: []byte("FbLoL")}
+		d.AddOption(&ouc)
+		return d
+	}
+}
+
+// WithArchType adds an arch type option to the packet
+func WithArchType(at ArchType) Modifier {
+	return func(d DHCPv6) DHCPv6 {
+		ao := OptClientArchType{ArchType: at}
+		d.AddOption(&ao)
+		return d
+	}
+}