improve udp driver based on other driver receiver thread implementati…

cflib/crtp/udpdriver.py

@@ -22,81 +22,278 @@
 #  GNU General Public License for more details.
 # You should have received a copy of the GNU General Public License
 # along with this program. If not, see <https://www.gnu.org/licenses/>.
-""" CRTP UDP Driver. Work either with the UDP server or with an UDP device
-See udpserver.py for the protocol"""
+"""
+Crazyflie UDP driver.
+
+This driver communicates with a Crazyflie (or simulator) over UDP using the CRTP
+protocol. It enables connecting to software-in-the-loop (SITL) simulations. 
+Scanning feature assumes a crazyflie server is running on port 19850-19859
+that will respond to a null CRTP packet with a valid CRTP packet.
+
+Wire Protocol
+-------------
+The UDP driver uses a simple wire protocol where each UDP datagram contains exactly
+one CRTP packet. The packet format is:
+
+    Byte 0:    CRTP header (port, channel, and reserved bits)
+    Bytes 1-N: CRTP payload data (0 to 30 bytes) set by CRTP_MAX_DATA_SIZE in firmware
+
+Total packet size: 1 to 31 bytes per UDP datagram
+
+The payload follows standard CRTP format as defined in the Crazyflie protocol
+specification. No additional framing, checksums, or encapsulation is added by
+the UDP driver.
+
+Scan Behavior
+-------------
+The scan_interface() method discovers available Crazyflie devices by probing
+UDP ports in sequence:
+
+Port Range: 19850 to 19859 (10 ports total, BASE_PORT + 0 through 9)
+Scan Address: Configurable via SCAN_ADDRESS environment variable (default: 127.0.0.1)
+Probe Packet: Single 0xFF byte (null CRTP packet)
+Timeout: 0.1 seconds per port
+
+Scan Procedure:
+    1. For each port in range:
+       - Send 0xFF probe packet to SCAN_ADDRESS:port
+       - Wait up to 0.1 seconds for response
+       - If any response received, device is present
+       - Add URI "udp://SCAN_ADDRESS:port" to results
+
+Environment Variables
+---------------------
+SCAN_ADDRESS: IP address to scan for Crazyflie devices
+              Default: 127.0.0.1
+              Example: export SCAN_ADDRESS=192.168.1.100
+
+              This is useful when the Crazyflie server and client run on different
+              hosts. The client will scan the specified address instead of localhost.
+
+Connection URI Format
+--------------------
+udp://<host>:<port>
+
+Examples:
+    udp://127.0.0.1:19850  - Local simulator on port 19850
+    udp://192.168.1.5:19850 - Remote device at 192.168.1.5
+"""
+
+# changelog:
+# - Complete rewrite to align with other CRTP driver implementations
+# - Added dedicated _UdpReceiveThread class for asynchronous packet reception
+# - Implemented functional scan_interface() that probes UDP ports 19850-19859
+# - Fixed send_packet() with null checks, proper error callbacks, and removed checksum
+# - Added proper socket cleanup in close() method
+# - Changed variable naming to align with other CRTP drivers and added docstrings
+# - Added environment variable SCAN_ADDRESS for scan_interface() to specify target IP address
+#   This is useful for server and clients running on different hosts
+
+import logging
+import os
 import queue
 import re
 import socket
 import struct
+import threading
 from urllib.parse import urlparse
 
-from .crtpdriver import CRTPDriver
 from .crtpstack import CRTPPacket
 from .exceptions import WrongUriType
+from cflib.crtp.crtpdriver import CRTPDriver
 
 __author__ = 'Bitcraze AB'
 __all__ = ['UdpDriver']
 
+logger = logging.getLogger(__name__)
+
+_BASE_PORT = 19850
+_NR_OF_PORTS_TO_SCAN = 10
+_SCAN_TIMEOUT = 0.1
+
 
 class UdpDriver(CRTPDriver):
+    """ Crazyflie UDP link driver """
 
     def __init__(self):
-        None
+        """ Create the link driver """
+        CRTPDriver.__init__(self)
+        self.socket = None
+        self.addr = None
+        self.uri = ''
+        self.link_error_callback = None
+        self.in_queue = None
+        self._thread = None
+        self.needs_resending = False
 
-    def connect(self, uri, linkQualityCallback, linkErrorCallback):
+    def connect(self, uri, radio_link_statistics_callback, link_error_callback):
+        """
+        Connect the link driver to a specified URI of the format:
+        udp://<host>:<port>
+
+        The callback for radio link statistics is not used by the UDP driver.
+        The callback from link_error_callback will be called when an error
+        occurs with an error message.
+        """
         if not re.search('^udp://', uri):
-            raise WrongUriType('Not an UDP URI')
+            raise WrongUriType('Not a UDP URI')
+
+        if self.socket is not None:
+            raise Exception('Link already open!')
+
+        self.uri = uri
+        self.link_error_callback = link_error_callback
 
         parse = urlparse(uri)
 
-        self.queue = queue.Queue()
+        # Prepare the inter-thread communication queue
+        self.in_queue = queue.Queue()
+
         self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
         self.addr = (parse.hostname, parse.port)
         self.socket.connect(self.addr)
 
-        self.socket.sendto('\xFF\x01\x01\x01'.encode(), self.addr)
+        # Launch the comm thread
+        self._thread = _UdpReceiveThread(self.socket, self.in_queue,
+                                         link_error_callback)
+        self._thread.start()
 
-    def receive_packet(self, time=0):
-        data, addr = self.socket.recvfrom(1024)
+    def receive_packet(self, wait=0):
+        """
+        Receive a packet though the link. This call is blocking but will
+        timeout and return None if a timeout is supplied.
+        """
+        if wait == 0:
+            try:
+                return self.in_queue.get(False)
+            except queue.Empty:
+                return None
+        elif wait < 0:
+            try:
+                return self.in_queue.get(True)
+            except queue.Empty:
+                return None
+        else:
+            try:
+                return self.in_queue.get(True, wait)
+            except queue.Empty:
+                return None
 
-        if data:
-            data = struct.unpack('B' * (len(data) - 1), data[0:len(data) - 1])
-            pk = CRTPPacket()
-            pk.port = data[0]
-            pk.data = data[1:]
-            return pk
+    def send_packet(self, pk):
+        """ Send the packet pk though the link """
+        if self.socket is None:
+            return
 
         try:
-            if time == 0:
-                return self.rxqueue.get(False)
-            elif time < 0:
-                while True:
-                    return self.rxqueue.get(True, 10)
-            else:
-                return self.rxqueue.get(True, time)
-        except queue.Empty:
-            return None
+            raw = (pk.header,) + struct.unpack('B' * len(pk.data), pk.data)
+            data = struct.pack('B' * len(raw), *raw)
+            self.socket.send(data)
+        except Exception as e:
+            if self.link_error_callback:
+                self.link_error_callback(
+                    'UdpDriver: Could not send packet to Crazyflie\n'
+                    'Exception: %s' % e)
 
-    def send_packet(self, pk):
-        raw = (pk.port,) + struct.unpack('B' * len(pk.data), pk.data)
+    def pause(self):
+        self._thread.stop()
+        self._thread = None
 
-        cksum = 0
-        for i in raw:
-            cksum += i
+    def restart(self):
+        if self._thread:
+            return
 
-        cksum %= 256
+        self._thread = _UdpReceiveThread(self.socket, self.in_queue,
+                                         self.link_error_callback)
+        self._thread.start()
 
-        data = ''.join(chr(v) for v in (raw + (cksum,)))
+    def close(self):
+        """ Close the link. """
+        # Stop the comm thread
+        if self._thread:
+            self._thread.stop()
 
-        # print tuple(data)
-        self.socket.sendto(data.encode(), self.addr)
+        # Close the UDP socket
+        try:
+            if self.socket:
+                self.socket.close()
+        except Exception as e:
+            logger.info('Could not close {}'.format(e))
+        self.socket = None
 
-    def close(self):
-        # Remove this from the server clients list
-        self.socket.sendto('\xFF\x01\x02\x02'.encode(), self.addr)
+        # Clear callbacks
+        self.link_error_callback = None
+
+    def get_status(self):
+        return 'No information available'
 
     def get_name(self):
         return 'udp'
 
-    def scan_interface(self, address):
-        return []
+    def scan_interface(self, address=None):
+        """ Scan interface for Crazyflies """
+        found = []
+        scan_address = os.getenv('SCAN_ADDRESS', '127.0.0.1')
+
+        for i in range(_NR_OF_PORTS_TO_SCAN):
+            port = _BASE_PORT + i
+            try:
+                s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+                s.settimeout(_SCAN_TIMEOUT)
+                s.connect((scan_address, port))
+                s.send(b'\xFF')  # Null CRTP packet as probe
+                s.recv(1024)
+                # Got a response, Crazyflie is available
+                s.close()
+                found.append(['udp://{}:{}'.format(scan_address, port), ''])
+            except socket.timeout:
+                s.close()
+            except Exception:
+                pass
+
+        return found
+
+
+# Receive thread
+class _UdpReceiveThread(threading.Thread):
+    """
+    UDP link receiver thread used to read data from the
+    UDP socket. """
+
+    def __init__(self, sock, inQueue, link_error_callback):
+        """ Create the object """
+        threading.Thread.__init__(self, name='UdpReceiveThread')
+        self._socket = sock
+        self._in_queue = inQueue
+        self._sp = False
+        self._link_error_callback = link_error_callback
+        self.daemon = True
+
+    def stop(self):
+        """ Stop the thread """
+        self._sp = True
+        try:
+            self.join()
+        except Exception:
+            pass
+
+    def run(self):
+        """ Run the receiver thread """
+        self._socket.settimeout(1.0)
+
+        while True:
+            if self._sp:
+                break
+            try:
+                packet = self._socket.recv(1024)
+                data = struct.unpack('B' * len(packet), packet)
+                if len(data) > 0:
+                    pk = CRTPPacket(header=data[0], data=data[1:])
+                    self._in_queue.put(pk)
+            except socket.timeout:
+                pass
+            except Exception as e:
+                import traceback
+
+                self._link_error_callback(
+                    'Error communicating with the Crazyflie\n'
+                    'Exception:%s\n\n%s' % (e, traceback.format_exc()))