diff options
Diffstat (limited to 'dhcpv4')
-rw-r--r-- | dhcpv4/dhcpv4_test.go | 4 | ||||
-rw-r--r-- | dhcpv4/modifiers.go | 46 | ||||
-rw-r--r-- | dhcpv4/modifiers_test.go | 57 |
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()) } |