Requested Options modifier for DHCPv4 (#147)

2 files changed, 35 insertions, 0 deletions

diff --git a/dhcpv4/modifiers.go b/dhcpv4/modifiers.go
index 3b6ce70..2a9c35b 100644
--- a/dhcpv4/modifiers.go
+++ b/dhcpv4/modifiers.go
@@ -47,3 +47,19 @@ func WithNetboot(d *DHCPv4) *DHCPv4 {
 	}
 	return d
 }
+
+// WithRequestedOptions adds requested options to the packet
+func WithRequestedOptions(optionCodes ...OptionCode) Modifier {
+	return func(d *DHCPv4) *DHCPv4 {
+		params := d.GetOneOption(OptionParameterRequestList)
+		if params == nil {
+			params = &OptParameterRequestList{}
+			d.AddOption(params)
+		}
+		opts := params.(*OptParameterRequestList)
+		for _, optionCode := range optionCodes {
+			opts.RequestedOpts = append(opts.RequestedOpts, optionCode)
+		}
+		return d
+	}
+}
diff --git a/dhcpv4/modifiers_test.go b/dhcpv4/modifiers_test.go
index 2f0e1ed..c40088c 100644
--- a/dhcpv4/modifiers_test.go
+++ b/dhcpv4/modifiers_test.go
@@ -67,3 +67,22 @@ func TestWithNetbootExistingBoth(t *testing.T) {
 	d = WithNetboot(d)
 	require.Equal(t, "Parameter Request List -> [Bootfile Name, TFTP Server Name]", d.options[0].String())
 }
+
+func TestWithRequestedOptions(t *testing.T) {
+	// Check if OptionParameterRequestList is created when not present
+	d, err := New()
+	require.NoError(t, err)
+	d = WithRequestedOptions(OptionFQDN)(d)
+	require.NotNil(t, d)
+	o := d.GetOneOption(OptionParameterRequestList)
+	require.NotNil(t, o)
+	opts := o.(*OptParameterRequestList)
+	require.ElementsMatch(t, opts.RequestedOpts, []OptionCode{OptionFQDN})
+	// Check if already set options are preserved
+	d = WithRequestedOptions(OptionHostName)(d)
+	require.NotNil(t, d)
+	o = d.GetOneOption(OptionParameterRequestList)
+	require.NotNil(t, o)
+	opts = o.(*OptParameterRequestList)
+	require.ElementsMatch(t, opts.RequestedOpts, []OptionCode{OptionFQDN, OptionHostName})
+}