From 210ae3190aad264ae3141155dbd4e4009644f405 Mon Sep 17 00:00:00 2001
From: Jason Ernst <ernstjason1@gmail.com>
Date: Sun, 8 Sep 2024 08:24:45 -0700
Subject: [PATCH] Added more ipv4 test coverage

---
 .../com/jasonernst/knet/ip/IPv4Header.kt      | 15 +++----
 .../com/jasonernst/knet/ip/IPv4HeaderTest.kt  | 40 +++++++++++++++++++
 2 files changed, 46 insertions(+), 9 deletions(-)

diff --git a/knet/src/main/kotlin/com/jasonernst/knet/ip/IPv4Header.kt b/knet/src/main/kotlin/com/jasonernst/knet/ip/IPv4Header.kt
index 86ccf50..baa7c2e 100644
--- a/knet/src/main/kotlin/com/jasonernst/knet/ip/IPv4Header.kt
+++ b/knet/src/main/kotlin/com/jasonernst/knet/ip/IPv4Header.kt
@@ -96,11 +96,6 @@ data class IPv4Header(
                 throw IllegalArgumentException("Invalid IPv4 header. IP version should be 4 but was $ipVersion")
             }
 
-            // ensure we have enough to to get IHL
-            if (stream.remaining() < 1) {
-                throw IllegalArgumentException("IPv4Header: stream too short to determine header length")
-            }
-
             // ensure we have enough capacity in the stream to parse out a full header
             val ihl: UByte = (versionAndHeaderLength.toInt() and 0x0F).toUByte()
             val headerAvailable = stream.limit() - start
@@ -131,12 +126,12 @@ data class IPv4Header(
             val destinationAddress = Inet4Address.getByAddress(destination) as Inet4Address
 
             // todo (compscidr): parse the options field instead of just dropping them
+            logger.debug("POS: ${stream.position()}, remaining: ${stream.remaining()}")
             if (ihl > 5u) {
-                // drop the IP option
-                for (i in 0u until (ihl - 5u)) {
-                    stream.int
-                }
+                logger.debug("Dropping IP options")
+                stream.position(stream.position() + ((ihl - 5u) * 4u).toInt())
             }
+            logger.debug("POS: ${stream.position()}, remaining: ${stream.remaining()}")
 
             return IPv4Header(
                 ihl = ihl,
@@ -157,6 +152,8 @@ data class IPv4Header(
     }
 
     override fun toByteArray(order: ByteOrder): ByteArray {
+        // note: this will reserve the space that was previously setup for options
+        // but they will be all zero'd out, not sure the impact of this
         val buffer = ByteBuffer.allocate((ihl * IP4_WORD_LENGTH).toInt())
         buffer.order(order)
 
diff --git a/knet/src/test/kotlin/com/jasonernst/knet/ip/IPv4HeaderTest.kt b/knet/src/test/kotlin/com/jasonernst/knet/ip/IPv4HeaderTest.kt
index 2dc6693..746f46f 100644
--- a/knet/src/test/kotlin/com/jasonernst/knet/ip/IPv4HeaderTest.kt
+++ b/knet/src/test/kotlin/com/jasonernst/knet/ip/IPv4HeaderTest.kt
@@ -1,6 +1,7 @@
 package com.jasonernst.knet.ip
 
 import com.jasonernst.icmp_common.Checksum
+import com.jasonernst.knet.PacketTooShortException
 import com.jasonernst.knet.ip.IPHeader.Companion.IP4_VERSION
 import com.jasonernst.knet.ip.IPv4Header.Companion.IP4_MIN_HEADER_LENGTH
 import com.jasonernst.knet.ip.IPv4Header.Companion.IP4_WORD_LENGTH
@@ -9,6 +10,7 @@ import com.jasonernst.knet.transport.tcp.options.TCPOptionEndOfOptionList
 import com.jasonernst.packetdumper.stringdumper.StringPacketDumper
 import org.junit.jupiter.api.Assertions.assertEquals
 import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.assertThrows
 import org.slf4j.LoggerFactory
 import java.net.InetAddress
 import java.nio.ByteBuffer
@@ -274,4 +276,42 @@ class IPv4HeaderTest {
         assertEquals(ipHeader2, parsedIpHeader2)
         assertEquals(tcpHeader2, parsedTCPHeader2)
     }
+
+    @Test
+    fun tooShortPacketTest() {
+        val buffer = ByteBuffer.allocate(0)
+        assertThrows<PacketTooShortException> {
+            IPv4Header.fromStream(buffer)
+        }
+    }
+
+    @Test
+    fun badIPVersion() {
+        val buffer = ByteBuffer.allocate(1)
+        buffer.put(0x00)
+        buffer.rewind()
+        assertThrows<IllegalArgumentException> {
+            IPv4Header.fromStream(buffer)
+        }
+    }
+
+    @Test fun tooShortForFullHeader() {
+        val buffer = ByteBuffer.allocate(2)
+        buffer.put(0x45)
+        buffer.put(0x00)
+        buffer.rewind()
+        assertThrows<PacketTooShortException> {
+            IPv4Header.fromStream(buffer)
+        }
+    }
+
+    // this is only a temp test until we properly implement ipv4 options
+    @Test fun testOptionDropping() {
+        val ipv4Packet = IPv4Header(ihl = 6u)
+        logger.debug("IPv4 packet: {}", ipv4Packet)
+        val buffer = ByteBuffer.wrap(ipv4Packet.toByteArray())
+        val parsedPacket = IPHeader.fromStream(buffer)
+        assertEquals(0u, buffer.remaining().toUInt())
+        assertEquals(ipv4Packet, parsedPacket)
+    }
 }