backported fixes for uds tunnel (timeout, stats accounting and stop mechanics)

tunnel-server/src/uds_tunnel/stats.py

@@ -105,6 +105,10 @@ class StatsManager:
         self.sent += size
         self.update()
 
+    def decrement_connections(self):
+        self.ns.current -= 1
+        self.ns.total -= 1
+
     @property
     def as_sent_counter(self) -> 'StatsSingleCounter':
         return StatsSingleCounter(self, False)

tunnel-server/src/uds_tunnel/tunnel.py

@@ -100,10 +100,11 @@ class TunnelProtocol(asyncio.Protocol):
         # Open Command has the ticket behind it
 
         if len(self.cmd) < consts.TICKET_LENGTH + consts.COMMAND_LENGTH:
-            # Reactivate timeout, will be deactivated on do_command
-            self.set_timeout(self.owner.cfg.command_timeout)
             return  # Wait for more data to complete OPEN command
 
+        # Clean timeout now, we have received all data
+        self.clean_timeout()
+
         # Ticket received, now process it with UDS
         ticket = self.cmd[consts.COMMAND_LENGTH :]
 
@@ -168,6 +169,9 @@ class TunnelProtocol(asyncio.Protocol):
         if len(self.cmd) < consts.PASSWORD_LENGTH + consts.COMMAND_LENGTH:
             return
 
+        self.stats_manager.decrement_connections()  # This connection does not count, it's just "stats"
+        # Clean timeout now, we have received all data
+        self.clean_timeout()
         logger.info('COMMAND: %s', self.cmd[: consts.COMMAND_LENGTH].decode())
 
         # Check valid source ip
@@ -225,8 +229,7 @@ class TunnelProtocol(asyncio.Protocol):
         if self.cmd == b'':
             logger.info('CONNECT FROM %s', self.pretty_source())
 
-        # Ensure we don't get a timeout
-        self.clean_timeout()
+        # We have at most self.owner.cfg.command_timeout seconds to receive the command and the ticket if needed
         self.cmd += data
 
         if len(self.cmd) >= consts.COMMAND_LENGTH:
@@ -254,8 +257,6 @@ class TunnelProtocol(asyncio.Protocol):
                 self.transport.write(consts.RESPONSE_ERROR_COMMAND)
                 self.close_connection()
                 return
-        else:
-            self.set_timeout(self.owner.cfg.command_timeout)
 
         # if not enough data to process command, wait for more
 

tunnel-server/src/udstunnel.py

@@ -189,7 +189,11 @@ async def tunnel_proc_async(
     logger.debug('Out of loop, stopping tasks: %s, running: %s', tasks, do_stop.is_set())
 
     # If any task is still running, cancel it
-    asyncio.gather(*tasks, return_exceptions=True).cancel()
+    for task in tasks:
+        try:
+            task.cancel()
+        except asyncio.CancelledError:
+            pass  # Ignore, we are stopping
     
     # for task in tasks:
     #    task.cancel()