17.12
This commit is contained in:
Mai Gillmann
@@ -0,0 +1,71 @@
|
||||
# -*- test-case-name: twisted.protocols.haproxy.test.test_parser -*-
|
||||
# Copyright (c) Twisted Matrix Laboratories.
|
||||
# See LICENSE for details.
|
||||
|
||||
"""
|
||||
Parser for 'haproxy:' string endpoint.
|
||||
"""
|
||||
|
||||
from zope.interface import implementer
|
||||
from twisted.plugin import IPlugin
|
||||
|
||||
from twisted.internet.endpoints import (
|
||||
quoteStringArgument, serverFromString, IStreamServerEndpointStringParser
|
||||
)
|
||||
from twisted.python.compat import iteritems
|
||||
|
||||
from . import proxyEndpoint
|
||||
|
||||
|
||||
def unparseEndpoint(args, kwargs):
|
||||
"""
|
||||
Un-parse the already-parsed args and kwargs back into endpoint syntax.
|
||||
|
||||
@param args: C{:}-separated arguments
|
||||
@type args: L{tuple} of native L{str}
|
||||
|
||||
@param kwargs: C{:} and then C{=}-separated keyword arguments
|
||||
|
||||
@type arguments: L{tuple} of native L{str}
|
||||
|
||||
@return: a string equivalent to the original format which this was parsed
|
||||
as.
|
||||
@rtype: native L{str}
|
||||
"""
|
||||
|
||||
description = ':'.join(
|
||||
[quoteStringArgument(str(arg)) for arg in args] +
|
||||
sorted(['%s=%s' % (quoteStringArgument(str(key)),
|
||||
quoteStringArgument(str(value)))
|
||||
for key, value in iteritems(kwargs)
|
||||
]))
|
||||
return description
|
||||
|
||||
|
||||
|
||||
@implementer(IPlugin, IStreamServerEndpointStringParser)
|
||||
class HAProxyServerParser(object):
|
||||
"""
|
||||
Stream server endpoint string parser for the HAProxyServerEndpoint type.
|
||||
|
||||
@ivar prefix: See L{IStreamServerEndpointStringParser.prefix}.
|
||||
"""
|
||||
prefix = "haproxy"
|
||||
|
||||
def parseStreamServer(self, reactor, *args, **kwargs):
|
||||
"""
|
||||
Parse a stream server endpoint from a reactor and string-only arguments
|
||||
and keyword arguments.
|
||||
|
||||
@param reactor: The reactor.
|
||||
|
||||
@param args: The parsed string arguments.
|
||||
|
||||
@param kwargs: The parsed keyword arguments.
|
||||
|
||||
@return: a stream server endpoint
|
||||
@rtype: L{IStreamServerEndpoint}
|
||||
"""
|
||||
subdescription = unparseEndpoint(args, kwargs)
|
||||
wrappedEndpoint = serverFromString(reactor, subdescription)
|
||||
return proxyEndpoint(wrappedEndpoint)
|
||||
@@ -0,0 +1,143 @@
|
||||
# -*- test-case-name: twisted.protocols.haproxy.test.test_v1parser -*-
|
||||
|
||||
# Copyright (c) Twisted Matrix Laboratories.
|
||||
# See LICENSE for details.
|
||||
|
||||
"""
|
||||
IProxyParser implementation for version one of the PROXY protocol.
|
||||
"""
|
||||
|
||||
from zope.interface import implementer
|
||||
from twisted.internet import address
|
||||
|
||||
from ._exceptions import (
|
||||
convertError, InvalidProxyHeader, InvalidNetworkProtocol,
|
||||
MissingAddressData
|
||||
)
|
||||
from . import _info
|
||||
from . import _interfaces
|
||||
|
||||
|
||||
|
||||
@implementer(_interfaces.IProxyParser)
|
||||
class V1Parser(object):
|
||||
"""
|
||||
PROXY protocol version one header parser.
|
||||
|
||||
Version one of the PROXY protocol is a human readable format represented
|
||||
by a single, newline delimited binary string that contains all of the
|
||||
relevant source and destination data.
|
||||
"""
|
||||
|
||||
PROXYSTR = b'PROXY'
|
||||
UNKNOWN_PROTO = b'UNKNOWN'
|
||||
TCP4_PROTO = b'TCP4'
|
||||
TCP6_PROTO = b'TCP6'
|
||||
ALLOWED_NET_PROTOS = (
|
||||
TCP4_PROTO,
|
||||
TCP6_PROTO,
|
||||
UNKNOWN_PROTO,
|
||||
)
|
||||
NEWLINE = b'\r\n'
|
||||
|
||||
def __init__(self):
|
||||
self.buffer = b''
|
||||
|
||||
|
||||
def feed(self, data):
|
||||
"""
|
||||
Consume a chunk of data and attempt to parse it.
|
||||
|
||||
@param data: A bytestring.
|
||||
@type data: L{bytes}
|
||||
|
||||
@return: A two-tuple containing, in order, a
|
||||
L{_interfaces.IProxyInfo} and any bytes fed to the
|
||||
parser that followed the end of the header. Both of these values
|
||||
are None until a complete header is parsed.
|
||||
|
||||
@raises InvalidProxyHeader: If the bytes fed to the parser create an
|
||||
invalid PROXY header.
|
||||
"""
|
||||
self.buffer += data
|
||||
if len(self.buffer) > 107 and self.NEWLINE not in self.buffer:
|
||||
raise InvalidProxyHeader()
|
||||
lines = (self.buffer).split(self.NEWLINE, 1)
|
||||
if not len(lines) > 1:
|
||||
return (None, None)
|
||||
self.buffer = b''
|
||||
remaining = lines.pop()
|
||||
header = lines.pop()
|
||||
info = self.parse(header)
|
||||
return (info, remaining)
|
||||
|
||||
|
||||
@classmethod
|
||||
def parse(cls, line):
|
||||
"""
|
||||
Parse a bytestring as a full PROXY protocol header line.
|
||||
|
||||
@param line: A bytestring that represents a valid HAProxy PROXY
|
||||
protocol header line.
|
||||
@type line: bytes
|
||||
|
||||
@return: A L{_interfaces.IProxyInfo} containing the parsed data.
|
||||
|
||||
@raises InvalidProxyHeader: If the bytestring does not represent a
|
||||
valid PROXY header.
|
||||
|
||||
@raises InvalidNetworkProtocol: When no protocol can be parsed or is
|
||||
not one of the allowed values.
|
||||
|
||||
@raises MissingAddressData: When the protocol is TCP* but the header
|
||||
does not contain a complete set of addresses and ports.
|
||||
"""
|
||||
originalLine = line
|
||||
proxyStr = None
|
||||
networkProtocol = None
|
||||
sourceAddr = None
|
||||
sourcePort = None
|
||||
destAddr = None
|
||||
destPort = None
|
||||
|
||||
with convertError(ValueError, InvalidProxyHeader):
|
||||
proxyStr, line = line.split(b' ', 1)
|
||||
|
||||
if proxyStr != cls.PROXYSTR:
|
||||
raise InvalidProxyHeader()
|
||||
|
||||
with convertError(ValueError, InvalidNetworkProtocol):
|
||||
networkProtocol, line = line.split(b' ', 1)
|
||||
|
||||
if networkProtocol not in cls.ALLOWED_NET_PROTOS:
|
||||
raise InvalidNetworkProtocol()
|
||||
|
||||
if networkProtocol == cls.UNKNOWN_PROTO:
|
||||
|
||||
return _info.ProxyInfo(originalLine, None, None)
|
||||
|
||||
with convertError(ValueError, MissingAddressData):
|
||||
sourceAddr, line = line.split(b' ', 1)
|
||||
|
||||
with convertError(ValueError, MissingAddressData):
|
||||
destAddr, line = line.split(b' ', 1)
|
||||
|
||||
with convertError(ValueError, MissingAddressData):
|
||||
sourcePort, line = line.split(b' ', 1)
|
||||
|
||||
with convertError(ValueError, MissingAddressData):
|
||||
destPort = line.split(b' ')[0]
|
||||
|
||||
if networkProtocol == cls.TCP4_PROTO:
|
||||
|
||||
return _info.ProxyInfo(
|
||||
originalLine,
|
||||
address.IPv4Address('TCP', sourceAddr, int(sourcePort)),
|
||||
address.IPv4Address('TCP', destAddr, int(destPort)),
|
||||
)
|
||||
|
||||
return _info.ProxyInfo(
|
||||
originalLine,
|
||||
address.IPv6Address('TCP', sourceAddr, int(sourcePort)),
|
||||
address.IPv6Address('TCP', destAddr, int(destPort)),
|
||||
)
|
||||
@@ -0,0 +1,106 @@
|
||||
# -*- 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 twisted.protocols import policies
|
||||
from twisted.internet import interfaces
|
||||
from twisted.internet.endpoints import _WrapperServerEndpoint
|
||||
|
||||
from ._exceptions import InvalidProxyHeader
|
||||
from ._v1parser import V1Parser
|
||||
from ._v2parser import V2Parser
|
||||
|
||||
|
||||
|
||||
class HAProxyProtocolWrapper(policies.ProtocolWrapper, object):
|
||||
"""
|
||||
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, wrappedProtocol):
|
||||
policies.ProtocolWrapper.__init__(self, factory, wrappedProtocol)
|
||||
self._proxyInfo = None
|
||||
self._parser = None
|
||||
|
||||
|
||||
def dataReceived(self, data):
|
||||
if self._proxyInfo is not None:
|
||||
return self.wrappedProtocol.dataReceived(data)
|
||||
|
||||
if self._parser is None:
|
||||
if (
|
||||
len(data) >= 16 and
|
||||
data[:12] == V2Parser.PREFIX and
|
||||
ord(data[12:13]) & 0b11110000 == 0x20
|
||||
):
|
||||
self._parser = V2Parser()
|
||||
elif len(data) >= 8 and data[:5] == V1Parser.PROXYSTR:
|
||||
self._parser = V1Parser()
|
||||
else:
|
||||
self.loseConnection()
|
||||
return None
|
||||
|
||||
try:
|
||||
self._proxyInfo, remaining = self._parser.feed(data)
|
||||
if remaining:
|
||||
self.wrappedProtocol.dataReceived(remaining)
|
||||
except InvalidProxyHeader:
|
||||
self.loseConnection()
|
||||
|
||||
|
||||
def getPeer(self):
|
||||
if self._proxyInfo and self._proxyInfo.source:
|
||||
return self._proxyInfo.source
|
||||
return self.transport.getPeer()
|
||||
|
||||
|
||||
def getHost(self):
|
||||
if self._proxyInfo and self._proxyInfo.destination:
|
||||
return self._proxyInfo.destination
|
||||
return self.transport.getHost()
|
||||
|
||||
|
||||
|
||||
class HAProxyWrappingFactory(policies.WrappingFactory):
|
||||
"""
|
||||
A Factory wrapper that adds PROXY protocol support to connections.
|
||||
"""
|
||||
protocol = HAProxyProtocolWrapper
|
||||
|
||||
def logPrefix(self):
|
||||
"""
|
||||
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 "%s (PROXY)" % (logPrefix,)
|
||||
|
||||
|
||||
|
||||
def proxyEndpoint(wrappedEndpoint):
|
||||
"""
|
||||
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)
|
||||
@@ -0,0 +1,7 @@
|
||||
# -*- test-case-name: twisted.protocols.haproxy.test -*-
|
||||
# Copyright (c) Twisted Matrix Laboratories.
|
||||
# See LICENSE for details.
|
||||
|
||||
"""
|
||||
Unit tests for L{twisted.protocols.haproxy}.
|
||||
"""
|
||||
Reference in New Issue
Block a user