~bzr-pqm/bzr/bzr.dev

# Copyright (C) 2006-2011 Canonical Ltd

#

# This program is free software; you can redistribute it and/or modify

# it under the terms of the GNU General Public License as published by

# the Free Software Foundation; either version 2 of the License, or

# (at your option) any later version.

#

# This program is distributed in the hope that it will be useful,

# but WITHOUT ANY WARRANTY; without even the implied warranty of

# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

# GNU General Public License for more details.

#

# You should have received a copy of the GNU General Public License

# along with this program; if not, write to the Free Software

# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA

"""Tests for smart transport"""

# all of this deals with byte strings so this is safe

from cStringIO import StringIO

import doctest

import os

import socket

import sys

import threading

import time

from testtools.matchers import DocTestMatches

import bzrlib

from bzrlib import (

        bzrdir,

        errors,

        osutils,

        tests,

        transport as _mod_transport,

        urlutils,

        )

from bzrlib.smart import (

        client,

        medium,

        message,

        protocol,

        request as _mod_request,

        server as _mod_server,

        vfs,

)

from bzrlib.tests import (

    features,

    test_smart,

    test_server,

    )

from bzrlib.transport import (

        http,

        local,

        memory,

        remote,

        ssh,

        )

def portable_socket_pair():

    """Return a pair of TCP sockets connected to each other.

    Unlike socket.socketpair, this should work on Windows.

    """

    listen_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

    listen_sock.bind(('127.0.0.1', 0))

    listen_sock.listen(1)

    client_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

    client_sock.connect(listen_sock.getsockname())

    server_sock, addr = listen_sock.accept()

    listen_sock.close()

    return server_sock, client_sock

class StringIOSSHVendor(object):

    """A SSH vendor that uses StringIO to buffer writes and answer reads."""

    def __init__(self, read_from, write_to):

        self.read_from = read_from

        self.write_to = write_to

        self.calls = []

    def connect_ssh(self, username, password, host, port, command):

        self.calls.append(('connect_ssh', username, password, host, port,

            command))

        return StringIOSSHConnection(self)

class StringIOSSHConnection(ssh.SSHConnection):

    """A SSH connection that uses StringIO to buffer writes and answer reads."""

    def __init__(self, vendor):

        self.vendor = vendor

    def close(self):

        self.vendor.calls.append(('close', ))

        self.vendor.read_from.close()

        self.vendor.write_to.close()

    def get_sock_or_pipes(self):

        return 'pipes', (self.vendor.read_from, self.vendor.write_to)

class _InvalidHostnameFeature(features.Feature):

    """Does 'non_existent.invalid' fail to resolve?

    RFC 2606 states that .invalid is reserved for invalid domain names, and

    also underscores are not a valid character in domain names.  Despite this,

    it's possible a badly misconfigured name server might decide to always

    return an address for any name, so this feature allows us to distinguish a

    broken system from a broken test.

    """

    def _probe(self):

        try:

            socket.gethostbyname('non_existent.invalid')

        except socket.gaierror:

            # The host name failed to resolve.  Good.

            return True

        else:

            return False

    def feature_name(self):

        return 'invalid hostname'

InvalidHostnameFeature = _InvalidHostnameFeature()

class SmartClientMediumTests(tests.TestCase):

    """Tests for SmartClientMedium.

    We should create a test scenario for this: we need a server module that

    construct the test-servers (like make_loopsocket_and_medium), and the list

    of SmartClientMedium classes to test.

    """

    def make_loopsocket_and_medium(self):

        """Create a loopback socket for testing, and a medium aimed at it."""

        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

        sock.bind(('127.0.0.1', 0))

        sock.listen(1)

        port = sock.getsockname()[1]

        client_medium = medium.SmartTCPClientMedium('127.0.0.1', port, 'base')

        return sock, client_medium

    def receive_bytes_on_server(self, sock, bytes):

        """Accept a connection on sock and read 3 bytes.

        The bytes are appended to the list bytes.

        :return: a Thread which is running to do the accept and recv.

        """

        def _receive_bytes_on_server():

            connection, address = sock.accept()

            bytes.append(osutils.recv_all(connection, 3))

            connection.close()

        t = threading.Thread(target=_receive_bytes_on_server)

        t.start()

        return t

    def test_construct_smart_simple_pipes_client_medium(self):

        # the SimplePipes client medium takes two pipes:

        # readable pipe, writeable pipe.

        # Constructing one should just save these and do nothing.

        # We test this by passing in None.

        client_medium = medium.SmartSimplePipesClientMedium(None, None, None)

    def test_simple_pipes_client_request_type(self):

        # SimplePipesClient should use SmartClientStreamMediumRequest's.

        client_medium = medium.SmartSimplePipesClientMedium(None, None, None)

        request = client_medium.get_request()

        self.assertIsInstance(request, medium.SmartClientStreamMediumRequest)

    def test_simple_pipes_client_get_concurrent_requests(self):

        # the simple_pipes client does not support pipelined requests:

        # but it does support serial requests: we construct one after

        # another is finished. This is a smoke test testing the integration

        # of the SmartClientStreamMediumRequest and the SmartClientStreamMedium

        # classes - as the sibling classes share this logic, they do not have

        # explicit tests for this.

        output = StringIO()

        client_medium = medium.SmartSimplePipesClientMedium(

            None, output, 'base')

        request = client_medium.get_request()

        request.finished_writing()

        request.finished_reading()

        request2 = client_medium.get_request()

        request2.finished_writing()

        request2.finished_reading()

    def test_simple_pipes_client__accept_bytes_writes_to_writable(self):

        # accept_bytes writes to the writeable pipe.

        output = StringIO()

        client_medium = medium.SmartSimplePipesClientMedium(

            None, output, 'base')

        client_medium._accept_bytes('abc')

        self.assertEqual('abc', output.getvalue())

    def test_simple_pipes_client_disconnect_does_nothing(self):

        # calling disconnect does nothing.

        input = StringIO()

        output = StringIO()

        client_medium = medium.SmartSimplePipesClientMedium(

            input, output, 'base')

        # send some bytes to ensure disconnecting after activity still does not

        # close.

        client_medium._accept_bytes('abc')

        client_medium.disconnect()

        self.assertFalse(input.closed)

        self.assertFalse(output.closed)

    def test_simple_pipes_client_accept_bytes_after_disconnect(self):

        # calling disconnect on the client does not alter the pipe that

        # accept_bytes writes to.

        input = StringIO()

        output = StringIO()

        client_medium = medium.SmartSimplePipesClientMedium(

            input, output, 'base')

        client_medium._accept_bytes('abc')

        client_medium.disconnect()

        client_medium._accept_bytes('abc')

        self.assertFalse(input.closed)

        self.assertFalse(output.closed)

        self.assertEqual('abcabc', output.getvalue())

    def test_simple_pipes_client_ignores_disconnect_when_not_connected(self):

        # Doing a disconnect on a new (and thus unconnected) SimplePipes medium

        # does nothing.

        client_medium = medium.SmartSimplePipesClientMedium(None, None, 'base')

        client_medium.disconnect()

    def test_simple_pipes_client_can_always_read(self):

        # SmartSimplePipesClientMedium is never disconnected, so read_bytes

        # always tries to read from the underlying pipe.

        input = StringIO('abcdef')

        client_medium = medium.SmartSimplePipesClientMedium(input, None, 'base')

        self.assertEqual('abc', client_medium.read_bytes(3))

        client_medium.disconnect()

        self.assertEqual('def', client_medium.read_bytes(3))

    def test_simple_pipes_client_supports__flush(self):

        # invoking _flush on a SimplePipesClient should flush the output

        # pipe. We test this by creating an output pipe that records

        # flush calls made to it.

        from StringIO import StringIO # get regular StringIO

        input = StringIO()

        output = StringIO()

        flush_calls = []

        def logging_flush(): flush_calls.append('flush')

        output.flush = logging_flush

        client_medium = medium.SmartSimplePipesClientMedium(

            input, output, 'base')

        # this call is here to ensure we only flush once, not on every

        # _accept_bytes call.

        client_medium._accept_bytes('abc')

        client_medium._flush()

        client_medium.disconnect()

        self.assertEqual(['flush'], flush_calls)

    def test_construct_smart_ssh_client_medium(self):

        # the SSH client medium takes:

        # host, port, username, password, vendor

        # Constructing one should just save these and do nothing.

        # we test this by creating a empty bound socket and constructing

        # a medium.

        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

        sock.bind(('127.0.0.1', 0))

        unopened_port = sock.getsockname()[1]

        # having vendor be invalid means that if it tries to connect via the

        # vendor it will blow up.

        ssh_params = medium.SSHParams('127.0.0.1', unopened_port, None, None)

        client_medium = medium.SmartSSHClientMedium(

            'base', ssh_params, "not a vendor")

        sock.close()

    def test_ssh_client_connects_on_first_use(self):

        # The only thing that initiates a connection from the medium is giving

        # it bytes.

        output = StringIO()

        vendor = StringIOSSHVendor(StringIO(), output)

        ssh_params = medium.SSHParams(

            'a hostname', 'a port', 'a username', 'a password', 'bzr')

        client_medium = medium.SmartSSHClientMedium('base', ssh_params, vendor)

        client_medium._accept_bytes('abc')

        self.assertEqual('abc', output.getvalue())

        self.assertEqual([('connect_ssh', 'a username', 'a password',

            'a hostname', 'a port',

            ['bzr', 'serve', '--inet', '--directory=/', '--allow-writes'])],

            vendor.calls)

    def test_ssh_client_changes_command_when_bzr_remote_path_passed(self):

        # The only thing that initiates a connection from the medium is giving

        # it bytes.

        output = StringIO()

        vendor = StringIOSSHVendor(StringIO(), output)

        ssh_params = medium.SSHParams(

            'a hostname', 'a port', 'a username', 'a password',

            bzr_remote_path='fugly')

        client_medium = medium.SmartSSHClientMedium('base', ssh_params, vendor)

        client_medium._accept_bytes('abc')

        self.assertEqual('abc', output.getvalue())

        self.assertEqual([('connect_ssh', 'a username', 'a password',

            'a hostname', 'a port',

            ['fugly', 'serve', '--inet', '--directory=/', '--allow-writes'])],

            vendor.calls)

    def test_ssh_client_disconnect_does_so(self):

        # calling disconnect should disconnect both the read_from and write_to

        # file-like object it from the ssh connection.

        input = StringIO()

        output = StringIO()

        vendor = StringIOSSHVendor(input, output)

        client_medium = medium.SmartSSHClientMedium(

            'base', medium.SSHParams('a hostname'), vendor)

        client_medium._accept_bytes('abc')

        client_medium.disconnect()

        self.assertTrue(input.closed)

        self.assertTrue(output.closed)

        self.assertEqual([

            ('connect_ssh', None, None, 'a hostname', None,

            ['bzr', 'serve', '--inet', '--directory=/', '--allow-writes']),

            ('close', ),

            ],

            vendor.calls)

    def test_ssh_client_disconnect_allows_reconnection(self):

        # calling disconnect on the client terminates the connection, but should

        # not prevent additional connections occuring.

        # we test this by initiating a second connection after doing a

        # disconnect.

        input = StringIO()

        output = StringIO()

        vendor = StringIOSSHVendor(input, output)

        client_medium = medium.SmartSSHClientMedium(

            'base', medium.SSHParams('a hostname'), vendor)

        client_medium._accept_bytes('abc')

        client_medium.disconnect()

        # the disconnect has closed output, so we need a new output for the

        # new connection to write to.

        input2 = StringIO()

        output2 = StringIO()

        vendor.read_from = input2

        vendor.write_to = output2

        client_medium._accept_bytes('abc')

        client_medium.disconnect()

        self.assertTrue(input.closed)

        self.assertTrue(output.closed)

        self.assertTrue(input2.closed)

        self.assertTrue(output2.closed)

        self.assertEqual([

            ('connect_ssh', None, None, 'a hostname', None,

            ['bzr', 'serve', '--inet', '--directory=/', '--allow-writes']),

            ('close', ),

            ('connect_ssh', None, None, 'a hostname', None,

            ['bzr', 'serve', '--inet', '--directory=/', '--allow-writes']),

            ('close', ),

            ],

            vendor.calls)

    def test_ssh_client_ignores_disconnect_when_not_connected(self):

        # Doing a disconnect on a new (and thus unconnected) SSH medium

        # does not fail.  It's ok to disconnect an unconnected medium.

        client_medium = medium.SmartSSHClientMedium(

            'base', medium.SSHParams(None))

        client_medium.disconnect()

    def test_ssh_client_raises_on_read_when_not_connected(self):

        # Doing a read on a new (and thus unconnected) SSH medium raises

        # MediumNotConnected.

        client_medium = medium.SmartSSHClientMedium(

            'base', medium.SSHParams(None))

        self.assertRaises(errors.MediumNotConnected, client_medium.read_bytes,

                          0)

        self.assertRaises(errors.MediumNotConnected, client_medium.read_bytes,

                          1)

    def test_ssh_client_supports__flush(self):

        # invoking _flush on a SSHClientMedium should flush the output

        # pipe. We test this by creating an output pipe that records

        # flush calls made to it.

        from StringIO import StringIO # get regular StringIO

        input = StringIO()

        output = StringIO()

        flush_calls = []

        def logging_flush(): flush_calls.append('flush')

        output.flush = logging_flush

        vendor = StringIOSSHVendor(input, output)

        client_medium = medium.SmartSSHClientMedium(

            'base', medium.SSHParams('a hostname'), vendor=vendor)

        # this call is here to ensure we only flush once, not on every

        # _accept_bytes call.

        client_medium._accept_bytes('abc')

        client_medium._flush()

        client_medium.disconnect()

        self.assertEqual(['flush'], flush_calls)

    def test_construct_smart_tcp_client_medium(self):

        # the TCP client medium takes a host and a port.  Constructing it won't

        # connect to anything.

        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

        sock.bind(('127.0.0.1', 0))

        unopened_port = sock.getsockname()[1]

        client_medium = medium.SmartTCPClientMedium(

            '127.0.0.1', unopened_port, 'base')

        sock.close()

    def test_tcp_client_connects_on_first_use(self):

        # The only thing that initiates a connection from the medium is giving

        # it bytes.

        sock, medium = self.make_loopsocket_and_medium()

        bytes = []

        t = self.receive_bytes_on_server(sock, bytes)

        medium.accept_bytes('abc')

        t.join()

        sock.close()

        self.assertEqual(['abc'], bytes)

    def test_tcp_client_disconnect_does_so(self):

        # calling disconnect on the client terminates the connection.

        # we test this by forcing a short read during a socket.MSG_WAITALL

        # call: write 2 bytes, try to read 3, and then the client disconnects.

        sock, medium = self.make_loopsocket_and_medium()

        bytes = []

        t = self.receive_bytes_on_server(sock, bytes)

        medium.accept_bytes('ab')

        medium.disconnect()

        t.join()

        sock.close()

        self.assertEqual(['ab'], bytes)

        # now disconnect again: this should not do anything, if disconnection

        # really did disconnect.

        medium.disconnect()

    def test_tcp_client_ignores_disconnect_when_not_connected(self):

        # Doing a disconnect on a new (and thus unconnected) TCP medium

        # does not fail.  It's ok to disconnect an unconnected medium.

        client_medium = medium.SmartTCPClientMedium(None, None, None)

        client_medium.disconnect()

    def test_tcp_client_raises_on_read_when_not_connected(self):

        # Doing a read on a new (and thus unconnected) TCP medium raises

        # MediumNotConnected.

        client_medium = medium.SmartTCPClientMedium(None, None, None)

        self.assertRaises(errors.MediumNotConnected, client_medium.read_bytes, 0)

        self.assertRaises(errors.MediumNotConnected, client_medium.read_bytes, 1)

    def test_tcp_client_supports__flush(self):

        # invoking _flush on a TCPClientMedium should do something useful.

        # RBC 20060922 not sure how to test/tell in this case.

        sock, medium = self.make_loopsocket_and_medium()

        bytes = []

        t = self.receive_bytes_on_server(sock, bytes)

        # try with nothing buffered

        medium._flush()

        medium._accept_bytes('ab')

        # and with something sent.

        medium._flush()

        medium.disconnect()

        t.join()

        sock.close()

        self.assertEqual(['ab'], bytes)

        # now disconnect again : this should not do anything, if disconnection

        # really did disconnect.

        medium.disconnect()

    def test_tcp_client_host_unknown_connection_error(self):

        self.requireFeature(InvalidHostnameFeature)

        client_medium = medium.SmartTCPClientMedium(

            'non_existent.invalid', 4155, 'base')

        self.assertRaises(

            errors.ConnectionError, client_medium._ensure_connection)

class TestSmartClientStreamMediumRequest(tests.TestCase):

    """Tests the for SmartClientStreamMediumRequest.

    SmartClientStreamMediumRequest is a helper for the three stream based

    mediums: TCP, SSH, SimplePipes, so we only test it once, and then test that

    those three mediums implement the interface it expects.

    """

    def test_accept_bytes_after_finished_writing_errors(self):

        # calling accept_bytes after calling finished_writing raises

        # WritingCompleted to prevent bad assumptions on stream environments

        # breaking the needs of message-based environments.

        output = StringIO()

        client_medium = medium.SmartSimplePipesClientMedium(

            None, output, 'base')

        request = medium.SmartClientStreamMediumRequest(client_medium)

        request.finished_writing()

        self.assertRaises(errors.WritingCompleted, request.accept_bytes, None)

    def test_accept_bytes(self):

        # accept bytes should invoke _accept_bytes on the stream medium.

        # we test this by using the SimplePipes medium - the most trivial one

        # and checking that the pipes get the data.

        input = StringIO()

        output = StringIO()

        client_medium = medium.SmartSimplePipesClientMedium(

            input, output, 'base')

        request = medium.SmartClientStreamMediumRequest(client_medium)

        request.accept_bytes('123')

        request.finished_writing()

        request.finished_reading()

        self.assertEqual('', input.getvalue())

        self.assertEqual('123', output.getvalue())

    def test_construct_sets_stream_request(self):

        # constructing a SmartClientStreamMediumRequest on a StreamMedium sets

        # the current request to the new SmartClientStreamMediumRequest

        output = StringIO()

        client_medium = medium.SmartSimplePipesClientMedium(

            None, output, 'base')

        request = medium.SmartClientStreamMediumRequest(client_medium)

        self.assertIs(client_medium._current_request, request)

    def test_construct_while_another_request_active_throws(self):

        # constructing a SmartClientStreamMediumRequest on a StreamMedium with

        # a non-None _current_request raises TooManyConcurrentRequests.

        output = StringIO()

        client_medium = medium.SmartSimplePipesClientMedium(

            None, output, 'base')

        client_medium._current_request = "a"

        self.assertRaises(errors.TooManyConcurrentRequests,

            medium.SmartClientStreamMediumRequest, client_medium)

    def test_finished_read_clears_current_request(self):

        # calling finished_reading clears the current request from the requests

        # medium

        output = StringIO()

        client_medium = medium.SmartSimplePipesClientMedium(

            None, output, 'base')

        request = medium.SmartClientStreamMediumRequest(client_medium)

        request.finished_writing()

        request.finished_reading()

        self.assertEqual(None, client_medium._current_request)

    def test_finished_read_before_finished_write_errors(self):

        # calling finished_reading before calling finished_writing triggers a

        # WritingNotComplete error.

        client_medium = medium.SmartSimplePipesClientMedium(

            None, None, 'base')

        request = medium.SmartClientStreamMediumRequest(client_medium)

        self.assertRaises(errors.WritingNotComplete, request.finished_reading)

    def test_read_bytes(self):

        # read bytes should invoke _read_bytes on the stream medium.

        # we test this by using the SimplePipes medium - the most trivial one

        # and checking that the data is supplied. Its possible that a

        # faulty implementation could poke at the pipe variables them selves,

        # but we trust that this will be caught as it will break the integration

        # smoke tests.

        input = StringIO('321')

        output = StringIO()

        client_medium = medium.SmartSimplePipesClientMedium(

            input, output, 'base')

        request = medium.SmartClientStreamMediumRequest(client_medium)

        request.finished_writing()

        self.assertEqual('321', request.read_bytes(3))

        request.finished_reading()

        self.assertEqual('', input.read())

        self.assertEqual('', output.getvalue())

    def test_read_bytes_before_finished_write_errors(self):

        # calling read_bytes before calling finished_writing triggers a

        # WritingNotComplete error because the Smart protocol is designed to be

        # compatible with strict message based protocols like HTTP where the

        # request cannot be submitted until the writing has completed.

        client_medium = medium.SmartSimplePipesClientMedium(None, None, 'base')

        request = medium.SmartClientStreamMediumRequest(client_medium)

        self.assertRaises(errors.WritingNotComplete, request.read_bytes, None)

    def test_read_bytes_after_finished_reading_errors(self):

        # calling read_bytes after calling finished_reading raises

        # ReadingCompleted to prevent bad assumptions on stream environments

        # breaking the needs of message-based environments.

        output = StringIO()

        client_medium = medium.SmartSimplePipesClientMedium(

            None, output, 'base')

        request = medium.SmartClientStreamMediumRequest(client_medium)

        request.finished_writing()

        request.finished_reading()

        self.assertRaises(errors.ReadingCompleted, request.read_bytes, None)

class RemoteTransportTests(test_smart.TestCaseWithSmartMedium):

    def test_plausible_url(self):

        self.assert_(self.get_url().startswith('bzr://'))

    def test_probe_transport(self):

        t = self.get_transport()

        self.assertIsInstance(t, remote.RemoteTransport)

    def test_get_medium_from_transport(self):

        """Remote transport has a medium always, which it can return."""

        t = self.get_transport()

        client_medium = t.get_smart_medium()

        self.assertIsInstance(client_medium, medium.SmartClientMedium)

class ErrorRaisingProtocol(object):

    def __init__(self, exception):

        self.exception = exception

    def next_read_size(self):

        raise self.exception

class SampleRequest(object):

    def __init__(self, expected_bytes):

        self.accepted_bytes = ''

        self._finished_reading = False

        self.expected_bytes = expected_bytes

        self.unused_data = ''

    def accept_bytes(self, bytes):

        self.accepted_bytes += bytes

        if self.accepted_bytes.startswith(self.expected_bytes):

            self._finished_reading = True

            self.unused_data = self.accepted_bytes[len(self.expected_bytes):]

    def next_read_size(self):

        if self._finished_reading:

            return 0

        else:

            return 1

class TestSmartServerStreamMedium(tests.TestCase):

    def setUp(self):

        super(TestSmartServerStreamMedium, self).setUp()

        self.overrideEnv('BZR_NO_SMART_VFS', None)

    def create_pipe_medium(self, to_server, from_server, transport,

                           timeout=4.0):

        """Create a new SmartServerPipeStreamMedium."""

        return medium.SmartServerPipeStreamMedium(to_server, from_server,

            transport, timeout=timeout)

    def create_pipe_context(self, to_server_bytes, transport):

        """Create a SmartServerSocketStreamMedium.

        This differes from create_pipe_medium, in that we initialize the

        request that is sent to the server, and return the StringIO class that

        will hold the response.

        """

        to_server = StringIO(to_server_bytes)

        from_server = StringIO()

        m = self.create_pipe_medium(to_server, from_server, transport)

        return m, from_server

    def create_socket_medium(self, server_sock, transport, timeout=4.0):

        """Initialize a new medium.SmartServerSocketStreamMedium."""

        return medium.SmartServerSocketStreamMedium(server_sock, transport,

            timeout=timeout)