diff options
Diffstat (limited to 'dhcpv4')
-rw-r--r-- | dhcpv4/modifiers.go | 16 | ||||
-rw-r--r-- | dhcpv4/modifiers_test.go | 19 |
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}) +} |