aboutsummaryrefslogtreecommitdiffstats
path: root/contrib/python/Twisted/py3/twisted/protocols/haproxy/_v2parser.py
blob: cfcf7c99bcc6260f43f6ec4b8a62d0bfc222a794 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
# -*- test-case-name: twisted.protocols.haproxy.test.test_v2parser -*-

# Copyright (c) Twisted Matrix Laboratories.
# See LICENSE for details.

"""
IProxyParser implementation for version two of the PROXY protocol.
"""

import binascii
import struct
from typing import Callable, Tuple, Type, Union

from zope.interface import implementer

from constantly import ValueConstant, Values
from typing_extensions import Literal

from twisted.internet import address
from twisted.python import compat
from . import _info, _interfaces
from ._exceptions import (
    InvalidNetworkProtocol,
    InvalidProxyHeader,
    MissingAddressData,
    convertError,
)


class NetFamily(Values):
    """
    Values for the 'family' field.
    """

    UNSPEC = ValueConstant(0x00)
    INET = ValueConstant(0x10)
    INET6 = ValueConstant(0x20)
    UNIX = ValueConstant(0x30)


class NetProtocol(Values):
    """
    Values for 'protocol' field.
    """

    UNSPEC = ValueConstant(0)
    STREAM = ValueConstant(1)
    DGRAM = ValueConstant(2)


_HIGH = 0b11110000
_LOW = 0b00001111
_LOCALCOMMAND = "LOCAL"
_PROXYCOMMAND = "PROXY"


@implementer(_interfaces.IProxyParser)
class V2Parser:
    """
    PROXY protocol version two header parser.

    Version two of the PROXY protocol is a binary format.
    """

    PREFIX = b"\x0D\x0A\x0D\x0A\x00\x0D\x0A\x51\x55\x49\x54\x0A"
    VERSIONS = [32]
    COMMANDS = {0: _LOCALCOMMAND, 1: _PROXYCOMMAND}
    ADDRESSFORMATS = {
        # TCP4
        17: "!4s4s2H",
        18: "!4s4s2H",
        # TCP6
        33: "!16s16s2H",
        34: "!16s16s2H",
        # UNIX
        49: "!108s108s",
        50: "!108s108s",
    }

    def __init__(self) -> None:
        self.buffer = b""

    def feed(
        self, data: bytes
    ) -> Union[Tuple[_info.ProxyInfo, bytes], Tuple[None, None]]:
        """
        Consume a chunk of data and attempt to parse it.

        @param data: A bytestring.
        @type data: 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) < 16:
            raise InvalidProxyHeader()

        size = struct.unpack("!H", self.buffer[14:16])[0] + 16
        if len(self.buffer) < size:
            return (None, None)

        header, remaining = self.buffer[:size], self.buffer[size:]
        self.buffer = b""
        info = self.parse(header)
        return (info, remaining)

    @staticmethod
    def _bytesToIPv4(bytestring: bytes) -> bytes:
        """
        Convert packed 32-bit IPv4 address bytes into a dotted-quad ASCII bytes
        representation of that address.

        @param bytestring: 4 octets representing an IPv4 address.
        @type bytestring: L{bytes}

        @return: a dotted-quad notation IPv4 address.
        @rtype: L{bytes}
        """
        return b".".join(
            ("%i" % (ord(b),)).encode("ascii") for b in compat.iterbytes(bytestring)
        )

    @staticmethod
    def _bytesToIPv6(bytestring: bytes) -> bytes:
        """
        Convert packed 128-bit IPv6 address bytes into a colon-separated ASCII
        bytes representation of that address.

        @param bytestring: 16 octets representing an IPv6 address.
        @type bytestring: L{bytes}

        @return: a dotted-quad notation IPv6 address.
        @rtype: L{bytes}
        """
        hexString = binascii.b2a_hex(bytestring)
        return b":".join(
            (f"{int(hexString[b : b + 4], 16):x}").encode("ascii")
            for b in range(0, 32, 4)
        )

    @classmethod
    def parse(cls, line: bytes) -> _info.ProxyInfo:
        """
        Parse a bytestring as a full PROXY protocol header.

        @param line: A bytestring that represents a valid HAProxy PROXY
            protocol version 2 header.
        @type line: bytes

        @return: A L{_interfaces.IProxyInfo} containing the
            parsed data.

        @raises InvalidProxyHeader: If the bytestring does not represent a
            valid PROXY header.
        """
        prefix = line[:12]
        addrInfo = None
        with convertError(IndexError, InvalidProxyHeader):
            # Use single value slices to ensure bytestring values are returned
            # instead of int in PY3.
            versionCommand = ord(line[12:13])
            familyProto = ord(line[13:14])

        if prefix != cls.PREFIX:
            raise InvalidProxyHeader()

        version, command = versionCommand & _HIGH, versionCommand & _LOW
        if version not in cls.VERSIONS or command not in cls.COMMANDS:
            raise InvalidProxyHeader()

        if cls.COMMANDS[command] == _LOCALCOMMAND:
            return _info.ProxyInfo(line, None, None)

        family, netproto = familyProto & _HIGH, familyProto & _LOW
        with convertError(ValueError, InvalidNetworkProtocol):
            family = NetFamily.lookupByValue(family)
            netproto = NetProtocol.lookupByValue(netproto)
        if family is NetFamily.UNSPEC or netproto is NetProtocol.UNSPEC:
            return _info.ProxyInfo(line, None, None)

        addressFormat = cls.ADDRESSFORMATS[familyProto]
        addrInfo = line[16 : 16 + struct.calcsize(addressFormat)]
        if family is NetFamily.UNIX:
            with convertError(struct.error, MissingAddressData):
                source, dest = struct.unpack(addressFormat, addrInfo)
            return _info.ProxyInfo(
                line,
                address.UNIXAddress(source.rstrip(b"\x00")),
                address.UNIXAddress(dest.rstrip(b"\x00")),
            )

        addrType: Union[Literal["TCP"], Literal["UDP"]] = "TCP"
        if netproto is NetProtocol.DGRAM:
            addrType = "UDP"
        addrCls: Union[
            Type[address.IPv4Address], Type[address.IPv6Address]
        ] = address.IPv4Address
        addrParser: Callable[[bytes], bytes] = cls._bytesToIPv4
        if family is NetFamily.INET6:
            addrCls = address.IPv6Address
            addrParser = cls._bytesToIPv6

        with convertError(struct.error, MissingAddressData):
            info = struct.unpack(addressFormat, addrInfo)
            source, dest, sPort, dPort = info

        return _info.ProxyInfo(
            line,
            addrCls(addrType, addrParser(source).decode(), sPort),
            addrCls(addrType, addrParser(dest).decode(), dPort),
        )