Added modifier for netbootv4 (#124)

3 files changed, 100 insertions, 7 deletions

diff --git a/dhcpv4/dhcpv4_test.go b/dhcpv4/dhcpv4_test.go
index 217c801..1bce05d 100644
--- a/dhcpv4/dhcpv4_test.go
+++ b/dhcpv4/dhcpv4_test.go
@@ -358,7 +358,7 @@ func TestDHCPv4RequestFromOfferWithModifier(t *testing.T) {
 	require.NoError(t, err)
 	offer.AddOption(&OptMessageType{MessageType: MessageTypeOffer})
 	offer.AddOption(&OptServerIdentifier{ServerID: net.IPv4(192, 168, 0, 1)})
-	userClass := WithUserClass([]byte("linuxboot"))
+	userClass := WithUserClass([]byte("linuxboot"), false)
 	req, err := RequestFromOffer(*offer, userClass)
 	require.NoError(t, err)
 	require.NotEqual(t, (*MessageType)(nil), *req.MessageType())
@@ -380,7 +380,7 @@ func TestNewReplyFromRequestWithModifier(t *testing.T) {
 	discover, err := New()
 	require.NoError(t, err)
 	discover.SetGatewayIPAddr(net.IPv4(192, 168, 0, 1))
-	userClass := WithUserClass([]byte("linuxboot"))
+	userClass := WithUserClass([]byte("linuxboot"), false)
 	reply, err := NewReplyFromRequest(discover, userClass)
 	require.NoError(t, err)
 	require.Equal(t, discover.TransactionID(), reply.TransactionID())
diff --git a/dhcpv4/modifiers.go b/dhcpv4/modifiers.go
index bc19219..3b6ce70 100644
--- a/dhcpv4/modifiers.go
+++ b/dhcpv4/modifiers.go
@@ -1,11 +1,49 @@
 package dhcpv4
 
-// WithUserClass adds a user class option to the packet
-func WithUserClass(uc []byte) Modifier {
+// WithUserClass adds a user class option to the packet.
+// The rfc parameter allows you to specify if the userclass should be
+// rfc compliant or not. More details in issue #113
+func WithUserClass(uc []byte, rfc bool) Modifier {
 	// TODO let the user specify multiple user classes
 	return func(d *DHCPv4) *DHCPv4 {
-		ouc := OptUserClass{UserClasses: [][]byte{uc}}
+		ouc := OptUserClass{
+			UserClasses: [][]byte{uc},
+			Rfc3004:     rfc,
+		}
 		d.AddOption(&ouc)
 		return d
 	}
-}
\ No newline at end of file
+}
+
+// WithNetboot adds bootfile URL and bootfile param options to a DHCPv4 packet.
+func WithNetboot(d *DHCPv4) *DHCPv4 {
+	params := d.GetOneOption(OptionParameterRequestList)
+
+	var (
+		OptParams                 *OptParameterRequestList
+		foundOptionTFTPServerName bool
+		foundOptionBootfileName   bool
+	)
+	if params != nil {
+		OptParams = params.(*OptParameterRequestList)
+		for _, option := range OptParams.RequestedOpts {
+			if option == OptionTFTPServerName {
+				foundOptionTFTPServerName = true
+			} else if option == OptionBootfileName {
+				foundOptionBootfileName = true
+			}
+		}
+		if !foundOptionTFTPServerName {
+			OptParams.RequestedOpts = append(OptParams.RequestedOpts, OptionTFTPServerName)
+		}
+		if !foundOptionBootfileName {
+			OptParams.RequestedOpts = append(OptParams.RequestedOpts, OptionBootfileName)
+		}
+	} else {
+		OptParams = &OptParameterRequestList{
+			RequestedOpts: []OptionCode{OptionTFTPServerName, OptionBootfileName},
+		}
+		d.AddOption(OptParams)
+	}
+	return d
+}
diff --git a/dhcpv4/modifiers_test.go b/dhcpv4/modifiers_test.go
index 2e249ee..2f0e1ed 100644
--- a/dhcpv4/modifiers_test.go
+++ b/dhcpv4/modifiers_test.go
@@ -8,7 +8,62 @@ import (
 
 func TestUserClassModifier(t *testing.T) {
 	d, _ := New()
-	userClass := WithUserClass([]byte("linuxboot"))
+	userClass := WithUserClass([]byte("linuxboot"), false)
 	d = userClass(d)
+	expected := []byte{
+		77, // OptionUserClass
+		9,  // length
+		'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't',
+	}
 	require.Equal(t, "User Class Information -> linuxboot", d.options[0].String())
+	require.Equal(t, expected, d.options[0].ToBytes())
+}
+
+func TestUserClassModifierRFC(t *testing.T) {
+	d, _ := New()
+	userClass := WithUserClass([]byte("linuxboot"), true)
+	d = userClass(d)
+	expected := []byte{
+		77, // OptionUserClass
+		10, // length
+		9, 'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't',
+	}
+	require.Equal(t, "User Class Information -> linuxboot", d.options[0].String())
+	require.Equal(t, expected, d.options[0].ToBytes())
+}
+
+func TestWithNetboot(t *testing.T) {
+	d, _ := New()
+	d = WithNetboot(d)
+	require.Equal(t, "Parameter Request List -> [TFTP Server Name, Bootfile Name]", d.options[0].String())
+}
+
+func TestWithNetbootExistingTFTP(t *testing.T) {
+	d, _ := New()
+	OptParams := &OptParameterRequestList{
+		RequestedOpts: []OptionCode{OptionTFTPServerName},
+	}
+	d.AddOption(OptParams)
+	d = WithNetboot(d)
+	require.Equal(t, "Parameter Request List -> [TFTP Server Name, Bootfile Name]", d.options[0].String())
+}
+
+func TestWithNetbootExistingBootfileName(t *testing.T) {
+	d, _ := New()
+	OptParams := &OptParameterRequestList{
+		RequestedOpts: []OptionCode{OptionBootfileName},
+	}
+	d.AddOption(OptParams)
+	d = WithNetboot(d)
+	require.Equal(t, "Parameter Request List -> [Bootfile Name, TFTP Server Name]", d.options[0].String())
+}
+
+func TestWithNetbootExistingBoth(t *testing.T) {
+	d, _ := New()
+	OptParams := &OptParameterRequestList{
+		RequestedOpts: []OptionCode{OptionBootfileName, OptionTFTPServerName},
+	}
+	d.AddOption(OptParams)
+	d = WithNetboot(d)
+	require.Equal(t, "Parameter Request List -> [Bootfile Name, TFTP Server Name]", d.options[0].String())
 }