Refactor DHCP receive as full packet always returned. Add support for 6100 UDP headers.

Author: BiffoBear

adafruit_wiznet5k/adafruit_wiznet5k.py

@@ -860,24 +860,19 @@ def read_udp(self, socket_num: int, length: int) -> Tuple[int, bytes]:
         """
         self._sock_num_in_range(socket_num)
         bytes_on_socket, bytes_read = 0, b""
-        # Parse the UDP Rx packet.
-        _, self._pbuff[:8] = self.socket_read(socket_num, 8)
-        try:
-            self.udp_from_ip[socket_num] = self._pbuff[:4]
-            self.udp_from_port[socket_num] = int.from_bytes(self._pbuff[4:6], "big")
-            udp_data_bytes = int.from_bytes(self._pbuff[6:8], "big")
-        except IndexError as err:
-            raise IndexError("Invalid UDP packet header.") from err
+        # Parse the UDP packet header.
+        header_length, self._pbuff[:8] = self.socket_read(socket_num, 8)
+        if header_length != 8:
+            raise ValueError("Invalid UDP header.")
+        data_length = self._chip_parse_udp_header(socket_num)
         # Read the UDP packet data.
-        if udp_data_bytes:
-            if udp_data_bytes <= length:
-                bytes_on_socket, bytes_read = self.socket_read(
-                    socket_num, udp_data_bytes
-                )
+        if data_length:
+            if data_length <= length:
+                bytes_on_socket, bytes_read = self.socket_read(socket_num, data_length)
             else:
                 bytes_on_socket, bytes_read = self.socket_read(socket_num, length)
                 # just consume the rest, it is lost to the higher layers
-                self.socket_read(socket_num, udp_data_bytes - length)
+                self.socket_read(socket_num, data_length - length)
         return bytes_on_socket, bytes_read
 
     def socket_write(
@@ -1340,3 +1335,21 @@ def _chip_socket_write(
                 self._write(dst_addr, 0x00, buffer[split_point:bytes_to_write])
             else:
                 self._write(dst_addr, 0x00, buffer[:bytes_to_write])
+
+    def _chip_parse_udp_header(self, socket_num) -> int:
+        """
+        Parse chip specific UDP header data for IPv4 packets.
+
+        Sets the source IPv4 address and port number and returns the UDP data length.
+
+        :return int: The UDP data length.
+        """
+        if self._chip_type in ("w5100s", "w5500"):
+            self.udp_from_ip[socket_num] = self._pbuff[:4]
+            self.udp_from_port[socket_num] = int.from_bytes(self._pbuff[4:6], "big")
+            return int.from_bytes(self._pbuff[6:], "big")
+        if self._chip_type == "w6100":
+            self.udp_from_ip[socket_num] = self._pbuff[3:7]
+            self.udp_from_port[socket_num] = int.from_bytes(self._pbuff[6:], "big")
+            return int.from_bytes(self._pbuff[:2], "big") & 0x07FF
+        raise ValueError("Unsupported chip type.")

adafruit_wiznet5k/adafruit_wiznet5k_dhcp.py

@@ -295,30 +295,18 @@ def _receive_dhcp_response(self, timeout: float) -> int:
         :returns int: The number of bytes stored in the global buffer.
         """
         debug_msg("Receiving a DHCP response.", self._debug)
-        # DHCP returns the query plus additional data. The query length is 236 bytes.
-        minimum_packet_length = 236
-        buffer = bytearray(b"")
-        bytes_read = 0
-        debug_msg("+ Beginning to receive…", self._debug)
-        while bytes_read < minimum_packet_length and time.monotonic() < timeout:
-            if self._eth.socket_available(self._wiz_sock, _SNMR_UDP):
-                x = self._eth.read_udp(self._wiz_sock, _BUFF_LENGTH - bytes_read)[1]
-                buffer.extend(x)
-                bytes_read = len(buffer)
-                debug_msg("+ Bytes read so far {}".format(bytes_read), self._debug)
-                debug_msg(x, self._debug)
-            if bytes_read == _BUFF_LENGTH:
-                break
-        debug_msg("Received {} bytes".format(bytes_read), self._debug)
-        if bytes_read < minimum_packet_length:
-            bytes_read = 0
-        else:
-            _BUFF[:bytes_read] = buffer
-            _BUFF[bytes_read:] = bytearray(_BUFF_LENGTH - bytes_read)
-        del buffer
-        gc.collect()
-        debug_msg(_BUFF[:bytes_read], self._debug)
-        return bytes_read
+        while time.monotonic() < timeout:
+            # DHCP returns the query plus additional data. The query length is 236 bytes.
+            if self._eth.socket_available(self._wiz_sock, _SNMR_UDP) > 236:
+                bytes_count, bytes_read = self._eth.read_udp(
+                    self._wiz_sock, _BUFF_LENGTH
+                )
+                _BUFF[:bytes_count] = bytes_read
+                debug_msg("Received {} bytes".format(bytes_count), self._debug)
+                del bytes_read
+                gc.collect()
+                return bytes_count
+        raise TimeoutError("No DHCP response received.")
 
     def _process_messaging_states(self, *, message_type: int):
         """