diff options
author | shmel1k <shmel1k@ydb.tech> | 2023-11-26 18:16:14 +0300 |
---|---|---|
committer | shmel1k <shmel1k@ydb.tech> | 2023-11-26 18:43:30 +0300 |
commit | b8cf9e88f4c5c64d9406af533d8948deb050d695 (patch) | |
tree | 218eb61fb3c3b96ec08b4d8cdfef383104a87d63 /contrib/python/Twisted/py3/twisted/protocols/haproxy/_wrapper.py | |
parent | 523f645a83a0ec97a0332dbc3863bb354c92a328 (diff) | |
download | ydb-b8cf9e88f4c5c64d9406af533d8948deb050d695.tar.gz |
add kikimr_configure
Diffstat (limited to 'contrib/python/Twisted/py3/twisted/protocols/haproxy/_wrapper.py')
-rw-r--r-- | contrib/python/Twisted/py3/twisted/protocols/haproxy/_wrapper.py | 109 |
1 files changed, 109 insertions, 0 deletions
diff --git a/contrib/python/Twisted/py3/twisted/protocols/haproxy/_wrapper.py b/contrib/python/Twisted/py3/twisted/protocols/haproxy/_wrapper.py new file mode 100644 index 0000000000..935dbfa9e2 --- /dev/null +++ b/contrib/python/Twisted/py3/twisted/protocols/haproxy/_wrapper.py @@ -0,0 +1,109 @@ +# -*- test-case-name: twisted.protocols.haproxy.test.test_wrapper -*- + +# Copyright (c) Twisted Matrix Laboratories. +# See LICENSE for details. + +""" +Protocol wrapper that provides HAProxy PROXY protocol support. +""" +from typing import Optional, Union + +from twisted.internet import interfaces +from twisted.internet.endpoints import _WrapperServerEndpoint +from twisted.protocols import policies +from . import _info +from ._exceptions import InvalidProxyHeader +from ._v1parser import V1Parser +from ._v2parser import V2Parser + + +class HAProxyProtocolWrapper(policies.ProtocolWrapper): + """ + A Protocol wrapper that provides HAProxy support. + + This protocol reads the PROXY stream header, v1 or v2, parses the provided + connection data, and modifies the behavior of getPeer and getHost to return + the data provided by the PROXY header. + """ + + def __init__( + self, factory: policies.WrappingFactory, wrappedProtocol: interfaces.IProtocol + ): + super().__init__(factory, wrappedProtocol) + self._proxyInfo: Optional[_info.ProxyInfo] = None + self._parser: Union[V2Parser, V1Parser, None] = None + + def dataReceived(self, data: bytes) -> None: + if self._proxyInfo is not None: + return self.wrappedProtocol.dataReceived(data) + + parser = self._parser + if parser is None: + if ( + len(data) >= 16 + and data[:12] == V2Parser.PREFIX + and ord(data[12:13]) & 0b11110000 == 0x20 + ): + self._parser = parser = V2Parser() + elif len(data) >= 8 and data[:5] == V1Parser.PROXYSTR: + self._parser = parser = V1Parser() + else: + self.loseConnection() + return None + + try: + self._proxyInfo, remaining = parser.feed(data) + if remaining: + self.wrappedProtocol.dataReceived(remaining) + except InvalidProxyHeader: + self.loseConnection() + + def getPeer(self) -> interfaces.IAddress: + if self._proxyInfo and self._proxyInfo.source: + return self._proxyInfo.source + assert self.transport + return self.transport.getPeer() + + def getHost(self) -> interfaces.IAddress: + if self._proxyInfo and self._proxyInfo.destination: + return self._proxyInfo.destination + assert self.transport + return self.transport.getHost() + + +class HAProxyWrappingFactory(policies.WrappingFactory): + """ + A Factory wrapper that adds PROXY protocol support to connections. + """ + + protocol = HAProxyProtocolWrapper + + def logPrefix(self) -> str: + """ + Annotate the wrapped factory's log prefix with some text indicating + the PROXY protocol is in use. + + @rtype: C{str} + """ + if interfaces.ILoggingContext.providedBy(self.wrappedFactory): + logPrefix = self.wrappedFactory.logPrefix() + else: + logPrefix = self.wrappedFactory.__class__.__name__ + return f"{logPrefix} (PROXY)" + + +def proxyEndpoint( + wrappedEndpoint: interfaces.IStreamServerEndpoint, +) -> _WrapperServerEndpoint: + """ + Wrap an endpoint with PROXY protocol support, so that the transport's + C{getHost} and C{getPeer} methods reflect the attributes of the proxied + connection rather than the underlying connection. + + @param wrappedEndpoint: The underlying listening endpoint. + @type wrappedEndpoint: L{IStreamServerEndpoint} + + @return: a new listening endpoint that speaks the PROXY protocol. + @rtype: L{IStreamServerEndpoint} + """ + return _WrapperServerEndpoint(wrappedEndpoint, HAProxyWrappingFactory) |