From 44bbada4ff277738e7bfbbe4dc6e29b6590e05f5 Mon Sep 17 00:00:00 2001
From: Nicolas Viennot <nicolas@viennot.biz>
Date: Sun, 12 May 2019 15:54:31 -0400
Subject: [PATCH] Use select() to avoid burning the CPU when reading from the
 UARTs

---
 btlejack/interface.py |  8 ++++++++
 btlejack/jobs.py      | 39 ++++++++++++++++++++-------------------
 btlejack/link.py      |  7 +++++++
 3 files changed, 35 insertions(+), 19 deletions(-)

diff --git a/btlejack/interface.py b/btlejack/interface.py
index bfc85bf..cc9ca1d 100755
--- a/btlejack/interface.py
+++ b/btlejack/interface.py
@@ -4,6 +4,8 @@
 Links the link layer to UI.
 """
 
+from btlejack.link import Link
+
 class AbstractInterface(object):
     """
     Sniffer/Hijacker abstract class.
@@ -112,6 +114,12 @@ def collab_channel_map(self):
         """
         self.mode = self.MODE_CCHM
 
+    @staticmethod
+    def wait_for_readeable_interfaces(interfaces, timeout=1.0):
+        links = [interface.link for interface in interfaces]
+        readables = Link.wait_for_readeable_links(links)
+        return [interface for interface in interfaces if interface.link in readables]
+
     def read_packet(self):
         """
         Read a packet from the underlying link.
diff --git a/btlejack/jobs.py b/btlejack/jobs.py
index 459de21..db07ed0 100755
--- a/btlejack/jobs.py
+++ b/btlejack/jobs.py
@@ -106,7 +106,7 @@ def sniff_connection(self, bd_address, channel=37):
         self.link.wait_packet(SniffConnReqResponse)
         super().sniff_connection()
 
-    def read_packet(self):
+    def read_packet_nowait(self):
         """
         Read packet if one is ready (asynchroous)
         """
@@ -116,6 +116,11 @@ def read_packet(self):
             packets.append(packet)
         return packets
 
+    def read_packet(self):
+        if AbstractInterface.wait_for_readeable_interfaces([self]):
+            return self.read_packet_nowait()
+        return []
+
 class MultiSnifferInterface(AbstractInterface):
     """
     Multi-sniffer interface class.
@@ -292,29 +297,25 @@ def read_packet(self):
         """
         Read packet(s) if one is ready (asynchroous)
         """
+        readable_links = AbstractInterface.wait_for_readeable_interfaces(self.interfaces)
+
         packets = []
         if (self.mode == self.MODE_RECOVER_CHM):
-            for link in self.interfaces:
-                if not link.is_idling():
-                    pkts = link.read_packet()
-                    if len(pkts) > 0:
-                        packets.extend(pkts)
-            return packets
+            for link in readable_links:
+                pkts = link.read_packet_nowait()
+                packets.extend(pkts)
         else:
             if self.active_link is None:
-                for link in self.interfaces:
-                    if not link.is_idling():
-                        pkts = link.read_packet()
-                        if len(pkts) > 0:
-                            self.active_link = link
-                            self.active_link.link.interface.timeout = 0.1
-                            packets.extend(pkts)
-                return packets
+                for link in readable_links:
+                    pkts = link.read_packet_nowait()
+                    if len(pkts) > 0:
+                        self.active_link = link
+                        self.active_link.link.interface.timeout = 0.1
+                        packets.extend(pkts)
             else:
-                pkts = self.active_link.read_packet()
-                if len(pkts) > 0:
-                    return pkts
-                return []
+                if self.active_link in readable_links:
+                    return self.active_link.read_packet_nowait()
+        return packets
 
     def send_packet(self, packet):
         """
diff --git a/btlejack/link.py b/btlejack/link.py
index 11a3a7b..015f851 100755
--- a/btlejack/link.py
+++ b/btlejack/link.py
@@ -6,6 +6,7 @@
 from serial.tools.list_ports import comports
 from struct import pack, unpack
 from threading import Lock
+from select import select
 from btlejack.packets import (Packet, PacketRegistry, ResetCommand,
     VersionCommand, ScanConnectionsCommand, RecoverCrcInitCommand,
     RecoverResponse, ResetResponse, VersionResponse, ScanConnectionsResponse,
@@ -57,6 +58,12 @@ def __init__(self, interface=None, baudrate=115200):
         self.interface = Serial(interface, baudrate, timeout=0)
         self.rx_buffer = bytes()
 
+    @staticmethod
+    def wait_for_readeable_links(links, timeout=1.0):
+        filenos = [link.interface.fileno() for link in links]
+        readables, _, _ = select(filenos, [], [], timeout)
+        return [link for link in links if link.interface.fileno() in readables]
+
     def set_timeout(self, timeout):
         """
         Set interface timeout.