contrib/tcpros: read field lengths in place instead of re-slicing

scapy/contrib/tcpros.py

@@ -25,6 +25,28 @@
 from scapy.config import conf
 
 
+def tcpros_iter_field_lengths(data, start=0):
+    """Iterate over the lengths of the ``[4-byte length + value]`` fields.
+
+    The fields are read in place with ``struct.unpack_from``: slicing the
+    buffer for every field copies the whole remainder, which makes the
+    traversal quadratic in the number of fields.
+
+    Args:
+        data: buffer holding the field list.
+        start: offset of the first field length in ``data``.
+
+    Yields:
+        the length of each field, in order.
+    """
+    end = len(data)
+    offset = start
+    while offset < end:
+        field_len = struct.unpack_from("<I", data, offset)[0]
+        yield field_len
+        offset += 4 + field_len
+
+
 class TCPROS(Packet):
     """
     TCPROS is a transport layer for ROS Messages and Services. It uses standard
@@ -97,15 +119,13 @@ def guess_payload_class(self, payload):
             # TCPROS encoding format:
             #  [4-byte field length + field=value ]*
             flag_encoding_format_subfields = True
-            while remain:
-                field_len_bytes = struct.unpack("<I", remain[:4])[0]
-                current = remain[4:4 + field_len_bytes]
-                remain = remain[4 + field_len_bytes:]
-
-                if int(field_len_bytes) != len(current):
+            offset = 4
+            for field_len in tcpros_iter_field_lengths(payload, 4):
+                if field_len > len(payload) - offset - 4:
                     # print("BREAKING - int(field_len_bytes) != len(current)")
                     flag_encoding_format_subfields = False
                     break
+                offset += 4 + field_len
 
         if (
             "callerid" in string_payload and
@@ -214,14 +234,9 @@ def pre_dissect(self, s):
             print("total_length: " + str(total_length))
             print("total_header_length: " + str(total_header_length))
         remain = raw(s)[4:total_header_length]
-        while remain:
-            field_len_bytes = struct.unpack("<I", remain[:4])[0]
-            remain = remain[4 + field_len_bytes:]
+        for field_len in tcpros_iter_field_lengths(remain):
             if conf.debug_dissector:
-                print("field_len_bytes: " + str(field_len_bytes))
-                current = remain[:4 + field_len_bytes]
-                print("current: " + str(current))
-                print("remain: " + str(remain))
+                print("field_len_bytes: " + str(field_len))
             self.nfields += 1
         return s
 
@@ -288,14 +303,9 @@ def pre_dissect(self, s):
             total_length = len(s)
             print("total_length: " + str(total_length))
             print("total_header_length: " + str(total_header_length))
-        while remain:
-            field_len_bytes = struct.unpack("<I", remain[:4])[0]
-            remain = remain[4 + field_len_bytes:]
+        for field_len in tcpros_iter_field_lengths(remain):
             if conf.debug_dissector:
-                print("field_len_bytes: " + str(field_len_bytes))
-                current = remain[:4 + field_len_bytes]
-                print("current: " + str(current))
-                print("remain: " + str(remain))
+                print("field_len_bytes: " + str(field_len))
             self.nfields_body += 1
         return s
 

test/contrib/tcpros.uts

@@ -56,3 +56,42 @@ assert(p[XMLRPCCall].version ==  b"<?xml version='1.0'?>\n")
 assert(p[XMLRPCCall].methodcall_opentag ==  b'<methodCall>\n')
 assert(p[XMLRPCCall].methodname ==  b'shutdown')
 assert(p[XMLRPCCall].params ==  b'<param>\n<value><string>/rosparam-92418</string></value>\n</param>\n<param>\n<value><string>BOOM</string></value>\n</param>\n')
+
++ Test TCPROS field list parsing
+= Test field lengths traversal
+from struct import pack
+from scapy.contrib.tcpros import tcpros_iter_field_lengths
+
+field_list = pack("<I", 3) + b"a=1" + pack("<I", 5) + b"bb=22" + pack("<I", 7) + b"ccc=333"
+assert(list(tcpros_iter_field_lengths(field_list)) == [3, 5, 7])
+
+= Test field lengths traversal from an offset
+assert(list(tcpros_iter_field_lengths(pack("<I", 42) + field_list, 4)) == [3, 5, 7])
+
+= Test header with many fields
+from scapy.config import conf
+
+saved_max_list_count = conf.max_list_count
+conf.max_list_count = 1024
+many_fields = [b"callerid=/test/%d" % i for i in range(300)]
+many_payload = b"".join(pack("<I", len(f)) + f for f in many_fields)
+many_header = TCPROS(pack("<I", len(many_payload)) + many_payload)[TCPROSHeader]
+conf.max_list_count = saved_max_list_count
+assert(many_header.nfields == 300)
+assert(len(many_header.list) == 300)
+assert(many_header.list[-1].field == many_fields[-1])
+
++ Test field list parsing with dissector debug output
+= Test header and body traversal with debug_dissector enabled
+from struct import pack
+from scapy.config import conf
+
+debug_fields = pack("<I", 3) + b"a=1" + pack("<I", 5) + b"bb=22"
+debug_payload = pack("<I", 4 + len(debug_fields)) + debug_fields
+saved_debug_dissector = conf.debug_dissector
+conf.debug_dissector = True
+try:
+    assert(TCPROSHeader(debug_payload).nfields == 2)
+    assert(TCPROSBody(debug_payload).nfields_body == 2)
+finally:
+    conf.debug_dissector = saved_debug_dissector