Add support for modifiers to dhcpv4 (#102)

6 files changed, 68 insertions, 5 deletions

diff --git a/dhcpv4/client.go b/dhcpv4/client.go
index 1ce6a5c..db7e71a 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
 
@@ -145,6 +145,9 @@ func (c *Client) Exchange(ifname string, discover *DHCPv4) ([]DHCPv4, error) {
 			return conversation, err
 		}
 	}
+	for _, mod := range modifiers {
+		discover = mod(discover)
+	}
 	conversation[0] = *discover
 
 	// Offer
@@ -155,7 +158,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 7965161..d670de2 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) {
@@ -184,7 +188,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,11 +218,14 @@ 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
 }
 
 // NewReplyFromRequest builds a DHCPv4 reply from a request.
-func NewReplyFromRequest(request *DHCPv4) (*DHCPv4, error) {
+func NewReplyFromRequest(request *DHCPv4, modifiers ...Modifier) (*DHCPv4, error) {
 	reply, err := New()
 	if err != nil {
 		return nil, err
@@ -231,6 +238,9 @@ func NewReplyFromRequest(request *DHCPv4) (*DHCPv4, error) {
 	reply.SetTransactionID(request.TransactionID())
 	reply.SetFlags(request.Flags())
 	reply.SetGatewayIPAddr(request.GatewayIPAddr())
+	for _, mod := range modifiers {
+		reply = mod(reply)
+	}
 	return reply, nil
 }
 
diff --git a/dhcpv4/dhcpv4_test.go b/dhcpv4/dhcpv4_test.go
index 55e082d..7e5f083 100644
--- a/dhcpv4/dhcpv4_test.go
+++ b/dhcpv4/dhcpv4_test.go
@@ -353,6 +353,19 @@ func TestDHCPv4RequestFromOffer(t *testing.T) {
 	require.Equal(t, MessageTypeRequest, *req.MessageType())
 }
 
+func TestDHCPv4RequestFromOfferWithModifier(t *testing.T) {
+	offer, err := New()
+	require.NoError(t, err)
+	offer.AddOption(&OptMessageType{MessageType: MessageTypeOffer})
+	offer.AddOption(&OptServerIdentifier{ServerID: net.IPv4(192, 168, 0, 1)})
+	userClass := WithUserClass([]byte("linuxboot"))
+	req, err := RequestFromOffer(*offer, userClass)
+	require.NoError(t, err)
+	require.NotEqual(t, (*MessageType)(nil), *req.MessageType())
+	require.Equal(t, MessageTypeRequest, *req.MessageType())
+	require.Equal(t, "OptUserClass{userclass=[linuxboot]}", req.options[3].String())
+}
+
 func TestNewReplyFromRequest(t *testing.T) {
 	discover, err := New()
 	require.NoError(t, err)
@@ -363,6 +376,18 @@ func TestNewReplyFromRequest(t *testing.T) {
 	require.Equal(t, discover.GatewayIPAddr(), reply.GatewayIPAddr())
 }
 
+func TestNewReplyFromRequestWithModifier(t *testing.T) {
+	discover, err := New()
+	require.NoError(t, err)
+	discover.SetGatewayIPAddr(net.IPv4(192, 168, 0, 1))
+	userClass := WithUserClass([]byte("linuxboot"))
+	reply, err := NewReplyFromRequest(discover, userClass)
+	require.NoError(t, err)
+	require.Equal(t, discover.TransactionID(), reply.TransactionID())
+	require.Equal(t, discover.GatewayIPAddr(), reply.GatewayIPAddr())
+	require.Equal(t, "OptUserClass{userclass=[linuxboot]}", reply.options[0].String())
+}
+
 func TestDHCPv4MessageTypeNil(t *testing.T) {
 	m, err := New()
 	require.NoError(t, err)
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))
 	}