add NewReplyFromRequest (#100)

3 files changed, 60 insertions, 9 deletions

diff --git a/dhcpv4/bsdp/bsdp.go b/dhcpv4/bsdp/bsdp.go
index 51ec20c..27de68e 100644
--- a/dhcpv4/bsdp/bsdp.go
+++ b/dhcpv4/bsdp/bsdp.go
@@ -73,7 +73,6 @@ func NewInformListForInterface(iface string, replyPort uint16) (*dhcpv4.DHCPv4,
 		return nil, err
 	}
 	d.AddOption(&dhcpv4.OptClassIdentifier{Identifier: vendorClassID})
-	d.AddOption(&dhcpv4.OptionGeneric{OptionCode: dhcpv4.OptionEnd})
 	return d, nil
 }
 
@@ -138,6 +137,5 @@ func InformSelectForAck(ack dhcpv4.DHCPv4, replyPort uint16, selectedImage BootI
 	})
 	d.AddOption(&dhcpv4.OptMessageType{MessageType: dhcpv4.MessageTypeInform})
 	d.AddOption(&OptVendorSpecificInformation{vendorOpts})
-	d.AddOption(&dhcpv4.OptionGeneric{OptionCode: dhcpv4.OptionEnd})
 	return d, nil
 }
diff --git a/dhcpv4/dhcpv4.go b/dhcpv4/dhcpv4.go
index a88be19..2d73003 100644
--- a/dhcpv4/dhcpv4.go
+++ b/dhcpv4/dhcpv4.go
@@ -112,6 +112,8 @@ func New() (*DHCPv4, error) {
 		return nil, err
 	}
 	d.options = options
+	// the End option has to be added explicitly
+	d.AddOption(&OptionGeneric{OptionCode: OptionEnd})
 	return &d, nil
 }
 
@@ -142,14 +144,12 @@ func NewDiscoveryForInterface(ifname string) (*DHCPv4, error) {
 			OptionDomainNameServer,
 		},
 	})
-	// the End option has to be added explicitly
-	d.AddOption(&OptionGeneric{OptionCode: OptionEnd})
 	return d, nil
 }
 
 // NewInformForInterface builds a new DHCPv4 Informational message with default
 // Ethernet HW type and the hardware address obtained from the specified
-// interface. It does NOT put a DHCP End option at the end.
+// interface.
 func NewInformForInterface(ifname string, needsBroadcast bool) (*DHCPv4, error) {
 	d, err := New()
 	if err != nil {
@@ -214,11 +214,26 @@ func RequestFromOffer(offer DHCPv4) (*DHCPv4, error) {
 	d.AddOption(&OptMessageType{MessageType: MessageTypeRequest})
 	d.AddOption(&OptRequestedIPAddress{RequestedAddr: offer.YourIPAddr()})
 	d.AddOption(&OptServerIdentifier{ServerID: serverIP})
-	// the End option has to be added explicitly
-	d.AddOption(&OptionGeneric{OptionCode: OptionEnd})
 	return d, nil
 }
 
+// NewReplyFromRequest builds a DHCPv4 reply from a request.
+func NewReplyFromRequest(request *DHCPv4) (*DHCPv4, error) {
+	reply, err := New()
+	if err != nil {
+		return nil, err
+	}
+	reply.SetOpcode(OpcodeBootReply)
+	reply.SetHwType(request.HwType())
+	reply.SetHwAddrLen(request.HwAddrLen())
+	hwaddr := request.ClientHwAddr()
+	reply.SetClientHwAddr(hwaddr[:])
+	reply.SetTransactionID(request.TransactionID())
+	reply.SetFlags(request.Flags())
+	reply.SetGatewayIPAddr(request.GatewayIPAddr())
+	return reply, nil
+}
+
 // FromBytes encodes the DHCPv4 packet into a sequence of bytes, and returns an
 // error if the packet is not valid.
 func FromBytes(data []byte) (*DHCPv4, error) {
@@ -556,9 +571,17 @@ func (d *DHCPv4) SetOptions(options []Option) {
 	d.options = options
 }
 
-// AddOption appends an option to the existing ones.
+// AddOption appends an option to the existing ones. If the last option is an
+// OptionEnd, it will be inserted before that. It does not deal with End
+// options that appead before the end, like in malformed packets.
 func (d *DHCPv4) AddOption(option Option) {
-	d.options = append(d.options, option)
+	if len(d.options) == 0 || d.options[len(d.options)-1].Code() != OptionEnd {
+		d.options = append(d.options, option)
+	} else {
+		end := d.options[len(d.options)-1]
+		d.options[len(d.options)-1] = option
+		d.options = append(d.options, end)
+	}
 }
 
 // MessageType returns the message type, trying to extract it from the
diff --git a/dhcpv4/dhcpv4_test.go b/dhcpv4/dhcpv4_test.go
index 9d44810..55e082d 100644
--- a/dhcpv4/dhcpv4_test.go
+++ b/dhcpv4/dhcpv4_test.go
@@ -290,6 +290,8 @@ func TestNewToBytes(t *testing.T) {
 	}
 	// Magic Cookie
 	expected = append(expected, MagicCookie...)
+	// End
+	expected = append(expected, 0xff)
 
 	d, err := New()
 	require.NoError(t, err)
@@ -322,6 +324,24 @@ func TestGetOption(t *testing.T) {
 	require.Equal(t, d.GetOneOption(OptionRouter), nil)
 }
 
+func TestAddOption(t *testing.T) {
+	d, err := New()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	hostnameOpt := &OptionGeneric{OptionCode: OptionHostName, Data: []byte("darkstar")}
+	bootFileOpt1 := &OptionGeneric{OptionCode: OptionBootfileName, Data: []byte("boot.img")}
+	bootFileOpt2 := &OptionGeneric{OptionCode: OptionBootfileName, Data: []byte("boot2.img")}
+	d.AddOption(hostnameOpt)
+	d.AddOption(bootFileOpt1)
+	d.AddOption(bootFileOpt2)
+
+	options := d.Options()
+	require.Equal(t, len(options), 4)
+	require.Equal(t, options[3].Code(), OptionEnd)
+}
+
 func TestDHCPv4RequestFromOffer(t *testing.T) {
 	offer, err := New()
 	require.NoError(t, err)
@@ -333,6 +353,16 @@ func TestDHCPv4RequestFromOffer(t *testing.T) {
 	require.Equal(t, MessageTypeRequest, *req.MessageType())
 }
 
+func TestNewReplyFromRequest(t *testing.T) {
+	discover, err := New()
+	require.NoError(t, err)
+	discover.SetGatewayIPAddr(net.IPv4(192, 168, 0, 1))
+	reply, err := NewReplyFromRequest(discover)
+	require.NoError(t, err)
+	require.Equal(t, discover.TransactionID(), reply.TransactionID())
+	require.Equal(t, discover.GatewayIPAddr(), reply.GatewayIPAddr())
+}
+
 func TestDHCPv4MessageTypeNil(t *testing.T) {
 	m, err := New()
 	require.NoError(t, err)