fix bug in IP marshal and add test to catch it

soypat · Sep 13, 2023 · d82c002 · d82c002
1 parent 009c067
commit d82c002
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 8 deletions.
diff --git a/internal/tcpctl/eth/headers.go b/internal/tcpctl/eth/headers.go
@@ -345,7 +345,7 @@ func DecodeIPv4Header(buf []byte) (iphdr IPv4Header) {
 // Put marshals the IPv4 frame onto buf. buf needs to be 20 bytes in length or Put panics.
 func (iphdr *IPv4Header) Put(buf []byte) {
 	_ = buf[19]
-	buf[0] = 4 & (iphdr.VersionAndIHL & 0xf0) // ignore set version.
+	buf[0] = (4 << 4) | (iphdr.VersionAndIHL & 0xf) // ignore set version.
 	buf[1] = iphdr.ToS
 	binary.BigEndian.PutUint16(buf[2:], iphdr.TotalLength)
 	binary.BigEndian.PutUint16(buf[4:], iphdr.ID)

diff --git a/internal/tcpctl/eth/headers_test.go b/internal/tcpctl/eth/headers_test.go
@@ -23,19 +23,26 @@ func TestUDPChecksum(t *testing.T) {
 		0x22, 0x3a, 0x20, 0x5b, 0x38, 0x31, 0x35, 0x32, 0x34, 0x36, 0x32, 0x30, 0x30, 0x30, 0x5d, 0x7d, // |": [8152462000]}|
 	}
 	// Process Ethernet header.
-	e := DecodeEthernetHeader(testUDPPacket[:14])
-	if !bytes.Equal(e.Destination[:], testUDPPacket[0:6]) {
+	ethData := testUDPPacket[:14]
+	e := DecodeEthernetHeader(ethData)
+	if !bytes.Equal(e.Destination[:], ethData[0:6]) {
 		t.Errorf("incorrect ethernet destination: %v", e.String())
 	}
-	if !bytes.Equal(e.Source[:], testUDPPacket[6:12]) {
+	if !bytes.Equal(e.Source[:], ethData[6:12]) {
 		t.Errorf("incorrect ethernet source: %v", e.String())
 	}
 	if e.AssertType() != EtherTypeIPv4 {
 		t.Errorf("incorrect ethertype: %v", e.String())
 	}
-
+	ethDataGot := make([]byte, len(ethData))
+	e.Put(ethDataGot)
+	if !bytes.Equal(ethData, ethDataGot) {
+		got := DecodeEthernetHeader(ethDataGot)
+		t.Error("ethernet marshal does not match original data", e.String(), got.String())
+	}
 	// Process IP header.
-	ip := DecodeIPv4Header(testUDPPacket[14:34])
+	ipData := testUDPPacket[14:34]
+	ip := DecodeIPv4Header(ipData)
 	if ip.Protocol != 17 {
 		t.Errorf("incorrect ip protocol: %v", ip.String())
 	}
@@ -45,9 +52,15 @@ func TestUDPChecksum(t *testing.T) {
 	if !bytes.Equal(ip.Destination[:], testUDPPacket[30:34]) {
 		t.Errorf("incorrect ip destination: %v", ip.String())
 	}
-
+	ipDataGot := make([]byte, len(ipData))
+	ip.Put(ipDataGot)
+	if !bytes.Equal(ipData, ipDataGot) {
+		got := DecodeIPv4Header(ipDataGot)
+		t.Error("ip marshal does not match original data", ip.String(), got.String())
+	}
 	// Process UDP header.
-	udp := DecodeUDPHeader(testUDPPacket[34 : 34+8])
+	udpData := testUDPPacket[34 : 34+8]
+	udp := DecodeUDPHeader(udpData)
 	if udp.SourcePort != 17500 {
 		t.Errorf("incorrect udp source port: %v", udp.String())
 	}
@@ -57,6 +70,12 @@ func TestUDPChecksum(t *testing.T) {
 	if udp.Length != 142 {
 		t.Errorf("incorrect udp length: %v", udp.String())
 	}
+	udpDataGot := make([]byte, len(udpData))
+	udp.Put(udpDataGot)
+	if !bytes.Equal(udpData, udpDataGot) {
+		got := DecodeUDPHeader(udpDataGot)
+		t.Error("udp marshal does not match original data", udp.String(), got.String())
+	}
 }
 
 func TestCRC791_oneshot(t *testing.T) {