added a packet encapsulation test that uses ip,nextheader and payload…

… together
compscidr · Sep 12, 2024 · 236fe90 · 236fe90
1 parent 3e9e824
commit 236fe90
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 3 deletions.
diff --git a/knet/src/main/kotlin/com/jasonernst/knet/Packet.kt b/knet/src/main/kotlin/com/jasonernst/knet/Packet.kt
@@ -23,10 +23,11 @@ data class Packet(
             val ipHeader = IPHeader.fromStream(stream)
             val nextHeader = NextHeader.fromStream(stream, ipHeader.protocol)
 
-            if (stream.remaining() < ipHeader.getPayloadLength().toInt()) {
-                throw PacketTooShortException("Packet too short to obtain entire payload")
+            val expectedRemaining = (ipHeader.getPayloadLength() - nextHeader.getHeaderLength()).toInt()
+            if (stream.remaining() < expectedRemaining) {
+                throw PacketTooShortException("Packet too short to obtain entire payload, have ${stream.remaining()}, expecting $expectedRemaining")
             }
-            val payload = ByteArray(ipHeader.getPayloadLength().toInt())
+            val payload = ByteArray(expectedRemaining)
             stream.get(payload)
             return Packet(ipHeader, nextHeader, payload)
         }

diff --git a/knet/src/test/kotlin/com/jasonernst/knet/EncapsulationTests.kt b/knet/src/test/kotlin/com/jasonernst/knet/EncapsulationTests.kt
@@ -148,4 +148,31 @@ class EncapsulationTests {
             )
         encapsulationTest(sourceAddress, destinationAddress, tcpHeader, payload)
     }
+
+    @Test
+    fun packetEncapsulationTest() {
+        val payload = "test".toByteArray()
+        val sourcePort = Random.Default.nextInt(2 * Short.MAX_VALUE - 1)
+        val sourceAddress = InetSocketAddress(InetAddress.getByName("::1"), sourcePort)
+        val destPort = Random.Default.nextInt(2 * Short.MAX_VALUE - 1)
+        val destinationAddress = InetSocketAddress(InetAddress.getByName("::2"), destPort)
+        val tcpHeader =
+            TCPHeader(
+                sourcePort = sourcePort.toUShort(),
+                destinationPort = destPort.toUShort(),
+                sequenceNumber = 100u,
+                acknowledgementNumber = 500u,
+                windowSize = 35000.toUShort(),
+            )
+        val ipHeader = IPHeader.createIPHeader(
+            sourceAddress.address,
+            destinationAddress.address,
+            IPType.TCP,
+            tcpHeader.getHeaderLength().toInt() + payload.size
+        )
+        val packet = Packet(ipHeader, tcpHeader, payload)
+        val stream = ByteBuffer.wrap(packet.toByteArray())
+        val parsedPacket = Packet.fromStream(stream)
+        assertEquals(packet, parsedPacket)
+    }
 }