From b5c839d7f28cc635376d5bee36de0e898ad8a1b8 Mon Sep 17 00:00:00 2001 From: Owen Mooney Date: Sun, 29 Jul 2018 14:14:31 +0100 Subject: Add support for modifiers to dhcpv4 --- dhcpv4/client.go | 10 +++++++--- dhcpv4/dhcpv4.go | 14 ++++++++++++-- dhcpv4/modifiers.go | 11 +++++++++++ dhcpv4/modifiers_test.go | 14 ++++++++++++++ dhcpv4/option_userclass.go | 2 +- 5 files changed, 45 insertions(+), 6 deletions(-) create mode 100644 dhcpv4/modifiers.go create mode 100644 dhcpv4/modifiers_test.go diff --git a/dhcpv4/client.go b/dhcpv4/client.go index 1ce6a5c..ce6f242 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 @@ -140,10 +140,14 @@ func (c *Client) Exchange(ifname string, discover *DHCPv4) ([]DHCPv4, error) { // Discover if discover == nil { - discover, err = NewDiscoveryForInterface(ifname) + discover, err = NewDiscoveryForInterface(ifname, modifiers...) if err != nil { return conversation, err } + } else { + for _, mod := range modifiers { + discover = mod(discover) + } } conversation[0] = *discover @@ -155,7 +159,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 2d73003..a74f675 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) { @@ -120,7 +124,7 @@ func New() (*DHCPv4, error) { // NewDiscoveryForInterface builds a new DHCPv4 Discovery message, with a default // Ethernet HW type and the hardware address obtained from the specified // interface. -func NewDiscoveryForInterface(ifname string) (*DHCPv4, error) { +func NewDiscoveryForInterface(ifname string, modifiers ...Modifier) (*DHCPv4, error) { d, err := New() if err != nil { return nil, err @@ -144,6 +148,9 @@ func NewDiscoveryForInterface(ifname string) (*DHCPv4, error) { OptionDomainNameServer, }, }) + for _, mod := range modifiers { + d = mod(d) + } return d, nil } @@ -184,7 +191,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,6 +221,9 @@ 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 } 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)) } -- cgit v1.2.3 From a6ecc33c23d061098e9a2bf4d9d2d4c0a710406b Mon Sep 17 00:00:00 2001 From: Owen Mooney Date: Sun, 29 Jul 2018 16:34:55 +0100 Subject: Add modifier support to NewReplyFromRequest --- dhcpv4/client.go | 9 ++++----- dhcpv4/dhcpv4.go | 10 +++++----- dhcpv4/dhcpv4_test.go | 25 +++++++++++++++++++++++++ 3 files changed, 34 insertions(+), 10 deletions(-) diff --git a/dhcpv4/client.go b/dhcpv4/client.go index ce6f242..db7e71a 100644 --- a/dhcpv4/client.go +++ b/dhcpv4/client.go @@ -140,14 +140,13 @@ func (c *Client) Exchange(ifname string, discover *DHCPv4, modifiers ...Modifier // Discover if discover == nil { - discover, err = NewDiscoveryForInterface(ifname, modifiers...) + discover, err = NewDiscoveryForInterface(ifname) if err != nil { return conversation, err } - } else { - for _, mod := range modifiers { - discover = mod(discover) - } + } + for _, mod := range modifiers { + discover = mod(discover) } conversation[0] = *discover diff --git a/dhcpv4/dhcpv4.go b/dhcpv4/dhcpv4.go index a74f675..e027f28 100644 --- a/dhcpv4/dhcpv4.go +++ b/dhcpv4/dhcpv4.go @@ -124,7 +124,7 @@ func New() (*DHCPv4, error) { // NewDiscoveryForInterface builds a new DHCPv4 Discovery message, with a default // Ethernet HW type and the hardware address obtained from the specified // interface. -func NewDiscoveryForInterface(ifname string, modifiers ...Modifier) (*DHCPv4, error) { +func NewDiscoveryForInterface(ifname string) (*DHCPv4, error) { d, err := New() if err != nil { return nil, err @@ -148,9 +148,6 @@ func NewDiscoveryForInterface(ifname string, modifiers ...Modifier) (*DHCPv4, er OptionDomainNameServer, }, }) - for _, mod := range modifiers { - d = mod(d) - } return d, nil } @@ -228,7 +225,7 @@ func RequestFromOffer(offer DHCPv4, modifiers ...Modifier) (*DHCPv4, error) { } // 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 @@ -241,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) -- cgit v1.2.3