/brz/remove-bazaar

To get this branch, use:
bzr branch http://gegoxaren.bato24.eu/bzr/brz/remove-bazaar

« back to all changes in this revision

Viewing changes to bzrlib/smart/protocol.py

Implement version 3 of the network protocol. (Andrew Bennetts)

Show diffs side-by-side

added added

removed removed

Lines of Context:
bzrlib/smart/protocol.py
20
20
 
21
21
import collections
22
22
from cStringIO import StringIO
 
23
import struct
 
24
import sys
23
25
import time
24
26
 
 
27
import bzrlib
25
28
from bzrlib import debug
26
29
from bzrlib import errors
27
 
from bzrlib.smart import request
 
30
from bzrlib.smart import message, request
28
31
from bzrlib.trace import log_exception_quietly, mutter
 
32
from bzrlib.util.bencode import bdecode, bencode
29
33
 
30
34
 
31
35
# Protocol version strings.  These are sent as prefixes of bzr requests and
34
38
REQUEST_VERSION_TWO = 'bzr request 2\n'
35
39
RESPONSE_VERSION_TWO = 'bzr response 2\n'
36
40
 
 
41
MESSAGE_VERSION_THREE = 'bzr message 3 (bzr 1.6)\n'
 
42
RESPONSE_VERSION_THREE = REQUEST_VERSION_THREE = MESSAGE_VERSION_THREE
 
43
 
37
44
 
38
45
def _recv_tuple(from_file):
39
46
    req_line = from_file.readline()
53
60
    return '\x01'.join(args) + '\n'
54
61
 
55
62
 
 
63
class Requester(object):
 
64
    """Abstract base class for an object that can issue requests on a smart
 
65
    medium.
 
66
    """
 
67
 
 
68
    def call(self, *args):
 
69
        """Make a remote call.
 
70
 
 
71
        :param args: the arguments of this call.
 
72
        """
 
73
        raise NotImplementedError(self.call)
 
74
 
 
75
    def call_with_body_bytes(self, args, body):
 
76
        """Make a remote call with a body.
 
77
 
 
78
        :param args: the arguments of this call.
 
79
        :type body: str
 
80
        :param body: the body to send with the request.
 
81
        """
 
82
        raise NotImplementedError(self.call_with_body_bytes)
 
83
 
 
84
    def call_with_body_readv_array(self, args, body):
 
85
        """Make a remote call with a readv array.
 
86
 
 
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.
 
90
        """
 
91
        raise NotImplementedError(self.call_with_body_readv_array)
 
92
 
 
93
    def set_headers(self, headers):
 
94
        raise NotImplementedError(self.set_headers)
 
95
 
 
96
 
56
97
class SmartProtocolBase(object):
57
98
    """Methods common to client and server"""
58
99
 
76
117
    def __init__(self, backing_transport, write_func, root_client_path='/'):
77
118
        self._backing_transport = backing_transport
78
119
        self._root_client_path = root_client_path
79
 
        self.excess_buffer = ''
 
120
        self.unused_data = ''
80
121
        self._finished = False
81
122
        self.in_buffer = ''
82
 
        self.has_dispatched = False
 
123
        self._has_dispatched = False
83
124
        self.request = None
84
125
        self._body_decoder = None
85
126
        self._write_func = write_func
92
133
        if not isinstance(bytes, str):
93
134
            raise ValueError(bytes)
94
135
        self.in_buffer += bytes
95
 
        if not self.has_dispatched:
 
136
        if not self._has_dispatched:
96
137
            if '\n' not in self.in_buffer:
97
138
                # no command line yet
98
139
                return
99
 
            self.has_dispatched = True
 
140
            self._has_dispatched = True
100
141
            try:
101
142
                first_line, self.in_buffer = self.in_buffer.split('\n', 1)
102
143
                first_line += '\n'
107
148
                self.request.dispatch_command(req_args[0], req_args[1:])
108
149
                if self.request.finished_reading:
109
150
                    # trivial request
110
 
                    self.excess_buffer = self.in_buffer
 
151
                    self.unused_data = self.in_buffer
111
152
                    self.in_buffer = ''
112
153
                    self._send_response(self.request.response)
113
154
            except KeyboardInterrupt:
114
155
                raise
 
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)
 
162
                return
115
163
            except Exception, exception:
116
164
                # everything else: pass to client, flush, and quit
117
165
                log_exception_quietly()
119
167
                    ('error', str(exception))))
120
168
                return
121
169
 
122
 
        if self.has_dispatched:
 
170
        if self._has_dispatched:
123
171
            if self._finished:
124
172
                # nothing to do.XXX: this routine should be a single state 
125
173
                # machine too.
126
 
                self.excess_buffer += self.in_buffer
 
174
                self.unused_data += self.in_buffer
127
175
                self.in_buffer = ''
128
176
                return
129
177
            if self._body_decoder is None:
138
186
                    raise AssertionError("no more body, request not finished")
139
187
            if self.request.response is not None:
140
188
                self._send_response(self.request.response)
141
 
                self.excess_buffer = self.in_buffer
 
189
                self.unused_data = self.in_buffer
142
190
                self.in_buffer = ''
143
191
            else:
144
192
                if self.request.finished_reading:
190
238
    This prefixes responses with the value of RESPONSE_VERSION_TWO.
191
239
    """
192
240
 
 
241
    response_marker = RESPONSE_VERSION_TWO
 
242
    request_marker = REQUEST_VERSION_TWO
 
243
 
193
244
    def _write_success_or_failure_prefix(self, response):
194
245
        """Write the protocol specific success/failure prefix."""
195
246
        if response.is_successful():
202
253
        
203
254
        Version two sends the value of RESPONSE_VERSION_TWO.
204
255
        """
205
 
        self._write_func(RESPONSE_VERSION_TWO)
 
256
        self._write_func(self.response_marker)
206
257
 
207
258
    def _send_response(self, response):
208
259
        """Send a smart server response down the output stream."""
245
296
                % chunk)
246
297
 
247
298
 
 
299
class _NeedMoreBytes(Exception):
 
300
    """Raise this inside a _StatefulDecoder to stop decoding until more bytes
 
301
    have been received.
 
302
    """
 
303
 
 
304
    def __init__(self, count=None):
 
305
        """Constructor.
 
306
 
 
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.
 
309
        """
 
310
        self.count = count
 
311
 
 
312
 
248
313
class _StatefulDecoder(object):
 
314
    """Base class for writing state machines to decode byte streams.
 
315
 
 
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.
 
320
 
 
321
    See ProtocolThreeDecoder for an example subclass.
 
322
    """
249
323
 
250
324
    def __init__(self):
251
325
        self.finished_reading = False
 
326
        self._in_buffer = ''
252
327
        self.unused_data = ''
253
328
        self.bytes_left = None
 
329
        self._number_needed_bytes = None
254
330
 
255
331
    def accept_bytes(self, bytes):
256
332
        """Decode as much of bytes as possible.
263
339
        """
264
340
        # accept_bytes is allowed to change the state
265
341
        current_state = self.state_accept
266
 
        self.state_accept(bytes)
267
 
        while current_state != self.state_accept:
268
 
            current_state = self.state_accept
269
 
            self.state_accept('')
 
342
        self._number_needed_bytes = None
 
343
        self._in_buffer += bytes
 
344
        try:
 
345
            # Run the function for the current state.
 
346
            self.state_accept()
 
347
            while current_state != self.state_accept:
 
348
                # The current state has changed.  Run the function for the new
 
349
                # current state, so that it can:
 
350
                #   - decode any unconsumed bytes left in a buffer, and
 
351
                #   - signal how many more bytes are expected (via raising
 
352
                #     _NeedMoreBytes).
 
353
                current_state = self.state_accept
 
354
                self.state_accept()
 
355
        except _NeedMoreBytes, e:
 
356
            self._number_needed_bytes = e.count
270
357
 
271
358
 
272
359
class ChunkedBodyDecoder(_StatefulDecoder):
279
366
    def __init__(self):
280
367
        _StatefulDecoder.__init__(self)
281
368
        self.state_accept = self._state_accept_expecting_header
282
 
        self._in_buffer = ''
283
369
        self.chunk_in_progress = None
284
370
        self.chunks = collections.deque()
285
371
        self.error = False
317
403
    def _extract_line(self):
318
404
        pos = self._in_buffer.find('\n')
319
405
        if pos == -1:
320
 
            # We haven't read a complete length prefix yet, so there's nothing
321
 
            # to do.
322
 
            return None
 
406
            # We haven't read a complete line yet, so request more bytes before
 
407
            # we continue.
 
408
            raise _NeedMoreBytes(1)
323
409
        line = self._in_buffer[:pos]
324
410
        # Trim the prefix (including '\n' delimiter) from the _in_buffer.
325
411
        self._in_buffer = self._in_buffer[pos+1:]
327
413
 
328
414
    def _finished(self):
329
415
        self.unused_data = self._in_buffer
330
 
        self._in_buffer = None
 
416
        self._in_buffer = ''
331
417
        self.state_accept = self._state_accept_reading_unused
332
418
        if self.error:
333
419
            error_args = tuple(self.error_in_progress)
335
421
            self.error_in_progress = None
336
422
        self.finished_reading = True
337
423
 
338
 
    def _state_accept_expecting_header(self, bytes):
339
 
        self._in_buffer += bytes
 
424
    def _state_accept_expecting_header(self):
340
425
        prefix = self._extract_line()
341
 
        if prefix is None:
342
 
            # We haven't read a complete length prefix yet, so there's nothing
343
 
            # to do.
344
 
            return
345
 
        elif prefix == 'chunked':
 
426
        if prefix == 'chunked':
346
427
            self.state_accept = self._state_accept_expecting_length
347
428
        else:
348
429
            raise errors.SmartProtocolError(
349
430
                'Bad chunked body header: "%s"' % (prefix,))
350
431
 
351
 
    def _state_accept_expecting_length(self, bytes):
352
 
        self._in_buffer += bytes
 
432
    def _state_accept_expecting_length(self):
353
433
        prefix = self._extract_line()
354
 
        if prefix is None:
355
 
            # We haven't read a complete length prefix yet, so there's nothing
356
 
            # to do.
357
 
            return
358
 
        elif prefix == 'ERR':
 
434
        if prefix == 'ERR':
359
435
            self.error = True
360
436
            self.error_in_progress = []
361
 
            self._state_accept_expecting_length('')
 
437
            self._state_accept_expecting_length()
362
438
            return
363
439
        elif prefix == 'END':
364
440
            # We've read the end-of-body marker.
371
447
            self.chunk_in_progress = ''
372
448
            self.state_accept = self._state_accept_reading_chunk
373
449
 
374
 
    def _state_accept_reading_chunk(self, bytes):
375
 
        self._in_buffer += bytes
 
450
    def _state_accept_reading_chunk(self):
376
451
        in_buffer_len = len(self._in_buffer)
377
452
        self.chunk_in_progress += self._in_buffer[:self.bytes_left]
378
453
        self._in_buffer = self._in_buffer[self.bytes_left:]
387
462
            self.chunk_in_progress = None
388
463
            self.state_accept = self._state_accept_expecting_length
389
464
        
390
 
    def _state_accept_reading_unused(self, bytes):
391
 
        self.unused_data += bytes
 
465
    def _state_accept_reading_unused(self):
 
466
        self.unused_data += self._in_buffer
 
467
        self._in_buffer = ''
392
468
 
393
469
 
394
470
class LengthPrefixedBodyDecoder(_StatefulDecoder):
398
474
        _StatefulDecoder.__init__(self)
399
475
        self.state_accept = self._state_accept_expecting_length
400
476
        self.state_read = self._state_read_no_data
401
 
        self._in_buffer = ''
 
477
        self._body = ''
402
478
        self._trailer_buffer = ''
403
479
    
404
480
    def next_read_size(self):
421
497
        """Return any pending data that has been decoded."""
422
498
        return self.state_read()
423
499
 
424
 
    def _state_accept_expecting_length(self, bytes):
425
 
        self._in_buffer += bytes
 
500
    def _state_accept_expecting_length(self):
426
501
        pos = self._in_buffer.find('\n')
427
502
        if pos == -1:
428
503
            return
429
504
        self.bytes_left = int(self._in_buffer[:pos])
430
505
        self._in_buffer = self._in_buffer[pos+1:]
 
506
        self.state_accept = self._state_accept_reading_body
 
507
        self.state_read = self._state_read_body_buffer
 
508
 
 
509
    def _state_accept_reading_body(self):
 
510
        self._body += self._in_buffer
431
511
        self.bytes_left -= len(self._in_buffer)
432
 
        self.state_accept = self._state_accept_reading_body
433
 
        self.state_read = self._state_read_in_buffer
434
 
 
435
 
    def _state_accept_reading_body(self, bytes):
436
 
        self._in_buffer += bytes
437
 
        self.bytes_left -= len(bytes)
 
512
        self._in_buffer = ''
438
513
        if self.bytes_left <= 0:
439
514
            # Finished with body
440
515
            if self.bytes_left != 0:
441
 
                self._trailer_buffer = self._in_buffer[self.bytes_left:]
442
 
                self._in_buffer = self._in_buffer[:self.bytes_left]
 
516
                self._trailer_buffer = self._body[self.bytes_left:]
 
517
                self._body = self._body[:self.bytes_left]
443
518
            self.bytes_left = None
444
519
            self.state_accept = self._state_accept_reading_trailer
445
520
        
446
 
    def _state_accept_reading_trailer(self, bytes):
447
 
        self._trailer_buffer += bytes
 
521
    def _state_accept_reading_trailer(self):
 
522
        self._trailer_buffer += self._in_buffer
 
523
        self._in_buffer = ''
448
524
        # TODO: what if the trailer does not match "done\n"?  Should this raise
449
525
        # a ProtocolViolation exception?
450
526
        if self._trailer_buffer.startswith('done\n'):
452
528
            self.state_accept = self._state_accept_reading_unused
453
529
            self.finished_reading = True
454
530
    
455
 
    def _state_accept_reading_unused(self, bytes):
456
 
        self.unused_data += bytes
 
531
    def _state_accept_reading_unused(self):
 
532
        self.unused_data += self._in_buffer
 
533
        self._in_buffer = ''
457
534
 
458
535
    def _state_read_no_data(self):
459
536
        return ''
460
537
 
461
 
    def _state_read_in_buffer(self):
462
 
        result = self._in_buffer
463
 
        self._in_buffer = ''
 
538
    def _state_read_body_buffer(self):
 
539
        result = self._body
 
540
        self._body = ''
464
541
        return result
465
542
 
466
543
 
467
 
class SmartClientRequestProtocolOne(SmartProtocolBase):
 
544
class SmartClientRequestProtocolOne(SmartProtocolBase, Requester,
 
545
                                    message.ResponseHandler):
468
546
    """The client-side protocol for smart version 1."""
469
547
 
470
548
    def __init__(self, request):
477
555
        self._body_buffer = None
478
556
        self._request_start_time = None
479
557
        self._last_verb = None
 
558
        self._headers = None
 
559
 
 
560
    def set_headers(self, headers):
 
561
        self._headers = dict(headers)
480
562
 
481
563
    def call(self, *args):
482
564
        if 'hpss' in debug.debug_flags:
536
618
        """
537
619
        self._request.finished_reading()
538
620
 
539
 
    def read_response_tuple(self, expect_body=False):
540
 
        """Read a response tuple from the wire.
541
 
 
542
 
        This should only be called once.
543
 
        """
 
621
    def _read_response_tuple(self):
544
622
        result = self._recv_tuple()
545
623
        if 'hpss' in debug.debug_flags:
546
624
            if self._request_start_time is not None:
550
628
                self._request_start_time = None
551
629
            else:
552
630
                mutter('   result:   %s', repr(result)[1:-1])
 
631
        return result
 
632
 
 
633
    def read_response_tuple(self, expect_body=False):
 
634
        """Read a response tuple from the wire.
 
635
 
 
636
        This should only be called once.
 
637
        """
 
638
        result = self._read_response_tuple()
553
639
        self._response_is_unknown_method(result)
 
640
        self._raise_args_if_error(result)
554
641
        if not expect_body:
555
642
            self._request.finished_reading()
556
643
        return result
557
644
 
 
645
    def _raise_args_if_error(self, result_tuple):
 
646
        # Later protocol versions have an explicit flag in the protocol to say
 
647
        # if an error response is "failed" or not.  In version 1 we don't have
 
648
        # that luxury.  So here is a complete list of errors that can be
 
649
        # returned in response to existing version 1 smart requests.  Responses
 
650
        # starting with these codes are always "failed" responses.
 
651
        v1_error_codes = [
 
652
            'norepository',
 
653
            'NoSuchFile',
 
654
            'FileExists',
 
655
            'DirectoryNotEmpty',
 
656
            'ShortReadvError',
 
657
            'UnicodeEncodeError',
 
658
            'UnicodeDecodeError',
 
659
            'ReadOnlyError',
 
660
            'nobranch',
 
661
            'NoSuchRevision',
 
662
            'nosuchrevision',
 
663
            'LockContention',
 
664
            'UnlockableTransport',
 
665
            'LockFailed',
 
666
            'TokenMismatch',
 
667
            'ReadError',
 
668
            'PermissionDenied',
 
669
            ]
 
670
        if result_tuple[0] in v1_error_codes:
 
671
            self._request.finished_reading()
 
672
            raise errors.ErrorFromSmartServer(result_tuple)
 
673
 
558
674
    def _response_is_unknown_method(self, result_tuple):
559
675
        """Raise UnexpectedSmartServerResponse if the response is an 'unknonwn
560
676
        method' response to the request.
644
760
    This prefixes the request with the value of REQUEST_VERSION_TWO.
645
761
    """
646
762
 
 
763
    response_marker = RESPONSE_VERSION_TWO
 
764
    request_marker = REQUEST_VERSION_TWO
 
765
 
647
766
    def read_response_tuple(self, expect_body=False):
648
767
        """Read a response tuple from the wire.
649
768
 
650
769
        This should only be called once.
651
770
        """
652
771
        version = self._request.read_line()
653
 
        if version != RESPONSE_VERSION_TWO:
654
 
            raise errors.SmartProtocolError('bad protocol marker %r' % version)
 
772
        if version != self.response_marker:
 
773
            self._request.finished_reading()
 
774
            raise errors.UnexpectedProtocolVersionMarker(version)
655
775
        response_status = self._recv_line()
656
 
        if response_status not in ('success\n', 'failed\n'):
 
776
        result = SmartClientRequestProtocolOne._read_response_tuple(self)
 
777
        self._response_is_unknown_method(result)
 
778
        if response_status == 'success\n':
 
779
            self.response_status = True
 
780
            if not expect_body:
 
781
                self._request.finished_reading()
 
782
            return result
 
783
        elif response_status == 'failed\n':
 
784
            self.response_status = False
 
785
            self._request.finished_reading()
 
786
            raise errors.ErrorFromSmartServer(result)
 
787
        else:
657
788
            raise errors.SmartProtocolError(
658
789
                'bad protocol status %r' % response_status)
659
 
        self.response_status = response_status == 'success\n'
660
 
        return SmartClientRequestProtocolOne.read_response_tuple(self, expect_body)
661
790
 
662
791
    def _write_protocol_version(self):
663
792
        """Write any prefixes this protocol requires.
664
793
        
665
794
        Version two sends the value of REQUEST_VERSION_TWO.
666
795
        """
667
 
        self._request.accept_bytes(REQUEST_VERSION_TWO)
 
796
        self._request.accept_bytes(self.request_marker)
668
797
 
669
798
    def read_streamed_body(self):
670
799
        """Read bytes from the body, decoding into a byte stream.
678
807
            bytes = self._request.read_bytes(bytes_wanted)
679
808
            _body_decoder.accept_bytes(bytes)
680
809
            for body_bytes in iter(_body_decoder.read_next_chunk, None):
681
 
                if 'hpss' in debug.debug_flags:
 
810
                if 'hpss' in debug.debug_flags and type(body_bytes) is str:
682
811
                    mutter('              %d byte chunk read',
683
812
                           len(body_bytes))
684
813
                yield body_bytes
685
814
        self._request.finished_reading()
686
815
 
 
816
 
 
817
def build_server_protocol_three(backing_transport, write_func,
 
818
                                root_client_path):
 
819
    request_handler = request.SmartServerRequestHandler(
 
820
        backing_transport, commands=request.request_handlers,
 
821
        root_client_path=root_client_path)
 
822
    responder = ProtocolThreeResponder(write_func)
 
823
    message_handler = message.ConventionalRequestHandler(request_handler, responder)
 
824
    return ProtocolThreeDecoder(message_handler)
 
825
 
 
826
 
 
827
class ProtocolThreeDecoder(_StatefulDecoder):
 
828
 
 
829
    response_marker = RESPONSE_VERSION_THREE
 
830
    request_marker = REQUEST_VERSION_THREE
 
831
 
 
832
    def __init__(self, message_handler, expect_version_marker=False):
 
833
        _StatefulDecoder.__init__(self)
 
834
        self._has_dispatched = False
 
835
        # Initial state
 
836
        if expect_version_marker:
 
837
            self.state_accept = self._state_accept_expecting_protocol_version
 
838
            # We're expecting at least the protocol version marker + some
 
839
            # headers.
 
840
            self._number_needed_bytes = len(MESSAGE_VERSION_THREE) + 4
 
841
        else:
 
842
            self.state_accept = self._state_accept_expecting_headers
 
843
            self._number_needed_bytes = 4
 
844
        self.decoding_failed = False
 
845
        self.request_handler = self.message_handler = message_handler
 
846
 
 
847
    def accept_bytes(self, bytes):
 
848
        self._number_needed_bytes = None
 
849
        try:
 
850
            _StatefulDecoder.accept_bytes(self, bytes)
 
851
        except KeyboardInterrupt:
 
852
            raise
 
853
        except errors.SmartMessageHandlerError, exception:
 
854
            # We do *not* set self.decoding_failed here.  The message handler
 
855
            # has raised an error, but the decoder is still able to parse bytes
 
856
            # and determine when this message ends.
 
857
            log_exception_quietly()
 
858
            self.message_handler.protocol_error(exception.exc_value)
 
859
            # The state machine is ready to continue decoding, but the
 
860
            # exception has interrupted the loop that runs the state machine.
 
861
            # So we call accept_bytes again to restart it.
 
862
            self.accept_bytes('')
 
863
        except Exception, exception:
 
864
            # The decoder itself has raised an exception.  We cannot continue
 
865
            # decoding.
 
866
            self.decoding_failed = True
 
867
            if isinstance(exception, errors.UnexpectedProtocolVersionMarker):
 
868
                # This happens during normal operation when the client tries a
 
869
                # protocol version the server doesn't understand, so no need to
 
870
                # log a traceback every time.
 
871
                # Note that this can only happen when
 
872
                # expect_version_marker=True, which is only the case on the
 
873
                # client side.
 
874
                pass
 
875
            else:
 
876
                log_exception_quietly()
 
877
            self.message_handler.protocol_error(exception)
 
878
 
 
879
    def _extract_length_prefixed_bytes(self):
 
880
        if len(self._in_buffer) < 4:
 
881
            # A length prefix by itself is 4 bytes, and we don't even have that
 
882
            # many yet.
 
883
            raise _NeedMoreBytes(4)
 
884
        (length,) = struct.unpack('!L', self._in_buffer[:4])
 
885
        end_of_bytes = 4 + length
 
886
        if len(self._in_buffer) < end_of_bytes:
 
887
            # We haven't yet read as many bytes as the length-prefix says there
 
888
            # are.
 
889
            raise _NeedMoreBytes(end_of_bytes)
 
890
        # Extract the bytes from the buffer.
 
891
        bytes = self._in_buffer[4:end_of_bytes]
 
892
        self._in_buffer = self._in_buffer[end_of_bytes:]
 
893
        return bytes
 
894
 
 
895
    def _extract_prefixed_bencoded_data(self):
 
896
        prefixed_bytes = self._extract_length_prefixed_bytes()
 
897
        try:
 
898
            decoded = bdecode(prefixed_bytes)
 
899
        except ValueError:
 
900
            raise errors.SmartProtocolError(
 
901
                'Bytes %r not bencoded' % (prefixed_bytes,))
 
902
        return decoded
 
903
 
 
904
    def _extract_single_byte(self):
 
905
        if self._in_buffer == '':
 
906
            # The buffer is empty
 
907
            raise _NeedMoreBytes(1)
 
908
        one_byte = self._in_buffer[0]
 
909
        self._in_buffer = self._in_buffer[1:]
 
910
        return one_byte
 
911
 
 
912
    def _state_accept_expecting_protocol_version(self):
 
913
        needed_bytes = len(MESSAGE_VERSION_THREE) - len(self._in_buffer)
 
914
        if needed_bytes > 0:
 
915
            # We don't have enough bytes to check if the protocol version
 
916
            # marker is right.  But we can check if it is already wrong by
 
917
            # checking that the start of MESSAGE_VERSION_THREE matches what
 
918
            # we've read so far.
 
919
            # [In fact, if the remote end isn't bzr we might never receive
 
920
            # len(MESSAGE_VERSION_THREE) bytes.  So if the bytes we have so far
 
921
            # are wrong then we should just raise immediately rather than
 
922
            # stall.]
 
923
            if not MESSAGE_VERSION_THREE.startswith(self._in_buffer):
 
924
                # We have enough bytes to know the protocol version is wrong
 
925
                raise errors.UnexpectedProtocolVersionMarker(self._in_buffer)
 
926
            raise _NeedMoreBytes(len(MESSAGE_VERSION_THREE))
 
927
        if not self._in_buffer.startswith(MESSAGE_VERSION_THREE):
 
928
            raise errors.UnexpectedProtocolVersionMarker(self._in_buffer)
 
929
        self._in_buffer = self._in_buffer[len(MESSAGE_VERSION_THREE):]
 
930
        self.state_accept = self._state_accept_expecting_headers
 
931
 
 
932
    def _state_accept_expecting_headers(self):
 
933
        decoded = self._extract_prefixed_bencoded_data()
 
934
        if type(decoded) is not dict:
 
935
            raise errors.SmartProtocolError(
 
936
                'Header object %r is not a dict' % (decoded,))
 
937
        self.state_accept = self._state_accept_expecting_message_part
 
938
        try:
 
939
            self.message_handler.headers_received(decoded)
 
940
        except:
 
941
            raise errors.SmartMessageHandlerError(sys.exc_info())
 
942
    
 
943
    def _state_accept_expecting_message_part(self):
 
944
        message_part_kind = self._extract_single_byte()
 
945
        if message_part_kind == 'o':
 
946
            self.state_accept = self._state_accept_expecting_one_byte
 
947
        elif message_part_kind == 's':
 
948
            self.state_accept = self._state_accept_expecting_structure
 
949
        elif message_part_kind == 'b':
 
950
            self.state_accept = self._state_accept_expecting_bytes
 
951
        elif message_part_kind == 'e':
 
952
            self.done()
 
953
        else:
 
954
            raise errors.SmartProtocolError(
 
955
                'Bad message kind byte: %r' % (message_part_kind,))
 
956
 
 
957
    def _state_accept_expecting_one_byte(self):
 
958
        byte = self._extract_single_byte()
 
959
        self.state_accept = self._state_accept_expecting_message_part
 
960
        try:
 
961
            self.message_handler.byte_part_received(byte)
 
962
        except:
 
963
            raise errors.SmartMessageHandlerError(sys.exc_info())
 
964
 
 
965
    def _state_accept_expecting_bytes(self):
 
966
        # XXX: this should not buffer whole message part, but instead deliver
 
967
        # the bytes as they arrive.
 
968
        prefixed_bytes = self._extract_length_prefixed_bytes()
 
969
        self.state_accept = self._state_accept_expecting_message_part
 
970
        try:
 
971
            self.message_handler.bytes_part_received(prefixed_bytes)
 
972
        except:
 
973
            raise errors.SmartMessageHandlerError(sys.exc_info())
 
974
 
 
975
    def _state_accept_expecting_structure(self):
 
976
        structure = self._extract_prefixed_bencoded_data()
 
977
        self.state_accept = self._state_accept_expecting_message_part
 
978
        try:
 
979
            self.message_handler.structure_part_received(structure)
 
980
        except:
 
981
            raise errors.SmartMessageHandlerError(sys.exc_info())
 
982
 
 
983
    def done(self):
 
984
        self.unused_data = self._in_buffer
 
985
        self._in_buffer = ''
 
986
        self.state_accept = self._state_accept_reading_unused
 
987
        try:
 
988
            self.message_handler.end_received()
 
989
        except:
 
990
            raise errors.SmartMessageHandlerError(sys.exc_info())
 
991
 
 
992
    def _state_accept_reading_unused(self):
 
993
        self.unused_data += self._in_buffer
 
994
        self._in_buffer = ''
 
995
 
 
996
    def next_read_size(self):
 
997
        if self.state_accept == self._state_accept_reading_unused:
 
998
            return 0
 
999
        elif self.decoding_failed:
 
1000
            # An exception occured while processing this message, probably from
 
1001
            # self.message_handler.  We're not sure that this state machine is
 
1002
            # in a consistent state, so just signal that we're done (i.e. give
 
1003
            # up).
 
1004
            return 0
 
1005
        else:
 
1006
            if self._number_needed_bytes is not None:
 
1007
                return self._number_needed_bytes - len(self._in_buffer)
 
1008
            else:
 
1009
                raise AssertionError("don't know how many bytes are expected!")
 
1010
 
 
1011
 
 
1012
class _ProtocolThreeEncoder(object):
 
1013
 
 
1014
    response_marker = request_marker = MESSAGE_VERSION_THREE
 
1015
 
 
1016
    def __init__(self, write_func):
 
1017
        self._write_func = write_func
 
1018
 
 
1019
    def _serialise_offsets(self, offsets):
 
1020
        """Serialise a readv offset list."""
 
1021
        txt = []
 
1022
        for start, length in offsets:
 
1023
            txt.append('%d,%d' % (start, length))
 
1024
        return '\n'.join(txt)
 
1025
        
 
1026
    def _write_protocol_version(self):
 
1027
        self._write_func(MESSAGE_VERSION_THREE)
 
1028
 
 
1029
    def _write_prefixed_bencode(self, structure):
 
1030
        bytes = bencode(structure)
 
1031
        self._write_func(struct.pack('!L', len(bytes)))
 
1032
        self._write_func(bytes)
 
1033
 
 
1034
    def _write_headers(self, headers):
 
1035
        self._write_prefixed_bencode(headers)
 
1036
 
 
1037
    def _write_structure(self, args):
 
1038
        self._write_func('s')
 
1039
        utf8_args = []
 
1040
        for arg in args:
 
1041
            if type(arg) is unicode:
 
1042
                utf8_args.append(arg.encode('utf8'))
 
1043
            else:
 
1044
                utf8_args.append(arg)
 
1045
        self._write_prefixed_bencode(utf8_args)
 
1046
 
 
1047
    def _write_end(self):
 
1048
        self._write_func('e')
 
1049
 
 
1050
    def _write_prefixed_body(self, bytes):
 
1051
        self._write_func('b')
 
1052
        self._write_func(struct.pack('!L', len(bytes)))
 
1053
        self._write_func(bytes)
 
1054
 
 
1055
    def _write_error_status(self):
 
1056
        self._write_func('oE')
 
1057
 
 
1058
    def _write_success_status(self):
 
1059
        self._write_func('oS')
 
1060
 
 
1061
 
 
1062
class ProtocolThreeResponder(_ProtocolThreeEncoder):
 
1063
 
 
1064
    def __init__(self, write_func):
 
1065
        _ProtocolThreeEncoder.__init__(self, write_func)
 
1066
        self.response_sent = False
 
1067
        self._headers = {'Software version': bzrlib.__version__}
 
1068
 
 
1069
    def send_error(self, exception):
 
1070
        if self.response_sent:
 
1071
            raise AssertionError(
 
1072
                "send_error(%s) called, but response already sent."
 
1073
                % (exception,))
 
1074
        if isinstance(exception, errors.UnknownSmartMethod):
 
1075
            failure = request.FailedSmartServerResponse(
 
1076
                ('UnknownMethod', exception.verb))
 
1077
            self.send_response(failure)
 
1078
            return
 
1079
        self.response_sent = True
 
1080
        self._write_protocol_version()
 
1081
        self._write_headers(self._headers)
 
1082
        self._write_error_status()
 
1083
        self._write_structure(('error', str(exception)))
 
1084
        self._write_end()
 
1085
 
 
1086
    def send_response(self, response):
 
1087
        if self.response_sent:
 
1088
            raise AssertionError(
 
1089
                "send_response(%r) called, but response already sent."
 
1090
                % (response,))
 
1091
        self.response_sent = True
 
1092
        self._write_protocol_version()
 
1093
        self._write_headers(self._headers)
 
1094
        if response.is_successful():
 
1095
            self._write_success_status()
 
1096
        else:
 
1097
            self._write_error_status()
 
1098
        self._write_structure(response.args)
 
1099
        if response.body is not None:
 
1100
            self._write_prefixed_body(response.body)
 
1101
        elif response.body_stream is not None:
 
1102
            for chunk in response.body_stream:
 
1103
                self._write_prefixed_body(chunk)
 
1104
        self._write_end()
 
1105
        
 
1106
 
 
1107
class ProtocolThreeRequester(_ProtocolThreeEncoder, Requester):
 
1108
 
 
1109
    def __init__(self, medium_request):
 
1110
        _ProtocolThreeEncoder.__init__(self, medium_request.accept_bytes)
 
1111
        self._medium_request = medium_request
 
1112
        self._headers = {}
 
1113
 
 
1114
    def set_headers(self, headers):
 
1115
        self._headers = headers.copy()
 
1116
        
 
1117
    def call(self, *args):
 
1118
        if 'hpss' in debug.debug_flags:
 
1119
            mutter('hpss call:   %s', repr(args)[1:-1])
 
1120
            base = getattr(self._medium_request._medium, 'base', None)
 
1121
            if base is not None:
 
1122
                mutter('             (to %s)', base)
 
1123
            self._request_start_time = time.time()
 
1124
        self._write_protocol_version()
 
1125
        self._write_headers(self._headers)
 
1126
        self._write_structure(args)
 
1127
        self._write_end()
 
1128
        self._medium_request.finished_writing()
 
1129
 
 
1130
    def call_with_body_bytes(self, args, body):
 
1131
        """Make a remote call of args with body bytes 'body'.
 
1132
 
 
1133
        After calling this, call read_response_tuple to find the result out.
 
1134
        """
 
1135
        if 'hpss' in debug.debug_flags:
 
1136
            mutter('hpss call w/body: %s (%r...)', repr(args)[1:-1], body[:20])
 
1137
            path = getattr(self._medium_request._medium, '_path', None)
 
1138
            if path is not None:
 
1139
                mutter('                  (to %s)', path)
 
1140
            mutter('              %d bytes', len(body))
 
1141
            self._request_start_time = time.time()
 
1142
        self._write_protocol_version()
 
1143
        self._write_headers(self._headers)
 
1144
        self._write_structure(args)
 
1145
        self._write_prefixed_body(body)
 
1146
        self._write_end()
 
1147
        self._medium_request.finished_writing()
 
1148
 
 
1149
    def call_with_body_readv_array(self, args, body):
 
1150
        """Make a remote call with a readv array.
 
1151
 
 
1152
        The body is encoded with one line per readv offset pair. The numbers in
 
1153
        each pair are separated by a comma, and no trailing \n is emitted.
 
1154
        """
 
1155
        if 'hpss' in debug.debug_flags:
 
1156
            mutter('hpss call w/readv: %s', repr(args)[1:-1])
 
1157
            path = getattr(self._medium_request._medium, '_path', None)
 
1158
            if path is not None:
 
1159
                mutter('                  (to %s)', path)
 
1160
            self._request_start_time = time.time()
 
1161
        self._write_protocol_version()
 
1162
        self._write_headers(self._headers)
 
1163
        self._write_structure(args)
 
1164
        readv_bytes = self._serialise_offsets(body)
 
1165
        if 'hpss' in debug.debug_flags:
 
1166
            mutter('              %d bytes in readv request', len(readv_bytes))
 
1167
        self._write_prefixed_body(readv_bytes)
 
1168
        self._write_end()
 
1169
        self._medium_request.finished_writing()
 
1170