Merge pull request #46 from icgood/callback

Simplify callback with asyncio.start_server

Author: icgood

.github/workflows/docker-build.yml

@@ -27,6 +27,8 @@ jobs:
         with:
           images: |
             ghcr.io/${{ github.repository }}
+          labels: |
+            org.opencontainers.image.source=https://github.com/icgood/proxy-protocol/tree/main/docker
           tags: |
             type=edge,branch=main
             type=pep440,pattern={{major}}

README.md

@@ -43,15 +43,13 @@ from proxyprotocol.sock import SocketInfo
 async def run(host: str, port: int) -> None:
     pp_detect = ProxyProtocolDetect()
     pp_reader = ProxyProtocolReader(pp_detect)
-    callback = partial(on_connection, pp_reader)
+    callback = reader.get_callback(on_connection)
     server = await asyncio.start_server(callback, host, port)
     async with server:
         await server.serve_forever()
 
-async def on_connection(pp_reader: ProxyProtocolReader,
-                        reader: StreamReader, writer: StreamWriter) -> None:
-    result = await pp_reader.read(reader)
-    info = SocketInfo.get(writer, result)
+async def on_connection(reader: StreamReader, writer: StreamWriter,
+                        info: SocketInfo) -> None:
     print(info.family, info.peername)
     # ... continue using connection
 ```

proxyprotocol/checksum.py

@@ -21,14 +21,19 @@ class Checksum(Protocol):
     """Provides typing compatible with the `crc32c.crc32c
     <https://github.com/ICRAR/crc32c#usage>`_ function, if it is installed.
 
-    Args:
-        val: The bytestring to checksum.
-        crc: The checksum of previous portions of data.
+    .. automethod:: __call__
 
     """
 
     __slots__: Sequence[str] = []
 
     @abstractmethod
     def __call__(self, val: bytes, crc: int = ...) -> int:
+        """The checksum function.
+
+        Args:
+            val: The bytestring to checksum.
+            crc: The checksum of previous portions of data.
+
+        """
         ...

proxyprotocol/reader.py

@@ -1,15 +1,25 @@
 
 from __future__ import annotations
 
-from typing_extensions import Final
+import asyncio
+from asyncio import StreamReader, StreamWriter
+from functools import partial
+from typing import Any, Awaitable, Callable, Coroutine, Union
+from typing_extensions import Final, TypeAlias
 
 from . import ProxyProtocol, ProxyProtocolIncompleteError, \
     ProxyProtocolWantRead
 from .result import ProxyResult, ProxyResultUnknown
+from .sock import SocketInfo
 from .typing import StreamReaderProtocol
 
 __all__ = ['ProxyProtocolReader']
 
+_Callback: TypeAlias = Callable[
+    [StreamReader, StreamWriter], Awaitable[None]]
+_WrappedCallback: TypeAlias = Callable[
+    [StreamReader, StreamWriter, SocketInfo], Coroutine[Any, Any, None]]
+
 
 class ProxyProtocolReader:
     """Read a PROXY protocol header from a stream.
@@ -47,7 +57,35 @@ async def read(self, reader: StreamReaderProtocol) -> ProxyResult:
                     return self.pp.unpack(view)
             except ProxyProtocolIncompleteError as exc:
                 want_read = exc.want_read
-            try:
-                data += await self._handle_want(reader, want_read)
-            except (EOFError, ConnectionResetError) as exc:
-                return ProxyResultUnknown(exc)
+            data += await self._handle_want(reader, want_read)
+
+    def get_callback(self, callback: _WrappedCallback,
+                     timeout: Union[None, int, float] = 3) -> _Callback:
+        """Get a callback object for use as the *client_connected_cb* argument
+        to :func:`asyncio.start_server`.
+
+        The returned callback will first read the PROXY protocol header before
+        starting the provided *callback* as a :class:`~asyncio.Task`. The
+        *callback* argument is similar to *client_connected_cb* but with an
+        additional positional argument -- the
+        :class:`~proxyprotocol.sock.SocketInfo` read from the header.
+
+        Args:
+            callback: Async function with arguments ``(reader, writer,
+                sock_info)`` called after successfully reading the header.
+            timeout: A timeout in seconds to allow for reading the header.
+
+        """
+        return partial(self._read_then_call, callback, timeout)
+
+    async def _read_then_call(self, callback: _WrappedCallback,
+                              timeout: Union[None, int, float],
+                              reader: StreamReader, writer: StreamWriter) \
+            -> None:
+        try:
+            result = await asyncio.wait_for(self.read(reader), timeout)
+        except Exception as exc:
+            writer.close()
+            result = ProxyResultUnknown(exc)
+        sock_info = SocketInfo.get(writer, result)
+        asyncio.create_task(callback(reader, writer, sock_info))

proxyprotocol/server/echo.py

@@ -11,7 +11,6 @@
 from contextlib import closing
 from functools import partial
 
-from .. import ProxyProtocol
 from ..reader import ProxyProtocolReader
 from ..sock import SocketInfo
 from . import Address
@@ -39,7 +38,8 @@ def main() -> int:
 
 async def run(address: Address) -> int:
     loop = asyncio.get_event_loop()
-    callback = partial(run_conn, address.pp)
+    reader = ProxyProtocolReader(address.pp)
+    callback = reader.get_callback(run_conn)
     server = await asyncio.start_server(
         callback, address.host, address.port or 0, ssl=address.ssl)
     async with server:
@@ -53,12 +53,9 @@ async def run(address: Address) -> int:
     return 0
 
 
-async def run_conn(pp: ProxyProtocol, reader: StreamReader,
-                   writer: StreamWriter) -> None:
-    pp_reader = ProxyProtocolReader(pp)
+async def run_conn(reader: StreamReader, writer: StreamWriter,
+                   sock_info: SocketInfo) -> None:
     with closing(writer):
-        result = await pp_reader.read(reader)
-        sock_info = SocketInfo.get(writer, result)
         _log.info('[%s] Connection received: %s',
                   sock_info.unique_id.hex(), sock_info)
         if sock_info.dnsbl is not None:

proxyprotocol/sock.py

@@ -5,9 +5,9 @@
 from abc import abstractmethod, ABCMeta
 from ipaddress import ip_address, IPv4Address, IPv6Address
 from socket import AddressFamily, SocketKind
-from typing import Union, Optional
+from typing import Dict, Optional, Union
 
-from .result import ProxyResult
+from .result import is_unknown, ProxyResult
 from .typing import SockAddr, Cipher, PeerCert, TransportProtocol
 
 __all__ = ['SocketInfo', 'SocketInfoProxy', 'SocketInfoLocal']
@@ -324,8 +324,13 @@ def dnsbl(self) -> Optional[str]:
         return self._result.tlv.ext.dnsbl
 
     def __repr__(self) -> str:
-        return f'<SocketInfoProxy peername={self.peername_str!r} ' \
-            f'sockname={self.sockname_str!r}>'
+        data: Dict[str, object] = {'peername': self.peername_str,
+                                   'sockname': self.sockname_str}
+        if is_unknown(self._result):
+            data['exc'] = self._result.exception
+        data_str = ''.join(f' {k}={v!r}' for k, v in data.items()
+                           if v is not None)
+        return f'<SocketInfoProxy{data_str}>'
 
 
 class SocketInfoLocal(SocketInfo):

pyproject.toml

@@ -25,7 +25,7 @@ build-backend = 'hatchling.build'
 
 [project]
 name = 'proxy-protocol'
-version = '0.9.1'
+version = '0.10.0'
 authors = [
     { name = 'Ian Good', email = 'ian@icgood.net' },
 ]
@@ -35,8 +35,6 @@ readme = { file = 'README.md', content-type = 'text/markdown' }
 requires-python = '~=3.8'
 classifiers = [
     'Development Status :: 3 - Alpha',
-    'Topic :: Communications :: Email :: Post-Office',
-    'Topic :: Communications :: Email :: Post-Office :: IMAP',
     'Intended Audience :: Developers',
     'Intended Audience :: Information Technology',
     'License :: OSI Approved :: MIT License',

test/test_sock.py

@@ -236,7 +236,7 @@ def test_str_ipv6(self) -> None:
                          "sockname='[::1]:10'>", repr(info))
 
     def test_str_unknown(self) -> None:
-        result = ProxyResultUnknown()
+        result = ProxyResultUnknown(RuntimeError())
         info = SocketInfo.get(self.transport, result)
-        self.assertEqual('<SocketInfoProxy peername=None sockname=None>',
+        self.assertEqual('<SocketInfoProxy exc=RuntimeError()>',
                          repr(info))