Merge pull request #42 from icgood/ref

Keep a reference to the connect task

Author: icgood

proxyprotocol/server/protocol.py

@@ -107,8 +107,9 @@ def proxy_data(self, data: bytes) -> None:
 
 class DownstreamProtocol(_Base):
 
-    __slots__ = ['loop', 'dnsbl', 'upstream', 'id', '_dnsbl_task', '_waiting',
-                 '_waiting_closed', '_upstream', '_upstream_factory']
+    __slots__ = ['loop', 'dnsbl', 'upstream', 'id', '_waiting',
+                 '_waiting_closed', '_upstream', '_upstream_factory',
+                 '_dnsbl_task', '_connect_task']
 
     def __init__(self, upstream_protocol: Type[UpstreamProtocol],
                  loop: AbstractEventLoop, buf_len: int, dnsbl: Dnsbl,
@@ -119,6 +120,7 @@ def __init__(self, upstream_protocol: Type[UpstreamProtocol],
         self.upstream: Final = upstream
         self.id: Final = uuid4().bytes
         self._dnsbl_task: Optional[Task[Optional[str]]] = None
+        self._connect_task: Optional[Task[_Connect]] = None
         self._waiting: Deque[bytes] = deque()
         self._waiting_closed = False
         self._upstream: Optional[UpstreamProtocol] = None
@@ -128,6 +130,8 @@ def __init__(self, upstream_protocol: Type[UpstreamProtocol],
     def _set_client(self, result: ProxyResult,
                     connect_task: Task[_Connect]) -> None:
         dnsbl_task = self._dnsbl_task
+        self._dnsbl_task = None
+        self._connect_task = None
         assert dnsbl_task is not None
         try:
             _, upstream = connect_task.result()
@@ -157,7 +161,7 @@ def connection_made(self, transport: BaseTransport) -> None:
                   self.id.hex(), self.sock_info)
         loop = self.loop
         self._dnsbl_task = loop.create_task(self.dnsbl.lookup(self.sock_info))
-        connect_task = loop.create_task(
+        self._connect_task = connect_task = loop.create_task(
             loop.create_connection(self._upstream_factory,
                                    self.upstream.host, self.upstream.port or 0,
                                    ssl=self.upstream.ssl))

proxyprotocol/version.py

@@ -2,7 +2,7 @@
 from __future__ import annotations
 
 from enum import Enum
-from typing import cast, Optional
+from typing import Optional
 
 from . import ProxyProtocol
 from .detect import ProxyProtocolDetect
@@ -44,4 +44,6 @@ def get(cls, name: Optional[str] = None) -> ProxyProtocol:
         """
         if not name:
             return cls.NOOP.value
-        return cast(ProxyProtocol, cls[name.upper()].value)
+        pp = cls[name.upper()].value
+        assert isinstance(pp, ProxyProtocol)
+        return pp

setup.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2021 Ian C. Good
+# Copyright (c) 2023 Ian C. Good
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
@@ -28,7 +28,7 @@
     license = f.read()
 
 setup(name='proxy-protocol',
-      version='0.9.0',
+      version='0.9.1',
       author='Ian Good',
       author_email='ian@icgood.net',
       description='PROXY protocol library with asyncio server implementation',