Add support for modifiers to dhcpv4

5 files changed, 45 insertions, 6 deletions

diff --git a/dhcpv4/client.go b/dhcpv4/client.go
index 1ce6a5c..ce6f242 100644
--- a/dhcpv4/client.go
+++ b/dhcpv4/client.go
@@ -124,7 +124,7 @@ func MakeListeningSocket(ifname string) (int, error) {
 // ordered as Discovery, Offer, Request and Acknowledge. In case of errors, an
 // error is returned, and the list of DHCPv4 objects will be shorted than 4,
 // containing all the sent and received DHCPv4 messages.
-func (c *Client) Exchange(ifname string, discover *DHCPv4) ([]DHCPv4, error) {
+func (c *Client) Exchange(ifname string, discover *DHCPv4, modifiers ...Modifier) ([]DHCPv4, error) {
 	conversation := make([]DHCPv4, 1)
 	var err error
 
@@ -140,10 +140,14 @@ func (c *Client) Exchange(ifname string, discover *DHCPv4) ([]DHCPv4, error) {
 
 	// Discover
 	if discover == nil {
-		discover, err = NewDiscoveryForInterface(ifname)
+		discover, err = NewDiscoveryForInterface(ifname, modifiers...)
 		if err != nil {
 			return conversation, err
 		}
+	} else {
+		for _, mod := range modifiers {
+			discover = mod(discover)
+		}
 	}
 	conversation[0] = *discover
 
@@ -155,7 +159,7 @@ func (c *Client) Exchange(ifname string, discover *DHCPv4) ([]DHCPv4, error) {
 	conversation = append(conversation, *offer)
 
 	// Request
-	request, err := RequestFromOffer(*offer)
+	request, err := RequestFromOffer(*offer, modifiers...)
 	if err != nil {
 		return conversation, err
 	}
diff --git a/dhcpv4/dhcpv4.go b/dhcpv4/dhcpv4.go
index 2d73003..a74f675 100644
--- a/dhcpv4/dhcpv4.go
+++ b/dhcpv4/dhcpv4.go
@@ -38,6 +38,10 @@ type DHCPv4 struct {
 	options        []Option
 }
 
+// Modifier defines the signature for functions that can modify DHCPv4
+// structures. This is used to simplify packet manipulation
+type Modifier func(d *DHCPv4) *DHCPv4
+
 // IPv4AddrsForInterface obtains the currently-configured, non-loopback IPv4
 // addresses for iface.
 func IPv4AddrsForInterface(iface *net.Interface) ([]net.IP, error) {
@@ -120,7 +124,7 @@ func New() (*DHCPv4, error) {
 // NewDiscoveryForInterface builds a new DHCPv4 Discovery message, with a default
 // Ethernet HW type and the hardware address obtained from the specified
 // interface.
-func NewDiscoveryForInterface(ifname string) (*DHCPv4, error) {
+func NewDiscoveryForInterface(ifname string, modifiers ...Modifier) (*DHCPv4, error) {
 	d, err := New()
 	if err != nil {
 		return nil, err
@@ -144,6 +148,9 @@ func NewDiscoveryForInterface(ifname string) (*DHCPv4, error) {
 			OptionDomainNameServer,
 		},
 	})
+	for _, mod := range modifiers {
+		d = mod(d)
+	}
 	return d, nil
 }
 
@@ -184,7 +191,7 @@ func NewInformForInterface(ifname string, needsBroadcast bool) (*DHCPv4, error)
 }
 
 // RequestFromOffer builds a DHCPv4 request from an offer.
-func RequestFromOffer(offer DHCPv4) (*DHCPv4, error) {
+func RequestFromOffer(offer DHCPv4, modifiers ...Modifier) (*DHCPv4, error) {
 	d, err := New()
 	if err != nil {
 		return nil, err
@@ -214,6 +221,9 @@ func RequestFromOffer(offer DHCPv4) (*DHCPv4, error) {
 	d.AddOption(&OptMessageType{MessageType: MessageTypeRequest})
 	d.AddOption(&OptRequestedIPAddress{RequestedAddr: offer.YourIPAddr()})
 	d.AddOption(&OptServerIdentifier{ServerID: serverIP})
+	for _, mod := range modifiers {
+		d = mod(d)
+	}
 	return d, nil
 }
 
diff --git a/dhcpv4/modifiers.go b/dhcpv4/modifiers.go
new file mode 100644
index 0000000..bc19219
--- /dev/null
+++ b/dhcpv4/modifiers.go
@@ -0,0 +1,11 @@
+package dhcpv4
+
+// WithUserClass adds a user class option to the packet
+func WithUserClass(uc []byte) Modifier {
+	// TODO let the user specify multiple user classes
+	return func(d *DHCPv4) *DHCPv4 {
+		ouc := OptUserClass{UserClasses: [][]byte{uc}}
+		d.AddOption(&ouc)
+		return d
+	}
+}
\ No newline at end of file
diff --git a/dhcpv4/modifiers_test.go b/dhcpv4/modifiers_test.go
new file mode 100644
index 0000000..415a4ea
--- /dev/null
+++ b/dhcpv4/modifiers_test.go
@@ -0,0 +1,14 @@
+package dhcpv4
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestUserClassModifier(t *testing.T) {
+	d, _ := New()
+	userClass := WithUserClass([]byte("linuxboot"))
+	d = userClass(d)
+	require.Equal(t, "OptUserClass{userclass=[linuxboot]}", d.options[0].String())
+}
\ No newline at end of file
diff --git a/dhcpv4/option_userclass.go b/dhcpv4/option_userclass.go
index 1505dbb..307a128 100644
--- a/dhcpv4/option_userclass.go
+++ b/dhcpv4/option_userclass.go
@@ -39,7 +39,7 @@ func (op *OptUserClass) Length() int {
 }
 
 func (op *OptUserClass) String() string {
-	ucStrings := make([]string, len(op.UserClasses))
+	ucStrings := make([]string, 0, len(op.UserClasses))
 	for _, uc := range op.UserClasses {
 		ucStrings = append(ucStrings, string(uc))
 	}