Modifying the writeIP function to always write IPv4 in 4 byte long format.

2 files changed, 12 insertions, 0 deletions

diff --git a/dhcpv4/dhcpv4.go b/dhcpv4/dhcpv4.go
index 57342d8..24e3ed0 100644
--- a/dhcpv4/dhcpv4.go
+++ b/dhcpv4/dhcpv4.go
@@ -429,6 +429,8 @@ func writeIP(b *uio.Lexer, ip net.IP) {
 	if ip == nil {
 		b.WriteBytes(zeros[:])
 	} else {
+		// Converting IP to 4 byte format
+		ip = ip.To4()
 		b.WriteBytes(ip[:net.IPv4len])
 	}
 }
diff --git a/dhcpv4/dhcpv4_test.go b/dhcpv4/dhcpv4_test.go
index 8ee0cdb..8d6ec54 100644
--- a/dhcpv4/dhcpv4_test.go
+++ b/dhcpv4/dhcpv4_test.go
@@ -7,6 +7,7 @@ import (
 
 	"github.com/insomniacslk/dhcp/iana"
 	"github.com/stretchr/testify/require"
+	"github.com/u-root/u-root/pkg/uio"
 )
 
 func TestGetExternalIPv4Addrs(t *testing.T) {
@@ -336,3 +337,12 @@ func TestSummary(t *testing.T) {
 		"    DHCP Message Type: INFORM\n"
 	require.Equal(t, want, packet.Summary())
 }
+
+func Test_withIP(t *testing.T) {
+	buff := uio.NewBigEndianBuffer(make([]byte, 0, 20))
+	// Converting a string into IP, ip1 will be in 16 byte format
+	ip1 := net.ParseIP("10.0.0.1")
+	writeIP(buff, ip1)
+	b := buff.Buffer
+	require.Equal(t, b.Len(), 4, "Testing no of bytes written by writeIP func")
+}