Merge pull request #105 from slimta/outbound-ipv6-flag

Add socket creator for IPv4 only connections

slimta/util/__init__.py

@@ -26,7 +26,10 @@
 
 from __future__ import absolute_import
 
-__all__ = ['validate_tls']
+from gevent import socket
+
+__all__ = ['validate_tls', 'build_ipv4_socket_creator',
+           'create_connection_ipv4']
 
 
 def validate_tls(tls, **overrides):
@@ -55,4 +58,46 @@ def validate_tls(tls, **overrides):
     return tls_copy
 
 
+def build_ipv4_socket_creator(only_ports=None):
+    """Returns a function that will act like
+    :py:func:`socket.create_connection` but only using IPv4 addresses. This
+    function can be used as the ``socket_creator`` argument to some classes
+    like :class:`~slimta.relay.smtp.mx.MxSmtpRelay`.
+
+    :param only_ports: If given, can be a list to limit which ports are
+                       restricted to IPv4. Connections to all other ports may
+                       be IPv6.
+
+    """
+    def socket_creator(*args, **kwargs):
+        return create_connection_ipv4(*args, only_ports=only_ports, **kwargs)
+    return socket_creator
+
+
+def create_connection_ipv4(address, timeout=None, source_address=None,
+                           only_ports=None):
+    """Attempts to mimick to :py:func:`socket.create_connection`, but
+    connections are only made to IPv4 addresses.
+
+    :param only_ports: If given, can be a list to limit which ports are
+                       restricted to IPv4. Connections to all other ports may
+                       be IPv6.
+
+    """
+    host, port = address
+    if only_ports and port not in only_ports:
+        return socket.create_connection(address, timeout, source_address)
+    last_exc = None
+    for res in socket.getaddrinfo(host, port, socket.AF_INET):
+        _, _, _, _, sockaddr = res
+        try:
+            return socket.create_connection(sockaddr,  timeout, source_address)
+        except socket.error as exc:
+            last_exc = exc
+    if last_exc is not None:
+        raise last_exc
+    else:
+        raise socket.error('getaddrinfo returns an empty list')
+
+
 # vim:et:fdm=marker:sts=4:sw=4:ts=4

test/test_slimta_util.py

@@ -0,0 +1,44 @@
+import unittest2 as unittest
+from mox3.mox import MoxTestBase, IgnoreArg
+from gevent import socket
+
+from slimta import util
+
+
+class TestIPv4SocketCreator(MoxTestBase):
+
+    def setUp(self):
+        super(TestIPv4SocketCreator, self).setUp()
+        self.mox.StubOutWithMock(socket, 'create_connection')
+        self.mox.StubOutWithMock(socket, 'getaddrinfo')
+        self.getaddrinfo = self.mox.CreateMock(socket.getaddrinfo)
+        self.socket_creator = util.build_ipv4_socket_creator([25])
+
+    def test_other_port(self):
+        socket.create_connection(('host', 443), 'timeout', 'source').AndReturn('socket')
+        self.mox.ReplayAll()
+        ret = self.socket_creator(('host', 443), 'timeout', 'source')
+        self.assertEqual('socket', ret)
+
+    def test_successful(self):
+        socket.getaddrinfo('host', 25, socket.AF_INET).AndReturn([(None, None, None, None, 'sockaddr')])
+        socket.create_connection('sockaddr', IgnoreArg(), IgnoreArg()).AndReturn('socket')
+        self.mox.ReplayAll()
+        ret = self.socket_creator(('host', 25), 'timeout', 'source')
+        self.assertEqual('socket', ret)
+
+    def test_error(self):
+        socket.getaddrinfo('host', 25, socket.AF_INET).AndReturn([(None, None, None, None, 'sockaddr')])
+        socket.create_connection('sockaddr', IgnoreArg(), IgnoreArg()).AndRaise(socket.error('error'))
+        self.mox.ReplayAll()
+        with self.assertRaises(socket.error):
+            self.socket_creator(('host', 25), 'timeout', 'source')
+
+    def test_no_addresses(self):
+        socket.getaddrinfo('host', 25, socket.AF_INET).AndReturn([])
+        self.mox.ReplayAll()
+        with self.assertRaises(socket.error):
+            self.socket_creator(('host', 25), 'timeout', 'source')
+
+
+# vim:et:fdm=marker:sts=4:sw=4:ts=4