fix(reassembly): IP reassembly reports only the octets it holds (truncated fragments, strict runs, zero-length datagrams)

pcapkit/foundation/reassembly/ip.py

@@ -37,7 +37,12 @@ class IP(ReassemblyBase[Packet[_AT], Datagram[_AT], BufferID, Buffer[_AT]], Gene
         strict: if :data:`True`, report a datagram that is not completely
             reassembled as the tuple of its received runs; otherwise as one
             contiguous payload, its holes zero-filled -- or, while the total
-            length is unknown, only the prefix up to the first hole
+            length is unknown, only the prefix up to the first hole. Receipt
+            is tracked in 8-octet blocks, so the partial last block of a
+            snaplen-truncated fragment is not counted as received: it is left
+            out of the runs and ends the prefix, but once the total length is
+            known the contiguous payload keeps its real octets, with only the
+            rest of the hole zero-filled
         store: if store reassembled datagram in memory, i.e.,
             :attr:`self._dtgram <pcapkit.foundation.reassembly.reassembly.Reassembly._dtgram>`
             (if not, datagram will be discarded after callback)
@@ -132,15 +137,24 @@ def reassembly(self, info: 'Packet[_AT]') -> 'None':
         # append packet index
         buf.index.append(info.num)
 
+        # The fragment's data length is what its header declares, but only the
+        # octets the capture actually holds can be written: a snaplen-truncated
+        # fragment carries fewer, and anything past the declared length is not
+        # fragment data. The data buffer is preallocated, so the write is also
+        # clipped to it -- its length never changes.
+        length = TL - IHL
+        held = max(min(len(info.payload), length, len(buf.datagram) - FO), 0)
+        payload = info.payload[:held]
+
         # put data into data buffer
         start = FO
-        stop = TL - IHL + FO
+        stop = FO + held
 
         # Find where this fragment disagrees with what the buffer already
         # holds *before* writing it -- ``buf.RCVBT`` and ``buf.TDL`` still
         # describe the state as every earlier fragment left it, which is
         # exactly what :meth:`_detect_conflicts` needs.
-        conflicts = self._detect_conflicts(buf.RCVBT, buf.TDL, buf.datagram, info.payload, start, stop)
+        conflicts = self._detect_conflicts(buf.RCVBT, buf.TDL, buf.datagram, payload, start, stop)
         if conflicts:
             buf.conflict.extend(conflicts)
 
@@ -150,23 +164,31 @@ def reassembly(self, info: 'Packet[_AT]') -> 'None':
         # payload always overwrites here; ``conflicts`` above is what records
         # that it *disagreed* with what it overwrote, which is the part RFC 791
         # leaves unrecorded.
-        buf.datagram[start:stop] = info.payload
+        buf.datagram[start:stop] = payload
 
-        # set RCVBT bits (in 8 octets)
+        # Set RCVBT bits (in 8 octets) for the blocks this fragment actually
+        # filled. A fragment held in full marks its last block even when it is
+        # short, as the final fragment's is; a truncated one leaves its partial
+        # last block clear, since the octets after the cut are a real hole.
         start = FO // 8
-        stop = FO // 8 + (TL - IHL + 7) // 8
-        buf.RCVBT[start:stop] = b'\x01' * (stop - start)
+        if held == length:
+            stop = (FO + held + 7) // 8
+        else:
+            stop = (FO + held) // 8
+        stop = min(stop, len(buf.RCVBT))
+        if stop > start:
+            buf.RCVBT[start:stop] = b'\x01' * (stop - start)
 
-        # get total data length (header excludes)
-        TDL = 0
+        # get total data length (header excludes); it stays ``-1`` until the
+        # fragment with MF clear has arrived, and ``0`` is a valid length
         if not MF:
-            TDL = TL - IHL + FO
-            buf.__update__(TDL=TDL)
+            buf.__update__(TDL=TL - IHL + FO)
+        TDL = buf.TDL
 
         # when datagram is reassembled in whole
         start = 0
         stop = (TDL + 7) // 8
-        if TDL and all(buf.RCVBT[start:stop]):
+        if TDL >= 0 and all(buf.RCVBT[start:stop]):
             self._dtgram.extend(
                 self.submit(self._buffer.pop(BUFID), bufid=BUFID, checked=True)
             )
@@ -270,7 +292,7 @@ def submit(self, buf: 'Buffer[_AT]', *, bufid: 'tuple[_AT, _AT, int, TransType]'
 
         start = 0
         stop = (TDL + 7) // 8
-        flag = checked or (TDL > 0 and all(RCVBT[start:stop]))
+        flag = checked or (TDL >= 0 and all(RCVBT[start:stop]))
         ret = []  # type: list[Datagram[_AT]]
 
         # How completely this datagram came out, and why it stopped. Derived once,
@@ -283,16 +305,21 @@ def submit(self, buf: 'Buffer[_AT]', *, bufid: 'tuple[_AT, _AT, int, TransType]'
         if not flag and self._flag_s:
             data = []  # type: list[bytes]
             byte = bytearray()
-            # extract received payload
+            # extract received payload; a run never extends past the total
+            # data length once it is known
+            limit = TDL if TDL >= 0 else len(datagram)
             for (bctr, bit) in enumerate(RCVBT):
                 if bit:     # received bit
                     this = bctr * 8
-                    that = this + 8
+                    that = min(this + 8, limit)
                     byte += datagram[this:that]
                 else:       # missing bit
                     if byte:    # strip empty payload
                         data.append(bytes(byte))
                     byte = bytearray()
+            # the last run may reach the end of the bit table
+            if byte:
+                data.append(bytes(byte))
             # strip empty packets
             if data or header:
                 packet = Datagram(
@@ -317,12 +344,14 @@ def submit(self, buf: 'Buffer[_AT]', *, bufid: 'tuple[_AT, _AT, int, TransType]'
             # fragment with **MF** clear has arrived -- until then it is still its
             # initial ``-1``. Which case this is decides how much of the buffer
             # there is to report.
-            if TDL > 0:
+            if TDL >= 0:
                 # The length is known. Report it, holes and all: the gaps read as
                 # zeros, exactly as they do in the TCP reassembler's loose mode.
                 # This is the reason ``strict=False`` exists -- a caller who
                 # wants the gaps *marked* rather than
-                # zero-filled uses ``strict=True`` and gets the runs.
+                # zero-filled uses ``strict=True`` and gets the runs. A truncated
+                # fragment's partial last block is not marked received, but the
+                # octets it did write are real and are reported here as they are.
                 stop = TDL
             else:
                 # The length is not known. Slicing ``datagram[:-1]`` here would
@@ -341,7 +370,9 @@ def submit(self, buf: 'Buffer[_AT]', *, bufid: 'tuple[_AT, _AT, int, TransType]'
                 # precisely because their offsets cannot be conveyed in a blob.
                 #
                 # ``RCVBT`` records receipt in 8-octet units, so the prefix ends at
-                # the first clear bit.
+                # the first clear bit -- which leaves out a truncated fragment's
+                # partial last block, octets that the known-length case above
+                # does report.
                 received = 0
                 for bit in RCVBT:
                     if not bit:

tests/foundation/reassembly/test_ip_reassembly_roundtrip_unit.py

@@ -0,0 +1,114 @@
+# -*- coding: utf-8 -*-
+"""IP reassembly reports exactly the octets it received.
+
+- GitHub issue #1357: a snaplen-truncated fragment shrank the data buffer and
+  marked its whole declared range received, so the datagram came out
+  ``COMPLETE`` with zeros in it and the rest shifted.
+- GitHub issue #1358: ``strict`` runs copied whole 8-octet blocks past the
+  total data length, and dropped a run that reached the end of the bit table.
+- GitHub issue #1359: a datagram whose total data length is ``0`` never
+  completed, because the length was tested for truthiness.
+
+Every case builds its fragments in memory and reads no capture. Classes are
+imported inside each test, after :func:`~tests._support.reimport_once_per_class`.
+
+"""
+
+from ipaddress import ip_address
+import unittest
+
+from tests._support import reimport_once_per_class
+
+#: A 2000-octet payload with no zero octet in it, so a zero-filled hole shows.
+FULL = (bytes(range(1, 256)) * 8)[:2000]
+
+
+class TestIPReassemblyRoundTrip(unittest.TestCase):
+    """Pin what the IP reassembler reports for truncated, sparse and empty datagrams."""
+
+    def setUp(self) -> None:
+        reimport_once_per_class(self)
+
+    @staticmethod
+    def _packet(num, fo, mf, payload, *, tl=None):
+        from pcapkit.const.reg.transtype import TransType
+        from pcapkit.foundation.reassembly.data.ip import Packet
+
+        tl = 20 + len(payload) if tl is None else tl
+        bufid = (ip_address('192.0.2.1'), ip_address('198.51.100.2'), 42, TransType.UDP)
+        return Packet(bufid, num, fo, 20, mf, tl, b'ip-header' if fo == 0 else b'',
+                      bytearray(payload), 1000.0)
+
+    @staticmethod
+    def _reassembler(strict):
+        from pcapkit.foundation.reassembly.ipv4 import IPv4
+
+        return IPv4(strict=strict)
+
+    def _truncated(self, order, strict):
+        frags = [(0, True, FULL[0:800]), (800, True, FULL[800:1600]), (1600, False, FULL[1600:])]
+        reasm = self._reassembler(strict)
+        for num, idx in enumerate(order, 1):
+            fo, mf, payload = frags[idx]
+            declared = len(payload)
+            if idx == 1:
+                payload = payload[:500]  # snaplen cut: 500 of 800 octets held
+            reasm(self._packet(num, fo, mf, payload, tl=20 + declared))
+        return reasm
+
+    def test_truncated_fragment_leaves_a_hole_strict(self) -> None:
+        for order in ([0, 1, 2], [2, 1, 0], [0, 2, 1]):
+            with self.subTest(order=order):
+                datagram, = self._truncated(order, strict=True).datagram
+                self.assertEqual(datagram.completed.name, 'PARTIAL')
+                # 800 + 500 = 1300 octets held, of which 162 whole blocks
+                self.assertEqual(datagram.payload, (FULL[:1296], FULL[1600:]))
+
+    def test_truncated_fragment_keeps_buffer_length_loose(self) -> None:
+        for order in ([0, 1, 2], [2, 1, 0], [0, 2, 1]):
+            with self.subTest(order=order):
+                reasm = self._truncated(order, strict=False)
+                datagram, = reasm.datagram
+                self.assertEqual(datagram.completed.name, 'PARTIAL')
+                self.assertEqual(len(datagram.payload), 2000)
+                self.assertEqual(datagram.payload[:1300], FULL[:1300])
+                self.assertEqual(datagram.payload[1300:1600], bytes(300))
+                self.assertEqual(datagram.payload[1600:], FULL[1600:])
+
+    def test_strict_run_clipped_at_total_data_length(self) -> None:
+        reasm = self._reassembler(strict=True)
+        reasm(self._packet(1, 0, True, b'A' * 1480))
+        reasm(self._packet(2, 2960, False, b'C' * 43))
+        datagram, = reasm.datagram
+        self.assertEqual(datagram.completed.name, 'PARTIAL')
+        self.assertEqual(datagram.payload, (b'A' * 1480, b'C' * 43))
+
+    def test_strict_run_at_end_of_bit_table_is_kept(self) -> None:
+        reasm = self._reassembler(strict=True)
+        reasm(self._packet(1, 0, True, b'A' * 8))
+        reasm(self._packet(2, 65520, False, b'Z' * 7))
+        datagram, = reasm.datagram
+        self.assertEqual(datagram.payload, (b'A' * 8, b'Z' * 7))
+
+    def test_zero_length_unfragmented_datagram_completes(self) -> None:
+        for strict in (True, False):
+            with self.subTest(strict=strict):
+                reasm = self._reassembler(strict)
+                reasm(self._packet(1, 0, False, b''))
+                self.assertEqual(dict(reasm._buffer), {})
+                datagram, = reasm.datagram
+                self.assertEqual(datagram.completed.name, 'COMPLETE')
+                self.assertEqual(datagram.payload, b'')
+
+    def test_completes_when_the_last_hole_fills_after_the_final_fragment(self) -> None:
+        reasm = self._reassembler(strict=True)
+        reasm(self._packet(1, 8, False, b'ijkl'))
+        reasm(self._packet(2, 0, True, b'abcdefgh'))
+        self.assertEqual(dict(reasm._buffer), {})
+        datagram, = reasm.datagram
+        self.assertEqual(datagram.completed.name, 'COMPLETE')
+        self.assertEqual(datagram.payload, b'abcdefghijkl')
+
+
+if __name__ == '__main__':
+    unittest.main()