Patched code sample

import struct

# The vulnerability was caused by uncontrolled recursion. The fix involves
# replacing the recursive function with an iterative one that uses a loop
# and tracks visited offsets to prevent infinite loops and excessive depth,
# which previously led to a RecursionError and CPU exhaustion.

# This custom exception is representative of the one used in the library.
class IncomingDecodeError(Exception):
    pass

# A simplified representation of the DNSIncoming class from the zeroconf library.
class PatchedDNSIncoming:
    def __init__(self, data: bytes) -> None:
        self.data = data

    def _decode_labels_at_offset(self, offset: int) -> list[bytes]:
        """
        Reads a domain name from the packet at a given offset.
        This is the patched, iterative version of the function.
        """
        labels: list[bytes] = []
        
        # --- START OF FIX ---
        # The fix replaces recursion with a loop and tracks visited offsets
        # to prevent pointer chains from causing a stack overflow.
        
        visited_offsets: set[int] = set()
        current_offset = offset

        while True:
            if current_offset in visited_offsets:
                raise IncomingDecodeError("DNS label loop detected")
            if len(visited_offsets) > 32: # Max pointer indirections
                raise IncomingDecodeError("DNS label pointer chain too deep")

            visited_offsets.add(current_offset)

            length = self.data[current_offset]
            current_offset += 1

            # Check if it's a pointer (first two bits are 11)
            if (length & 0b11000000) == 0b11000000:
                # It's a pointer, so we jump to a new offset.
                # The 'continue' ensures we restart the loop from the new offset
                # instead of making a recursive call.
                if current_offset >= len(self.data):
                    raise IncomingDecodeError("DNS label pointer is truncated")
                pointer_offset = ((length & 0b00111111) << 8) | self.data[current_offset]
                current_offset = pointer_offset
                continue # The core of the iterative fix
            
            # --- END OF FIX ---

            # Check if it's the end of the name (zero-length label)
            elif length == 0:
                break # Name decoding is complete

            # It's a normal label
            else:
                if current_offset + length > len(self.data):
                    raise IncomingDecodeError("DNS label is truncated")
                label = self.data[current_offset : current_offset + length]
                labels.append(label)
                current_offset += length
        
        return labels

    def read_name_from_offset(self, offset: int) -> str:
        """Helper to decode and format the name."""
        labels = self._decode_labels_at_offset(offset)
        return '.'.join(label.decode('utf-8') for label in labels) + '.'

# --- Example Usage ---

# A legitimate packet with a single pointer
# 0: \x04test\x07example\x03com\x00
# 17: \x03www\xc0\x00 -> points back to 'test.example.com'
packet_ok = b'\x04test\x07example\x03com\x00\x03www\xc0\x00'
msg_ok = PatchedDNSIncoming(packet_ok)
# This will correctly decode 'www.test.example.com.'
# print(f"Decoded OK: {msg_ok.read_name_from_offset(17)}")

# Malicious packet with a pointer loop (\xc0\x00 -> \xc0\x00)
# A vulnerable recursive implementation would crash here.
# The patched version detects the loop and raises an error gracefully.
packet_loop = b'\xc0\x00'
msg_loop = PatchedDNSIncoming(packet_loop)
try:
    msg_loop.read_name_from_offset(0)
except IncomingDecodeError as e:
    # This demonstrates the fix is working by catching the attack.
    # print(f"Caught attack (loop): {e}")
    pass

# Malicious packet with a long pointer chain
# 0: \xc0\x02 -> 2: \xc0\x04 -> 4: \xc0\x06 ...
packet_chain = b''.join([struct.pack('!H', 0b11000000_00000000 | (i + 2)) for i in range(0, 40, 2)])
msg_chain = PatchedDNSIncoming(packet_chain)
try:
    msg_chain.read_name_from_offset(0)
except IncomingDecodeError as e:
    # This demonstrates the fix is working by catching the deep chain.
    # print(f"Caught attack (deep chain): {e}")
    pass