#!/usr/bin/env python
"""Tests of non-blocking connection states.

The goal is to build a reliable procedure for non-blocking connects.
These are tricky; see <http://cr.yp.to/docs/connect.html>.

Scott Lamb <slamb@slamb.org>
"""

from util import *
import select

__version__ = '$Id$'

class NonblockingConnectTests(unittest.TestCase):
    def setUp(self):
        """Setup for connecting or success tests."""
        self.listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.listener.listen(1)
        self.client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.address = TestAddress()

        # Print the listener port number for use with tcpdump/Ethereal.
        sys.stderr.write('(%d) ' % (self.listener.getsockname()[1]))
        sys.stderr.flush()

        # Ensure the connection will not succeed immediately.
        self.address.drop(self.listener)

        # Put it in the connecting state.
        self.client.setblocking(0)
        try:
            self.client.connect(self.listener.getsockname())
        except socket.error, (err, errstr):
            self.assertErrorEquals(errno.EINPROGRESS, err)
        else:
            self.fail('Connection should not have succeeded immediately')

    def tearDown(self):
        self.client = None
        self.listener = None
        self.address = None

    def _connect(self):
        """Makes the connection succeed."""
        self.address = None
        l = [self.client]
        (rlist, wlist, xlist) = select.select(l, l, l)
        self.assertTrue(self.client in rlist or self.client in wlist)
        self.assertEquals(xlist, [])

        # Ensure a second connection attempt will fail.
        self.server, _ = self.listener.accept()
        self.listener = None

    def _fail(self):
        """Makes the connection fail with ECONNREFUSED."""
        self.listener = None # shut it down
        self.address.open()
        l = [self.client]
        (rlist, wlist, xlist) = select.select(l, l, l)
        self.assertTrue(self.client in rlist or self.client in wlist)
        self.assertEquals(xlist, [])

    @firewall_test
    def testGetpeernameWhileConnecting(self):
        """getpeername while connecting yields ENOTCONN."""
        try:
            rv = self.client.getpeername()
        except socket.error, (err, errstr):
            self.assertErrorEquals(errno.ENOTCONN, err)
        else:
            self.fail("getpeername while connecting returned %s (listener is %s)"
                      % (rv, self.listener.getsockname()))

    @firewall_test
    def testGetErrorWhileConnecting(self):
        """getsockopt(SO_ERROR) while connecting yields EINPROGRESS."""
        self.assertErrorEquals(errno.EINPROGRESS,
            self.client.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR))

    @firewall_test
    def testReadZeroWhileConnecting(self):
        """read(0) while connecting yields EINPROGRESS."""
        try:
            self.client.recv(0)
        except socket.error, (err, errstr):
            self.assertErrorEquals(errno.EINPROGRESS, err)
        else:
            self.fail('No error reading while connecting')

    @firewall_test
    def testReadOneWhileConnecting(self):
        """read(1) while connecting yields EINPROGRESS."""
        try:
            self.client.recv(1)
        except socket.error, (err, errstr):
            self.assertErrorEquals(errno.EINPROGRESS, err)
        else:
            self.fail('No error reading while connecting')

    @firewall_test
    def testWriteZeroWhileConnecting(self):
        """write(0) while connecting yields ENOTCONN."""
        try:
            self.client.send('')
        except socket.error, (err, errstr):
            self.assertErrorEquals(errno.ENOTCONN, err)
        else:
            self.fail('No error writing while connecting')

    @firewall_test
    def testWriteOneWhileConnecting(self):
        """write(1) while connecting yields ENOTCONN."""
        try:
            self.client.send('x')
        except socket.error, (err, errstr):
            self.assertErrorEquals(errno.ENOTCONN, err)
        else:
            self.fail('No error writing while connecting')

    @firewall_test
    def testConnectWhileConnecting(self):
        """connect while connecting yields ENOTCONN."""
        try:
            self.client.connect(self.listener.getpeername())
        except socket.error, (err, errstr):
            self.assertErrorEquals(errno.ENOTCONN, err)
        else:
            self.fail('No error connecting while connecting')

    @firewall_test
    def testGetpeernameAfterConnected(self):
        """getpeername after connected yields correct address."""
        self.listener_addr = self.listener.getsockname()
        self._connect()

        # Only check ports, since listener_addr[0] will be 0.0.0.0.
        self.assertEquals(self.listener_addr[1], self.client.getpeername()[1])

    @firewall_test
    def testGetErrorAfterConnected(self):
        """getsockopt(SO_ERROR) after connected yields 0."""
        self._connect()
        self.assertErrorEquals(0,
            self.client.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR))

    @firewall_test
    def testReadZeroAfterConnected(self):
        """read(0) after connected yields EWOULDBLOCK."""
        self._connect()
        try:
            result = self.client.recv(0)
        except socket.error, (err, errstr):
            self.assertErrorEquals(errno.EWOULDBLOCK, err)
        else:
            self.fail('Nope. Instead it returned %r' % result)

    @firewall_test
    def testReadZeroWithData(self):
        """read(0) with data yields 0."""
        self._connect()
        self.assertEquals(1, self.server.send('x'))
        self.assertEquals('', self.client.recv(0))

    @firewall_test
    def testReadOneAfterConnected(self):
        """read(1) after connected yields EWOULDBLOCK."""
        self._connect()
        try:
            self.client.recv(1)
        except socket.error, (err, errstr):
            self.assertErrorEquals(errno.EWOULDBLOCK, err)
        else:
            self.fail("No error reading when nothing's been sent")

    @firewall_test
    def testReadOneWithData(self):
        """read(1) with data yields 1."""
        self._connect()
        self.assertEquals(1, self.server.send('x'))
        self.assertEquals('x', self.client.recv(1))

    @firewall_test
    def testWriteZeroAfterConnected(self):
        """write(0) after connected succeeds."""
        self._connect()
        self.assertEquals(0, self.client.send(''))

    @firewall_test
    def testWriteOneAfterConnected(self):
        """write(1) after connected succeeds."""
        self._connect()
        self.assertEquals(1, self.client.send('x'))

    @firewall_test
    def testConnectAfterConnected(self):
        """connect after connected yields EISCONN or succeeds."""
        listener_addr = self.listener.getsockname()
        self._connect()
        try:
            self.client.connect(listener_addr)
        except socket.error, (err, errstr):
            self.assertErrorEquals(errno.EISCONN, err) # okay; Darwin 8.2.0
        else:
            pass # okay; Linux kernel-2.6.9-5.EL

    @firewall_test
    def testGetpeernameAfterRefused(self):
        """getpeername after refused yields ENOTCONN."""
        self._fail()
        try:
            self.client.getpeername()
        except socket.error, (err, errstr):
            self.assertErrorEquals(errno.ENOTCONN, err)

    @firewall_test
    def testReadOneAfterRefused(self):
        """read(1) after refused yields ECONNREFUSED."""
        self._fail()
        try:
            self.client.recv(1)
        except socket.error, (err, errstr):
            self.assertErrorEquals(errno.ECONNREFUSED, err)

    @firewall_test
    def testReadZeroAfterRefused(self):
        """read(0) after refused yields ECONNREFUSED."""
        self._fail()
        try:
            self.client.recv(0)
        except socket.error, (err, errstr):
            self.assertErrorEquals(errno.ECONNREFUSED, err)

    @firewall_test
    def testWriteZeroAfterRefused(self):
        """write(0) after refused yields ECONNREFUSED."""
        self._fail()
        try:
            self.client.send('')
        except socket.error, (err, errstr):
            self.assertErrorEquals(errno.ECONNREFUSED, err)

    @firewall_test
    def testWriteOneAfterRefused(self):
        """write(1) after refused yields ECONNREFUSED."""
        self._fail()
        try:
            self.client.send('x')
        except socket.error, (err, errstr):
            self.assertErrorEquals(errno.ECONNREFUSED, err)

    @firewall_test
    def testGetErrorAfterRefused(self):
        """getsockopt(SO_ERROR) after refused yields ECONNREFUSED."""
        self._fail()
        self.assertErrorEquals(errno.ECONNREFUSED,
            self.client.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR))

    @firewall_test
    def testConnectAfterRefused(self):
        """connect after refused yields ECONNREFUSED."""
        listener_addr = self.listener.getsockname()
        self._fail()

        # DJB claims this sometimes makes another connection.
        # Fail if so.
        self.listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.listener.bind(listener_addr)
        self.listener.listen(1)

        try:
            self.client.connect(listener_addr)
        except socket.error, (err, errstr):
            self.assertErrorEquals(errno.ECONNREFUSED, err)
        else:
            self.fail("Made another connection attempt!")

if __name__ == '__main__':
    unittest.main()