1
# Copyright (C) 2006, 2007 Canonical Ltd
 
 
3
# This program is free software; you can redistribute it and/or modify
 
 
4
# it under the terms of the GNU General Public License as published by
 
 
5
# the Free Software Foundation; either version 2 of the License, or
 
 
6
# (at your option) any later version.
 
 
8
# This program is distributed in the hope that it will be useful,
 
 
9
# but WITHOUT ANY WARRANTY; without even the implied warranty of
 
 
10
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 
 
11
# GNU General Public License for more details.
 
 
13
# You should have received a copy of the GNU General Public License
 
 
14
# along with this program; if not, write to the Free Software
 
 
15
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 
 
17
"""Wire-level encoding and decoding of requests and responses for the smart
 
 
22
from cStringIO import StringIO
 
 
28
from bzrlib import debug
 
 
29
from bzrlib import errors
 
 
30
from bzrlib.smart import message, request
 
 
31
from bzrlib.trace import log_exception_quietly, mutter
 
 
32
from bzrlib.util.bencode import bdecode, bencode
 
 
35
# Protocol version strings.  These are sent as prefixes of bzr requests and
 
 
36
# responses to identify the protocol version being used. (There are no version
 
 
37
# one strings because that version doesn't send any).
 
 
38
REQUEST_VERSION_TWO = 'bzr request 2\n'
 
 
39
RESPONSE_VERSION_TWO = 'bzr response 2\n'
 
 
41
MESSAGE_VERSION_THREE = 'bzr message 3 (bzr 1.6)\n'
 
 
42
RESPONSE_VERSION_THREE = REQUEST_VERSION_THREE = MESSAGE_VERSION_THREE
 
 
45
def _recv_tuple(from_file):
 
 
46
    req_line = from_file.readline()
 
 
47
    return _decode_tuple(req_line)
 
 
50
def _decode_tuple(req_line):
 
 
51
    if req_line is None or req_line == '':
 
 
53
    if req_line[-1] != '\n':
 
 
54
        raise errors.SmartProtocolError("request %r not terminated" % req_line)
 
 
55
    return tuple(req_line[:-1].split('\x01'))
 
 
58
def _encode_tuple(args):
 
 
59
    """Encode the tuple args to a bytestream."""
 
 
60
    return '\x01'.join(args) + '\n'
 
 
63
class Requester(object):
 
 
64
    """Abstract base class for an object that can issue requests on a smart
 
 
68
    def call(self, *args):
 
 
69
        """Make a remote call.
 
 
71
        :param args: the arguments of this call.
 
 
73
        raise NotImplementedError(self.call)
 
 
75
    def call_with_body_bytes(self, args, body):
 
 
76
        """Make a remote call with a body.
 
 
78
        :param args: the arguments of this call.
 
 
80
        :param body: the body to send with the request.
 
 
82
        raise NotImplementedError(self.call_with_body_bytes)
 
 
84
    def call_with_body_readv_array(self, args, body):
 
 
85
        """Make a remote call with a readv array.
 
 
87
        :param args: the arguments of this call.
 
 
88
        :type body: iterable of (start, length) tuples.
 
 
89
        :param body: the readv ranges to send with this request.
 
 
91
        raise NotImplementedError(self.call_with_body_readv_array)
 
 
93
    def set_headers(self, headers):
 
 
94
        raise NotImplementedError(self.set_headers)
 
 
97
class SmartProtocolBase(object):
 
 
98
    """Methods common to client and server"""
 
 
100
    # TODO: this only actually accomodates a single block; possibly should
 
 
101
    # support multiple chunks?
 
 
102
    def _encode_bulk_data(self, body):
 
 
103
        """Encode body as a bulk data chunk."""
 
 
104
        return ''.join(('%d\n' % len(body), body, 'done\n'))
 
 
106
    def _serialise_offsets(self, offsets):
 
 
107
        """Serialise a readv offset list."""
 
 
109
        for start, length in offsets:
 
 
110
            txt.append('%d,%d' % (start, length))
 
 
111
        return '\n'.join(txt)
 
 
114
class SmartServerRequestProtocolOne(SmartProtocolBase):
 
 
115
    """Server-side encoding and decoding logic for smart version 1."""
 
 
117
    def __init__(self, backing_transport, write_func, root_client_path='/'):
 
 
118
        self._backing_transport = backing_transport
 
 
119
        self._root_client_path = root_client_path
 
 
120
        self.unused_data = ''
 
 
121
        self._finished = False
 
 
123
        self._has_dispatched = False
 
 
125
        self._body_decoder = None
 
 
126
        self._write_func = write_func
 
 
128
    def accept_bytes(self, bytes):
 
 
129
        """Take bytes, and advance the internal state machine appropriately.
 
 
131
        :param bytes: must be a byte string
 
 
133
        if not isinstance(bytes, str):
 
 
134
            raise ValueError(bytes)
 
 
135
        self.in_buffer += bytes
 
 
136
        if not self._has_dispatched:
 
 
137
            if '\n' not in self.in_buffer:
 
 
138
                # no command line yet
 
 
140
            self._has_dispatched = True
 
 
142
                first_line, self.in_buffer = self.in_buffer.split('\n', 1)
 
 
144
                req_args = _decode_tuple(first_line)
 
 
145
                self.request = request.SmartServerRequestHandler(
 
 
146
                    self._backing_transport, commands=request.request_handlers,
 
 
147
                    root_client_path=self._root_client_path)
 
 
148
                self.request.dispatch_command(req_args[0], req_args[1:])
 
 
149
                if self.request.finished_reading:
 
 
151
                    self.unused_data = self.in_buffer
 
 
153
                    self._send_response(self.request.response)
 
 
154
            except KeyboardInterrupt:
 
 
156
            except errors.UnknownSmartMethod, err:
 
 
157
                protocol_error = errors.SmartProtocolError(
 
 
158
                    "bad request %r" % (err.verb,))
 
 
159
                failure = request.FailedSmartServerResponse(
 
 
160
                    ('error', str(protocol_error)))
 
 
161
                self._send_response(failure)
 
 
163
            except Exception, exception:
 
 
164
                # everything else: pass to client, flush, and quit
 
 
165
                log_exception_quietly()
 
 
166
                self._send_response(request.FailedSmartServerResponse(
 
 
167
                    ('error', str(exception))))
 
 
170
        if self._has_dispatched:
 
 
172
                # nothing to do.XXX: this routine should be a single state 
 
 
174
                self.unused_data += self.in_buffer
 
 
177
            if self._body_decoder is None:
 
 
178
                self._body_decoder = LengthPrefixedBodyDecoder()
 
 
179
            self._body_decoder.accept_bytes(self.in_buffer)
 
 
180
            self.in_buffer = self._body_decoder.unused_data
 
 
181
            body_data = self._body_decoder.read_pending_data()
 
 
182
            self.request.accept_body(body_data)
 
 
183
            if self._body_decoder.finished_reading:
 
 
184
                self.request.end_of_body()
 
 
185
                if not self.request.finished_reading:
 
 
186
                    raise AssertionError("no more body, request not finished")
 
 
187
            if self.request.response is not None:
 
 
188
                self._send_response(self.request.response)
 
 
189
                self.unused_data = self.in_buffer
 
 
192
                if self.request.finished_reading:
 
 
193
                    raise AssertionError(
 
 
194
                        "no response and we have finished reading.")
 
 
196
    def _send_response(self, response):
 
 
197
        """Send a smart server response down the output stream."""
 
 
199
            raise AssertionError('response already sent')
 
 
202
        self._finished = True
 
 
203
        self._write_protocol_version()
 
 
204
        self._write_success_or_failure_prefix(response)
 
 
205
        self._write_func(_encode_tuple(args))
 
 
207
            if not isinstance(body, str):
 
 
208
                raise ValueError(body)
 
 
209
            bytes = self._encode_bulk_data(body)
 
 
210
            self._write_func(bytes)
 
 
212
    def _write_protocol_version(self):
 
 
213
        """Write any prefixes this protocol requires.
 
 
215
        Version one doesn't send protocol versions.
 
 
218
    def _write_success_or_failure_prefix(self, response):
 
 
219
        """Write the protocol specific success/failure prefix.
 
 
221
        For SmartServerRequestProtocolOne this is omitted but we
 
 
222
        call is_successful to ensure that the response is valid.
 
 
224
        response.is_successful()
 
 
226
    def next_read_size(self):
 
 
229
        if self._body_decoder is None:
 
 
232
            return self._body_decoder.next_read_size()
 
 
235
class SmartServerRequestProtocolTwo(SmartServerRequestProtocolOne):
 
 
236
    r"""Version two of the server side of the smart protocol.
 
 
238
    This prefixes responses with the value of RESPONSE_VERSION_TWO.
 
 
241
    response_marker = RESPONSE_VERSION_TWO
 
 
242
    request_marker = REQUEST_VERSION_TWO
 
 
244
    def _write_success_or_failure_prefix(self, response):
 
 
245
        """Write the protocol specific success/failure prefix."""
 
 
246
        if response.is_successful():
 
 
247
            self._write_func('success\n')
 
 
249
            self._write_func('failed\n')
 
 
251
    def _write_protocol_version(self):
 
 
252
        r"""Write any prefixes this protocol requires.
 
 
254
        Version two sends the value of RESPONSE_VERSION_TWO.
 
 
256
        self._write_func(self.response_marker)
 
 
258
    def _send_response(self, response):
 
 
259
        """Send a smart server response down the output stream."""
 
 
261
            raise AssertionError('response already sent')
 
 
262
        self._finished = True
 
 
263
        self._write_protocol_version()
 
 
264
        self._write_success_or_failure_prefix(response)
 
 
265
        self._write_func(_encode_tuple(response.args))
 
 
266
        if response.body is not None:
 
 
267
            if not isinstance(response.body, str):
 
 
268
                raise AssertionError('body must be a str')
 
 
269
            if not (response.body_stream is None):
 
 
270
                raise AssertionError(
 
 
271
                    'body_stream and body cannot both be set')
 
 
272
            bytes = self._encode_bulk_data(response.body)
 
 
273
            self._write_func(bytes)
 
 
274
        elif response.body_stream is not None:
 
 
275
            _send_stream(response.body_stream, self._write_func)
 
 
278
def _send_stream(stream, write_func):
 
 
279
    write_func('chunked\n')
 
 
280
    _send_chunks(stream, write_func)
 
 
284
def _send_chunks(stream, write_func):
 
 
286
        if isinstance(chunk, str):
 
 
287
            bytes = "%x\n%s" % (len(chunk), chunk)
 
 
289
        elif isinstance(chunk, request.FailedSmartServerResponse):
 
 
291
            _send_chunks(chunk.args, write_func)
 
 
294
            raise errors.BzrError(
 
 
295
                'Chunks must be str or FailedSmartServerResponse, got %r'
 
 
299
class _NeedMoreBytes(Exception):
 
 
300
    """Raise this inside a _StatefulDecoder to stop decoding until more bytes
 
 
304
    def __init__(self, count=None):
 
 
307
        :param count: the total number of bytes needed by the current state.
 
 
308
            May be None if the number of bytes needed is unknown.
 
 
313
class _StatefulDecoder(object):
 
 
314
    """Base class for writing state machines to decode byte streams.
 
 
316
    Subclasses should provide a self.state_accept attribute that accepts bytes
 
 
317
    and, if appropriate, updates self.state_accept to a different function.
 
 
318
    accept_bytes will call state_accept as often as necessary to make sure the
 
 
319
    state machine has progressed as far as possible before it returns.
 
 
321
    See ProtocolThreeDecoder for an example subclass.
 
 
325
        self.finished_reading = False
 
 
326
        self._in_buffer_list = []
 
 
327
        self._in_buffer_len = 0
 
 
328
        self.unused_data = ''
 
 
329
        self.bytes_left = None
 
 
330
        self._number_needed_bytes = None
 
 
332
    def _get_in_buffer(self):
 
 
333
        if len(self._in_buffer_list) == 1:
 
 
334
            return self._in_buffer_list[0]
 
 
335
        in_buffer = ''.join(self._in_buffer_list)
 
 
336
        if len(in_buffer) != self._in_buffer_len:
 
 
337
            raise AssertionError(
 
 
338
                "Length of buffer did not match expected value: %s != %s"
 
 
339
                % self._in_buffer_len, len(in_buffer))
 
 
340
        self._in_buffer_list = [in_buffer]
 
 
343
    def _get_in_bytes(self, count):
 
 
344
        """Grab X bytes from the input_buffer.
 
 
346
        Callers should have already checked that self._in_buffer_len is >
 
 
347
        count. Note, this does not consume the bytes from the buffer. The
 
 
348
        caller will still need to call _get_in_buffer() and then
 
 
349
        _set_in_buffer() if they actually need to consume the bytes.
 
 
351
        # check if we can yield the bytes from just the first entry in our list
 
 
352
        if len(self._in_buffer_list) == 0:
 
 
353
            raise AssertionError('Callers must be sure we have buffered bytes'
 
 
354
                ' before calling _get_in_bytes')
 
 
355
        if len(self._in_buffer_list[0]) > count:
 
 
356
            return self._in_buffer_list[0][:count]
 
 
357
        # We can't yield it from the first buffer, so collapse all buffers, and
 
 
359
        in_buf = self._get_in_buffer()
 
 
360
        return in_buf[:count]
 
 
362
    def _set_in_buffer(self, new_buf):
 
 
363
        if new_buf is not None:
 
 
364
            self._in_buffer_list = [new_buf]
 
 
365
            self._in_buffer_len = len(new_buf)
 
 
367
            self._in_buffer_list = []
 
 
368
            self._in_buffer_len = 0
 
 
370
    def accept_bytes(self, bytes):
 
 
371
        """Decode as much of bytes as possible.
 
 
373
        If 'bytes' contains too much data it will be appended to
 
 
376
        finished_reading will be set when no more data is required.  Further
 
 
377
        data will be appended to self.unused_data.
 
 
379
        # accept_bytes is allowed to change the state
 
 
380
        self._number_needed_bytes = None
 
 
381
        # lsprof puts a very large amount of time on this specific call for
 
 
383
        self._in_buffer_list.append(bytes)
 
 
384
        self._in_buffer_len += len(bytes)
 
 
386
            # Run the function for the current state.
 
 
387
            current_state = self.state_accept
 
 
389
            while current_state != self.state_accept:
 
 
390
                # The current state has changed.  Run the function for the new
 
 
391
                # current state, so that it can:
 
 
392
                #   - decode any unconsumed bytes left in a buffer, and
 
 
393
                #   - signal how many more bytes are expected (via raising
 
 
395
                current_state = self.state_accept
 
 
397
        except _NeedMoreBytes, e:
 
 
398
            self._number_needed_bytes = e.count
 
 
401
class ChunkedBodyDecoder(_StatefulDecoder):
 
 
402
    """Decoder for chunked body data.
 
 
404
    This is very similar the HTTP's chunked encoding.  See the description of
 
 
405
    streamed body data in `doc/developers/network-protocol.txt` for details.
 
 
409
        _StatefulDecoder.__init__(self)
 
 
410
        self.state_accept = self._state_accept_expecting_header
 
 
411
        self.chunk_in_progress = None
 
 
412
        self.chunks = collections.deque()
 
 
414
        self.error_in_progress = None
 
 
416
    def next_read_size(self):
 
 
417
        # Note: the shortest possible chunk is 2 bytes: '0\n', and the
 
 
418
        # end-of-body marker is 4 bytes: 'END\n'.
 
 
419
        if self.state_accept == self._state_accept_reading_chunk:
 
 
420
            # We're expecting more chunk content.  So we're expecting at least
 
 
421
            # the rest of this chunk plus an END chunk.
 
 
422
            return self.bytes_left + 4
 
 
423
        elif self.state_accept == self._state_accept_expecting_length:
 
 
424
            if self._in_buffer_len == 0:
 
 
425
                # We're expecting a chunk length.  There's at least two bytes
 
 
426
                # left: a digit plus '\n'.
 
 
429
                # We're in the middle of reading a chunk length.  So there's at
 
 
430
                # least one byte left, the '\n' that terminates the length.
 
 
432
        elif self.state_accept == self._state_accept_reading_unused:
 
 
434
        elif self.state_accept == self._state_accept_expecting_header:
 
 
435
            return max(0, len('chunked\n') - self._in_buffer_len)
 
 
437
            raise AssertionError("Impossible state: %r" % (self.state_accept,))
 
 
439
    def read_next_chunk(self):
 
 
441
            return self.chunks.popleft()
 
 
445
    def _extract_line(self):
 
 
446
        in_buf = self._get_in_buffer()
 
 
447
        pos = in_buf.find('\n')
 
 
449
            # We haven't read a complete line yet, so request more bytes before
 
 
451
            raise _NeedMoreBytes(1)
 
 
453
        # Trim the prefix (including '\n' delimiter) from the _in_buffer.
 
 
454
        self._set_in_buffer(in_buf[pos+1:])
 
 
458
        self.unused_data = self._get_in_buffer()
 
 
459
        self._in_buffer_list = []
 
 
460
        self._in_buffer_len = 0
 
 
461
        self.state_accept = self._state_accept_reading_unused
 
 
463
            error_args = tuple(self.error_in_progress)
 
 
464
            self.chunks.append(request.FailedSmartServerResponse(error_args))
 
 
465
            self.error_in_progress = None
 
 
466
        self.finished_reading = True
 
 
468
    def _state_accept_expecting_header(self):
 
 
469
        prefix = self._extract_line()
 
 
470
        if prefix == 'chunked':
 
 
471
            self.state_accept = self._state_accept_expecting_length
 
 
473
            raise errors.SmartProtocolError(
 
 
474
                'Bad chunked body header: "%s"' % (prefix,))
 
 
476
    def _state_accept_expecting_length(self):
 
 
477
        prefix = self._extract_line()
 
 
480
            self.error_in_progress = []
 
 
481
            self._state_accept_expecting_length()
 
 
483
        elif prefix == 'END':
 
 
484
            # We've read the end-of-body marker.
 
 
485
            # Any further bytes are unused data, including the bytes left in
 
 
490
            self.bytes_left = int(prefix, 16)
 
 
491
            self.chunk_in_progress = ''
 
 
492
            self.state_accept = self._state_accept_reading_chunk
 
 
494
    def _state_accept_reading_chunk(self):
 
 
495
        in_buf = self._get_in_buffer()
 
 
496
        in_buffer_len = len(in_buf)
 
 
497
        self.chunk_in_progress += in_buf[:self.bytes_left]
 
 
498
        self._set_in_buffer(in_buf[self.bytes_left:])
 
 
499
        self.bytes_left -= in_buffer_len
 
 
500
        if self.bytes_left <= 0:
 
 
501
            # Finished with chunk
 
 
502
            self.bytes_left = None
 
 
504
                self.error_in_progress.append(self.chunk_in_progress)
 
 
506
                self.chunks.append(self.chunk_in_progress)
 
 
507
            self.chunk_in_progress = None
 
 
508
            self.state_accept = self._state_accept_expecting_length
 
 
510
    def _state_accept_reading_unused(self):
 
 
511
        self.unused_data += self._get_in_buffer()
 
 
512
        self._in_buffer_list = []
 
 
515
class LengthPrefixedBodyDecoder(_StatefulDecoder):
 
 
516
    """Decodes the length-prefixed bulk data."""
 
 
519
        _StatefulDecoder.__init__(self)
 
 
520
        self.state_accept = self._state_accept_expecting_length
 
 
521
        self.state_read = self._state_read_no_data
 
 
523
        self._trailer_buffer = ''
 
 
525
    def next_read_size(self):
 
 
526
        if self.bytes_left is not None:
 
 
527
            # Ideally we want to read all the remainder of the body and the
 
 
529
            return self.bytes_left + 5
 
 
530
        elif self.state_accept == self._state_accept_reading_trailer:
 
 
531
            # Just the trailer left
 
 
532
            return 5 - len(self._trailer_buffer)
 
 
533
        elif self.state_accept == self._state_accept_expecting_length:
 
 
534
            # There's still at least 6 bytes left ('\n' to end the length, plus
 
 
538
            # Reading excess data.  Either way, 1 byte at a time is fine.
 
 
541
    def read_pending_data(self):
 
 
542
        """Return any pending data that has been decoded."""
 
 
543
        return self.state_read()
 
 
545
    def _state_accept_expecting_length(self):
 
 
546
        in_buf = self._get_in_buffer()
 
 
547
        pos = in_buf.find('\n')
 
 
550
        self.bytes_left = int(in_buf[:pos])
 
 
551
        self._set_in_buffer(in_buf[pos+1:])
 
 
552
        self.state_accept = self._state_accept_reading_body
 
 
553
        self.state_read = self._state_read_body_buffer
 
 
555
    def _state_accept_reading_body(self):
 
 
556
        in_buf = self._get_in_buffer()
 
 
558
        self.bytes_left -= len(in_buf)
 
 
559
        self._set_in_buffer(None)
 
 
560
        if self.bytes_left <= 0:
 
 
562
            if self.bytes_left != 0:
 
 
563
                self._trailer_buffer = self._body[self.bytes_left:]
 
 
564
                self._body = self._body[:self.bytes_left]
 
 
565
            self.bytes_left = None
 
 
566
            self.state_accept = self._state_accept_reading_trailer
 
 
568
    def _state_accept_reading_trailer(self):
 
 
569
        self._trailer_buffer += self._get_in_buffer()
 
 
570
        self._set_in_buffer(None)
 
 
571
        # TODO: what if the trailer does not match "done\n"?  Should this raise
 
 
572
        # a ProtocolViolation exception?
 
 
573
        if self._trailer_buffer.startswith('done\n'):
 
 
574
            self.unused_data = self._trailer_buffer[len('done\n'):]
 
 
575
            self.state_accept = self._state_accept_reading_unused
 
 
576
            self.finished_reading = True
 
 
578
    def _state_accept_reading_unused(self):
 
 
579
        self.unused_data += self._get_in_buffer()
 
 
580
        self._set_in_buffer(None)
 
 
582
    def _state_read_no_data(self):
 
 
585
    def _state_read_body_buffer(self):
 
 
591
class SmartClientRequestProtocolOne(SmartProtocolBase, Requester,
 
 
592
                                    message.ResponseHandler):
 
 
593
    """The client-side protocol for smart version 1."""
 
 
595
    def __init__(self, request):
 
 
596
        """Construct a SmartClientRequestProtocolOne.
 
 
598
        :param request: A SmartClientMediumRequest to serialise onto and
 
 
601
        self._request = request
 
 
602
        self._body_buffer = None
 
 
603
        self._request_start_time = None
 
 
604
        self._last_verb = None
 
 
607
    def set_headers(self, headers):
 
 
608
        self._headers = dict(headers)
 
 
610
    def call(self, *args):
 
 
611
        if 'hpss' in debug.debug_flags:
 
 
612
            mutter('hpss call:   %s', repr(args)[1:-1])
 
 
613
            if getattr(self._request._medium, 'base', None) is not None:
 
 
614
                mutter('             (to %s)', self._request._medium.base)
 
 
615
            self._request_start_time = time.time()
 
 
616
        self._write_args(args)
 
 
617
        self._request.finished_writing()
 
 
618
        self._last_verb = args[0]
 
 
620
    def call_with_body_bytes(self, args, body):
 
 
621
        """Make a remote call of args with body bytes 'body'.
 
 
623
        After calling this, call read_response_tuple to find the result out.
 
 
625
        if 'hpss' in debug.debug_flags:
 
 
626
            mutter('hpss call w/body: %s (%r...)', repr(args)[1:-1], body[:20])
 
 
627
            if getattr(self._request._medium, '_path', None) is not None:
 
 
628
                mutter('                  (to %s)', self._request._medium._path)
 
 
629
            mutter('              %d bytes', len(body))
 
 
630
            self._request_start_time = time.time()
 
 
631
            if 'hpssdetail' in debug.debug_flags:
 
 
632
                mutter('hpss body content: %s', body)
 
 
633
        self._write_args(args)
 
 
634
        bytes = self._encode_bulk_data(body)
 
 
635
        self._request.accept_bytes(bytes)
 
 
636
        self._request.finished_writing()
 
 
637
        self._last_verb = args[0]
 
 
639
    def call_with_body_readv_array(self, args, body):
 
 
640
        """Make a remote call with a readv array.
 
 
642
        The body is encoded with one line per readv offset pair. The numbers in
 
 
643
        each pair are separated by a comma, and no trailing \n is emitted.
 
 
645
        if 'hpss' in debug.debug_flags:
 
 
646
            mutter('hpss call w/readv: %s', repr(args)[1:-1])
 
 
647
            if getattr(self._request._medium, '_path', None) is not None:
 
 
648
                mutter('                  (to %s)', self._request._medium._path)
 
 
649
            self._request_start_time = time.time()
 
 
650
        self._write_args(args)
 
 
651
        readv_bytes = self._serialise_offsets(body)
 
 
652
        bytes = self._encode_bulk_data(readv_bytes)
 
 
653
        self._request.accept_bytes(bytes)
 
 
654
        self._request.finished_writing()
 
 
655
        if 'hpss' in debug.debug_flags:
 
 
656
            mutter('              %d bytes in readv request', len(readv_bytes))
 
 
657
        self._last_verb = args[0]
 
 
659
    def cancel_read_body(self):
 
 
660
        """After expecting a body, a response code may indicate one otherwise.
 
 
662
        This method lets the domain client inform the protocol that no body
 
 
663
        will be transmitted. This is a terminal method: after calling it the
 
 
664
        protocol is not able to be used further.
 
 
666
        self._request.finished_reading()
 
 
668
    def _read_response_tuple(self):
 
 
669
        result = self._recv_tuple()
 
 
670
        if 'hpss' in debug.debug_flags:
 
 
671
            if self._request_start_time is not None:
 
 
672
                mutter('   result:   %6.3fs  %s',
 
 
673
                       time.time() - self._request_start_time,
 
 
675
                self._request_start_time = None
 
 
677
                mutter('   result:   %s', repr(result)[1:-1])
 
 
680
    def read_response_tuple(self, expect_body=False):
 
 
681
        """Read a response tuple from the wire.
 
 
683
        This should only be called once.
 
 
685
        result = self._read_response_tuple()
 
 
686
        self._response_is_unknown_method(result)
 
 
687
        self._raise_args_if_error(result)
 
 
689
            self._request.finished_reading()
 
 
692
    def _raise_args_if_error(self, result_tuple):
 
 
693
        # Later protocol versions have an explicit flag in the protocol to say
 
 
694
        # if an error response is "failed" or not.  In version 1 we don't have
 
 
695
        # that luxury.  So here is a complete list of errors that can be
 
 
696
        # returned in response to existing version 1 smart requests.  Responses
 
 
697
        # starting with these codes are always "failed" responses.
 
 
704
            'UnicodeEncodeError',
 
 
705
            'UnicodeDecodeError',
 
 
711
            'UnlockableTransport',
 
 
717
        if result_tuple[0] in v1_error_codes:
 
 
718
            self._request.finished_reading()
 
 
719
            raise errors.ErrorFromSmartServer(result_tuple)
 
 
721
    def _response_is_unknown_method(self, result_tuple):
 
 
722
        """Raise UnexpectedSmartServerResponse if the response is an 'unknonwn
 
 
723
        method' response to the request.
 
 
725
        :param response: The response from a smart client call_expecting_body
 
 
727
        :param verb: The verb used in that call.
 
 
728
        :raises: UnexpectedSmartServerResponse
 
 
730
        if (result_tuple == ('error', "Generic bzr smart protocol error: "
 
 
731
                "bad request '%s'" % self._last_verb) or
 
 
732
              result_tuple == ('error', "Generic bzr smart protocol error: "
 
 
733
                "bad request u'%s'" % self._last_verb)):
 
 
734
            # The response will have no body, so we've finished reading.
 
 
735
            self._request.finished_reading()
 
 
736
            raise errors.UnknownSmartMethod(self._last_verb)
 
 
738
    def read_body_bytes(self, count=-1):
 
 
739
        """Read bytes from the body, decoding into a byte stream.
 
 
741
        We read all bytes at once to ensure we've checked the trailer for 
 
 
742
        errors, and then feed the buffer back as read_body_bytes is called.
 
 
744
        if self._body_buffer is not None:
 
 
745
            return self._body_buffer.read(count)
 
 
746
        _body_decoder = LengthPrefixedBodyDecoder()
 
 
748
        while not _body_decoder.finished_reading:
 
 
749
            bytes = self._request.read_bytes(_body_decoder.next_read_size())
 
 
751
                # end of file encountered reading from server
 
 
752
                raise errors.ConnectionReset(
 
 
753
                    "Connection lost while reading response body.")
 
 
754
            _body_decoder.accept_bytes(bytes)
 
 
755
        self._request.finished_reading()
 
 
756
        self._body_buffer = StringIO(_body_decoder.read_pending_data())
 
 
757
        # XXX: TODO check the trailer result.
 
 
758
        if 'hpss' in debug.debug_flags:
 
 
759
            mutter('              %d body bytes read',
 
 
760
                   len(self._body_buffer.getvalue()))
 
 
761
        return self._body_buffer.read(count)
 
 
763
    def _recv_tuple(self):
 
 
764
        """Receive a tuple from the medium request."""
 
 
765
        return _decode_tuple(self._request.read_line())
 
 
767
    def query_version(self):
 
 
768
        """Return protocol version number of the server."""
 
 
770
        resp = self.read_response_tuple()
 
 
771
        if resp == ('ok', '1'):
 
 
773
        elif resp == ('ok', '2'):
 
 
776
            raise errors.SmartProtocolError("bad response %r" % (resp,))
 
 
778
    def _write_args(self, args):
 
 
779
        self._write_protocol_version()
 
 
780
        bytes = _encode_tuple(args)
 
 
781
        self._request.accept_bytes(bytes)
 
 
783
    def _write_protocol_version(self):
 
 
784
        """Write any prefixes this protocol requires.
 
 
786
        Version one doesn't send protocol versions.
 
 
790
class SmartClientRequestProtocolTwo(SmartClientRequestProtocolOne):
 
 
791
    """Version two of the client side of the smart protocol.
 
 
793
    This prefixes the request with the value of REQUEST_VERSION_TWO.
 
 
796
    response_marker = RESPONSE_VERSION_TWO
 
 
797
    request_marker = REQUEST_VERSION_TWO
 
 
799
    def read_response_tuple(self, expect_body=False):
 
 
800
        """Read a response tuple from the wire.
 
 
802
        This should only be called once.
 
 
804
        version = self._request.read_line()
 
 
805
        if version != self.response_marker:
 
 
806
            self._request.finished_reading()
 
 
807
            raise errors.UnexpectedProtocolVersionMarker(version)
 
 
808
        response_status = self._request.read_line()
 
 
809
        result = SmartClientRequestProtocolOne._read_response_tuple(self)
 
 
810
        self._response_is_unknown_method(result)
 
 
811
        if response_status == 'success\n':
 
 
812
            self.response_status = True
 
 
814
                self._request.finished_reading()
 
 
816
        elif response_status == 'failed\n':
 
 
817
            self.response_status = False
 
 
818
            self._request.finished_reading()
 
 
819
            raise errors.ErrorFromSmartServer(result)
 
 
821
            raise errors.SmartProtocolError(
 
 
822
                'bad protocol status %r' % response_status)
 
 
824
    def _write_protocol_version(self):
 
 
825
        """Write any prefixes this protocol requires.
 
 
827
        Version two sends the value of REQUEST_VERSION_TWO.
 
 
829
        self._request.accept_bytes(self.request_marker)
 
 
831
    def read_streamed_body(self):
 
 
832
        """Read bytes from the body, decoding into a byte stream.
 
 
834
        # Read no more than 64k at a time so that we don't risk error 10055 (no
 
 
835
        # buffer space available) on Windows.
 
 
836
        _body_decoder = ChunkedBodyDecoder()
 
 
837
        while not _body_decoder.finished_reading:
 
 
838
            bytes = self._request.read_bytes(_body_decoder.next_read_size())
 
 
840
                # end of file encountered reading from server
 
 
841
                raise errors.ConnectionReset(
 
 
842
                    "Connection lost while reading streamed body.")
 
 
843
            _body_decoder.accept_bytes(bytes)
 
 
844
            for body_bytes in iter(_body_decoder.read_next_chunk, None):
 
 
845
                if 'hpss' in debug.debug_flags and type(body_bytes) is str:
 
 
846
                    mutter('              %d byte chunk read',
 
 
849
        self._request.finished_reading()
 
 
852
def build_server_protocol_three(backing_transport, write_func,
 
 
854
    request_handler = request.SmartServerRequestHandler(
 
 
855
        backing_transport, commands=request.request_handlers,
 
 
856
        root_client_path=root_client_path)
 
 
857
    responder = ProtocolThreeResponder(write_func)
 
 
858
    message_handler = message.ConventionalRequestHandler(request_handler, responder)
 
 
859
    return ProtocolThreeDecoder(message_handler)
 
 
862
class ProtocolThreeDecoder(_StatefulDecoder):
 
 
864
    response_marker = RESPONSE_VERSION_THREE
 
 
865
    request_marker = REQUEST_VERSION_THREE
 
 
867
    def __init__(self, message_handler, expect_version_marker=False):
 
 
868
        _StatefulDecoder.__init__(self)
 
 
869
        self._has_dispatched = False
 
 
871
        if expect_version_marker:
 
 
872
            self.state_accept = self._state_accept_expecting_protocol_version
 
 
873
            # We're expecting at least the protocol version marker + some
 
 
875
            self._number_needed_bytes = len(MESSAGE_VERSION_THREE) + 4
 
 
877
            self.state_accept = self._state_accept_expecting_headers
 
 
878
            self._number_needed_bytes = 4
 
 
879
        self.decoding_failed = False
 
 
880
        self.request_handler = self.message_handler = message_handler
 
 
882
    def accept_bytes(self, bytes):
 
 
883
        self._number_needed_bytes = None
 
 
885
            _StatefulDecoder.accept_bytes(self, bytes)
 
 
886
        except KeyboardInterrupt:
 
 
888
        except errors.SmartMessageHandlerError, exception:
 
 
889
            # We do *not* set self.decoding_failed here.  The message handler
 
 
890
            # has raised an error, but the decoder is still able to parse bytes
 
 
891
            # and determine when this message ends.
 
 
892
            log_exception_quietly()
 
 
893
            self.message_handler.protocol_error(exception.exc_value)
 
 
894
            # The state machine is ready to continue decoding, but the
 
 
895
            # exception has interrupted the loop that runs the state machine.
 
 
896
            # So we call accept_bytes again to restart it.
 
 
897
            self.accept_bytes('')
 
 
898
        except Exception, exception:
 
 
899
            # The decoder itself has raised an exception.  We cannot continue
 
 
901
            self.decoding_failed = True
 
 
902
            if isinstance(exception, errors.UnexpectedProtocolVersionMarker):
 
 
903
                # This happens during normal operation when the client tries a
 
 
904
                # protocol version the server doesn't understand, so no need to
 
 
905
                # log a traceback every time.
 
 
906
                # Note that this can only happen when
 
 
907
                # expect_version_marker=True, which is only the case on the
 
 
911
                log_exception_quietly()
 
 
912
            self.message_handler.protocol_error(exception)
 
 
914
    def _extract_length_prefixed_bytes(self):
 
 
915
        if self._in_buffer_len < 4:
 
 
916
            # A length prefix by itself is 4 bytes, and we don't even have that
 
 
918
            raise _NeedMoreBytes(4)
 
 
919
        (length,) = struct.unpack('!L', self._get_in_bytes(4))
 
 
920
        end_of_bytes = 4 + length
 
 
921
        if self._in_buffer_len < end_of_bytes:
 
 
922
            # We haven't yet read as many bytes as the length-prefix says there
 
 
924
            raise _NeedMoreBytes(end_of_bytes)
 
 
925
        # Extract the bytes from the buffer.
 
 
926
        in_buf = self._get_in_buffer()
 
 
927
        bytes = in_buf[4:end_of_bytes]
 
 
928
        self._set_in_buffer(in_buf[end_of_bytes:])
 
 
931
    def _extract_prefixed_bencoded_data(self):
 
 
932
        prefixed_bytes = self._extract_length_prefixed_bytes()
 
 
934
            decoded = bdecode(prefixed_bytes)
 
 
936
            raise errors.SmartProtocolError(
 
 
937
                'Bytes %r not bencoded' % (prefixed_bytes,))
 
 
940
    def _extract_single_byte(self):
 
 
941
        if self._in_buffer_len == 0:
 
 
942
            # The buffer is empty
 
 
943
            raise _NeedMoreBytes(1)
 
 
944
        in_buf = self._get_in_buffer()
 
 
946
        self._set_in_buffer(in_buf[1:])
 
 
949
    def _state_accept_expecting_protocol_version(self):
 
 
950
        needed_bytes = len(MESSAGE_VERSION_THREE) - self._in_buffer_len
 
 
951
        in_buf = self._get_in_buffer()
 
 
953
            # We don't have enough bytes to check if the protocol version
 
 
954
            # marker is right.  But we can check if it is already wrong by
 
 
955
            # checking that the start of MESSAGE_VERSION_THREE matches what
 
 
957
            # [In fact, if the remote end isn't bzr we might never receive
 
 
958
            # len(MESSAGE_VERSION_THREE) bytes.  So if the bytes we have so far
 
 
959
            # are wrong then we should just raise immediately rather than
 
 
961
            if not MESSAGE_VERSION_THREE.startswith(in_buf):
 
 
962
                # We have enough bytes to know the protocol version is wrong
 
 
963
                raise errors.UnexpectedProtocolVersionMarker(in_buf)
 
 
964
            raise _NeedMoreBytes(len(MESSAGE_VERSION_THREE))
 
 
965
        if not in_buf.startswith(MESSAGE_VERSION_THREE):
 
 
966
            raise errors.UnexpectedProtocolVersionMarker(in_buf)
 
 
967
        self._set_in_buffer(in_buf[len(MESSAGE_VERSION_THREE):])
 
 
968
        self.state_accept = self._state_accept_expecting_headers
 
 
970
    def _state_accept_expecting_headers(self):
 
 
971
        decoded = self._extract_prefixed_bencoded_data()
 
 
972
        if type(decoded) is not dict:
 
 
973
            raise errors.SmartProtocolError(
 
 
974
                'Header object %r is not a dict' % (decoded,))
 
 
975
        self.state_accept = self._state_accept_expecting_message_part
 
 
977
            self.message_handler.headers_received(decoded)
 
 
979
            raise errors.SmartMessageHandlerError(sys.exc_info())
 
 
981
    def _state_accept_expecting_message_part(self):
 
 
982
        message_part_kind = self._extract_single_byte()
 
 
983
        if message_part_kind == 'o':
 
 
984
            self.state_accept = self._state_accept_expecting_one_byte
 
 
985
        elif message_part_kind == 's':
 
 
986
            self.state_accept = self._state_accept_expecting_structure
 
 
987
        elif message_part_kind == 'b':
 
 
988
            self.state_accept = self._state_accept_expecting_bytes
 
 
989
        elif message_part_kind == 'e':
 
 
992
            raise errors.SmartProtocolError(
 
 
993
                'Bad message kind byte: %r' % (message_part_kind,))
 
 
995
    def _state_accept_expecting_one_byte(self):
 
 
996
        byte = self._extract_single_byte()
 
 
997
        self.state_accept = self._state_accept_expecting_message_part
 
 
999
            self.message_handler.byte_part_received(byte)
 
 
1001
            raise errors.SmartMessageHandlerError(sys.exc_info())
 
 
1003
    def _state_accept_expecting_bytes(self):
 
 
1004
        # XXX: this should not buffer whole message part, but instead deliver
 
 
1005
        # the bytes as they arrive.
 
 
1006
        prefixed_bytes = self._extract_length_prefixed_bytes()
 
 
1007
        self.state_accept = self._state_accept_expecting_message_part
 
 
1009
            self.message_handler.bytes_part_received(prefixed_bytes)
 
 
1011
            raise errors.SmartMessageHandlerError(sys.exc_info())
 
 
1013
    def _state_accept_expecting_structure(self):
 
 
1014
        structure = self._extract_prefixed_bencoded_data()
 
 
1015
        self.state_accept = self._state_accept_expecting_message_part
 
 
1017
            self.message_handler.structure_part_received(structure)
 
 
1019
            raise errors.SmartMessageHandlerError(sys.exc_info())
 
 
1022
        self.unused_data = self._get_in_buffer()
 
 
1023
        self._set_in_buffer(None)
 
 
1024
        self.state_accept = self._state_accept_reading_unused
 
 
1026
            self.message_handler.end_received()
 
 
1028
            raise errors.SmartMessageHandlerError(sys.exc_info())
 
 
1030
    def _state_accept_reading_unused(self):
 
 
1031
        self.unused_data = self._get_in_buffer()
 
 
1032
        self._set_in_buffer(None)
 
 
1034
    def next_read_size(self):
 
 
1035
        if self.state_accept == self._state_accept_reading_unused:
 
 
1037
        elif self.decoding_failed:
 
 
1038
            # An exception occured while processing this message, probably from
 
 
1039
            # self.message_handler.  We're not sure that this state machine is
 
 
1040
            # in a consistent state, so just signal that we're done (i.e. give
 
 
1044
            if self._number_needed_bytes is not None:
 
 
1045
                return self._number_needed_bytes - self._in_buffer_len
 
 
1047
                raise AssertionError("don't know how many bytes are expected!")
 
 
1050
class _ProtocolThreeEncoder(object):
 
 
1052
    response_marker = request_marker = MESSAGE_VERSION_THREE
 
 
1054
    def __init__(self, write_func):
 
 
1056
        self._real_write_func = write_func
 
 
1058
    def _write_func(self, bytes):
 
 
1063
            self._real_write_func(self._buf)
 
 
1066
    def _serialise_offsets(self, offsets):
 
 
1067
        """Serialise a readv offset list."""
 
 
1069
        for start, length in offsets:
 
 
1070
            txt.append('%d,%d' % (start, length))
 
 
1071
        return '\n'.join(txt)
 
 
1073
    def _write_protocol_version(self):
 
 
1074
        self._write_func(MESSAGE_VERSION_THREE)
 
 
1076
    def _write_prefixed_bencode(self, structure):
 
 
1077
        bytes = bencode(structure)
 
 
1078
        self._write_func(struct.pack('!L', len(bytes)))
 
 
1079
        self._write_func(bytes)
 
 
1081
    def _write_headers(self, headers):
 
 
1082
        self._write_prefixed_bencode(headers)
 
 
1084
    def _write_structure(self, args):
 
 
1085
        self._write_func('s')
 
 
1088
            if type(arg) is unicode:
 
 
1089
                utf8_args.append(arg.encode('utf8'))
 
 
1091
                utf8_args.append(arg)
 
 
1092
        self._write_prefixed_bencode(utf8_args)
 
 
1094
    def _write_end(self):
 
 
1095
        self._write_func('e')
 
 
1098
    def _write_prefixed_body(self, bytes):
 
 
1099
        self._write_func('b')
 
 
1100
        self._write_func(struct.pack('!L', len(bytes)))
 
 
1101
        self._write_func(bytes)
 
 
1103
    def _write_error_status(self):
 
 
1104
        self._write_func('oE')
 
 
1106
    def _write_success_status(self):
 
 
1107
        self._write_func('oS')
 
 
1110
class ProtocolThreeResponder(_ProtocolThreeEncoder):
 
 
1112
    def __init__(self, write_func):
 
 
1113
        _ProtocolThreeEncoder.__init__(self, write_func)
 
 
1114
        self.response_sent = False
 
 
1115
        self._headers = {'Software version': bzrlib.__version__}
 
 
1117
    def send_error(self, exception):
 
 
1118
        if self.response_sent:
 
 
1119
            raise AssertionError(
 
 
1120
                "send_error(%s) called, but response already sent."
 
 
1122
        if isinstance(exception, errors.UnknownSmartMethod):
 
 
1123
            failure = request.FailedSmartServerResponse(
 
 
1124
                ('UnknownMethod', exception.verb))
 
 
1125
            self.send_response(failure)
 
 
1127
        self.response_sent = True
 
 
1128
        self._write_protocol_version()
 
 
1129
        self._write_headers(self._headers)
 
 
1130
        self._write_error_status()
 
 
1131
        self._write_structure(('error', str(exception)))
 
 
1134
    def send_response(self, response):
 
 
1135
        if self.response_sent:
 
 
1136
            raise AssertionError(
 
 
1137
                "send_response(%r) called, but response already sent."
 
 
1139
        self.response_sent = True
 
 
1140
        self._write_protocol_version()
 
 
1141
        self._write_headers(self._headers)
 
 
1142
        if response.is_successful():
 
 
1143
            self._write_success_status()
 
 
1145
            self._write_error_status()
 
 
1146
        self._write_structure(response.args)
 
 
1147
        if response.body is not None:
 
 
1148
            self._write_prefixed_body(response.body)
 
 
1149
        elif response.body_stream is not None:
 
 
1150
            for chunk in response.body_stream:
 
 
1151
                self._write_prefixed_body(chunk)
 
 
1156
class ProtocolThreeRequester(_ProtocolThreeEncoder, Requester):
 
 
1158
    def __init__(self, medium_request):
 
 
1159
        _ProtocolThreeEncoder.__init__(self, medium_request.accept_bytes)
 
 
1160
        self._medium_request = medium_request
 
 
1163
    def set_headers(self, headers):
 
 
1164
        self._headers = headers.copy()
 
 
1166
    def call(self, *args):
 
 
1167
        if 'hpss' in debug.debug_flags:
 
 
1168
            mutter('hpss call:   %s', repr(args)[1:-1])
 
 
1169
            base = getattr(self._medium_request._medium, 'base', None)
 
 
1170
            if base is not None:
 
 
1171
                mutter('             (to %s)', base)
 
 
1172
            self._request_start_time = time.time()
 
 
1173
        self._write_protocol_version()
 
 
1174
        self._write_headers(self._headers)
 
 
1175
        self._write_structure(args)
 
 
1177
        self._medium_request.finished_writing()
 
 
1179
    def call_with_body_bytes(self, args, body):
 
 
1180
        """Make a remote call of args with body bytes 'body'.
 
 
1182
        After calling this, call read_response_tuple to find the result out.
 
 
1184
        if 'hpss' in debug.debug_flags:
 
 
1185
            mutter('hpss call w/body: %s (%r...)', repr(args)[1:-1], body[:20])
 
 
1186
            path = getattr(self._medium_request._medium, '_path', None)
 
 
1187
            if path is not None:
 
 
1188
                mutter('                  (to %s)', path)
 
 
1189
            mutter('              %d bytes', len(body))
 
 
1190
            self._request_start_time = time.time()
 
 
1191
        self._write_protocol_version()
 
 
1192
        self._write_headers(self._headers)
 
 
1193
        self._write_structure(args)
 
 
1194
        self._write_prefixed_body(body)
 
 
1196
        self._medium_request.finished_writing()
 
 
1198
    def call_with_body_readv_array(self, args, body):
 
 
1199
        """Make a remote call with a readv array.
 
 
1201
        The body is encoded with one line per readv offset pair. The numbers in
 
 
1202
        each pair are separated by a comma, and no trailing \n is emitted.
 
 
1204
        if 'hpss' in debug.debug_flags:
 
 
1205
            mutter('hpss call w/readv: %s', repr(args)[1:-1])
 
 
1206
            path = getattr(self._medium_request._medium, '_path', None)
 
 
1207
            if path is not None:
 
 
1208
                mutter('                  (to %s)', path)
 
 
1209
            self._request_start_time = time.time()
 
 
1210
        self._write_protocol_version()
 
 
1211
        self._write_headers(self._headers)
 
 
1212
        self._write_structure(args)
 
 
1213
        readv_bytes = self._serialise_offsets(body)
 
 
1214
        if 'hpss' in debug.debug_flags:
 
 
1215
            mutter('              %d bytes in readv request', len(readv_bytes))
 
 
1216
        self._write_prefixed_body(readv_bytes)
 
 
1218
        self._medium_request.finished_writing()