/brz/remove-bazaar

# Copyright (C) 2006-2010 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

"""Wire-level encoding and decoding of requests and responses for the smart

client and server.

"""

from __future__ import absolute_import

import collections

import struct

import sys

try:

    import _thread

except ImportError:

    import thread as _thread

import time

import breezy

from .. import (

    debug,

    errors,

    osutils,

    )

from ..sixish import (

    BytesIO,

    reraise,

)

from . import message, request

from ..trace import log_exception_quietly, mutter

from ..bencode import bdecode_as_tuple, bencode

# Protocol version strings.  These are sent as prefixes of bzr requests and

# responses to identify the protocol version being used. (There are no version

# one strings because that version doesn't send any).

REQUEST_VERSION_TWO = 'bzr request 2\n'

RESPONSE_VERSION_TWO = 'bzr response 2\n'

MESSAGE_VERSION_THREE = 'bzr message 3 (bzr 1.6)\n'

RESPONSE_VERSION_THREE = REQUEST_VERSION_THREE = MESSAGE_VERSION_THREE

def _recv_tuple(from_file):

    req_line = from_file.readline()

    return _decode_tuple(req_line)

def _decode_tuple(req_line):

    if req_line is None or req_line == '':

        return None

    if req_line[-1] != '\n':

        raise errors.SmartProtocolError("request %r not terminated" % req_line)

    return tuple(req_line[:-1].split('\x01'))

def _encode_tuple(args):

    """Encode the tuple args to a bytestream."""

    joined = '\x01'.join(args) + '\n'

    if isinstance(joined, unicode):

        # XXX: We should fix things so this never happens!  -AJB, 20100304

        mutter('response args contain unicode, should be only bytes: %r',

               joined)

        joined = joined.encode('ascii')

    return joined

class Requester(object):

    """Abstract base class for an object that can issue requests on a smart

    medium.

    """

    def call(self, *args):

        """Make a remote call.

        :param args: the arguments of this call.

        """

        raise NotImplementedError(self.call)

    def call_with_body_bytes(self, args, body):

        """Make a remote call with a body.

        :param args: the arguments of this call.

        :type body: str

        :param body: the body to send with the request.

        """

        raise NotImplementedError(self.call_with_body_bytes)

    def call_with_body_readv_array(self, args, body):

        """Make a remote call with a readv array.

        :param args: the arguments of this call.

        :type body: iterable of (start, length) tuples.

        :param body: the readv ranges to send with this request.

        """

        raise NotImplementedError(self.call_with_body_readv_array)

    def set_headers(self, headers):

        raise NotImplementedError(self.set_headers)

class SmartProtocolBase(object):

    """Methods common to client and server"""

    # TODO: this only actually accomodates a single block; possibly should

    # support multiple chunks?

    def _encode_bulk_data(self, body):

        """Encode body as a bulk data chunk."""

        return ''.join(('%d\n' % len(body), body, 'done\n'))

    def _serialise_offsets(self, offsets):

        """Serialise a readv offset list."""

        txt = []

        for start, length in offsets:

            txt.append('%d,%d' % (start, length))

        return '\n'.join(txt)

class SmartServerRequestProtocolOne(SmartProtocolBase):

    """Server-side encoding and decoding logic for smart version 1."""

    def __init__(self, backing_transport, write_func, root_client_path='/',

            jail_root=None):

        self._backing_transport = backing_transport

        self._root_client_path = root_client_path

        self._jail_root = jail_root

        self.unused_data = ''

        self._finished = False

        self.in_buffer = ''

        self._has_dispatched = False

        self.request = None

        self._body_decoder = None

        self._write_func = write_func

    def accept_bytes(self, bytes):

        """Take bytes, and advance the internal state machine appropriately.

        :param bytes: must be a byte string

        """

        if not isinstance(bytes, str):

            raise ValueError(bytes)

        self.in_buffer += bytes

        if not self._has_dispatched:

            if '\n' not in self.in_buffer:

                # no command line yet

                return

            self._has_dispatched = True

            try:

                first_line, self.in_buffer = self.in_buffer.split('\n', 1)

                first_line += '\n'

                req_args = _decode_tuple(first_line)

                self.request = request.SmartServerRequestHandler(

                    self._backing_transport, commands=request.request_handlers,

                    root_client_path=self._root_client_path,

                    jail_root=self._jail_root)

                self.request.args_received(req_args)

                if self.request.finished_reading:

                    # trivial request

                    self.unused_data = self.in_buffer

                    self.in_buffer = ''

                    self._send_response(self.request.response)

            except KeyboardInterrupt:

                raise

            except errors.UnknownSmartMethod as err:

                protocol_error = errors.SmartProtocolError(

                    "bad request %r" % (err.verb,))

                failure = request.FailedSmartServerResponse(

                    ('error', str(protocol_error)))

                self._send_response(failure)

                return

            except Exception as exception:

                # everything else: pass to client, flush, and quit

                log_exception_quietly()

                self._send_response(request.FailedSmartServerResponse(

                    ('error', str(exception))))

                return

        if self._has_dispatched:

            if self._finished:

                # nothing to do.XXX: this routine should be a single state

                # machine too.

                self.unused_data += self.in_buffer

                self.in_buffer = ''

                return

            if self._body_decoder is None:

                self._body_decoder = LengthPrefixedBodyDecoder()

            self._body_decoder.accept_bytes(self.in_buffer)

            self.in_buffer = self._body_decoder.unused_data

            body_data = self._body_decoder.read_pending_data()

            self.request.accept_body(body_data)

            if self._body_decoder.finished_reading:

                self.request.end_of_body()

                if not self.request.finished_reading:

                    raise AssertionError("no more body, request not finished")

            if self.request.response is not None:

                self._send_response(self.request.response)

                self.unused_data = self.in_buffer

                self.in_buffer = ''

            else:

                if self.request.finished_reading:

                    raise AssertionError(

                        "no response and we have finished reading.")

    def _send_response(self, response):

        """Send a smart server response down the output stream."""

        if self._finished:

            raise AssertionError('response already sent')

        args = response.args

        body = response.body

        self._finished = True

        self._write_protocol_version()

        self._write_success_or_failure_prefix(response)

        self._write_func(_encode_tuple(args))

        if body is not None:

            if not isinstance(body, str):

                raise ValueError(body)

            bytes = self._encode_bulk_data(body)

            self._write_func(bytes)

    def _write_protocol_version(self):

        """Write any prefixes this protocol requires.

        Version one doesn't send protocol versions.

        """

    def _write_success_or_failure_prefix(self, response):

        """Write the protocol specific success/failure prefix.

        For SmartServerRequestProtocolOne this is omitted but we

        call is_successful to ensure that the response is valid.

        """

        response.is_successful()

    def next_read_size(self):

        if self._finished:

            return 0

        if self._body_decoder is None:

            return 1

        else:

            return self._body_decoder.next_read_size()

class SmartServerRequestProtocolTwo(SmartServerRequestProtocolOne):

    r"""Version two of the server side of the smart protocol.

    This prefixes responses with the value of RESPONSE_VERSION_TWO.

    """

    response_marker = RESPONSE_VERSION_TWO

    request_marker = REQUEST_VERSION_TWO

    def _write_success_or_failure_prefix(self, response):

        """Write the protocol specific success/failure prefix."""

        if response.is_successful():

            self._write_func('success\n')

        else:

            self._write_func('failed\n')

    def _write_protocol_version(self):

        r"""Write any prefixes this protocol requires.

        Version two sends the value of RESPONSE_VERSION_TWO.

        """

        self._write_func(self.response_marker)

    def _send_response(self, response):

        """Send a smart server response down the output stream."""

        if (self._finished):

            raise AssertionError('response already sent')

        self._finished = True

        self._write_protocol_version()

        self._write_success_or_failure_prefix(response)

        self._write_func(_encode_tuple(response.args))

        if response.body is not None:

            if not isinstance(response.body, str):

                raise AssertionError('body must be a str')

            if not (response.body_stream is None):

                raise AssertionError(

                    'body_stream and body cannot both be set')

            bytes = self._encode_bulk_data(response.body)

            self._write_func(bytes)

        elif response.body_stream is not None:

            _send_stream(response.body_stream, self._write_func)

def _send_stream(stream, write_func):

    write_func('chunked\n')

    _send_chunks(stream, write_func)

    write_func('END\n')

def _send_chunks(stream, write_func):

    for chunk in stream:

        if isinstance(chunk, str):

            bytes = "%x\n%s" % (len(chunk), chunk)

            write_func(bytes)

        elif isinstance(chunk, request.FailedSmartServerResponse):

            write_func('ERR\n')

            _send_chunks(chunk.args, write_func)

            return

        else:

            raise errors.BzrError(

                'Chunks must be str or FailedSmartServerResponse, got %r'

                % chunk)

class _NeedMoreBytes(Exception):

    """Raise this inside a _StatefulDecoder to stop decoding until more bytes

    have been received.

    """

    def __init__(self, count=None):

        """Constructor.

        :param count: the total number of bytes needed by the current state.

            May be None if the number of bytes needed is unknown.

        """

        self.count = count

class _StatefulDecoder(object):

    """Base class for writing state machines to decode byte streams.

    Subclasses should provide a self.state_accept attribute that accepts bytes

    and, if appropriate, updates self.state_accept to a different function.

    accept_bytes will call state_accept as often as necessary to make sure the

    state machine has progressed as far as possible before it returns.

    See ProtocolThreeDecoder for an example subclass.

    """

    def __init__(self):

        self.finished_reading = False

        self._in_buffer_list = []

        self._in_buffer_len = 0

        self.unused_data = ''

        self.bytes_left = None

        self._number_needed_bytes = None

    def _get_in_buffer(self):

        if len(self._in_buffer_list) == 1:

            return self._in_buffer_list[0]

        in_buffer = ''.join(self._in_buffer_list)

        if len(in_buffer) != self._in_buffer_len:

            raise AssertionError(

                "Length of buffer did not match expected value: %s != %s"

                % self._in_buffer_len, len(in_buffer))

        self._in_buffer_list = [in_buffer]

        return in_buffer

    def _get_in_bytes(self, count):

        """Grab X bytes from the input_buffer.

        Callers should have already checked that self._in_buffer_len is >

        count. Note, this does not consume the bytes from the buffer. The

        caller will still need to call _get_in_buffer() and then

        _set_in_buffer() if they actually need to consume the bytes.

        """

        # check if we can yield the bytes from just the first entry in our list

        if len(self._in_buffer_list) == 0:

            raise AssertionError('Callers must be sure we have buffered bytes'

                ' before calling _get_in_bytes')

        if len(self._in_buffer_list[0]) > count:

            return self._in_buffer_list[0][:count]

        # We can't yield it from the first buffer, so collapse all buffers, and

        # yield it from that

        in_buf = self._get_in_buffer()

        return in_buf[:count]

    def _set_in_buffer(self, new_buf):

        if new_buf is not None:

            self._in_buffer_list = [new_buf]

            self._in_buffer_len = len(new_buf)

        else:

            self._in_buffer_list = []

            self._in_buffer_len = 0

    def accept_bytes(self, bytes):

        """Decode as much of bytes as possible.

        If 'bytes' contains too much data it will be appended to

        self.unused_data.

        finished_reading will be set when no more data is required.  Further

        data will be appended to self.unused_data.

        """

        # accept_bytes is allowed to change the state

        self._number_needed_bytes = None

        # lsprof puts a very large amount of time on this specific call for

        # large readv arrays

        self._in_buffer_list.append(bytes)

        self._in_buffer_len += len(bytes)

        try:

            # Run the function for the current state.

            current_state = self.state_accept

            self.state_accept()

            while current_state != self.state_accept:

                # The current state has changed.  Run the function for the new

                # current state, so that it can:

                #   - decode any unconsumed bytes left in a buffer, and

                #   - signal how many more bytes are expected (via raising

                #     _NeedMoreBytes).

                current_state = self.state_accept

                self.state_accept()

        except _NeedMoreBytes as e:

            self._number_needed_bytes = e.count

class ChunkedBodyDecoder(_StatefulDecoder):

    """Decoder for chunked body data.

    This is very similar the HTTP's chunked encoding.  See the description of

    streamed body data in `doc/developers/network-protocol.txt` for details.

    """

    def __init__(self):

        _StatefulDecoder.__init__(self)

        self.state_accept = self._state_accept_expecting_header

        self.chunk_in_progress = None

        self.chunks = collections.deque()

        self.error = False

        self.error_in_progress = None

    def next_read_size(self):

        # Note: the shortest possible chunk is 2 bytes: '0\n', and the

        # end-of-body marker is 4 bytes: 'END\n'.

        if self.state_accept == self._state_accept_reading_chunk:

            # We're expecting more chunk content.  So we're expecting at least

            # the rest of this chunk plus an END chunk.

            return self.bytes_left + 4

        elif self.state_accept == self._state_accept_expecting_length:

            if self._in_buffer_len == 0:

                # We're expecting a chunk length.  There's at least two bytes

                # left: a digit plus '\n'.

                return 2

            else:

                # We're in the middle of reading a chunk length.  So there's at

                # least one byte left, the '\n' that terminates the length.

                return 1

        elif self.state_accept == self._state_accept_reading_unused:

            return 1

        elif self.state_accept == self._state_accept_expecting_header:

            return max(0, len('chunked\n') - self._in_buffer_len)

        else:

            raise AssertionError("Impossible state: %r" % (self.state_accept,))

    def read_next_chunk(self):

        try:

            return self.chunks.popleft()

        except IndexError:

            return None

    def _extract_line(self):

        in_buf = self._get_in_buffer()

        pos = in_buf.find('\n')

        if pos == -1:

            # We haven't read a complete line yet, so request more bytes before

            # we continue.

            raise _NeedMoreBytes(1)

        line = in_buf[:pos]

        # Trim the prefix (including '\n' delimiter) from the _in_buffer.

        self._set_in_buffer(in_buf[pos+1:])

        return line

    def _finished(self):

        self.unused_data = self._get_in_buffer()

        self._in_buffer_list = []

        self._in_buffer_len = 0

        self.state_accept = self._state_accept_reading_unused

        if self.error:

            error_args = tuple(self.error_in_progress)

            self.chunks.append(request.FailedSmartServerResponse(error_args))

            self.error_in_progress = None

        self.finished_reading = True

    def _state_accept_expecting_header(self):

        prefix = self._extract_line()

        if prefix == 'chunked':

            self.state_accept = self._state_accept_expecting_length

        else:

            raise errors.SmartProtocolError(

                'Bad chunked body header: "%s"' % (prefix,))

    def _state_accept_expecting_length(self):

        prefix = self._extract_line()

        if prefix == 'ERR':

            self.error = True

            self.error_in_progress = []

            self._state_accept_expecting_length()

            return

        elif prefix == 'END':

            # We've read the end-of-body marker.

            # Any further bytes are unused data, including the bytes left in

            # the _in_buffer.

            self._finished()

            return

        else:

            self.bytes_left = int(prefix, 16)

            self.chunk_in_progress = ''

            self.state_accept = self._state_accept_reading_chunk

    def _state_accept_reading_chunk(self):

        in_buf = self._get_in_buffer()

        in_buffer_len = len(in_buf)

        self.chunk_in_progress += in_buf[:self.bytes_left]

        self._set_in_buffer(in_buf[self.bytes_left:])

        self.bytes_left -= in_buffer_len

        if self.bytes_left <= 0:

            # Finished with chunk

            self.bytes_left = None

            if self.error:

                self.error_in_progress.append(self.chunk_in_progress)

            else:

                self.chunks.append(self.chunk_in_progress)

            self.chunk_in_progress = None

            self.state_accept = self._state_accept_expecting_length

    def _state_accept_reading_unused(self):

        self.unused_data += self._get_in_buffer()

        self._in_buffer_list = []

class LengthPrefixedBodyDecoder(_StatefulDecoder):

    """Decodes the length-prefixed bulk data."""

    def __init__(self):

        _StatefulDecoder.__init__(self)

        self.state_accept = self._state_accept_expecting_length

        self.state_read = self._state_read_no_data

        self._body = ''

        self._trailer_buffer = ''

    def next_read_size(self):

        if self.bytes_left is not None:

            # Ideally we want to read all the remainder of the body and the

            # trailer in one go.

            return self.bytes_left + 5

        elif self.state_accept == self._state_accept_reading_trailer:

            # Just the trailer left

            return 5 - len(self._trailer_buffer)

        elif self.state_accept == self._state_accept_expecting_length:

            # There's still at least 6 bytes left ('\n' to end the length, plus

            # 'done\n').

            return 6

        else:

            # Reading excess data.  Either way, 1 byte at a time is fine.

            return 1

    def read_pending_data(self):

        """Return any pending data that has been decoded."""

        return self.state_read()

    def _state_accept_expecting_length(self):

        in_buf = self._get_in_buffer()

        pos = in_buf.find('\n')

        if pos == -1:

            return

        self.bytes_left = int(in_buf[:pos])

        self._set_in_buffer(in_buf[pos+1:])

        self.state_accept = self._state_accept_reading_body

        self.state_read = self._state_read_body_buffer

    def _state_accept_reading_body(self):

        in_buf = self._get_in_buffer()

        self._body += in_buf

        self.bytes_left -= len(in_buf)

        self._set_in_buffer(None)

        if self.bytes_left <= 0:

            # Finished with body

            if self.bytes_left != 0:

                self._trailer_buffer = self._body[self.bytes_left:]

                self._body = self._body[:self.bytes_left]

            self.bytes_left = None

            self.state_accept = self._state_accept_reading_trailer

    def _state_accept_reading_trailer(self):

        self._trailer_buffer += self._get_in_buffer()

        self._set_in_buffer(None)

        # TODO: what if the trailer does not match "done\n"?  Should this raise

        # a ProtocolViolation exception?

        if self._trailer_buffer.startswith('done\n'):

            self.unused_data = self._trailer_buffer[len('done\n'):]

            self.state_accept = self._state_accept_reading_unused

            self.finished_reading = True

    def _state_accept_reading_unused(self):

        self.unused_data += self._get_in_buffer()

        self._set_in_buffer(None)

    def _state_read_no_data(self):

        return ''

    def _state_read_body_buffer(self):

        result = self._body

        self._body = ''

        return result

class SmartClientRequestProtocolOne(SmartProtocolBase, Requester,

                                    message.ResponseHandler):

    """The client-side protocol for smart version 1."""

    def __init__(self, request):

        """Construct a SmartClientRequestProtocolOne.

        :param request: A SmartClientMediumRequest to serialise onto and

            deserialise from.

        """

        self._request = request

        self._body_buffer = None

        self._request_start_time = None

        self._last_verb = None

        self._headers = None

    def set_headers(self, headers):

        self._headers = dict(headers)

    def call(self, *args):

        if 'hpss' in debug.debug_flags:

            mutter('hpss call:   %s', repr(args)[1:-1])

            if getattr(self._request._medium, 'base', None) is not None:

                mutter('             (to %s)', self._request._medium.base)

            self._request_start_time = osutils.timer_func()

        self._write_args(args)

        self._request.finished_writing()

        self._last_verb = args[0]

    def call_with_body_bytes(self, args, body):

        """Make a remote call of args with body bytes 'body'.

        After calling this, call read_response_tuple to find the result out.

        """

        if 'hpss' in debug.debug_flags:

            mutter('hpss call w/body: %s (%r...)', repr(args)[1:-1], body[:20])

            if getattr(self._request._medium, '_path', None) is not None:

                mutter('                  (to %s)', self._request._medium._path)

            mutter('              %d bytes', len(body))

            self._request_start_time = osutils.timer_func()