Fix tcp proxy for allow keep alive connections (also allowed proxy database connections)

Author: AMEST

src/__pycache__/docker_service.cpython-39.pyc

@@ -1,5 +1,5 @@
 from socketserver import BaseRequestHandler, TCPServer, ThreadingMixIn
-from socket import socket, AF_INET, SOCK_STREAM, gethostbyname, error
+from socket import socket, AF_INET, SOCK_STREAM, gethostbyname, error, timeout
 import threading
 import logging
 from enum import Enum, auto
@@ -21,6 +21,7 @@ class TcpProxySockHandler(BaseRequestHandler):
     PROXY_HOST = None
     PROXY_PORT = None
     Logger = logging.getLogger("TcpProxyHandler")
+    CONNECTION_ACTIVE = True
 
     def handle(self):
         # self.request is the TCP socket connected to the client
@@ -30,46 +31,38 @@ def handle(self):
         # Try to connect to the server and send data
         try:
             sock.connect((self.PROXY_HOST, self.PROXY_PORT))
-            self.transferData(self.request, sock)
-            # Receive data from the server
-            self.Logger.info("Receive data: %s <- %s:%s" %(self.client_address[0], self.PROXY_HOST, self.PROXY_PORT))
-            self.transferData(sock, self.request)
+
+            requestThread = threading.Thread(target=self.transfer_data, args=(self.request, sock))
+            requestThread.daemon = True
+            requestThread.start()
+
+            self.transfer_data(sock, self.request)
+
         finally:
             sock.close()
+            self.CONNECTION_ACTIVE = False
+            self.Logger.info("Connection %s -> %s:%s closed" %(self.client_address[0], self.PROXY_HOST, self.PROXY_PORT))
 
-    def transferData(self, sock1 : socket, sock2 : socket, max_read_try = 10, receive_message_length = 2048):
-        read_try = 0
-        last_try = False
-        sock1.setblocking(0)
+    def transfer_data(self, sock1 : socket, sock2 : socket, receive_message_length = 2048):
+        #setup socket receive timeout
+        sock1.settimeout(5.0)
         while True:
             try:
+                if not self.CONNECTION_ACTIVE or self.server.SHUTDOWN_REQUESTED:
+                    return
                 msg = sock1.recv(receive_message_length)
                 if not msg:
-                    return
+                    break
                 sock2.sendall(msg)
-
-                read_try = 0
-                if last_try and len(msg) > 0:
-                    last_try = False
-                if len(msg) < (receive_message_length / 2):
-                    if last_try:
-                        return
-                    else:
-                        last_try = True
-            except error as e:
-                if e.errno != errno.EAGAIN and e.errno != errno.EWOULDBLOCK:
-                    raise e
-                if read_try > max_read_try:
-                    return
-                read_try += 1
-                time.sleep(0.1)
-                continue
+            except timeout:
+                pass
+        self.CONNECTION_ACTIVE = False
 
 class ThreadedTCPServer(ThreadingMixIn, TCPServer):
-    pass
+    SHUTDOWN_REQUESTED = False
 
 class ProxyServer():
-    
+
     def __init__(self, proxy_host, proxy_port, type=ProxyType.TCP):
         if(type != ProxyType.TCP):
             raise NotImplementedError("Only TCP Proxy implemented")
@@ -93,5 +86,6 @@ def start(self):
         self.logger.info("Starting %s proxy to %s:%s", self.type, self.proxy_host, self.proxy_port)
 
     def stop(self):
+        self.server.SHUTDOWN_REQUESTED = True
         self.server.shutdown()
         self.logger.info("Stopping %s proxy to %s:%s", self.type, self.proxy_host, self.proxy_port)