Merge pull request #62 from tom24/udpfixes

Fixed issues with UDP datagram receives

adafruit_wiznet5k/adafruit_wiznet5k.py

@@ -122,14 +122,12 @@
 W5100_MAX_SOCK_NUM = const(0x04)
 SOCKET_INVALID = const(255)
 
-# UDP socket struct.
-UDP_SOCK = {"bytes_remaining": 0, "remote_ip": 0, "remote_port": 0}
 
 # Source ports in use
 SRC_PORTS = [0] * W5200_W5500_MAX_SOCK_NUM
 
 
-class WIZNET5K:  # pylint: disable=too-many-public-methods
+class WIZNET5K:  # pylint: disable=too-many-public-methods, too-many-instance-attributes
     """Interface for WIZNET5K module.
 
     :param ~busio.SPI spi_bus: The SPI bus the Wiznet module is connected to.
@@ -181,7 +179,11 @@ def __init__(
         # Set MAC address
         self.mac_address = mac
         self.src_port = 0
-        self._dns = 0
+        self._dns = (0, 0, 0, 0)
+        # udp related
+        self.udp_datasize = [0] * self.max_sockets
+        self.udp_from_ip = [b"\x00\x00\x00\x00"] * self.max_sockets
+        self.udp_from_port = [0] * self.max_sockets
 
         # First, wait link status is on
         # to avoid the code during DHCP, socket listen, connect ...
@@ -495,11 +497,6 @@ def write(self, addr, callback, data):
                     bus_device.write(bytes([data[i]]))  # pylint: disable=no-member
 
     # Socket-Register API
-    def udp_remaining(self):
-        """Returns amount of bytes remaining in a udp socket."""
-        if self._debug:
-            print("* UDP Bytes Remaining: ", UDP_SOCK["bytes_remaining"])
-        return UDP_SOCK["bytes_remaining"]
 
     def socket_available(self, socket_num, sock_type=SNMR_TCP):
         """Returns the amount of bytes to be read from the socket.
@@ -520,16 +517,16 @@ def socket_available(self, socket_num, sock_type=SNMR_TCP):
         if sock_type == SNMR_TCP:
             return res
         if res > 0:
-            if UDP_SOCK["bytes_remaining"]:
-                return UDP_SOCK["bytes_remaining"]
+            if self.udp_datasize[socket_num]:
+                return self.udp_datasize[socket_num]
             # parse the udp rx packet
             # read the first 8 header bytes
             ret, self._pbuff = self.socket_read(socket_num, 8)
             if ret > 0:
-                UDP_SOCK["remote_ip"] = self._pbuff[:4]
-                UDP_SOCK["remote_port"] = (self._pbuff[4] << 8) + self._pbuff[5]
-                UDP_SOCK["bytes_remaining"] = (self._pbuff[6] << 8) + self._pbuff[7]
-                ret = UDP_SOCK["bytes_remaining"]
+                self.udp_from_ip[socket_num] = self._pbuff[:4]
+                self.udp_from_port[socket_num] = (self._pbuff[4] << 8) + self._pbuff[5]
+                self.udp_datasize[socket_num] = (self._pbuff[6] << 8) + self._pbuff[7]
+                ret = self.udp_datasize[socket_num]
                 return ret
         return 0
 
@@ -573,7 +570,7 @@ def socket_connect(self, socket_num, dest, port, conn_mode=SNMR_TCP):
                 if self.socket_status(socket_num)[0] == SNSR_SOCK_CLOSED:
                     raise RuntimeError("Failed to establish connection.")
         elif conn_mode == SNMR_UDP:
-            UDP_SOCK["bytes_remaining"] = 0
+            self.udp_datasize[socket_num] = 0
         return 1
 
     def _send_socket_cmd(self, socket, cmd):
@@ -765,14 +762,15 @@ def socket_read(self, socket_num, length):
         return ret, resp
 
     def read_udp(self, socket_num, length):
-        """Read UDP socket's remaining bytes."""
-        if UDP_SOCK["bytes_remaining"] > 0:
-            if UDP_SOCK["bytes_remaining"] <= length:
-                ret, resp = self.socket_read(socket_num, UDP_SOCK["bytes_remaining"])
+        """Read UDP socket's current message bytes."""
+        if self.udp_datasize[socket_num] > 0:
+            if self.udp_datasize[socket_num] <= length:
+                ret, resp = self.socket_read(socket_num, self.udp_datasize[socket_num])
             else:
                 ret, resp = self.socket_read(socket_num, length)
-            if ret > 0:
-                UDP_SOCK["bytes_remaining"] -= ret
+                # just consume the rest, it is lost to the higher layers
+                self.socket_read(socket_num, self.udp_datasize[socket_num] - length)
+            self.udp_datasize[socket_num] = 0
             return ret, resp
         return -1
 

adafruit_wiznet5k/adafruit_wiznet5k_socket.py

@@ -294,6 +294,7 @@ def recv(self, bufsize=0, flags=0):  # pylint: disable=too-many-branches
                         ]
                     elif self._sock_type == SOCK_DGRAM:
                         self._buffer += _the_interface.read_udp(self.socknum, avail)[1]
+                        break
                 else:
                     break
             gc.collect()
@@ -315,6 +316,7 @@ def recv(self, bufsize=0, flags=0):  # pylint: disable=too-many-branches
                     )[1]
                 elif self._sock_type == SOCK_DGRAM:
                     recv = _the_interface.read_udp(self.socknum, min(to_read, avail))[1]
+                    to_read = len(recv)  # only get this dgram
                 recv = bytes(recv)
                 received.append(recv)
                 to_read -= len(recv)
@@ -333,6 +335,26 @@ def recv(self, bufsize=0, flags=0):  # pylint: disable=too-many-branches
         gc.collect()
         return ret
 
+    def embed_recv(self, bufsize=0, flags=0):  # pylint: disable=too-many-branches
+        """Reads some bytes from the connected remote address and then return recv().
+        :param int bufsize: Maximum number of bytes to receive.
+        :param int flags: ignored, present for compatibility.
+        """
+        # print("Socket read", bufsize)
+        ret = None
+        avail = self.available()
+        if avail:
+            if self._sock_type == SOCK_STREAM:
+                self._buffer += _the_interface.socket_read(self.socknum, avail)[1]
+            elif self._sock_type == SOCK_DGRAM:
+                self._buffer += _the_interface.read_udp(self.socknum, avail)[1]
+        gc.collect()
+        ret = self._buffer
+        # print("RET ptr:", id(ret), id(self._buffer))
+        self._buffer = b""
+        gc.collect()
+        return ret
+
     def recvfrom(self, bufsize=0, flags=0):
         """Reads some bytes from the connected remote address.
 
@@ -343,8 +365,8 @@ def recvfrom(self, bufsize=0, flags=0):
         return (
             self.recv(bufsize),
             (
-                _the_interface.remote_ip(self.socknum),
-                _the_interface.remote_port(self.socknum),
+                _the_interface.pretty_ip(_the_interface.udp_from_ip[self.socknum]),
+                _the_interface.udp_from_port[self.socknum],
             ),
         )