/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

  • Committer: John Arbash Meinel
  • Date: 2008-06-16 17:11:49 UTC
  • mfrom: (3496 +trunk)
  • mto: This revision was merged to the branch mainline in revision 3497.
  • Revision ID: john@arbash-meinel.com-20080616171149-q4cryeayazevsjlk
merge dev, resolve

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()
41
48
 
42
49
 
43
50
def _decode_tuple(req_line):
44
 
    if req_line == None or req_line == '':
 
51
    if req_line is None or req_line == '':
45
52
        return None
46
53
    if req_line[-1] != '\n':
47
54
        raise errors.SmartProtocolError("request %r not terminated" % req_line)
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
 
73
114
class SmartServerRequestProtocolOne(SmartProtocolBase):
74
115
    """Server-side encoding and decoding logic for smart version 1."""
75
116
    
76
 
    def __init__(self, backing_transport, write_func):
 
117
    def __init__(self, backing_transport, write_func, root_client_path='/'):
77
118
        self._backing_transport = backing_transport
78
 
        self.excess_buffer = ''
 
119
        self._root_client_path = root_client_path
 
120
        self.unused_data = ''
79
121
        self._finished = False
80
122
        self.in_buffer = ''
81
 
        self.has_dispatched = False
 
123
        self._has_dispatched = False
82
124
        self.request = None
83
125
        self._body_decoder = None
84
126
        self._write_func = write_func
88
130
        
89
131
        :param bytes: must be a byte string
90
132
        """
91
 
        assert isinstance(bytes, str)
 
133
        if not isinstance(bytes, str):
 
134
            raise ValueError(bytes)
92
135
        self.in_buffer += bytes
93
 
        if not self.has_dispatched:
 
136
        if not self._has_dispatched:
94
137
            if '\n' not in self.in_buffer:
95
138
                # no command line yet
96
139
                return
97
 
            self.has_dispatched = True
 
140
            self._has_dispatched = True
98
141
            try:
99
142
                first_line, self.in_buffer = self.in_buffer.split('\n', 1)
100
143
                first_line += '\n'
101
144
                req_args = _decode_tuple(first_line)
102
145
                self.request = request.SmartServerRequestHandler(
103
 
                    self._backing_transport, commands=request.request_handlers)
 
146
                    self._backing_transport, commands=request.request_handlers,
 
147
                    root_client_path=self._root_client_path)
104
148
                self.request.dispatch_command(req_args[0], req_args[1:])
105
149
                if self.request.finished_reading:
106
150
                    # trivial request
107
 
                    self.excess_buffer = self.in_buffer
 
151
                    self.unused_data = self.in_buffer
108
152
                    self.in_buffer = ''
109
153
                    self._send_response(self.request.response)
110
154
            except KeyboardInterrupt:
111
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
112
163
            except Exception, exception:
113
164
                # everything else: pass to client, flush, and quit
114
165
                log_exception_quietly()
116
167
                    ('error', str(exception))))
117
168
                return
118
169
 
119
 
        if self.has_dispatched:
 
170
        if self._has_dispatched:
120
171
            if self._finished:
121
172
                # nothing to do.XXX: this routine should be a single state 
122
173
                # machine too.
123
 
                self.excess_buffer += self.in_buffer
 
174
                self.unused_data += self.in_buffer
124
175
                self.in_buffer = ''
125
176
                return
126
177
            if self._body_decoder is None:
131
182
            self.request.accept_body(body_data)
132
183
            if self._body_decoder.finished_reading:
133
184
                self.request.end_of_body()
134
 
                assert self.request.finished_reading, \
135
 
                    "no more body, request not finished"
 
185
                if not self.request.finished_reading:
 
186
                    raise AssertionError("no more body, request not finished")
136
187
            if self.request.response is not None:
137
188
                self._send_response(self.request.response)
138
 
                self.excess_buffer = self.in_buffer
 
189
                self.unused_data = self.in_buffer
139
190
                self.in_buffer = ''
140
191
            else:
141
 
                assert not self.request.finished_reading, \
142
 
                    "no response and we have finished reading."
 
192
                if self.request.finished_reading:
 
193
                    raise AssertionError(
 
194
                        "no response and we have finished reading.")
143
195
 
144
196
    def _send_response(self, response):
145
197
        """Send a smart server response down the output stream."""
146
 
        assert not self._finished, 'response already sent'
 
198
        if self._finished:
 
199
            raise AssertionError('response already sent')
147
200
        args = response.args
148
201
        body = response.body
149
202
        self._finished = True
151
204
        self._write_success_or_failure_prefix(response)
152
205
        self._write_func(_encode_tuple(args))
153
206
        if body is not None:
154
 
            assert isinstance(body, str), 'body must be a str'
 
207
            if not isinstance(body, str):
 
208
                raise ValueError(body)
155
209
            bytes = self._encode_bulk_data(body)
156
210
            self._write_func(bytes)
157
211
 
184
238
    This prefixes responses with the value of RESPONSE_VERSION_TWO.
185
239
    """
186
240
 
 
241
    response_marker = RESPONSE_VERSION_TWO
 
242
    request_marker = REQUEST_VERSION_TWO
 
243
 
187
244
    def _write_success_or_failure_prefix(self, response):
188
245
        """Write the protocol specific success/failure prefix."""
189
246
        if response.is_successful():
196
253
        
197
254
        Version two sends the value of RESPONSE_VERSION_TWO.
198
255
        """
199
 
        self._write_func(RESPONSE_VERSION_TWO)
 
256
        self._write_func(self.response_marker)
200
257
 
201
258
    def _send_response(self, response):
202
259
        """Send a smart server response down the output stream."""
203
 
        assert not self._finished, 'response already sent'
 
260
        if (self._finished):
 
261
            raise AssertionError('response already sent')
204
262
        self._finished = True
205
263
        self._write_protocol_version()
206
264
        self._write_success_or_failure_prefix(response)
207
265
        self._write_func(_encode_tuple(response.args))
208
266
        if response.body is not None:
209
 
            assert isinstance(response.body, str), 'body must be a str'
210
 
            assert response.body_stream is None, (
211
 
                'body_stream and body cannot both be set')
 
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')
212
272
            bytes = self._encode_bulk_data(response.body)
213
273
            self._write_func(bytes)
214
274
        elif response.body_stream is not None:
236
296
                % chunk)
237
297
 
238
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
 
239
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
    """
240
323
 
241
324
    def __init__(self):
242
325
        self.finished_reading = False
 
326
        self._in_buffer = ''
243
327
        self.unused_data = ''
244
328
        self.bytes_left = None
 
329
        self._number_needed_bytes = None
245
330
 
246
331
    def accept_bytes(self, bytes):
247
332
        """Decode as much of bytes as possible.
254
339
        """
255
340
        # accept_bytes is allowed to change the state
256
341
        current_state = self.state_accept
257
 
        self.state_accept(bytes)
258
 
        while current_state != self.state_accept:
259
 
            current_state = self.state_accept
260
 
            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
261
357
 
262
358
 
263
359
class ChunkedBodyDecoder(_StatefulDecoder):
270
366
    def __init__(self):
271
367
        _StatefulDecoder.__init__(self)
272
368
        self.state_accept = self._state_accept_expecting_header
273
 
        self._in_buffer = ''
274
369
        self.chunk_in_progress = None
275
370
        self.chunks = collections.deque()
276
371
        self.error = False
308
403
    def _extract_line(self):
309
404
        pos = self._in_buffer.find('\n')
310
405
        if pos == -1:
311
 
            # We haven't read a complete length prefix yet, so there's nothing
312
 
            # to do.
313
 
            return None
 
406
            # We haven't read a complete line yet, so request more bytes before
 
407
            # we continue.
 
408
            raise _NeedMoreBytes(1)
314
409
        line = self._in_buffer[:pos]
315
410
        # Trim the prefix (including '\n' delimiter) from the _in_buffer.
316
411
        self._in_buffer = self._in_buffer[pos+1:]
318
413
 
319
414
    def _finished(self):
320
415
        self.unused_data = self._in_buffer
321
 
        self._in_buffer = None
 
416
        self._in_buffer = ''
322
417
        self.state_accept = self._state_accept_reading_unused
323
418
        if self.error:
324
419
            error_args = tuple(self.error_in_progress)
326
421
            self.error_in_progress = None
327
422
        self.finished_reading = True
328
423
 
329
 
    def _state_accept_expecting_header(self, bytes):
330
 
        self._in_buffer += bytes
 
424
    def _state_accept_expecting_header(self):
331
425
        prefix = self._extract_line()
332
 
        if prefix is None:
333
 
            # We haven't read a complete length prefix yet, so there's nothing
334
 
            # to do.
335
 
            return
336
 
        elif prefix == 'chunked':
 
426
        if prefix == 'chunked':
337
427
            self.state_accept = self._state_accept_expecting_length
338
428
        else:
339
429
            raise errors.SmartProtocolError(
340
430
                'Bad chunked body header: "%s"' % (prefix,))
341
431
 
342
 
    def _state_accept_expecting_length(self, bytes):
343
 
        self._in_buffer += bytes
 
432
    def _state_accept_expecting_length(self):
344
433
        prefix = self._extract_line()
345
 
        if prefix is None:
346
 
            # We haven't read a complete length prefix yet, so there's nothing
347
 
            # to do.
348
 
            return
349
 
        elif prefix == 'ERR':
 
434
        if prefix == 'ERR':
350
435
            self.error = True
351
436
            self.error_in_progress = []
352
 
            self._state_accept_expecting_length('')
 
437
            self._state_accept_expecting_length()
353
438
            return
354
439
        elif prefix == 'END':
355
440
            # We've read the end-of-body marker.
362
447
            self.chunk_in_progress = ''
363
448
            self.state_accept = self._state_accept_reading_chunk
364
449
 
365
 
    def _state_accept_reading_chunk(self, bytes):
366
 
        self._in_buffer += bytes
 
450
    def _state_accept_reading_chunk(self):
367
451
        in_buffer_len = len(self._in_buffer)
368
452
        self.chunk_in_progress += self._in_buffer[:self.bytes_left]
369
453
        self._in_buffer = self._in_buffer[self.bytes_left:]
378
462
            self.chunk_in_progress = None
379
463
            self.state_accept = self._state_accept_expecting_length
380
464
        
381
 
    def _state_accept_reading_unused(self, bytes):
382
 
        self.unused_data += bytes
 
465
    def _state_accept_reading_unused(self):
 
466
        self.unused_data += self._in_buffer
 
467
        self._in_buffer = ''
383
468
 
384
469
 
385
470
class LengthPrefixedBodyDecoder(_StatefulDecoder):
389
474
        _StatefulDecoder.__init__(self)
390
475
        self.state_accept = self._state_accept_expecting_length
391
476
        self.state_read = self._state_read_no_data
392
 
        self._in_buffer = ''
 
477
        self._body = ''
393
478
        self._trailer_buffer = ''
394
479
    
395
480
    def next_read_size(self):
412
497
        """Return any pending data that has been decoded."""
413
498
        return self.state_read()
414
499
 
415
 
    def _state_accept_expecting_length(self, bytes):
416
 
        self._in_buffer += bytes
 
500
    def _state_accept_expecting_length(self):
417
501
        pos = self._in_buffer.find('\n')
418
502
        if pos == -1:
419
503
            return
420
504
        self.bytes_left = int(self._in_buffer[:pos])
421
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
422
511
        self.bytes_left -= len(self._in_buffer)
423
 
        self.state_accept = self._state_accept_reading_body
424
 
        self.state_read = self._state_read_in_buffer
425
 
 
426
 
    def _state_accept_reading_body(self, bytes):
427
 
        self._in_buffer += bytes
428
 
        self.bytes_left -= len(bytes)
 
512
        self._in_buffer = ''
429
513
        if self.bytes_left <= 0:
430
514
            # Finished with body
431
515
            if self.bytes_left != 0:
432
 
                self._trailer_buffer = self._in_buffer[self.bytes_left:]
433
 
                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]
434
518
            self.bytes_left = None
435
519
            self.state_accept = self._state_accept_reading_trailer
436
520
        
437
 
    def _state_accept_reading_trailer(self, bytes):
438
 
        self._trailer_buffer += bytes
 
521
    def _state_accept_reading_trailer(self):
 
522
        self._trailer_buffer += self._in_buffer
 
523
        self._in_buffer = ''
439
524
        # TODO: what if the trailer does not match "done\n"?  Should this raise
440
525
        # a ProtocolViolation exception?
441
526
        if self._trailer_buffer.startswith('done\n'):
443
528
            self.state_accept = self._state_accept_reading_unused
444
529
            self.finished_reading = True
445
530
    
446
 
    def _state_accept_reading_unused(self, bytes):
447
 
        self.unused_data += bytes
 
531
    def _state_accept_reading_unused(self):
 
532
        self.unused_data += self._in_buffer
 
533
        self._in_buffer = ''
448
534
 
449
535
    def _state_read_no_data(self):
450
536
        return ''
451
537
 
452
 
    def _state_read_in_buffer(self):
453
 
        result = self._in_buffer
454
 
        self._in_buffer = ''
 
538
    def _state_read_body_buffer(self):
 
539
        result = self._body
 
540
        self._body = ''
455
541
        return result
456
542
 
457
543
 
458
 
class SmartClientRequestProtocolOne(SmartProtocolBase):
 
544
class SmartClientRequestProtocolOne(SmartProtocolBase, Requester,
 
545
                                    message.ResponseHandler):
459
546
    """The client-side protocol for smart version 1."""
460
547
 
461
548
    def __init__(self, request):
467
554
        self._request = request
468
555
        self._body_buffer = None
469
556
        self._request_start_time = None
 
557
        self._last_verb = None
 
558
        self._headers = None
 
559
 
 
560
    def set_headers(self, headers):
 
561
        self._headers = dict(headers)
470
562
 
471
563
    def call(self, *args):
472
564
        if 'hpss' in debug.debug_flags:
476
568
            self._request_start_time = time.time()
477
569
        self._write_args(args)
478
570
        self._request.finished_writing()
 
571
        self._last_verb = args[0]
479
572
 
480
573
    def call_with_body_bytes(self, args, body):
481
574
        """Make a remote call of args with body bytes 'body'.
494
587
        bytes = self._encode_bulk_data(body)
495
588
        self._request.accept_bytes(bytes)
496
589
        self._request.finished_writing()
 
590
        self._last_verb = args[0]
497
591
 
498
592
    def call_with_body_readv_array(self, args, body):
499
593
        """Make a remote call with a readv array.
513
607
        self._request.finished_writing()
514
608
        if 'hpss' in debug.debug_flags:
515
609
            mutter('              %d bytes in readv request', len(readv_bytes))
 
610
        self._last_verb = args[0]
516
611
 
517
612
    def cancel_read_body(self):
518
613
        """After expecting a body, a response code may indicate one otherwise.
523
618
        """
524
619
        self._request.finished_reading()
525
620
 
526
 
    def read_response_tuple(self, expect_body=False):
527
 
        """Read a response tuple from the wire.
528
 
 
529
 
        This should only be called once.
530
 
        """
 
621
    def _read_response_tuple(self):
531
622
        result = self._recv_tuple()
532
623
        if 'hpss' in debug.debug_flags:
533
624
            if self._request_start_time is not None:
537
628
                self._request_start_time = None
538
629
            else:
539
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()
 
639
        self._response_is_unknown_method(result)
 
640
        self._raise_args_if_error(result)
540
641
        if not expect_body:
541
642
            self._request.finished_reading()
542
643
        return result
543
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
 
 
674
    def _response_is_unknown_method(self, result_tuple):
 
675
        """Raise UnexpectedSmartServerResponse if the response is an 'unknonwn
 
676
        method' response to the request.
 
677
        
 
678
        :param response: The response from a smart client call_expecting_body
 
679
            call.
 
680
        :param verb: The verb used in that call.
 
681
        :raises: UnexpectedSmartServerResponse
 
682
        """
 
683
        if (result_tuple == ('error', "Generic bzr smart protocol error: "
 
684
                "bad request '%s'" % self._last_verb) or
 
685
              result_tuple == ('error', "Generic bzr smart protocol error: "
 
686
                "bad request u'%s'" % self._last_verb)):
 
687
            # The response will have no body, so we've finished reading.
 
688
            self._request.finished_reading()
 
689
            raise errors.UnknownSmartMethod(self._last_verb)
 
690
        
544
691
    def read_body_bytes(self, count=-1):
545
692
        """Read bytes from the body, decoding into a byte stream.
546
693
        
557
704
        while not _body_decoder.finished_reading:
558
705
            bytes_wanted = min(_body_decoder.next_read_size(), max_read)
559
706
            bytes = self._request.read_bytes(bytes_wanted)
 
707
            if bytes == '':
 
708
                # end of file encountered reading from server
 
709
                raise errors.ConnectionReset(
 
710
                    "Connection lost while reading response body.")
560
711
            _body_decoder.accept_bytes(bytes)
561
712
        self._request.finished_reading()
562
713
        self._body_buffer = StringIO(_body_decoder.read_pending_data())
613
764
    This prefixes the request with the value of REQUEST_VERSION_TWO.
614
765
    """
615
766
 
 
767
    response_marker = RESPONSE_VERSION_TWO
 
768
    request_marker = REQUEST_VERSION_TWO
 
769
 
616
770
    def read_response_tuple(self, expect_body=False):
617
771
        """Read a response tuple from the wire.
618
772
 
619
773
        This should only be called once.
620
774
        """
621
775
        version = self._request.read_line()
622
 
        if version != RESPONSE_VERSION_TWO:
623
 
            raise errors.SmartProtocolError('bad protocol marker %r' % version)
 
776
        if version != self.response_marker:
 
777
            self._request.finished_reading()
 
778
            raise errors.UnexpectedProtocolVersionMarker(version)
624
779
        response_status = self._recv_line()
625
 
        if response_status not in ('success\n', 'failed\n'):
 
780
        result = SmartClientRequestProtocolOne._read_response_tuple(self)
 
781
        self._response_is_unknown_method(result)
 
782
        if response_status == 'success\n':
 
783
            self.response_status = True
 
784
            if not expect_body:
 
785
                self._request.finished_reading()
 
786
            return result
 
787
        elif response_status == 'failed\n':
 
788
            self.response_status = False
 
789
            self._request.finished_reading()
 
790
            raise errors.ErrorFromSmartServer(result)
 
791
        else:
626
792
            raise errors.SmartProtocolError(
627
793
                'bad protocol status %r' % response_status)
628
 
        self.response_status = response_status == 'success\n'
629
 
        return SmartClientRequestProtocolOne.read_response_tuple(self, expect_body)
630
794
 
631
795
    def _write_protocol_version(self):
632
796
        """Write any prefixes this protocol requires.
633
797
        
634
798
        Version two sends the value of REQUEST_VERSION_TWO.
635
799
        """
636
 
        self._request.accept_bytes(REQUEST_VERSION_TWO)
 
800
        self._request.accept_bytes(self.request_marker)
637
801
 
638
802
    def read_streamed_body(self):
639
803
        """Read bytes from the body, decoding into a byte stream.
645
809
        while not _body_decoder.finished_reading:
646
810
            bytes_wanted = min(_body_decoder.next_read_size(), max_read)
647
811
            bytes = self._request.read_bytes(bytes_wanted)
 
812
            if bytes == '':
 
813
                # end of file encountered reading from server
 
814
                raise errors.ConnectionReset(
 
815
                    "Connection lost while reading streamed body.")
648
816
            _body_decoder.accept_bytes(bytes)
649
817
            for body_bytes in iter(_body_decoder.read_next_chunk, None):
650
 
                if 'hpss' in debug.debug_flags:
 
818
                if 'hpss' in debug.debug_flags and type(body_bytes) is str:
651
819
                    mutter('              %d byte chunk read',
652
820
                           len(body_bytes))
653
821
                yield body_bytes
654
822
        self._request.finished_reading()
655
823
 
 
824
 
 
825
def build_server_protocol_three(backing_transport, write_func,
 
826
                                root_client_path):
 
827
    request_handler = request.SmartServerRequestHandler(
 
828
        backing_transport, commands=request.request_handlers,
 
829
        root_client_path=root_client_path)
 
830
    responder = ProtocolThreeResponder(write_func)
 
831
    message_handler = message.ConventionalRequestHandler(request_handler, responder)
 
832
    return ProtocolThreeDecoder(message_handler)
 
833
 
 
834
 
 
835
class ProtocolThreeDecoder(_StatefulDecoder):
 
836
 
 
837
    response_marker = RESPONSE_VERSION_THREE
 
838
    request_marker = REQUEST_VERSION_THREE
 
839
 
 
840
    def __init__(self, message_handler, expect_version_marker=False):
 
841
        _StatefulDecoder.__init__(self)
 
842
        self._has_dispatched = False
 
843
        # Initial state
 
844
        if expect_version_marker:
 
845
            self.state_accept = self._state_accept_expecting_protocol_version
 
846
            # We're expecting at least the protocol version marker + some
 
847
            # headers.
 
848
            self._number_needed_bytes = len(MESSAGE_VERSION_THREE) + 4
 
849
        else:
 
850
            self.state_accept = self._state_accept_expecting_headers
 
851
            self._number_needed_bytes = 4
 
852
        self.decoding_failed = False
 
853
        self.request_handler = self.message_handler = message_handler
 
854
 
 
855
    def accept_bytes(self, bytes):
 
856
        self._number_needed_bytes = None
 
857
        try:
 
858
            _StatefulDecoder.accept_bytes(self, bytes)
 
859
        except KeyboardInterrupt:
 
860
            raise
 
861
        except errors.SmartMessageHandlerError, exception:
 
862
            # We do *not* set self.decoding_failed here.  The message handler
 
863
            # has raised an error, but the decoder is still able to parse bytes
 
864
            # and determine when this message ends.
 
865
            log_exception_quietly()
 
866
            self.message_handler.protocol_error(exception.exc_value)
 
867
            # The state machine is ready to continue decoding, but the
 
868
            # exception has interrupted the loop that runs the state machine.
 
869
            # So we call accept_bytes again to restart it.
 
870
            self.accept_bytes('')
 
871
        except Exception, exception:
 
872
            # The decoder itself has raised an exception.  We cannot continue
 
873
            # decoding.
 
874
            self.decoding_failed = True
 
875
            if isinstance(exception, errors.UnexpectedProtocolVersionMarker):
 
876
                # This happens during normal operation when the client tries a
 
877
                # protocol version the server doesn't understand, so no need to
 
878
                # log a traceback every time.
 
879
                # Note that this can only happen when
 
880
                # expect_version_marker=True, which is only the case on the
 
881
                # client side.
 
882
                pass
 
883
            else:
 
884
                log_exception_quietly()
 
885
            self.message_handler.protocol_error(exception)
 
886
 
 
887
    def _extract_length_prefixed_bytes(self):
 
888
        if len(self._in_buffer) < 4:
 
889
            # A length prefix by itself is 4 bytes, and we don't even have that
 
890
            # many yet.
 
891
            raise _NeedMoreBytes(4)
 
892
        (length,) = struct.unpack('!L', self._in_buffer[:4])
 
893
        end_of_bytes = 4 + length
 
894
        if len(self._in_buffer) < end_of_bytes:
 
895
            # We haven't yet read as many bytes as the length-prefix says there
 
896
            # are.
 
897
            raise _NeedMoreBytes(end_of_bytes)
 
898
        # Extract the bytes from the buffer.
 
899
        bytes = self._in_buffer[4:end_of_bytes]
 
900
        self._in_buffer = self._in_buffer[end_of_bytes:]
 
901
        return bytes
 
902
 
 
903
    def _extract_prefixed_bencoded_data(self):
 
904
        prefixed_bytes = self._extract_length_prefixed_bytes()
 
905
        try:
 
906
            decoded = bdecode(prefixed_bytes)
 
907
        except ValueError:
 
908
            raise errors.SmartProtocolError(
 
909
                'Bytes %r not bencoded' % (prefixed_bytes,))
 
910
        return decoded
 
911
 
 
912
    def _extract_single_byte(self):
 
913
        if self._in_buffer == '':
 
914
            # The buffer is empty
 
915
            raise _NeedMoreBytes(1)
 
916
        one_byte = self._in_buffer[0]
 
917
        self._in_buffer = self._in_buffer[1:]
 
918
        return one_byte
 
919
 
 
920
    def _state_accept_expecting_protocol_version(self):
 
921
        needed_bytes = len(MESSAGE_VERSION_THREE) - len(self._in_buffer)
 
922
        if needed_bytes > 0:
 
923
            # We don't have enough bytes to check if the protocol version
 
924
            # marker is right.  But we can check if it is already wrong by
 
925
            # checking that the start of MESSAGE_VERSION_THREE matches what
 
926
            # we've read so far.
 
927
            # [In fact, if the remote end isn't bzr we might never receive
 
928
            # len(MESSAGE_VERSION_THREE) bytes.  So if the bytes we have so far
 
929
            # are wrong then we should just raise immediately rather than
 
930
            # stall.]
 
931
            if not MESSAGE_VERSION_THREE.startswith(self._in_buffer):
 
932
                # We have enough bytes to know the protocol version is wrong
 
933
                raise errors.UnexpectedProtocolVersionMarker(self._in_buffer)
 
934
            raise _NeedMoreBytes(len(MESSAGE_VERSION_THREE))
 
935
        if not self._in_buffer.startswith(MESSAGE_VERSION_THREE):
 
936
            raise errors.UnexpectedProtocolVersionMarker(self._in_buffer)
 
937
        self._in_buffer = self._in_buffer[len(MESSAGE_VERSION_THREE):]
 
938
        self.state_accept = self._state_accept_expecting_headers
 
939
 
 
940
    def _state_accept_expecting_headers(self):
 
941
        decoded = self._extract_prefixed_bencoded_data()
 
942
        if type(decoded) is not dict:
 
943
            raise errors.SmartProtocolError(
 
944
                'Header object %r is not a dict' % (decoded,))
 
945
        self.state_accept = self._state_accept_expecting_message_part
 
946
        try:
 
947
            self.message_handler.headers_received(decoded)
 
948
        except:
 
949
            raise errors.SmartMessageHandlerError(sys.exc_info())
 
950
    
 
951
    def _state_accept_expecting_message_part(self):
 
952
        message_part_kind = self._extract_single_byte()
 
953
        if message_part_kind == 'o':
 
954
            self.state_accept = self._state_accept_expecting_one_byte
 
955
        elif message_part_kind == 's':
 
956
            self.state_accept = self._state_accept_expecting_structure
 
957
        elif message_part_kind == 'b':
 
958
            self.state_accept = self._state_accept_expecting_bytes
 
959
        elif message_part_kind == 'e':
 
960
            self.done()
 
961
        else:
 
962
            raise errors.SmartProtocolError(
 
963
                'Bad message kind byte: %r' % (message_part_kind,))
 
964
 
 
965
    def _state_accept_expecting_one_byte(self):
 
966
        byte = self._extract_single_byte()
 
967
        self.state_accept = self._state_accept_expecting_message_part
 
968
        try:
 
969
            self.message_handler.byte_part_received(byte)
 
970
        except:
 
971
            raise errors.SmartMessageHandlerError(sys.exc_info())
 
972
 
 
973
    def _state_accept_expecting_bytes(self):
 
974
        # XXX: this should not buffer whole message part, but instead deliver
 
975
        # the bytes as they arrive.
 
976
        prefixed_bytes = self._extract_length_prefixed_bytes()
 
977
        self.state_accept = self._state_accept_expecting_message_part
 
978
        try:
 
979
            self.message_handler.bytes_part_received(prefixed_bytes)
 
980
        except:
 
981
            raise errors.SmartMessageHandlerError(sys.exc_info())
 
982
 
 
983
    def _state_accept_expecting_structure(self):
 
984
        structure = self._extract_prefixed_bencoded_data()
 
985
        self.state_accept = self._state_accept_expecting_message_part
 
986
        try:
 
987
            self.message_handler.structure_part_received(structure)
 
988
        except:
 
989
            raise errors.SmartMessageHandlerError(sys.exc_info())
 
990
 
 
991
    def done(self):
 
992
        self.unused_data = self._in_buffer
 
993
        self._in_buffer = ''
 
994
        self.state_accept = self._state_accept_reading_unused
 
995
        try:
 
996
            self.message_handler.end_received()
 
997
        except:
 
998
            raise errors.SmartMessageHandlerError(sys.exc_info())
 
999
 
 
1000
    def _state_accept_reading_unused(self):
 
1001
        self.unused_data += self._in_buffer
 
1002
        self._in_buffer = ''
 
1003
 
 
1004
    def next_read_size(self):
 
1005
        if self.state_accept == self._state_accept_reading_unused:
 
1006
            return 0
 
1007
        elif self.decoding_failed:
 
1008
            # An exception occured while processing this message, probably from
 
1009
            # self.message_handler.  We're not sure that this state machine is
 
1010
            # in a consistent state, so just signal that we're done (i.e. give
 
1011
            # up).
 
1012
            return 0
 
1013
        else:
 
1014
            if self._number_needed_bytes is not None:
 
1015
                return self._number_needed_bytes - len(self._in_buffer)
 
1016
            else:
 
1017
                raise AssertionError("don't know how many bytes are expected!")
 
1018
 
 
1019
 
 
1020
class _ProtocolThreeEncoder(object):
 
1021
 
 
1022
    response_marker = request_marker = MESSAGE_VERSION_THREE
 
1023
 
 
1024
    def __init__(self, write_func):
 
1025
        self._buf = ''
 
1026
        self._real_write_func = write_func
 
1027
 
 
1028
    def _write_func(self, bytes):
 
1029
        self._buf += bytes
 
1030
 
 
1031
    def flush(self):
 
1032
        if self._buf:
 
1033
            self._real_write_func(self._buf)
 
1034
            self._buf = ''
 
1035
 
 
1036
    def _serialise_offsets(self, offsets):
 
1037
        """Serialise a readv offset list."""
 
1038
        txt = []
 
1039
        for start, length in offsets:
 
1040
            txt.append('%d,%d' % (start, length))
 
1041
        return '\n'.join(txt)
 
1042
        
 
1043
    def _write_protocol_version(self):
 
1044
        self._write_func(MESSAGE_VERSION_THREE)
 
1045
 
 
1046
    def _write_prefixed_bencode(self, structure):
 
1047
        bytes = bencode(structure)
 
1048
        self._write_func(struct.pack('!L', len(bytes)))
 
1049
        self._write_func(bytes)
 
1050
 
 
1051
    def _write_headers(self, headers):
 
1052
        self._write_prefixed_bencode(headers)
 
1053
 
 
1054
    def _write_structure(self, args):
 
1055
        self._write_func('s')
 
1056
        utf8_args = []
 
1057
        for arg in args:
 
1058
            if type(arg) is unicode:
 
1059
                utf8_args.append(arg.encode('utf8'))
 
1060
            else:
 
1061
                utf8_args.append(arg)
 
1062
        self._write_prefixed_bencode(utf8_args)
 
1063
 
 
1064
    def _write_end(self):
 
1065
        self._write_func('e')
 
1066
        self.flush()
 
1067
 
 
1068
    def _write_prefixed_body(self, bytes):
 
1069
        self._write_func('b')
 
1070
        self._write_func(struct.pack('!L', len(bytes)))
 
1071
        self._write_func(bytes)
 
1072
 
 
1073
    def _write_error_status(self):
 
1074
        self._write_func('oE')
 
1075
 
 
1076
    def _write_success_status(self):
 
1077
        self._write_func('oS')
 
1078
 
 
1079
 
 
1080
class ProtocolThreeResponder(_ProtocolThreeEncoder):
 
1081
 
 
1082
    def __init__(self, write_func):
 
1083
        _ProtocolThreeEncoder.__init__(self, write_func)
 
1084
        self.response_sent = False
 
1085
        self._headers = {'Software version': bzrlib.__version__}
 
1086
 
 
1087
    def send_error(self, exception):
 
1088
        if self.response_sent:
 
1089
            raise AssertionError(
 
1090
                "send_error(%s) called, but response already sent."
 
1091
                % (exception,))
 
1092
        if isinstance(exception, errors.UnknownSmartMethod):
 
1093
            failure = request.FailedSmartServerResponse(
 
1094
                ('UnknownMethod', exception.verb))
 
1095
            self.send_response(failure)
 
1096
            return
 
1097
        self.response_sent = True
 
1098
        self._write_protocol_version()
 
1099
        self._write_headers(self._headers)
 
1100
        self._write_error_status()
 
1101
        self._write_structure(('error', str(exception)))
 
1102
        self._write_end()
 
1103
 
 
1104
    def send_response(self, response):
 
1105
        if self.response_sent:
 
1106
            raise AssertionError(
 
1107
                "send_response(%r) called, but response already sent."
 
1108
                % (response,))
 
1109
        self.response_sent = True
 
1110
        self._write_protocol_version()
 
1111
        self._write_headers(self._headers)
 
1112
        if response.is_successful():
 
1113
            self._write_success_status()
 
1114
        else:
 
1115
            self._write_error_status()
 
1116
        self._write_structure(response.args)
 
1117
        if response.body is not None:
 
1118
            self._write_prefixed_body(response.body)
 
1119
        elif response.body_stream is not None:
 
1120
            for chunk in response.body_stream:
 
1121
                self._write_prefixed_body(chunk)
 
1122
                self.flush()
 
1123
        self._write_end()
 
1124
        
 
1125
 
 
1126
class ProtocolThreeRequester(_ProtocolThreeEncoder, Requester):
 
1127
 
 
1128
    def __init__(self, medium_request):
 
1129
        _ProtocolThreeEncoder.__init__(self, medium_request.accept_bytes)
 
1130
        self._medium_request = medium_request
 
1131
        self._headers = {}
 
1132
 
 
1133
    def set_headers(self, headers):
 
1134
        self._headers = headers.copy()
 
1135
        
 
1136
    def call(self, *args):
 
1137
        if 'hpss' in debug.debug_flags:
 
1138
            mutter('hpss call:   %s', repr(args)[1:-1])
 
1139
            base = getattr(self._medium_request._medium, 'base', None)
 
1140
            if base is not None:
 
1141
                mutter('             (to %s)', base)
 
1142
            self._request_start_time = time.time()
 
1143
        self._write_protocol_version()
 
1144
        self._write_headers(self._headers)
 
1145
        self._write_structure(args)
 
1146
        self._write_end()
 
1147
        self._medium_request.finished_writing()
 
1148
 
 
1149
    def call_with_body_bytes(self, args, body):
 
1150
        """Make a remote call of args with body bytes 'body'.
 
1151
 
 
1152
        After calling this, call read_response_tuple to find the result out.
 
1153
        """
 
1154
        if 'hpss' in debug.debug_flags:
 
1155
            mutter('hpss call w/body: %s (%r...)', repr(args)[1:-1], body[:20])
 
1156
            path = getattr(self._medium_request._medium, '_path', None)
 
1157
            if path is not None:
 
1158
                mutter('                  (to %s)', path)
 
1159
            mutter('              %d bytes', len(body))
 
1160
            self._request_start_time = time.time()
 
1161
        self._write_protocol_version()
 
1162
        self._write_headers(self._headers)
 
1163
        self._write_structure(args)
 
1164
        self._write_prefixed_body(body)
 
1165
        self._write_end()
 
1166
        self._medium_request.finished_writing()
 
1167
 
 
1168
    def call_with_body_readv_array(self, args, body):
 
1169
        """Make a remote call with a readv array.
 
1170
 
 
1171
        The body is encoded with one line per readv offset pair. The numbers in
 
1172
        each pair are separated by a comma, and no trailing \n is emitted.
 
1173
        """
 
1174
        if 'hpss' in debug.debug_flags:
 
1175
            mutter('hpss call w/readv: %s', repr(args)[1:-1])
 
1176
            path = getattr(self._medium_request._medium, '_path', None)
 
1177
            if path is not None:
 
1178
                mutter('                  (to %s)', path)
 
1179
            self._request_start_time = time.time()
 
1180
        self._write_protocol_version()
 
1181
        self._write_headers(self._headers)
 
1182
        self._write_structure(args)
 
1183
        readv_bytes = self._serialise_offsets(body)
 
1184
        if 'hpss' in debug.debug_flags:
 
1185
            mutter('              %d bytes in readv request', len(readv_bytes))
 
1186
        self._write_prefixed_body(readv_bytes)
 
1187
        self._write_end()
 
1188
        self._medium_request.finished_writing()
 
1189