Fix proxy protocol usage

setup.py

@@ -23,7 +23,7 @@
 
 
 setup(name='python-slimta',
-      version='4.0.2',
+      version='4.0.3',
       author='Ian Good',
       author_email='icgood@gmail.com',
       description='Lightweight, asynchronous SMTP libraries.',

slimta/edge/smtp.py

@@ -248,7 +248,7 @@ def handle(self, socket, address):
         smtp_server = None
         try:
             handlers = SmtpSession(address, self.validator_class, self.handoff)
-            smtp_server = Server(socket, handlers, self.auth,
+            smtp_server = Server(socket, handlers, address, self.auth,
                                  self.context, self.tls_immediately,
                                  command_timeout=self.command_timeout,
                                  data_timeout=self.data_timeout)

slimta/smtp/io.py

@@ -72,6 +72,8 @@ def close(self):
         if self.encrypted:
             try:
                 self.socket.unwrap()
+            except ValueError:
+                pass
             except SSLWantReadError:
                 pass
             except socket_error as e:

slimta/smtp/server.py

@@ -75,6 +75,7 @@ class Server(object):
                      corresponding SMTP commands are received. These methods
                      can modify the |Reply| before the command response is
                      sent.
+    :param address: The address of the connected client.
     :param auth: If True, enable authentication with default mechanisms. May
                  also be given as a list of SASL mechanism names to support,
                  e.g. ``['PLAIN', 'LOGIN', 'CRAM-MD5']``.
@@ -90,13 +91,13 @@ class Server(object):
 
     """
 
-    def __init__(self, socket, handlers, auth=False,
+    def __init__(self, socket, handlers, address=None, auth=False,
                  context=None, tls_immediately=False,
                  command_timeout=None, data_timeout=None):
         self.handlers = handlers
         self.extensions = Extensions()
 
-        self.io = IO(socket)
+        self.io = IO(socket, address)
 
         self.bannered = False
         self.have_mailfrom = None

slimta/util/proxyproto.py

@@ -32,6 +32,7 @@
 
 from gevent import socket
 
+from slimta.edge import EdgeServer
 from slimta.logging import getSocketLogger
 
 __all__ = ['ProxyProtocol', 'ProxyProtocolV1', 'ProxyProtocolV2']
@@ -157,6 +158,22 @@ def process_pp_v1(cls, sock, initial):
         log.recv(sock, line)
         return cls.parse_pp_line(line)
 
+    @classmethod
+    def mixin(cls, edge):
+        """Dynamically mix-in the :class:`ProxyProtocolV1` class as a base of
+        the given edge object. Use with caution.
+
+        :param edge: the edge object to use proxy protocol on.
+        :type edge: :class:`~slimta.edge.EdgeServer`
+        :raises: ValueError
+
+        """
+        if not isinstance(edge, EdgeServer):
+            raise ValueError(edge)
+        old_class = edge.__class__
+        new_class_name = cls.__name__ + old_class.__name__
+        edge.__class__ = type(new_class_name, (cls, old_class), {})
+
     def handle(self, sock, addr):
         """Intercepts calls to :meth:`~slimta.edge.EdgeServer.handle`, reads
         the proxy protocol header, and then resumes the original call.
@@ -209,32 +226,32 @@ def __read_pp_data(cls, sock, length, initial):
     def __parse_pp_data(cls, data):
         assert data[0:12] == b'\r\n\r\n\x00\r\nQUIT\n', \
             'Invalid proxy protocol v2 signature'
-        assert data[13] & 0xf0 == 0x20, 'Invalid proxy protocol version'
+        assert data[12] & 0xf0 == 0x20, 'Invalid proxy protocol version'
         command = cls.__commands.get(data[12] & 0x0f)
         family = cls.__families.get(data[13] & 0xf0)
         protocol = cls.__protocols.get(data[13] & 0x0f)
-        addr_len = struct.unpack('<H', data[14:16])[0]
+        addr_len = struct.unpack('!H', data[14:16])[0]
         return command, family, protocol, addr_len
 
     @classmethod
     def __parse_pp_addresses(cls, family, addr_data):
         if family == socket.AF_INET:
             src_ip, dst_ip, src_port, dst_port = \
-                struct.unpack('<4s4sHH', addr_data)
+                struct.unpack('!4s4sHH', addr_data)
             src_addr = (socket.inet_ntop(family, src_ip), src_port)
             dst_addr = (socket.inet_ntop(family, dst_ip), dst_port)
             return src_addr, dst_addr
         elif family == socket.AF_INET6:
             src_ip, dst_ip, src_port, dst_port = \
-                struct.unpack('<16s16sHH', addr_data)
+                struct.unpack('!16s16sHH', addr_data)
             src_addr = (socket.inet_ntop(family, src_ip), src_port)
             dst_addr = (socket.inet_ntop(family, dst_ip), dst_port)
             return src_addr, dst_addr
         elif family == socket.AF_UNIX:
-            src_addr, dst_addr = struct.unpack('<108s108s', addr_data)
+            src_addr, dst_addr = struct.unpack('!108s108s', addr_data)
             return src_addr.rstrip(b'\x00'), dst_addr.rstrip(b'\x00')
         else:
-            return unknown_pp_source_address,  unknown_pp_dest_address
+            return unknown_pp_source_address, unknown_pp_dest_address
 
     @classmethod
     def process_pp_v2(cls, sock, initial):
@@ -249,6 +266,22 @@ def process_pp_v2(cls, sock, initial):
         except struct.error:
             raise AssertionError('Invalid proxy protocol data')
 
+    @classmethod
+    def mixin(cls, edge):
+        """Dynamically mix-in the :class:`ProxyProtocolV2` class as a base of
+        the given edge object. Use with caution.
+
+        :param edge: the edge object to use proxy protocol on.
+        :type edge: :class:`~slimta.edge.EdgeServer`
+        :raises: ValueError
+
+        """
+        if not isinstance(edge, EdgeServer):
+            raise ValueError(edge)
+        old_class = edge.__class__
+        new_class_name = cls.__name__ + old_class.__name__
+        edge.__class__ = type(new_class_name, (cls, old_class), {})
+
     def handle(self, sock, addr):
         """Intercepts calls to :meth:`~slimta.edge.EdgeServer.handle`, reads
         the proxy protocol header, and then resumes the original call.
@@ -291,6 +324,22 @@ def __read_pp_initial(cls, sock):
             read = memoryview(buf)[0:len(read)+read_n].tobytes()
         return read
 
+    @classmethod
+    def mixin(cls, edge):
+        """Dynamically mix-in the :class:`ProxyProtocol` class as a base of the
+        given edge object. Use with caution.
+
+        :param edge: the edge object to use proxy protocol on.
+        :type edge: :class:`~slimta.edge.EdgeServer`
+        :raises: ValueError
+
+        """
+        if not isinstance(edge, EdgeServer):
+            raise ValueError(edge)
+        old_class = edge.__class__
+        new_class_name = cls.__name__ + old_class.__name__
+        edge.__class__ = type(new_class_name, (cls, old_class), {})
+
     def handle(self, sock, addr):
         """Intercepts calls to :meth:`~slimta.edge.EdgeServer.handle`, reads
         the proxy protocol header, and then resumes the original call.

test/test_slimta_util_proxyproto.py

@@ -202,11 +202,11 @@ def test_parse_pp_data(self):
         self.assertEqual(61680, addr_len)
 
     def test_parse_pp_addresses(self):
-        data = bytearray(b'\x00\x00\x00\x00\x7f\x00\x00\x01\x00\x00\x19\x00')
+        data = bytearray(b'\x00\x00\x00\x00\x7f\x00\x00\x01\x00\x00\x00\x19')
         src_addr, dst_addr = self.pp._ProxyProtocolV2__parse_pp_addresses(socket.AF_INET, data)
         self.assertEqual(('0.0.0.0', 0), src_addr)
         self.assertEqual(('127.0.0.1', 25), dst_addr)
-        data = bytearray((b'\x00'*15 + b'\x01')*2 + b'\x00\x00\x19\x00')
+        data = bytearray((b'\x00'*15 + b'\x01')*2 + b'\x00\x00\x00\x19')
         src_addr, dst_addr = self.pp._ProxyProtocolV2__parse_pp_addresses(socket.AF_INET6, data)
         self.assertEqual(('::1', 0), src_addr)
         self.assertEqual(('::1', 25), dst_addr)