Merge pull request #29 from icgood/short

Make sure short-lived connections are proxied

Author: icgood

proxyprotocol/server/__init__.py

@@ -26,6 +26,8 @@ class Address:
 
     """
 
+    __slots__ = ['url', 'query', 'server', '_ssl']
+
     def __init__(self, addr: str, *, server: bool = False) -> None:
         super().__init__()
         url = urlsplit(addr)

proxyprotocol/server/protocol.py

@@ -3,13 +3,13 @@
 
 import logging
 from abc import abstractmethod, ABCMeta
-from asyncio import Task, AbstractEventLoop, CancelledError
+from asyncio import Task, AbstractEventLoop
 from asyncio.protocols import BufferedProtocol, BaseProtocol
 from asyncio.transports import Transport, BaseTransport
 from collections import deque
 from functools import partial
 from socket import AddressFamily, SocketKind
-from typing import Any, Type, Optional, Tuple, Deque, Set
+from typing import Type, Optional, Tuple, Deque
 from typing_extensions import Final
 from uuid import uuid4
 
@@ -27,6 +27,8 @@
 
 class _Base(BufferedProtocol, metaclass=ABCMeta):
 
+    __slots__ = ['_paused', '_buf', '_queue', '_transport', '_sock_info']
+
     def __init__(self, buf_len: int) -> None:
         super().__init__()
         self._paused = False
@@ -41,15 +43,10 @@ def sock_info(self) -> SocketInfo:
         assert sock_info is not None
         return sock_info
 
-    @property
-    def is_open(self) -> bool:
-        return self._sock_info is not None
-
     def close(self) -> None:
         if self._transport is not None:
             self._transport.close()
             self._transport = None
-            self._sock_info = None
 
     def write(self, data: memoryview) -> None:
         transport = self._transport
@@ -97,6 +94,9 @@ def proxy_data(self, data: memoryview) -> None:
 
 class DownstreamProtocol(_Base):
 
+    __slots__ = ['loop', 'dnsbl', 'upstream', 'id', '_dnsbl_task', '_waiting',
+                 '_waiting_closed', '_upstream', '_upstream_factory']
+
     def __init__(self, upstream_protocol: Type[UpstreamProtocol],
                  loop: AbstractEventLoop, buf_len: int, dnsbl: Dnsbl,
                  upstream: Address) -> None:
@@ -105,56 +105,57 @@ def __init__(self, upstream_protocol: Type[UpstreamProtocol],
         self.dnsbl: Final = dnsbl
         self.upstream: Final = upstream
         self.id: Final = uuid4().bytes
+        self._dnsbl_task: Optional[Task[Optional[str]]] = None
         self._waiting: Deque[memoryview] = deque()
-        self._tasks: Set[Task[Any]] = set()
+        self._waiting_closed = False
         self._upstream: Optional[UpstreamProtocol] = None
         self._upstream_factory = partial(upstream_protocol, self, buf_len,
                                          upstream.pp)
 
-    def close(self) -> None:
-        super().close()
-        for task in self._tasks:
-            task.cancel()
-        if self._upstream is not None:
-            upstream = self._upstream
-            self._upstream = None
-            upstream.close()
-
     def _set_client(self, connect_task: Task[_Connect]) -> None:
+        dnsbl_task = self._dnsbl_task
+        assert dnsbl_task is not None
         try:
             _, upstream = connect_task.result()
-        except CancelledError:
-            pass  # Connection was never established
         except OSError:
             self.close()
+            dnsbl_task.cancel()
             _log.exception('[%s] Connection failed: %s',
                            self.id.hex(), self.upstream)
         else:
             assert isinstance(upstream, UpstreamProtocol)
             self._upstream = upstream
-            waiting = self._waiting
-            while waiting:
-                data = waiting.popleft()
-                upstream.write(data)
+            dnsbl_task.add_done_callback(self._send_initial)
+
+    def _send_initial(self, dnsbl_task: Task[Optional[str]]) -> None:
+        upstream = self._upstream
+        assert upstream is not None
+        upstream.write_header(dnsbl_task.result())
+        waiting = self._waiting
+        while waiting:
+            data = waiting.popleft()
+            upstream.write(data)
+        if self._waiting_closed:
+            upstream.close()
 
     def connection_made(self, transport: BaseTransport) -> None:
         super().connection_made(transport)
         _log.info('[%s] Downstream connection received: %s',
                   self.id.hex(), self.sock_info)
         loop = self.loop
-        dnsbl_task = loop.create_task(self.dnsbl.lookup(self.sock_info))
-        self._tasks.add(dnsbl_task)
-        dnsbl_task.add_done_callback(self._tasks.discard)
+        self._dnsbl_task = loop.create_task(self.dnsbl.lookup(self.sock_info))
         connect_task = loop.create_task(
-            loop.create_connection(partial(self._upstream_factory, dnsbl_task),
+            loop.create_connection(self._upstream_factory,
                                    self.upstream.host, self.upstream.port or 0,
                                    ssl=self.upstream.ssl))
-        self._tasks.add(connect_task)
-        connect_task.add_done_callback(self._tasks.discard)
         connect_task.add_done_callback(self._set_client)
 
     def connection_lost(self, exc: Optional[Exception]) -> None:
         super().connection_lost(exc)
+        if self._upstream is None:
+            self._waiting_closed = True
+        else:
+            self._upstream.close()
         _log.info('[%s] Downstream connection closed', self.id.hex())
 
     def proxy_data(self, data: memoryview) -> None:
@@ -167,39 +168,34 @@ def proxy_data(self, data: memoryview) -> None:
 
 class UpstreamProtocol(_Base):
 
+    __slots__ = ['pp', 'downstream']
+
     def __init__(self, downstream: DownstreamProtocol, buf_len: int,
-                 pp: ProxyProtocol, dnsbl_task: Task[Optional[str]]) -> None:
+                 pp: ProxyProtocol) -> None:
         super().__init__(buf_len)
         self.pp: Final = pp
         self.downstream: Final = downstream
-        self.dnsbl_task: Final = dnsbl_task
 
     def close(self) -> None:
         super().close()
         self.downstream.close()
 
-    def build_pp_header(self) -> bytes:
+    def build_pp_header(self, dnsbl: Optional[str]) -> bytes:
         sock_info = self.downstream.sock_info
         sock = sock_info.socket
         ssl_object = sock_info.transport.get_extra_info('ssl_object')
         try:
             protocol: Optional[SocketKind] = SocketKind(sock.proto)
         except ValueError:
             protocol = None
-        dnsbl = self.dnsbl_task.result()
-        return self.pp.build(sock.getpeername(), sock.getsockname(),
+        return self.pp.build(sock_info.peername, sock_info.sockname,
                              family=AddressFamily(sock.family),
                              protocol=protocol, unique_id=self.downstream.id,
                              ssl=ssl_object, dnsbl=dnsbl)
 
-    def connection_made(self, transport: BaseTransport) -> None:
-        super().connection_made(transport)
-        self.dnsbl_task.add_done_callback(self._write_header)
-
-    def _write_header(self, task: Task[Any]) -> None:
-        if self.downstream.is_open:
-            header = self.build_pp_header()
-            self.write(memoryview(header))
+    def write_header(self, dnsbl: Optional[str]) -> None:
+        header = self.build_pp_header(dnsbl)
+        self.write(memoryview(header))
 
     def proxy_data(self, data: memoryview) -> None:
         self.downstream.write(data)