diff --git a/pcapkit/foundation/traceflow/tcp.py b/pcapkit/foundation/traceflow/tcp.py index d61474c02d..103ba81acd 100644 --- a/pcapkit/foundation/traceflow/tcp.py +++ b/pcapkit/foundation/traceflow/tcp.py @@ -140,6 +140,9 @@ class TCP(TraceFlowBase[BufferID, Buffer[_AT], Index, Packet[_AT]], Generic[_AT] #: and cached so that :meth:`_make_segment` does not import per packet. Set #: only when ``analyse`` is on, which is the only time it is read. _reasm_packet: 'type[Any]' + #: Every flow label this tracer has handed out, created with the first + #: flow and read by :meth:`_unique_label` to keep labels unique. + _labels: 'set[str]' ########################################################################## # Defaults. @@ -231,6 +234,12 @@ def trace(self, packet: 'Packet[_AT]', *, output: 'bool' = False) -> 'Dumper | s reverse half of a bidirectional conversation joins that flow rather than minting a label of its own. + A label is unique within the tracer. When a flow would repeat the + label of an earlier one -- the same endpoints opening a new flow at + the same capture timestamp -- it is suffixed with ``-1``, ``-2`` and + so on, so that each flow keeps its own output file; see + :meth:`_unique_label`. + """ # clear cache self.__cached__['submit'] = None @@ -267,6 +276,7 @@ def trace(self, packet: 'Packet[_AT]', *, output: 'bool' = False) -> 'Dumper | s # ``:`` is a path separator on Windows and a drive separator # elsewhere in the tooling, so an IPv6 label cannot carry it label = label.replace(':', '.') + label = self._unique_label(label) logger.debug('new TCP flow %s', label) self._buffer[BUFID] = Buffer( fpout=self._foutio(fname=f'{self._fproot}/{label}{self._fdpext or ""}', protocol=packet.protocol, @@ -322,6 +332,33 @@ def trace(self, packet: 'Packet[_AT]', *, output: 'bool' = False) -> 'Dumper | s # return label or output object return fpout if output else label + def _unique_label(self, label: 'str') -> 'str': + """Make a new flow's label unique within this tracer. + + Arguments: + label: the label built from the packet that opens the flow + + Returns: + ``label`` itself if no earlier flow carries it; otherwise ``label`` + suffixed with the first of ``-1``, ``-2``, ... that is still unused. + + The label names the flow's output file, which is opened -- and so + truncated -- when the flow starts. Two flows with one label would + therefore share one file, and the later would erase the earlier: the + same endpoints opening a second flow at the same capture timestamp, as a + unidirectional flow resumed after its FIN or a new connection reusing a + bidirectional flow's endpoints can. A label that does not collide is + returned unchanged. + + """ + used = self.__dict__.setdefault('_labels', set()) # type: set[str] + unique, count = label, 0 + while unique in used: + count += 1 + unique = f'{label}-{count}' + used.add(unique) + return unique + def _make_reassembly(self) -> 'Optional[TCP_Reassembly]': """Build the reassembler a new flow will feed, if analysis was asked for. diff --git a/tests/foundation/traceflow/test_traceflow_label_unique_unit.py b/tests/foundation/traceflow/test_traceflow_label_unique_unit.py new file mode 100644 index 0000000000..cf09a3c37e --- /dev/null +++ b/tests/foundation/traceflow/test_traceflow_label_unique_unit.py @@ -0,0 +1,103 @@ +# -*- coding: utf-8 -*- +"""TCP flow labels are unique, so no flow truncates another's output file. + +GitHub issue #1360: the label is built from the packet that opens a flow, and the +output file is named after it and opened -- truncating it -- when the flow starts. +Two flows opened on the same endpoints at the same capture timestamp got one label +and one file, and the later flow erased the earlier one's frames. That happens +unidirectionally when a direction resumes after its FIN, and bidirectionally when a +new connection reuses a torn-down conversation's endpoints. + +Every case builds its packets in memory and writes to a temporary directory. + +""" + +import importlib.util +from ipaddress import ip_address +import os +import tempfile +import unittest + +from tests._support import reimport_once_per_class + +HAS_RUNTIME = all(importlib.util.find_spec(name) is not None + for name in ('tbtrim', 'aenum', 'chardet', 'dictdumper')) + +#: The label both colliding flows used to share. +LABEL = '192.0.2.1_12345-198.51.100.2_443-300.0' + + +@unittest.skipUnless(HAS_RUNTIME, 'runtime dependencies not installed') +class TestTraceFlowLabelUnique(unittest.TestCase): + """Pin one label, and one file, per traced flow.""" + + def setUp(self) -> None: + reimport_once_per_class(self) + + @staticmethod + def _packet(index: int, *, syn: bool = False, fin: bool = False, reverse: bool = False, + timestamp: float = 300.0) -> 'object': + from pcapkit.const.reg.linktype import LinkType + from pcapkit.foundation.traceflow.data.tcp import Packet + + near, far = (ip_address('192.0.2.1'), 12345), (ip_address('198.51.100.2'), 443) + (src, srcport), (dst, dstport) = (far, near) if reverse else (near, far) + return Packet(LinkType.ETHERNET, index, {'frame': index}, syn, fin, False, src, dst, + srcport, dstport, timestamp, 0, 0, b'h', bytearray(b'')) + + def _trace(self, bidirectional: bool, packets: 'list[object]') -> 'tuple[list, dict]': + from pcapkit.foundation.traceflow.tcp import TCP + + with tempfile.TemporaryDirectory() as tempdir: + trace = TCP(tempdir, 'json', bidirectional=bidirectional) + for packet in packets: + trace.dump(packet) + trace.finish() + flows = [(flow.label, flow.index, flow.fpout) for flow in trace.index] + files = {} + for name in sorted(os.listdir(tempdir)): + with open(os.path.join(tempdir, name), encoding='utf-8') as file: + text = file.read() + files[name] = tuple(i for i in range(1, 5) if f'Frame {i}' in text) + flows = [(label, index, os.path.basename(fpout)) for label, index, fpout in flows] + return flows, files + + def test_unidirectional_flow_resumed_at_the_same_timestamp(self) -> None: + flows, files = self._trace(False, [self._packet(1, fin=True), self._packet(2)]) + + self.assertEqual(flows, [(LABEL, (1,), f'{LABEL}.json'), + (f'{LABEL}-1', (2,), f'{LABEL}-1.json')]) + self.assertEqual(files, {f'{LABEL}-1.json': (2,), f'{LABEL}.json': (1,)}) + + def test_bidirectional_connection_reusing_endpoints_at_the_same_timestamp(self) -> None: + flows, files = self._trace(True, [ + self._packet(1, syn=True), + self._packet(2, fin=True), + self._packet(3, fin=True, reverse=True), + self._packet(4, syn=True), # a new connection on the torn-down endpoints + ]) + + self.assertEqual(flows, [(LABEL, (1, 2, 3), f'{LABEL}.json'), + (f'{LABEL}-1', (4,), f'{LABEL}-1.json')]) + self.assertEqual(files, {f'{LABEL}-1.json': (4,), f'{LABEL}.json': (1, 2, 3)}) + + def test_every_repeat_gets_the_next_free_suffix(self) -> None: + flows, files = self._trace(False, [self._packet(i, fin=True) for i in range(1, 5)]) + + labels = [LABEL, f'{LABEL}-1', f'{LABEL}-2', f'{LABEL}-3'] + self.assertEqual([flow[0] for flow in flows], labels) + self.assertEqual(files, {f'{label}.json': (i,) for i, label in enumerate(labels, 1)}) + + def test_non_colliding_labels_are_unchanged(self) -> None: + flows, _ = self._trace(False, [self._packet(1, fin=True), self._packet(2, timestamp=301.0), + self._packet(3, reverse=True)]) + + self.assertEqual([flow[0] for flow in flows], [ + LABEL, + '192.0.2.1_12345-198.51.100.2_443-301.0', + '198.51.100.2_443-192.0.2.1_12345-300.0', + ]) + + +if __name__ == '__main__': + unittest.main()