Merge pull request #18 from icgood/pp-query

Specify PROXY protocol version in query params

Author: icgood

README.md

@@ -87,6 +87,10 @@ together can demonstrate the process end-to-end: use `proxyprotocol-server`
 to proxy connections with a PROXY protocol header to `proxyprotocol-echo`,
 which then displays the original connection information.
 
+The `hostname:port` arguments used by both types of servers are parsed by the
+[`Address`][8] class, which allows for customization of SSL/TLS and PROXY
+protocol versions.
+
 ### Echo Server
 
 The `proxyprotocol-echo` server expects inbound connections to provide a PROXY
@@ -106,7 +110,7 @@ header to indicate the original connection information.
 
 ```bash
 proxyprotocol-server --help
-proxyprotocol-server v2 --service localhost:10000 localhost:10007
+proxyprotocol-server --service localhost:10000 localhost:10007
 ```
 
 ## Development and Testing
@@ -157,3 +161,4 @@ hinting to the extent possible and common in the rest of the codebase.
 [5]: https://docs.python.org/3/library/venv.html
 [6]: https://www.python.org/dev/peps/pep-0484/
 [7]: http://mypy-lang.org/
+[8]: https://icgood.github.io/proxy-protocol/proxyprotocol.html#proxyprotocol.server.Address

doc/source/proxyprotocol.rst

@@ -59,3 +59,9 @@
 
 .. automodule:: proxyprotocol.sock
    :members:
+
+``proxyprotocol.server``
+------------------------
+
+.. automodule:: proxyprotocol.server
+   :members:

proxyprotocol/server/__init__.py

@@ -4,6 +4,9 @@
 from typing import Optional
 from typing_extensions import Final
 
+from .. import ProxyProtocol
+from ..version import ProxyProtocolVersion
+
 __all__ = ['Address']
 
 
@@ -12,6 +15,7 @@ class Address:
 
     * ``HOST``
     * ``HOST:PORT``
+    * ``HOST:PORT?pp=v1``
     * ``ssl://HOST:PORT`` (outbound addresses only)
     * ``ssl://HOST:PORT?cert=/path/to/cert.pem``
     * ``ssl://HOST:PORT?cert=cert.pem&key=privkey.pem&verify=CERT_REQUIRED``
@@ -25,7 +29,7 @@ class Address:
     def __init__(self, addr: str, *, server: bool = False) -> None:
         super().__init__()
         url = urlsplit(addr)
-        if not url.scheme and not url.netloc:
+        if not url.scheme or not url.netloc:
             url = urlsplit('//' + addr)
         if url.query:
             query = parse_qs(url.query)
@@ -46,9 +50,15 @@ def port(self) -> Optional[int]:
         """The port parsed from the address."""
         return self.url.port or None
 
+    @property
+    def pp(self) -> ProxyProtocol:
+        """The PROXY protocol implementation."""
+        pp_version = self.query.get('pp', [''])[-1] or 'detect'
+        return ProxyProtocolVersion.get(pp_version)
+
     @property
     def ssl(self) -> Optional[SSLContext]:
-        """The :class:`~ssl.SSLContext` to use on the address."""
+        """The: class:`~ssl.SSLContext` to use on the address."""
         if self.url.scheme == 'ssl':
             if self._ssl is None:
                 if self.server:

proxyprotocol/server/echo.py

@@ -8,11 +8,11 @@
 import sys
 from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
 from asyncio import CancelledError, StreamReader, StreamWriter
+from contextlib import closing
 from functools import partial
 
 from .. import ProxyProtocol
 from ..sock import SocketInfo
-from ..version import ProxyProtocolVersion
 from . import Address
 
 __all__ = ['main']
@@ -23,9 +23,6 @@
 def main() -> int:
     parser = ArgumentParser(description=__doc__,
                             formatter_class=ArgumentDefaultsHelpFormatter)
-    parser.add_argument('type', default='detect', nargs='?',
-                        choices=[v.name.lower() for v in ProxyProtocolVersion],
-                        help='the PROXY protocol version')
     parser.add_argument('address', metavar='HOST:PORT',
                         type=partial(Address, server=True),
                         nargs='?', default=':10007',
@@ -36,13 +33,12 @@ def main() -> int:
         level=logging.INFO,
         format='%(asctime)-15s %(name)s %(message)s')
 
-    pp = ProxyProtocolVersion.get(args.type)
-    return asyncio.run(run(pp, args.address))
+    return asyncio.run(run(args.address))
 
 
-async def run(pp: ProxyProtocol, address: Address) -> int:
+async def run(address: Address) -> int:
     loop = asyncio.get_event_loop()
-    callback = partial(run_conn, pp)
+    callback = partial(run_conn, address.pp)
     server = await asyncio.start_server(
         callback, address.host, address.port or 0, ssl=address.ssl)
     async with server:
@@ -58,20 +54,21 @@ async def run(pp: ProxyProtocol, address: Address) -> int:
 
 async def run_conn(pp: ProxyProtocol, reader: StreamReader,
                    writer: StreamWriter) -> None:
-    result = await pp.read(reader)
-    sock_info = SocketInfo(writer, result)
-    _log.info('[%s] Connection received: %s',
-              sock_info.unique_id.hex(), sock_info)
-    try:
-        while True:
-            line = await reader.readline()
-            if not line:
-                break
-            writer.write(line)
-    except IOError:
-        pass
-    finally:
-        _log.info('[%s] Connection lost', sock_info.unique_id.hex())
+    with closing(writer):
+        result = await pp.read(reader)
+        sock_info = SocketInfo(writer, result)
+        _log.info('[%s] Connection received: %s',
+                  sock_info.unique_id.hex(), sock_info)
+        try:
+            while True:
+                line = await reader.readline()
+                if not line:
+                    break
+                writer.write(line)
+        except IOError:
+            pass
+        finally:
+            _log.info('[%s] Connection lost', sock_info.unique_id.hex())
 
 
 if __name__ == '__main__':

proxyprotocol/server/main.py

@@ -14,8 +14,6 @@
 from contextlib import AsyncExitStack
 from functools import partial
 
-from .. import ProxyProtocol
-from ..version import ProxyProtocolVersion
 from . import Address
 from .protocol import DownstreamProtocol, UpstreamProtocol
 
@@ -25,16 +23,14 @@
 def main() -> int:
     parser = ArgumentParser(description=__doc__,
                             formatter_class=ArgumentDefaultsHelpFormatter)
-    parser.add_argument('--service', nargs=2, metavar='HOST:PORT', default=[],
+    parser.add_argument('--service', nargs=2, default=[],
+                        metavar=('SOURCE:PORT', 'DEST:PORT'),
                         action='append', dest='services',
                         help='source and destination of a service')
     parser.add_argument('--buf-len', metavar='BYTES', default=262144, type=int,
                         help='size of the read buffer')
     parser.add_argument('-q', '--quiet', action='store_true',
                         help='show only upstream connection errors')
-    parser.add_argument('type', default='detect', nargs='?',
-                        choices=[v.name.lower() for v in ProxyProtocolVersion],
-                        help='the PROXY protocol version')
     args = parser.parse_args()
 
     if not args.services:
@@ -44,17 +40,16 @@ def main() -> int:
         level=logging.ERROR if args.quiet else logging.INFO,
         format='%(asctime)-15s %(name)s %(message)s')
 
-    pp = ProxyProtocolVersion.get(args.type)
-    return asyncio.run(run(pp, args))
+    return asyncio.run(run(args))
 
 
-async def run(pp: ProxyProtocol, args: Namespace) -> int:
+async def run(args: Namespace) -> int:
     loop = asyncio.get_running_loop()
     services = [(Address(source, server=True), Address(dest))
                 for (source, dest) in args.services]
     buf_len: int = args.buf_len
     new_server = partial(DownstreamProtocol, UpstreamProtocol,
-                         pp, loop, buf_len)
+                         loop, buf_len)
     servers = [
         await loop.create_server(partial(new_server, dest),
                                  source.host, source.port or 0,

proxyprotocol/server/protocol.py

@@ -26,9 +26,8 @@
 
 class _Base(BufferedProtocol, metaclass=ABCMeta):
 
-    def __init__(self, pp: ProxyProtocol, buf_len: int) -> None:
+    def __init__(self, buf_len: int) -> None:
         super().__init__()
-        self.pp: Final = pp
         self._paused = False
         self._buf: bytearray = bytearray(buf_len)
         self._queue: Deque[bytes] = deque()
@@ -94,16 +93,17 @@ def proxy_data(self, data: memoryview) -> None:
 class DownstreamProtocol(_Base):
 
     def __init__(self, upstream_protocol: Type[UpstreamProtocol],
-                 pp: ProxyProtocol, loop: AbstractEventLoop,
-                 buf_len: int, upstream: Address) -> None:
-        super().__init__(pp, buf_len)
+                 loop: AbstractEventLoop, buf_len: int,
+                 upstream: Address) -> None:
+        super().__init__(buf_len)
         self.loop: Final = loop
         self.upstream: Final = upstream
         self.id: Final = uuid4().bytes
         self._waiting: Deque[memoryview] = deque()
         self._connect: Optional[Task[Any]] = None
         self._upstream: Optional[UpstreamProtocol] = None
-        self._upstream_factory = partial(upstream_protocol, self)
+        self._upstream_factory = partial(upstream_protocol, self, buf_len,
+                                         upstream.pp)
 
     def close(self) -> None:
         super().close()
@@ -158,8 +158,10 @@ def proxy_data(self, data: memoryview) -> None:
 
 class UpstreamProtocol(_Base):
 
-    def __init__(self, downstream: DownstreamProtocol) -> None:
-        super().__init__(downstream.pp, len(downstream._buf))
+    def __init__(self, downstream: DownstreamProtocol, buf_len: int,
+                 pp: ProxyProtocol) -> None:
+        super().__init__(buf_len)
+        self.pp: Final = pp
         self.downstream: Final = downstream
 
     def close(self) -> None:

setup.py

@@ -28,7 +28,7 @@
     license = f.read()
 
 setup(name='proxy-protocol',
-      version='0.5.6',
+      version='0.5.7',
       author='Ian Good',
       author_email='ian@icgood.net',
       description='PROXY protocol library with asyncio server implementation',