Implement Haproxy PROXY protocol v1 and v2

Signed-off-by: Frank Villaro-Dixon <frank.villaro@infomaniak.com>

Author: pascal-dubois

pyftpdlib/handlers.py

@@ -1234,6 +1234,7 @@ def __init__(self, conn, server, ioloop=None):
         self.remote_ip = ""
         self.remote_port = ""
         self.started = time.time()
+        self.proxy_proto_obj = self.server.proxy_proto_obj
 
         # private session attributes
         self._last_response = ""
@@ -1294,6 +1295,15 @@ def __init__(self, conn, server, ioloop=None):
                 self.handle_error()
             return
         else:
+            if self.server.proxy_proto_enabled and self.proxy_proto_obj and \
+                    (self.proxy_proto_obj.trusted or self.server.proxy_proto_allow_untrusted):
+                # Override the socket values if PROXY protocol data is
+                # available and comes from a trusted proxy
+                if self.proxy_proto_obj.remote_ip:
+                    self.remote_ip = self.proxy_proto_obj.remote_ip
+                if self.proxy_proto_obj.remote_port:
+                    self.remote_port = self.proxy_proto_obj.remote_port
+
             self.log("FTP session opened (connect)")
 
         # try to handle urgent data inline

pyftpdlib/ioloop.py

@@ -76,6 +76,7 @@ def handle_accepted(self, sock, addr):
 from .log import debug
 from .log import is_logging_configured
 from .log import logger
+from .proxy_proto import ProxyProtocol
 
 
 timer = getattr(time, 'monotonic', time.time)
@@ -427,9 +428,9 @@ def unregister(self, fd):
             del self.socket_map[fd]
         except KeyError:
             debug("call: unregister(); fd was no longer in socket_map", self)
-        for l in (self._r, self._w):
+        for loop in (self._r, self._w):
             try:
-                l.remove(fd)
+                loop.remove(fd)
             except ValueError:
                 pass
 
@@ -977,8 +978,29 @@ def add_channel(self, map=None, events=None):
 class Acceptor(AsyncChat):
     """Same as base AsyncChat and supposed to be used to
     accept new connections.
+
+    All relevant PROXY protocol information is stored in class attributes
+    described below.
+
+     - (bool) proxy_proto_enabled:
+        enable the use of PROXY protocol (defaults to False).
+
+     - (list) proxy_proto_trusted_nets:
+        the IP networks of the proxies you want to trust. Use a /32 (or /128) network
+        mask if you want to declare a single IP (defaults to []).
+
+     - (bool) proxy_proto_allow_untrusted:
+        whether or not to parse untrusted proxies headers (defaults to True).
+
+     - (instance) proxy_proto_obj:
+        the ProxyProtocol instance populated with header's information
     """
 
+    proxy_proto_enabled = False
+    proxy_proto_trusted_nets = []
+    proxy_proto_allow_untrusted = True
+    proxy_proto_obj = None
+
     def add_channel(self, map=None, events=None):
         AsyncChat.add_channel(self, map=map, events=self.ioloop.READ)
 
@@ -1045,6 +1067,21 @@ def handle_accept(self):
                 debug("call: handle_accept(); accept() returned ECONNABORTED",
                       self)
         else:
+            if self.proxy_proto_enabled:
+                # Retrieve a populated PROXY protocol object. If no exception
+                # is raised the returned object is considered valid
+                try:
+                    ProxyProtocol.trusted_networks = self.proxy_proto_trusted_nets
+                    ProxyProtocol.allow_untrusted = self.proxy_proto_allow_untrusted
+                    self.proxy_proto_obj = ProxyProtocol.create(sock)
+                except Exception as e:
+                    logger.error("proxy: {}".format(e))
+                    return
+
+                if self.proxy_proto_obj.trusted or self.proxy_proto_allow_untrusted:
+                    # Could result in a (None, None) tuple
+                    addr = (self.proxy_proto_obj.remote_ip, self.proxy_proto_obj.remote_port)
+
             # sometimes addr == None instead of (ip, port) (see issue 104)
             if addr is not None:
                 self.handle_accepted(sock, addr)