From a30b9f1a62f683e46be4b9b5cea4337e212b53c9 Mon Sep 17 00:00:00 2001
From: Sunpy <sunpy@osufx.com>
Date: Mon, 30 Mar 2020 00:59:53 +0200
Subject: [PATCH] Rewrote timeout system

---
 sampy/client/__init__.py | 22 ++++++++++++++++++++--
 sampy/client/base.py     | 21 ---------------------
 sampy/client/player.py   |  2 +-
 sampy/client/query.py    |  2 +-
 sampy/server.py          |  7 +++++++
 5 files changed, 29 insertions(+), 25 deletions(-)

diff --git a/sampy/client/__init__.py b/sampy/client/__init__.py
index e058fe3..88a7dae 100644
--- a/sampy/client/__init__.py
+++ b/sampy/client/__init__.py
@@ -1,5 +1,7 @@
 import struct
 import socket
+import asyncio
+from time import time
 
 from . import base
 from . import query
@@ -9,6 +11,8 @@ STATE_UNKNOWN   = (0, base.BaseClient)
 STATE_QUERY     = (1, query.QueryClient)
 STATE_PLAYER    = (2, player.PlayerClient)
 
+TIMEOUT = 10 # assume connection is closed after 10 seconds of inactivity (change this to a higher value so you dont timeout while debugging might be a good idea)
+
 class Client:
     def __init__(self, server: "__ServerInstance__", ip: str, port: int):
         self.server = server
@@ -16,19 +20,33 @@ class Client:
         self.port = port
 
         self.set_state(STATE_UNKNOWN)
+
+        self.last_active = time()
+        self.keep_alive_task = asyncio.create_task( self.keep_alive() )
+        self.connected = True # keep_alive will set this to False if connection has not been interacted with for a while (allowing server loop to remove their reference)
     
     def set_state(self, state: tuple):
+        #self.keep_alive_task.cancel()
         self.state = state
         self.client = self.state[1](self.server, self.ip, self.port)
 
     async def on_packet(self, packet: bytes):
+        self.last_active = time()
+
         if self.state == STATE_UNKNOWN:
             # We are currently unaware if this is a player client or query client, but we got a packet that will be our check to know
             if packet.startswith(b"SAMP"):
-                await self.client.on_state_change()
                 self.set_state(STATE_QUERY)
             else:
-                await self.client.on_state_change()
                 self.set_state(STATE_PLAYER)
 
         await self.client.on_packet(packet)
+
+    async def keep_alive(self): # Maybe bad name for this method as it rather checks if connection is dropped
+        while True:
+            timestamp = time()
+            if self.last_active + TIMEOUT - timestamp < 0:
+                self.connected = False
+                return
+            
+            await asyncio.sleep(self.last_active + TIMEOUT - timestamp)
\ No newline at end of file
diff --git a/sampy/client/base.py b/sampy/client/base.py
index 02524d7..59bda42 100644
--- a/sampy/client/base.py
+++ b/sampy/client/base.py
@@ -1,13 +1,9 @@
 import socket
 import struct
-import asyncio
-from time import time
 
 import logging
 logger = logging.getLogger(__name__)
 
-TIMEOUT = 30 # assume connection is closed after 30 seconds if inactivity
-
 class BaseClient:
     def __init__(self, server: "__ServerInstance__", ip: str, port: int):
         self.server = server
@@ -15,27 +11,10 @@ class BaseClient:
         self.port = port
 
         self.ip_uint, = struct.unpack(b"<I", bytes(int(x) for x in self.ip.split(".")))
-
-        self.last_active = time()
-        self.keep_alive_task = asyncio.create_task( self.keep_alive() )
-        self.connected = True # keep_alive will set this to False if connection has not been interacted with for a while (allowing server loop to remove their reference)
     
     async def on_packet(self, packet: bytes):
         logger.debug("on_packet(%s)" % packet)
-        self.last_active = time()
     
     async def send(self, packet: bytes):
         sock: socket.socket = self.server.socket
         sock.sendto(packet, (self.ip, self.port))
-
-    async def on_state_change(self): # Stop the keep alive task as the whole class is being replaced, TODO: Find out if this could be automated with a magic method
-        self.keep_alive_task.cancel()
-
-    async def keep_alive(self):
-        while True:
-            timestamp = time()
-            if self.last_active + TIMEOUT - timestamp < 0:
-                self.connected = False
-                return
-            
-            await asyncio.sleep(self.last_active + TIMEOUT - timestamp)
diff --git a/sampy/client/player.py b/sampy/client/player.py
index d98110e..1b26df8 100644
--- a/sampy/client/player.py
+++ b/sampy/client/player.py
@@ -11,4 +11,4 @@ class PlayerClient(BaseClient):
         logger.debug("Client resolved to PlayerClient")
     
     async def on_packet(self, packet: bytes):
-        logger.debug("on_packet(%s)" % packet)
\ No newline at end of file
+        logger.debug("on_packet(%s)" % packet)
diff --git a/sampy/client/query.py b/sampy/client/query.py
index 61ef514..9993b86 100644
--- a/sampy/client/query.py
+++ b/sampy/client/query.py
@@ -22,7 +22,7 @@ class QueryClient(BaseClient):
         }
     
     async def on_packet(self, packet: bytes):
-        await super().on_packet(packet)
+        logger.debug("on_packet(%s)" % packet)
 
         if len(packet) <= 10: # Invalid
             return
diff --git a/sampy/server.py b/sampy/server.py
index f5a3f18..1132eda 100644
--- a/sampy/server.py
+++ b/sampy/server.py
@@ -50,4 +50,11 @@ class Server:
                     self.clients[addr] = Client(self, ip, port)
                 await self.clients[addr].on_packet(data)
 
+            disconnected = [c for c in self.clients.values() if c.connected == False]
+            for c in disconnected: # Remove dead connections
+                addr = (c.ip, c.port)
+                if addr in self.clients:
+                    del self.clients[addr]
+                    logger.debug("free(%s)" % c)
+
             await asyncio.sleep(0)