Fix reconnects and multiple packets

Author: tannewt

adafruit_ble_file_transfer.py

@@ -118,6 +118,7 @@ def __init__(self, service):
         self._service = service
 
     def _write(self, buffer):
+        # print("write", binascii.hexlify(buffer))
         sent = 0
         while sent < len(buffer):
             remaining = len(buffer) - sent
@@ -127,21 +128,22 @@ def _write(self, buffer):
 
     def _readinto(self, buffer):
         read = 0
+        long_buffer = bytearray(512)
         # Read back how much we can write
         while read == 0:
             try:
                 read = self._service.raw.readinto(buffer)
-            except ValueError as error:
-                print(error)
-                long_buffer = bytearray(512)
+            except ValueError:
                 read = self._service.raw.readinto(long_buffer)
-                print("long packet", long_buffer[:read])
+        buffer[:read] = long_buffer[:read]
         return read
 
     def read(self, path, *, offset=0):
         """Returns the contents of the file at the given path starting at the given offset"""
+        # pylint: disable=too-many-locals
         path = path.encode("utf-8")
         chunk_size = CHUNK_SIZE
+        start_offset = offset
         encoded = (
             struct.pack(
                 "<BxHII", FileTransferService.READ, len(path), offset, chunk_size
@@ -150,30 +152,45 @@ def read(self, path, *, offset=0):
         )
         self._write(encoded)
         b = bytearray(struct.calcsize("<BBxxIII") + chunk_size)
-        current_offset = offset
+        current_offset = start_offset
+        chunk_done = True
         content_length = None
-        buf = None
+        chunk_end = 0
+        # pylint: disable=unsupported-assignment-operation
+        buf = []
         data_header_size = struct.calcsize("<BBxxIII")
         while content_length is None or current_offset < content_length:
             read = self._readinto(b)
-            (
-                cmd,
-                status,
-                current_offset,
-                content_length,
-                _,
-            ) = struct.unpack_from("<BBxxIII", b)
-            if cmd != FileTransferService.READ_DATA:
-                raise ProtocolError("Incorrect reply")
-            if status != FileTransferService.OK:
-                raise ValueError("Missing file")
-            if buf is None:
-                buf = bytearray(content_length - offset)
-            out_offset = current_offset - offset
-            buf[out_offset : out_offset + (read - data_header_size)] = b[
-                data_header_size:read
-            ]
-            current_offset += read - data_header_size
+            if chunk_done:
+                (
+                    cmd,
+                    status,
+                    content_offset,
+                    content_length,
+                    chunk_length,
+                ) = struct.unpack_from("<BBxxIII", b)
+                chunk_end = content_offset + chunk_length
+                if cmd != FileTransferService.READ_DATA:
+                    print("error:", b)
+                    raise ProtocolError("Incorrect reply")
+                if status != FileTransferService.OK:
+                    raise ValueError("Missing file")
+                if not buf:
+                    buf = bytearray(content_length - start_offset)
+                out_offset = current_offset - start_offset
+                buf[out_offset : out_offset + (read - data_header_size)] = b[
+                    data_header_size:read
+                ]
+                current_offset += read - data_header_size
+            else:
+                out_offset = current_offset - start_offset
+                buf[out_offset : out_offset + read] = b[:read]
+                current_offset += read
+
+            chunk_done = current_offset == chunk_end
+            if not chunk_done:
+                continue
+
             chunk_size = min(CHUNK_SIZE, content_length - current_offset)
             if chunk_size == 0:
                 break
@@ -206,6 +223,7 @@ def write(self, path, contents, *, offset=0):
             self._readinto(b)
             cmd, status, current_offset, free_space = struct.unpack("<BBxxII", b)
             if status != FileTransferService.OK:
+                print("write error", status)
                 raise RuntimeError()
             if (
                 cmd != FileTransferService.WRITE_PACING
@@ -256,6 +274,7 @@ def mkdir(self, path):
 
     def listdir(self, path):
         """Returns a list of tuples, one tuple for each file or directory in the given path"""
+        # pylint: disable=too-many-locals
         paths = []
         path = path.encode("utf-8")
         encoded = struct.pack("<BxH", FileTransferService.LISTDIR, len(path)) + path
@@ -264,31 +283,42 @@ def listdir(self, path):
         i = 0
         total = 10  # starting value that will be replaced by the first response
         header_size = struct.calcsize("<BBHIIII")
+        path_length = 0
+        encoded_path = b""
         while i < total:
             read = self._readinto(b)
             offset = 0
+            file_size = 0
+            flags = 0
             while offset < read:
-                (
-                    cmd,
-                    status,
-                    path_length,
-                    i,
-                    total,
-                    flags,
-                    file_size,
-                ) = struct.unpack_from("<BBHIIII", b, offset=offset)
-                if cmd != FileTransferService.LISTDIR_ENTRY:
-                    raise ProtocolError()
-                if i >= total:
-                    break
-                path = str(
-                    b[offset + header_size : offset + header_size + path_length],
-                    "utf-8",
-                )
-                paths.append((path, file_size, flags))
-                offset += header_size + path_length
-                if status != FileTransferService.OK:
-                    break
+                if len(encoded_path) == path_length:
+                    if path_length > 0:
+                        path = str(
+                            encoded_path,
+                            "utf-8",
+                        )
+                        paths.append((path, file_size, flags))
+                    (
+                        cmd,
+                        status,
+                        path_length,
+                        i,
+                        total,
+                        flags,
+                        file_size,
+                    ) = struct.unpack_from("<BBHIIII", b, offset=offset)
+                    offset += header_size
+                    encoded_path = b""
+                    if cmd != FileTransferService.LISTDIR_ENTRY:
+                        raise ProtocolError()
+                    if status != FileTransferService.OK:
+                        break
+                    if i >= total:
+                        break
+
+                path_read = min(path_length - len(encoded_path), read - offset)
+                encoded_path += b[offset : offset + path_read]
+                offset += path_read
         return paths
 
     def delete(self, path):

examples/ble_file_transfer_simpletest.py

@@ -5,55 +5,95 @@
 Used with ble_uart_echo_test.py. Transmits "echo" to the UARTService and receives it back.
 """
 
+import binascii
 import random
 import time
 
 from adafruit_ble import BLERadio
-from adafruit_ble.advertising.standard import ProvideServicesAdvertisement
+from adafruit_ble.advertising.standard import (
+    ProvideServicesAdvertisement,
+    Advertisement,
+)
 import adafruit_ble_file_transfer
 
 
 def _write(client, filename, contents, *, offset=0):
+    # pylint: disable=redefined-outer-name
     start = time.monotonic()
-    client.write(filename, contents, offset=offset)
-    duration = time.monotonic() - start
+    try:
+        client.write(filename, contents, offset=offset)
+        duration = time.monotonic() - start
+        client = wait_for_reconnect()
+    except RuntimeError:
+        print("write failed. is usb connected?")
+        return client
     print("wrote", filename, "at rate", len(contents) / duration, "B/s")
+    return client
 
 
 def _read(client, filename, *, offset=0):
+    # pylint: disable=redefined-outer-name
     start = time.monotonic()
-    contents = client.read(filename, offset=offset)
-    duration = time.monotonic() - start
+    try:
+        contents = client.read(filename, offset=offset)
+        duration = time.monotonic() - start
+    except ValueError:
+        print("missing file:", filename)
+        return b""
     print("read", filename, "at rate", len(contents) / duration, "B/s")
     return contents
 
 
 ble = BLERadio()
+
+peer_address = None
+
+
+def wait_for_reconnect():
+    print("waiting for disconnect")
+    while ble.connected:
+        pass
+    print("reconnecting to", peer_address)
+    new_connection = ble.connect(peer_address)
+    print("reconnected")
+    if not new_connection.paired:
+        print("pairing")
+        new_connection.pair()
+    new_service = new_connection[adafruit_ble_file_transfer.FileTransferService]
+    new_client = adafruit_ble_file_transfer.FileTransferClient(new_service)
+    print("sleeping")
+    time.sleep(2)
+    return new_client
+
+
 # ble._adapter.erase_bonding()
 # print("erased")
 while True:
     try:
         while ble.connected:
             for connection in ble.connections:
-                print(
-                    "services", connection._bleio_connection.discover_remote_services()
-                )
+                # pylint: disable=redefined-outer-name
                 if adafruit_ble_file_transfer.FileTransferService not in connection:
                     continue
                 if not connection.paired:
                     print("pairing")
                     connection.pair()
+                print("paired")
                 service = connection[adafruit_ble_file_transfer.FileTransferService]
                 client = adafruit_ble_file_transfer.FileTransferClient(service)
-                _write(client, "/hello.txt", "Hello world".encode("utf-8"))
+                client = _write(client, "/hello.txt", "Hello world".encode("utf-8"))
+                time.sleep(1)
                 c = _read(client, "/hello.txt")
                 print(len(c), c)
-                client.mkdir("/world/")
+                try:
+                    client.mkdir("/world/")
+                except ValueError:
+                    print("path exists or isn't valid")
                 print(client.listdir("/world/"))
-                _write(client, "/world/hi.txt", "Hi world".encode("utf-8"))
+                client = _write(client, "/world/hi.txt", "Hi world".encode("utf-8"))
 
                 hello_world = "Hello world".encode("utf-8")
-                _write(client, "/world/hello.txt", hello_world)
+                client = _write(client, "/world/hello.txt", hello_world)
                 c = _read(client, "/world/hello.txt")
                 print(c)
 
@@ -62,39 +102,60 @@ def _read(client, filename, *, offset=0):
                 c = _read(client, "/world/hello.txt", offset=hello)
                 print(c)
 
-                _write(client, "/world/hello.txt", "offsets!", offset=hello)
+                client = _write(
+                    client, "/world/hello.txt", "offsets!".encode("utf-8"), offset=hello
+                )
                 c = _read(client, "/world/hello.txt", offset=0)
                 print(c)
 
                 # Test deleting
                 print(client.listdir("/world/"))
-                client.delete("/world/hello.txt")
+                try:
+                    client.delete("/world/hello.txt")
+                except ValueError:
+                    print("exception correctly raised")
+
                 try:
                     client.delete("/world/")  # should raise an exception
                 except ValueError:
                     print("exception correctly raised")
                 print(client.listdir("/world/"))
-                client.delete("/world/hi.txt")
-                client.delete("/world/")
+                try:
+                    client.delete("/world/hi.txt")
+                except ValueError:
+                    print("missing /world/hi.txt")
+                try:
+                    client.delete("/world/")
+                except ValueError:
+                    print("cannot delete /world/")
                 print(client.listdir("/"))
 
                 large_1k = bytearray(1024)
-                for i in range(len(large_1k)):
+                for i, _ in enumerate(large_1k):
                     large_1k[i] = random.randint(0, 255)
-                _write(client, "/random.txt", large_1k)
-                time.sleep(0.1)
+                client = _write(client, "/random.txt", large_1k)
                 contents = _read(client, "/random.txt")
                 if large_1k != contents:
-                    print("large contents don't match!")
+                    print(binascii.hexlify(large_1k))
+                    print(binascii.hexlify(contents))
+                    raise RuntimeError("large contents don't match!")
             time.sleep(20)
-    except ConnectionError:
+    except ConnectionError as e:
         pass
 
     print("disconnected, scanning")
-    for advertisement in ble.start_scan(ProvideServicesAdvertisement, timeout=1):
-        if adafruit_ble_file_transfer.FileTransferService not in advertisement.services:
+    for advertisement in ble.start_scan(
+        ProvideServicesAdvertisement, Advertisement, timeout=1
+    ):
+        # print(advertisement.address, advertisement.address.type)
+        if (
+            not hasattr(advertisement, "services")
+            or adafruit_ble_file_transfer.FileTransferService
+            not in advertisement.services
+        ):
             continue
         ble.connect(advertisement)
-        print("connected")
+        peer_address = advertisement.address
+        print("connected to", advertisement.address)
         break
     ble.stop_scan()