Merge pull request #20 from icgood/want-read

Make parsing independent of transport

Author: icgood

README.md

@@ -37,18 +37,20 @@ from functools import partial
 
 from proxyprotocol.base import ProxyProtocol
 from proxyprotocol.detect import ProxyProtocolDetect
+from proxyprotocol.reader import ProxyProtocolReader
 from proxyprotocol.socket import SocketInfo
 
 async def run(host: str, port: int) -> None:
-    pp = ProxyProtocolDetect()
-    callback = partial(on_connection, pp)
+    pp_detect = ProxyProtocolDetect()
+    pp_reader = ProxyProtocolReader(pp_detect)
+    callback = partial(on_connection, pp_reader)
     server = await asyncio.start_server(callback, host, port)
     async with server:
         await server.serve_forever()
 
-async def on_connection(pp: ProxyProtocolDetect,
+async def on_connection(pp_reader: ProxyProtocolReader,
                         reader: StreamReader, writer: StreamWriter) -> None:
-    result = await pp.read(reader)
+    result = await pp_reader.read(reader)
     info = SocketInfo(writer, result)
     print(info.family, info.peername)
     # ... continue using connection
@@ -139,10 +141,6 @@ $ mypy --strict proxyprotocol test
 $ flake8 proxyprotocol test
 ```
 
-A py.test run executes both unit and integration tests. The integration tests
-use mocked sockets to simulate the sending and receiving of commands and
-responses, and are kept in the `test/server/` subdirectory.
-
 ### Type Hinting
 
 This project makes heavy use of Python's [type hinting][6] system, with the

doc/source/proxyprotocol.rst

@@ -60,6 +60,12 @@
 .. automodule:: proxyprotocol.sock
    :members:
 
+``proxyprotocol.reader``
+------------------------
+
+.. automodule:: proxyprotocol.reader
+   :members:
+
 ``proxyprotocol.server``
 ------------------------
 

proxyprotocol/__init__.py

@@ -5,12 +5,13 @@
 from socket import AddressFamily, SocketKind
 from ssl import SSLSocket, SSLObject
 from typing import Any, Union, Optional, Sequence
+from typing_extensions import Final
 
 from .tlv import ProxyProtocolTLV
-from .typing import Address, StreamReaderProtocol
+from .typing import Address
 
-__all__ = ['__version__', 'ProxyProtocolError', 'ProxyProtocolResult',
-           'ProxyProtocol']
+__all__ = ['__version__', 'ProxyProtocolError', 'ProxyProtocolWantRead',
+           'ProxyProtocolResult', 'ProxyProtocol']
 
 #: The package version string.
 __version__: str = pkg_resources.require('proxy-protocol')[0].version
@@ -26,7 +27,30 @@ class ProxyProtocolError(ValueError):
         and closed.
 
     """
-    pass
+
+    __slots__: Sequence[str] = []
+
+
+class ProxyProtocolWantRead(Exception):
+    """Thrown when the PROXY protocol header cannot be parsed because the
+    provided data is not enough to be parsed. Additional data must be read
+    before parsing can succeed or fail.
+
+    Either *want_bytes* or *want_line* must be given, but not both.
+
+    Args:
+        want_bytes: Number of bytes needed before parsing may proceed.
+        want_line: Additional data should be read until the end of a line.
+
+    """
+
+    __slots__ = ['want_bytes', 'want_line']
+
+    def __init__(self, want_bytes: Optional[int] = None, *,
+                 want_line: bool = False) -> None:
+        super().__init__('Additional data needed')
+        self.want_bytes: Final = want_bytes
+        self.want_line: Final = want_line
 
 
 class ProxyProtocolResult(metaclass=ABCMeta):
@@ -100,16 +124,16 @@ def is_valid(self, signature: bytes) -> bool:
         ...
 
     @abstractmethod
-    async def read(self, reader: StreamReaderProtocol, *,
-                   signature: bytes = b'') -> ProxyProtocolResult:
-        """Read a PROXY protocol header from the given stream and return
+    def parse(self, data: bytes) -> ProxyProtocolResult:
+        """Parse a PROXY protocol header from the given bytestring and return
         information about the original connection.
 
         Args:
-            reader: The input stream.
-            signature: Any data that has already been read from the stream.
+            data: The bytestring read for the header thus far.
 
         Raises:
+            :exc:`ProxyProtocolWantRead`: The header was incomplete and must
+                be extended with additional bytes or lines to finish parsing.
             :exc:`ProxyProtocolError`: The header failed to parse due to a
                 syntax error or unsupported format.
             :exc:`ValueError`: Malformed or out-of-range data was encountered

proxyprotocol/detect.py

@@ -3,9 +3,9 @@
 from ssl import SSLSocket, SSLObject
 from typing import Union, Optional
 
-from . import ProxyProtocolError, ProxyProtocolResult, ProxyProtocol
-from .result import ProxyProtocolResultUnknown
-from .typing import Address, StreamReaderProtocol
+from . import ProxyProtocolError, ProxyProtocolWantRead, \
+    ProxyProtocolResult, ProxyProtocol
+from .typing import Address
 from .v1 import ProxyProtocolV1
 from .v2 import ProxyProtocolV2
 
@@ -31,15 +31,11 @@ def __init__(self, *versions: ProxyProtocol) -> None:
     def is_valid(self, signature: bytes) -> bool:
         return any(v.is_valid(signature) for v in self.versions)
 
-    async def read(self, reader: StreamReaderProtocol, *,
-                   signature: bytes = b'') \
-            -> ProxyProtocolResult:  # pragma: no cover
-        try:
-            signature += await reader.readexactly(8 - len(signature))
-        except (EOFError, ConnectionResetError) as exc:
-            return ProxyProtocolResultUnknown(exc)
-        pp = self.choose_version(signature)
-        return await pp.read(reader, signature=signature)
+    def parse(self, data: bytes) -> ProxyProtocolResult:
+        if len(data) < 8:
+            raise ProxyProtocolWantRead(8 - len(data))
+        pp = self.choose_version(data[0:8])
+        return pp.parse(data)
 
     def choose_version(self, signature: bytes) -> ProxyProtocol:
         """Choose the PROXY protocol version based on the 8-byte signature.

proxyprotocol/noop.py

@@ -5,7 +5,7 @@
 
 from . import ProxyProtocol
 from .result import ProxyProtocolResultLocal
-from .typing import Address, StreamReaderProtocol
+from .typing import Address
 
 __all__ = ['ProxyProtocolNoop']
 
@@ -24,9 +24,7 @@ def is_valid(self, signature: bytes) -> NoReturn:
         # This implementation may not be detected
         raise NotImplementedError()
 
-    async def read(self, reader: StreamReaderProtocol, *,
-                   signature: bytes = b'') \
-            -> ProxyProtocolResultLocal:  # pragma: no cover
+    def parse(self, data: bytes) -> ProxyProtocolResultLocal:
         return ProxyProtocolResultLocal()
 
     def build(self, source: Address, dest: Address, *, family: AddressFamily,

proxyprotocol/reader.py

@@ -0,0 +1,54 @@
+
+from typing_extensions import Final
+
+from . import ProxyProtocolWantRead, ProxyProtocol, ProxyProtocolResult
+from .result import ProxyProtocolResultUnknown
+from .typing import StreamReaderProtocol
+
+__all__ = ['ProxyProtocolReader']
+
+
+class ProxyProtocolReader:
+    """Read a PROXY protocol header from a stream.
+
+    Args:
+        pp: The PROXY protocol implementation.
+
+    """
+
+    def __init__(self, pp: ProxyProtocol) -> None:
+        super().__init__()
+        self.pp: Final = pp
+
+    def _parse(self, data: bytearray) -> ProxyProtocolResult:
+        view = memoryview(data)
+        try:
+            return self.pp.parse(view)
+        finally:
+            view.release()
+
+    async def _handle_want(self, reader: StreamReaderProtocol,
+                           want_read: ProxyProtocolWantRead) -> bytes:
+        if want_read.want_bytes is not None:
+            return await reader.readexactly(want_read.want_bytes)
+        elif want_read.want_line:
+            return await reader.readline()
+        raise ValueError() from want_read
+
+    async def read(self, reader: StreamReaderProtocol) -> ProxyProtocolResult:
+        """Read a complete PROXY protocol header from the input stream and
+        return the result.
+
+        Args:
+            reader: The input stream.
+
+        """
+        data = bytearray()
+        while True:
+            try:
+                return self._parse(data)
+            except ProxyProtocolWantRead as want_read:
+                try:
+                    data += await self._handle_want(reader, want_read)
+                except (EOFError, ConnectionResetError) as exc:
+                    return ProxyProtocolResultUnknown(exc)

proxyprotocol/server/echo.py

@@ -12,6 +12,7 @@
 from functools import partial
 
 from .. import ProxyProtocol
+from ..reader import ProxyProtocolReader
 from ..sock import SocketInfo
 from . import Address
 
@@ -54,8 +55,9 @@ async def run(address: Address) -> int:
 
 async def run_conn(pp: ProxyProtocol, reader: StreamReader,
                    writer: StreamWriter) -> None:
+    pp_reader = ProxyProtocolReader(pp)
     with closing(writer):
-        result = await pp.read(reader)
+        result = await pp_reader.read(reader)
         sock_info = SocketInfo(writer, result)
         _log.info('[%s] Connection received: %s',
                   sock_info.unique_id.hex(), sock_info)