← Back to branch summary

~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/smart/protocol.py

  • Committer: Canonical.com Patch Queue Manager
  • Date: 2009-04-09 20:23:07 UTC
  • mfrom: (4265.1.4 bbc-merge)
  • Revision ID: pqm@pqm.ubuntu.com-20090409202307-n0depb16qepoe21o
(jam) Change _fetch_uses_deltas = False for CHK repos until we can
        write a better fix.

Show diffs side-by-side

added added

removed removed

Lines of Context:
bzrlib/smart/protocol.py
12
12
#
13
13
# You should have received a copy of the GNU General Public License
14
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
 
15
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
16
16
 
17
17
"""Wire-level encoding and decoding of requests and responses for the smart
18
18
client and server.
19
19
"""
20
20
 
21
 
 
 
21
import collections
22
22
from cStringIO import StringIO
 
23
import struct
 
24
import sys
 
25
import time
23
26
 
 
27
import bzrlib
 
28
from bzrlib import debug
24
29
from bzrlib import errors
25
 
from bzrlib.smart import request
 
30
from bzrlib.smart import message, request
 
31
from bzrlib.trace import log_exception_quietly, mutter
 
32
from bzrlib.util.bencode import bdecode_as_tuple, bencode
 
33
 
 
34
 
 
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'
 
40
 
 
41
MESSAGE_VERSION_THREE = 'bzr message 3 (bzr 1.6)\n'
 
42
RESPONSE_VERSION_THREE = REQUEST_VERSION_THREE = MESSAGE_VERSION_THREE
26
43
 
27
44
 
28
45
def _recv_tuple(from_file):
31
48
 
32
49
 
33
50
def _decode_tuple(req_line):
34
 
    if req_line == None or req_line == '':
 
51
    if req_line is None or req_line == '':
35
52
        return None
36
53
    if req_line[-1] != '\n':
37
54
        raise errors.SmartProtocolError("request %r not terminated" % req_line)
43
60
    return '\x01'.join(args) + '\n'
44
61
 
45
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
 
46
97
class SmartProtocolBase(object):
47
98
    """Methods common to client and server"""
48
99
 
58
109
        for start, length in offsets:
59
110
            txt.append('%d,%d' % (start, length))
60
111
        return '\n'.join(txt)
61
 
        
 
112
 
62
113
 
63
114
class SmartServerRequestProtocolOne(SmartProtocolBase):
64
115
    """Server-side encoding and decoding logic for smart version 1."""
65
 
    
66
 
    def __init__(self, backing_transport, write_func):
 
116
 
 
117
    def __init__(self, backing_transport, write_func, root_client_path='/'):
67
118
        self._backing_transport = backing_transport
68
 
        self.excess_buffer = ''
 
119
        self._root_client_path = root_client_path
 
120
        self.unused_data = ''
69
121
        self._finished = False
70
122
        self.in_buffer = ''
71
 
        self.has_dispatched = False
 
123
        self._has_dispatched = False
72
124
        self.request = None
73
125
        self._body_decoder = None
74
126
        self._write_func = write_func
75
127
 
76
128
    def accept_bytes(self, bytes):
77
129
        """Take bytes, and advance the internal state machine appropriately.
78
 
        
 
130
 
79
131
        :param bytes: must be a byte string
80
132
        """
81
 
        assert isinstance(bytes, str)
 
133
        if not isinstance(bytes, str):
 
134
            raise ValueError(bytes)
82
135
        self.in_buffer += bytes
83
 
        if not self.has_dispatched:
 
136
        if not self._has_dispatched:
84
137
            if '\n' not in self.in_buffer:
85
138
                # no command line yet
86
139
                return
87
 
            self.has_dispatched = True
 
140
            self._has_dispatched = True
88
141
            try:
89
142
                first_line, self.in_buffer = self.in_buffer.split('\n', 1)
90
143
                first_line += '\n'
91
144
                req_args = _decode_tuple(first_line)
92
145
                self.request = request.SmartServerRequestHandler(
93
 
                    self._backing_transport, commands=request.request_handlers)
 
146
                    self._backing_transport, commands=request.request_handlers,
 
147
                    root_client_path=self._root_client_path)
94
148
                self.request.dispatch_command(req_args[0], req_args[1:])
95
149
                if self.request.finished_reading:
96
150
                    # trivial request
97
 
                    self.excess_buffer = self.in_buffer
 
151
                    self.unused_data = self.in_buffer
98
152
                    self.in_buffer = ''
99
 
                    self._send_response(self.request.response.args,
100
 
                        self.request.response.body)
 
153
                    self._send_response(self.request.response)
101
154
            except KeyboardInterrupt:
102
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
103
163
            except Exception, exception:
104
164
                # everything else: pass to client, flush, and quit
105
 
                self._send_response(('error', str(exception)))
 
165
                log_exception_quietly()
 
166
                self._send_response(request.FailedSmartServerResponse(
 
167
                    ('error', str(exception))))
106
168
                return
107
169
 
108
 
        if self.has_dispatched:
 
170
        if self._has_dispatched:
109
171
            if self._finished:
110
 
                # nothing to do.XXX: this routine should be a single state 
 
172
                # nothing to do.XXX: this routine should be a single state
111
173
                # machine too.
112
 
                self.excess_buffer += self.in_buffer
 
174
                self.unused_data += self.in_buffer
113
175
                self.in_buffer = ''
114
176
                return
115
177
            if self._body_decoder is None:
120
182
            self.request.accept_body(body_data)
121
183
            if self._body_decoder.finished_reading:
122
184
                self.request.end_of_body()
123
 
                assert self.request.finished_reading, \
124
 
                    "no more body, request not finished"
 
185
                if not self.request.finished_reading:
 
186
                    raise AssertionError("no more body, request not finished")
125
187
            if self.request.response is not None:
126
 
                self._send_response(self.request.response.args,
127
 
                    self.request.response.body)
128
 
                self.excess_buffer = self.in_buffer
 
188
                self._send_response(self.request.response)
 
189
                self.unused_data = self.in_buffer
129
190
                self.in_buffer = ''
130
191
            else:
131
 
                assert not self.request.finished_reading, \
132
 
                    "no response and we have finished reading."
 
192
                if self.request.finished_reading:
 
193
                    raise AssertionError(
 
194
                        "no response and we have finished reading.")
133
195
 
134
 
    def _send_response(self, args, body=None):
 
196
    def _send_response(self, response):
135
197
        """Send a smart server response down the output stream."""
136
 
        assert not self._finished, 'response already sent'
 
198
        if self._finished:
 
199
            raise AssertionError('response already sent')
 
200
        args = response.args
 
201
        body = response.body
137
202
        self._finished = True
 
203
        self._write_protocol_version()
 
204
        self._write_success_or_failure_prefix(response)
138
205
        self._write_func(_encode_tuple(args))
139
206
        if body is not None:
140
 
            assert isinstance(body, str), 'body must be a str'
 
207
            if not isinstance(body, str):
 
208
                raise ValueError(body)
141
209
            bytes = self._encode_bulk_data(body)
142
210
            self._write_func(bytes)
143
211
 
 
212
    def _write_protocol_version(self):
 
213
        """Write any prefixes this protocol requires.
 
214
 
 
215
        Version one doesn't send protocol versions.
 
216
        """
 
217
 
 
218
    def _write_success_or_failure_prefix(self, response):
 
219
        """Write the protocol specific success/failure prefix.
 
220
 
 
221
        For SmartServerRequestProtocolOne this is omitted but we
 
222
        call is_successful to ensure that the response is valid.
 
223
        """
 
224
        response.is_successful()
 
225
 
144
226
    def next_read_size(self):
145
227
        if self._finished:
146
228
            return 0
150
232
            return self._body_decoder.next_read_size()
151
233
 
152
234
 
153
 
class LengthPrefixedBodyDecoder(object):
154
 
    """Decodes the length-prefixed bulk data."""
155
 
    
 
235
class SmartServerRequestProtocolTwo(SmartServerRequestProtocolOne):
 
236
    r"""Version two of the server side of the smart protocol.
 
237
 
 
238
    This prefixes responses with the value of RESPONSE_VERSION_TWO.
 
239
    """
 
240
 
 
241
    response_marker = RESPONSE_VERSION_TWO
 
242
    request_marker = REQUEST_VERSION_TWO
 
243
 
 
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')
 
248
        else:
 
249
            self._write_func('failed\n')
 
250
 
 
251
    def _write_protocol_version(self):
 
252
        r"""Write any prefixes this protocol requires.
 
253
 
 
254
        Version two sends the value of RESPONSE_VERSION_TWO.
 
255
        """
 
256
        self._write_func(self.response_marker)
 
257
 
 
258
    def _send_response(self, response):
 
259
        """Send a smart server response down the output stream."""
 
260
        if (self._finished):
 
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)
 
276
 
 
277
 
 
278
def _send_stream(stream, write_func):
 
279
    write_func('chunked\n')
 
280
    _send_chunks(stream, write_func)
 
281
    write_func('END\n')
 
282
 
 
283
 
 
284
def _send_chunks(stream, write_func):
 
285
    for chunk in stream:
 
286
        if isinstance(chunk, str):
 
287
            bytes = "%x\n%s" % (len(chunk), chunk)
 
288
            write_func(bytes)
 
289
        elif isinstance(chunk, request.FailedSmartServerResponse):
 
290
            write_func('ERR\n')
 
291
            _send_chunks(chunk.args, write_func)
 
292
            return
 
293
        else:
 
294
            raise errors.BzrError(
 
295
                'Chunks must be str or FailedSmartServerResponse, got %r'
 
296
                % chunk)
 
297
 
 
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
 
 
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
    """
 
323
 
156
324
    def __init__(self):
 
325
        self.finished_reading = False
 
326
        self._in_buffer_list = []
 
327
        self._in_buffer_len = 0
 
328
        self.unused_data = ''
157
329
        self.bytes_left = None
158
 
        self.finished_reading = False
159
 
        self.unused_data = ''
160
 
        self.state_accept = self._state_accept_expecting_length
161
 
        self.state_read = self._state_read_no_data
162
 
        self._in_buffer = ''
163
 
        self._trailer_buffer = ''
164
 
    
 
330
        self._number_needed_bytes = None
 
331
 
 
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]
 
341
        return in_buffer
 
342
 
 
343
    def _get_in_bytes(self, count):
 
344
        """Grab X bytes from the input_buffer.
 
345
 
 
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.
 
350
        """
 
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
 
358
        # yield it from that
 
359
        in_buf = self._get_in_buffer()
 
360
        return in_buf[:count]
 
361
 
 
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)
 
366
        else:
 
367
            self._in_buffer_list = []
 
368
            self._in_buffer_len = 0
 
369
 
165
370
    def accept_bytes(self, bytes):
166
371
        """Decode as much of bytes as possible.
167
372
 
172
377
        data will be appended to self.unused_data.
173
378
        """
174
379
        # accept_bytes is allowed to change the state
175
 
        current_state = self.state_accept
176
 
        self.state_accept(bytes)
177
 
        while current_state != self.state_accept:
 
380
        self._number_needed_bytes = None
 
381
        # lsprof puts a very large amount of time on this specific call for
 
382
        # large readv arrays
 
383
        self._in_buffer_list.append(bytes)
 
384
        self._in_buffer_len += len(bytes)
 
385
        try:
 
386
            # Run the function for the current state.
178
387
            current_state = self.state_accept
179
 
            self.state_accept('')
 
388
            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
 
394
                #     _NeedMoreBytes).
 
395
                current_state = self.state_accept
 
396
                self.state_accept()
 
397
        except _NeedMoreBytes, e:
 
398
            self._number_needed_bytes = e.count
 
399
 
 
400
 
 
401
class ChunkedBodyDecoder(_StatefulDecoder):
 
402
    """Decoder for chunked body data.
 
403
 
 
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.
 
406
    """
 
407
 
 
408
    def __init__(self):
 
409
        _StatefulDecoder.__init__(self)
 
410
        self.state_accept = self._state_accept_expecting_header
 
411
        self.chunk_in_progress = None
 
412
        self.chunks = collections.deque()
 
413
        self.error = False
 
414
        self.error_in_progress = None
 
415
 
 
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'.
 
427
                return 2
 
428
            else:
 
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.
 
431
                return 1
 
432
        elif self.state_accept == self._state_accept_reading_unused:
 
433
            return 1
 
434
        elif self.state_accept == self._state_accept_expecting_header:
 
435
            return max(0, len('chunked\n') - self._in_buffer_len)
 
436
        else:
 
437
            raise AssertionError("Impossible state: %r" % (self.state_accept,))
 
438
 
 
439
    def read_next_chunk(self):
 
440
        try:
 
441
            return self.chunks.popleft()
 
442
        except IndexError:
 
443
            return None
 
444
 
 
445
    def _extract_line(self):
 
446
        in_buf = self._get_in_buffer()
 
447
        pos = in_buf.find('\n')
 
448
        if pos == -1:
 
449
            # We haven't read a complete line yet, so request more bytes before
 
450
            # we continue.
 
451
            raise _NeedMoreBytes(1)
 
452
        line = in_buf[:pos]
 
453
        # Trim the prefix (including '\n' delimiter) from the _in_buffer.
 
454
        self._set_in_buffer(in_buf[pos+1:])
 
455
        return line
 
456
 
 
457
    def _finished(self):
 
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
 
462
        if self.error:
 
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
 
467
 
 
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
 
472
        else:
 
473
            raise errors.SmartProtocolError(
 
474
                'Bad chunked body header: "%s"' % (prefix,))
 
475
 
 
476
    def _state_accept_expecting_length(self):
 
477
        prefix = self._extract_line()
 
478
        if prefix == 'ERR':
 
479
            self.error = True
 
480
            self.error_in_progress = []
 
481
            self._state_accept_expecting_length()
 
482
            return
 
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
 
486
            # the _in_buffer.
 
487
            self._finished()
 
488
            return
 
489
        else:
 
490
            self.bytes_left = int(prefix, 16)
 
491
            self.chunk_in_progress = ''
 
492
            self.state_accept = self._state_accept_reading_chunk
 
493
 
 
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
 
503
            if self.error:
 
504
                self.error_in_progress.append(self.chunk_in_progress)
 
505
            else:
 
506
                self.chunks.append(self.chunk_in_progress)
 
507
            self.chunk_in_progress = None
 
508
            self.state_accept = self._state_accept_expecting_length
 
509
 
 
510
    def _state_accept_reading_unused(self):
 
511
        self.unused_data += self._get_in_buffer()
 
512
        self._in_buffer_list = []
 
513
 
 
514
 
 
515
class LengthPrefixedBodyDecoder(_StatefulDecoder):
 
516
    """Decodes the length-prefixed bulk data."""
 
517
 
 
518
    def __init__(self):
 
519
        _StatefulDecoder.__init__(self)
 
520
        self.state_accept = self._state_accept_expecting_length
 
521
        self.state_read = self._state_read_no_data
 
522
        self._body = ''
 
523
        self._trailer_buffer = ''
180
524
 
181
525
    def next_read_size(self):
182
526
        if self.bytes_left is not None:
193
537
        else:
194
538
            # Reading excess data.  Either way, 1 byte at a time is fine.
195
539
            return 1
196
 
        
 
540
 
197
541
    def read_pending_data(self):
198
542
        """Return any pending data that has been decoded."""
199
543
        return self.state_read()
200
544
 
201
 
    def _state_accept_expecting_length(self, bytes):
202
 
        self._in_buffer += bytes
203
 
        pos = self._in_buffer.find('\n')
 
545
    def _state_accept_expecting_length(self):
 
546
        in_buf = self._get_in_buffer()
 
547
        pos = in_buf.find('\n')
204
548
        if pos == -1:
205
549
            return
206
 
        self.bytes_left = int(self._in_buffer[:pos])
207
 
        self._in_buffer = self._in_buffer[pos+1:]
208
 
        self.bytes_left -= len(self._in_buffer)
 
550
        self.bytes_left = int(in_buf[:pos])
 
551
        self._set_in_buffer(in_buf[pos+1:])
209
552
        self.state_accept = self._state_accept_reading_body
210
 
        self.state_read = self._state_read_in_buffer
 
553
        self.state_read = self._state_read_body_buffer
211
554
 
212
 
    def _state_accept_reading_body(self, bytes):
213
 
        self._in_buffer += bytes
214
 
        self.bytes_left -= len(bytes)
 
555
    def _state_accept_reading_body(self):
 
556
        in_buf = self._get_in_buffer()
 
557
        self._body += in_buf
 
558
        self.bytes_left -= len(in_buf)
 
559
        self._set_in_buffer(None)
215
560
        if self.bytes_left <= 0:
216
561
            # Finished with body
217
562
            if self.bytes_left != 0:
218
 
                self._trailer_buffer = self._in_buffer[self.bytes_left:]
219
 
                self._in_buffer = self._in_buffer[:self.bytes_left]
 
563
                self._trailer_buffer = self._body[self.bytes_left:]
 
564
                self._body = self._body[:self.bytes_left]
220
565
            self.bytes_left = None
221
566
            self.state_accept = self._state_accept_reading_trailer
222
 
        
223
 
    def _state_accept_reading_trailer(self, bytes):
224
 
        self._trailer_buffer += bytes
 
567
 
 
568
    def _state_accept_reading_trailer(self):
 
569
        self._trailer_buffer += self._get_in_buffer()
 
570
        self._set_in_buffer(None)
225
571
        # TODO: what if the trailer does not match "done\n"?  Should this raise
226
572
        # a ProtocolViolation exception?
227
573
        if self._trailer_buffer.startswith('done\n'):
228
574
            self.unused_data = self._trailer_buffer[len('done\n'):]
229
575
            self.state_accept = self._state_accept_reading_unused
230
576
            self.finished_reading = True
231
 
    
232
 
    def _state_accept_reading_unused(self, bytes):
233
 
        self.unused_data += bytes
 
577
 
 
578
    def _state_accept_reading_unused(self):
 
579
        self.unused_data += self._get_in_buffer()
 
580
        self._set_in_buffer(None)
234
581
 
235
582
    def _state_read_no_data(self):
236
583
        return ''
237
584
 
238
 
    def _state_read_in_buffer(self):
239
 
        result = self._in_buffer
240
 
        self._in_buffer = ''
 
585
    def _state_read_body_buffer(self):
 
586
        result = self._body
 
587
        self._body = ''
241
588
        return result
242
589
 
243
590
 
244
 
class SmartClientRequestProtocolOne(SmartProtocolBase):
 
591
class SmartClientRequestProtocolOne(SmartProtocolBase, Requester,
 
592
                                    message.ResponseHandler):
245
593
    """The client-side protocol for smart version 1."""
246
594
 
247
595
    def __init__(self, request):
252
600
        """
253
601
        self._request = request
254
602
        self._body_buffer = None
 
603
        self._request_start_time = None
 
604
        self._last_verb = None
 
605
        self._headers = None
 
606
 
 
607
    def set_headers(self, headers):
 
608
        self._headers = dict(headers)
255
609
 
256
610
    def call(self, *args):
257
 
        bytes = _encode_tuple(args)
258
 
        self._request.accept_bytes(bytes)
 
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)
259
617
        self._request.finished_writing()
 
618
        self._last_verb = args[0]
260
619
 
261
620
    def call_with_body_bytes(self, args, body):
262
621
        """Make a remote call of args with body bytes 'body'.
263
622
 
264
623
        After calling this, call read_response_tuple to find the result out.
265
624
        """
266
 
        bytes = _encode_tuple(args)
267
 
        self._request.accept_bytes(bytes)
 
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)
268
634
        bytes = self._encode_bulk_data(body)
269
635
        self._request.accept_bytes(bytes)
270
636
        self._request.finished_writing()
 
637
        self._last_verb = args[0]
271
638
 
272
639
    def call_with_body_readv_array(self, args, body):
273
640
        """Make a remote call with a readv array.
275
642
        The body is encoded with one line per readv offset pair. The numbers in
276
643
        each pair are separated by a comma, and no trailing \n is emitted.
277
644
        """
278
 
        bytes = _encode_tuple(args)
279
 
        self._request.accept_bytes(bytes)
 
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)
280
651
        readv_bytes = self._serialise_offsets(body)
281
652
        bytes = self._encode_bulk_data(readv_bytes)
282
653
        self._request.accept_bytes(bytes)
283
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]
 
658
 
 
659
    def call_with_body_stream(self, args, stream):
 
660
        # Protocols v1 and v2 don't support body streams.  So it's safe to
 
661
        # assume that a v1/v2 server doesn't support whatever method we're
 
662
        # trying to call with a body stream.
 
663
        self._request.finished_writing()
 
664
        self._request.finished_reading()
 
665
        raise errors.UnknownSmartMethod(args[0])
284
666
 
285
667
    def cancel_read_body(self):
286
668
        """After expecting a body, a response code may indicate one otherwise.
291
673
        """
292
674
        self._request.finished_reading()
293
675
 
294
 
    def read_response_tuple(self, expect_body=False):
295
 
        """Read a response tuple from the wire.
296
 
 
297
 
        This should only be called once.
298
 
        """
 
676
    def _read_response_tuple(self):
299
677
        result = self._recv_tuple()
 
678
        if 'hpss' in debug.debug_flags:
 
679
            if self._request_start_time is not None:
 
680
                mutter('   result:   %6.3fs  %s',
 
681
                       time.time() - self._request_start_time,
 
682
                       repr(result)[1:-1])
 
683
                self._request_start_time = None
 
684
            else:
 
685
                mutter('   result:   %s', repr(result)[1:-1])
 
686
        return result
 
687
 
 
688
    def read_response_tuple(self, expect_body=False):
 
689
        """Read a response tuple from the wire.
 
690
 
 
691
        This should only be called once.
 
692
        """
 
693
        result = self._read_response_tuple()
 
694
        self._response_is_unknown_method(result)
 
695
        self._raise_args_if_error(result)
300
696
        if not expect_body:
301
697
            self._request.finished_reading()
302
698
        return result
303
699
 
 
700
    def _raise_args_if_error(self, result_tuple):
 
701
        # Later protocol versions have an explicit flag in the protocol to say
 
702
        # if an error response is "failed" or not.  In version 1 we don't have
 
703
        # that luxury.  So here is a complete list of errors that can be
 
704
        # returned in response to existing version 1 smart requests.  Responses
 
705
        # starting with these codes are always "failed" responses.
 
706
        v1_error_codes = [
 
707
            'norepository',
 
708
            'NoSuchFile',
 
709
            'FileExists',
 
710
            'DirectoryNotEmpty',
 
711
            'ShortReadvError',
 
712
            'UnicodeEncodeError',
 
713
            'UnicodeDecodeError',
 
714
            'ReadOnlyError',
 
715
            'nobranch',
 
716
            'NoSuchRevision',
 
717
            'nosuchrevision',
 
718
            'LockContention',
 
719
            'UnlockableTransport',
 
720
            'LockFailed',
 
721
            'TokenMismatch',
 
722
            'ReadError',
 
723
            'PermissionDenied',
 
724
            ]
 
725
        if result_tuple[0] in v1_error_codes:
 
726
            self._request.finished_reading()
 
727
            raise errors.ErrorFromSmartServer(result_tuple)
 
728
 
 
729
    def _response_is_unknown_method(self, result_tuple):
 
730
        """Raise UnexpectedSmartServerResponse if the response is an 'unknonwn
 
731
        method' response to the request.
 
732
 
 
733
        :param response: The response from a smart client call_expecting_body
 
734
            call.
 
735
        :param verb: The verb used in that call.
 
736
        :raises: UnexpectedSmartServerResponse
 
737
        """
 
738
        if (result_tuple == ('error', "Generic bzr smart protocol error: "
 
739
                "bad request '%s'" % self._last_verb) or
 
740
              result_tuple == ('error', "Generic bzr smart protocol error: "
 
741
                "bad request u'%s'" % self._last_verb)):
 
742
            # The response will have no body, so we've finished reading.
 
743
            self._request.finished_reading()
 
744
            raise errors.UnknownSmartMethod(self._last_verb)
 
745
 
304
746
    def read_body_bytes(self, count=-1):
305
747
        """Read bytes from the body, decoding into a byte stream.
306
 
        
307
 
        We read all bytes at once to ensure we've checked the trailer for 
 
748
 
 
749
        We read all bytes at once to ensure we've checked the trailer for
308
750
        errors, and then feed the buffer back as read_body_bytes is called.
309
751
        """
310
752
        if self._body_buffer is not None:
312
754
        _body_decoder = LengthPrefixedBodyDecoder()
313
755
 
314
756
        while not _body_decoder.finished_reading:
315
 
            bytes_wanted = _body_decoder.next_read_size()
316
 
            bytes = self._request.read_bytes(bytes_wanted)
 
757
            bytes = self._request.read_bytes(_body_decoder.next_read_size())
 
758
            if bytes == '':
 
759
                # end of file encountered reading from server
 
760
                raise errors.ConnectionReset(
 
761
                    "Connection lost while reading response body.")
317
762
            _body_decoder.accept_bytes(bytes)
318
763
        self._request.finished_reading()
319
764
        self._body_buffer = StringIO(_body_decoder.read_pending_data())
320
765
        # XXX: TODO check the trailer result.
 
766
        if 'hpss' in debug.debug_flags:
 
767
            mutter('              %d body bytes read',
 
768
                   len(self._body_buffer.getvalue()))
321
769
        return self._body_buffer.read(count)
322
770
 
323
771
    def _recv_tuple(self):
324
772
        """Receive a tuple from the medium request."""
325
 
        line = ''
326
 
        while not line or line[-1] != '\n':
327
 
            # TODO: this is inefficient - but tuples are short.
328
 
            new_char = self._request.read_bytes(1)
329
 
            line += new_char
330
 
            assert new_char != '', "end of file reading from server."
331
 
        return _decode_tuple(line)
 
773
        return _decode_tuple(self._request.read_line())
332
774
 
333
775
    def query_version(self):
334
776
        """Return protocol version number of the server."""
336
778
        resp = self.read_response_tuple()
337
779
        if resp == ('ok', '1'):
338
780
            return 1
 
781
        elif resp == ('ok', '2'):
 
782
            return 2
339
783
        else:
340
784
            raise errors.SmartProtocolError("bad response %r" % (resp,))
341
785
 
342
 
 
 
786
    def _write_args(self, args):
 
787
        self._write_protocol_version()
 
788
        bytes = _encode_tuple(args)
 
789
        self._request.accept_bytes(bytes)
 
790
 
 
791
    def _write_protocol_version(self):
 
792
        """Write any prefixes this protocol requires.
 
793
 
 
794
        Version one doesn't send protocol versions.
 
795
        """
 
796
 
 
797
 
 
798
class SmartClientRequestProtocolTwo(SmartClientRequestProtocolOne):
 
799
    """Version two of the client side of the smart protocol.
 
800
 
 
801
    This prefixes the request with the value of REQUEST_VERSION_TWO.
 
802
    """
 
803
 
 
804
    response_marker = RESPONSE_VERSION_TWO
 
805
    request_marker = REQUEST_VERSION_TWO
 
806
 
 
807
    def read_response_tuple(self, expect_body=False):
 
808
        """Read a response tuple from the wire.
 
809
 
 
810
        This should only be called once.
 
811
        """
 
812
        version = self._request.read_line()
 
813
        if version != self.response_marker:
 
814
            self._request.finished_reading()
 
815
            raise errors.UnexpectedProtocolVersionMarker(version)
 
816
        response_status = self._request.read_line()
 
817
        result = SmartClientRequestProtocolOne._read_response_tuple(self)
 
818
        self._response_is_unknown_method(result)
 
819
        if response_status == 'success\n':
 
820
            self.response_status = True
 
821
            if not expect_body:
 
822
                self._request.finished_reading()
 
823
            return result
 
824
        elif response_status == 'failed\n':
 
825
            self.response_status = False
 
826
            self._request.finished_reading()
 
827
            raise errors.ErrorFromSmartServer(result)
 
828
        else:
 
829
            raise errors.SmartProtocolError(
 
830
                'bad protocol status %r' % response_status)
 
831
 
 
832
    def _write_protocol_version(self):
 
833
        """Write any prefixes this protocol requires.
 
834
 
 
835
        Version two sends the value of REQUEST_VERSION_TWO.
 
836
        """
 
837
        self._request.accept_bytes(self.request_marker)
 
838
 
 
839
    def read_streamed_body(self):
 
840
        """Read bytes from the body, decoding into a byte stream.
 
841
        """
 
842
        # Read no more than 64k at a time so that we don't risk error 10055 (no
 
843
        # buffer space available) on Windows.
 
844
        _body_decoder = ChunkedBodyDecoder()
 
845
        while not _body_decoder.finished_reading:
 
846
            bytes = self._request.read_bytes(_body_decoder.next_read_size())
 
847
            if bytes == '':
 
848
                # end of file encountered reading from server
 
849
                raise errors.ConnectionReset(
 
850
                    "Connection lost while reading streamed body.")
 
851
            _body_decoder.accept_bytes(bytes)
 
852
            for body_bytes in iter(_body_decoder.read_next_chunk, None):
 
853
                if 'hpss' in debug.debug_flags and type(body_bytes) is str:
 
854
                    mutter('              %d byte chunk read',
 
855
                           len(body_bytes))
 
856
                yield body_bytes
 
857
        self._request.finished_reading()
 
858
 
 
859
 
 
860
def build_server_protocol_three(backing_transport, write_func,
 
861
                                root_client_path):
 
862
    request_handler = request.SmartServerRequestHandler(
 
863
        backing_transport, commands=request.request_handlers,
 
864
        root_client_path=root_client_path)
 
865
    responder = ProtocolThreeResponder(write_func)
 
866
    message_handler = message.ConventionalRequestHandler(request_handler, responder)
 
867
    return ProtocolThreeDecoder(message_handler)
 
868
 
 
869
 
 
870
class ProtocolThreeDecoder(_StatefulDecoder):
 
871
 
 
872
    response_marker = RESPONSE_VERSION_THREE
 
873
    request_marker = REQUEST_VERSION_THREE
 
874
 
 
875
    def __init__(self, message_handler, expect_version_marker=False):
 
876
        _StatefulDecoder.__init__(self)
 
877
        self._has_dispatched = False
 
878
        # Initial state
 
879
        if expect_version_marker:
 
880
            self.state_accept = self._state_accept_expecting_protocol_version
 
881
            # We're expecting at least the protocol version marker + some
 
882
            # headers.
 
883
            self._number_needed_bytes = len(MESSAGE_VERSION_THREE) + 4
 
884
        else:
 
885
            self.state_accept = self._state_accept_expecting_headers
 
886
            self._number_needed_bytes = 4
 
887
        self.decoding_failed = False
 
888
        self.request_handler = self.message_handler = message_handler
 
889
 
 
890
    def accept_bytes(self, bytes):
 
891
        self._number_needed_bytes = None
 
892
        try:
 
893
            _StatefulDecoder.accept_bytes(self, bytes)
 
894
        except KeyboardInterrupt:
 
895
            raise
 
896
        except errors.SmartMessageHandlerError, exception:
 
897
            # We do *not* set self.decoding_failed here.  The message handler
 
898
            # has raised an error, but the decoder is still able to parse bytes
 
899
            # and determine when this message ends.
 
900
            log_exception_quietly()
 
901
            self.message_handler.protocol_error(exception.exc_value)
 
902
            # The state machine is ready to continue decoding, but the
 
903
            # exception has interrupted the loop that runs the state machine.
 
904
            # So we call accept_bytes again to restart it.
 
905
            self.accept_bytes('')
 
906
        except Exception, exception:
 
907
            # The decoder itself has raised an exception.  We cannot continue
 
908
            # decoding.
 
909
            self.decoding_failed = True
 
910
            if isinstance(exception, errors.UnexpectedProtocolVersionMarker):
 
911
                # This happens during normal operation when the client tries a
 
912
                # protocol version the server doesn't understand, so no need to
 
913
                # log a traceback every time.
 
914
                # Note that this can only happen when
 
915
                # expect_version_marker=True, which is only the case on the
 
916
                # client side.
 
917
                pass
 
918
            else:
 
919
                log_exception_quietly()
 
920
            self.message_handler.protocol_error(exception)
 
921
 
 
922
    def _extract_length_prefixed_bytes(self):
 
923
        if self._in_buffer_len < 4:
 
924
            # A length prefix by itself is 4 bytes, and we don't even have that
 
925
            # many yet.
 
926
            raise _NeedMoreBytes(4)
 
927
        (length,) = struct.unpack('!L', self._get_in_bytes(4))
 
928
        end_of_bytes = 4 + length
 
929
        if self._in_buffer_len < end_of_bytes:
 
930
            # We haven't yet read as many bytes as the length-prefix says there
 
931
            # are.
 
932
            raise _NeedMoreBytes(end_of_bytes)
 
933
        # Extract the bytes from the buffer.
 
934
        in_buf = self._get_in_buffer()
 
935
        bytes = in_buf[4:end_of_bytes]
 
936
        self._set_in_buffer(in_buf[end_of_bytes:])
 
937
        return bytes
 
938
 
 
939
    def _extract_prefixed_bencoded_data(self):
 
940
        prefixed_bytes = self._extract_length_prefixed_bytes()
 
941
        try:
 
942
            decoded = bdecode_as_tuple(prefixed_bytes)
 
943
        except ValueError:
 
944
            raise errors.SmartProtocolError(
 
945
                'Bytes %r not bencoded' % (prefixed_bytes,))
 
946
        return decoded
 
947
 
 
948
    def _extract_single_byte(self):
 
949
        if self._in_buffer_len == 0:
 
950
            # The buffer is empty
 
951
            raise _NeedMoreBytes(1)
 
952
        in_buf = self._get_in_buffer()
 
953
        one_byte = in_buf[0]
 
954
        self._set_in_buffer(in_buf[1:])
 
955
        return one_byte
 
956
 
 
957
    def _state_accept_expecting_protocol_version(self):
 
958
        needed_bytes = len(MESSAGE_VERSION_THREE) - self._in_buffer_len
 
959
        in_buf = self._get_in_buffer()
 
960
        if needed_bytes > 0:
 
961
            # We don't have enough bytes to check if the protocol version
 
962
            # marker is right.  But we can check if it is already wrong by
 
963
            # checking that the start of MESSAGE_VERSION_THREE matches what
 
964
            # we've read so far.
 
965
            # [In fact, if the remote end isn't bzr we might never receive
 
966
            # len(MESSAGE_VERSION_THREE) bytes.  So if the bytes we have so far
 
967
            # are wrong then we should just raise immediately rather than
 
968
            # stall.]
 
969
            if not MESSAGE_VERSION_THREE.startswith(in_buf):
 
970
                # We have enough bytes to know the protocol version is wrong
 
971
                raise errors.UnexpectedProtocolVersionMarker(in_buf)
 
972
            raise _NeedMoreBytes(len(MESSAGE_VERSION_THREE))
 
973
        if not in_buf.startswith(MESSAGE_VERSION_THREE):
 
974
            raise errors.UnexpectedProtocolVersionMarker(in_buf)
 
975
        self._set_in_buffer(in_buf[len(MESSAGE_VERSION_THREE):])
 
976
        self.state_accept = self._state_accept_expecting_headers
 
977
 
 
978
    def _state_accept_expecting_headers(self):
 
979
        decoded = self._extract_prefixed_bencoded_data()
 
980
        if type(decoded) is not dict:
 
981
            raise errors.SmartProtocolError(
 
982
                'Header object %r is not a dict' % (decoded,))
 
983
        self.state_accept = self._state_accept_expecting_message_part
 
984
        try:
 
985
            self.message_handler.headers_received(decoded)
 
986
        except:
 
987
            raise errors.SmartMessageHandlerError(sys.exc_info())
 
988
 
 
989
    def _state_accept_expecting_message_part(self):
 
990
        message_part_kind = self._extract_single_byte()
 
991
        if message_part_kind == 'o':
 
992
            self.state_accept = self._state_accept_expecting_one_byte
 
993
        elif message_part_kind == 's':
 
994
            self.state_accept = self._state_accept_expecting_structure
 
995
        elif message_part_kind == 'b':
 
996
            self.state_accept = self._state_accept_expecting_bytes
 
997
        elif message_part_kind == 'e':
 
998
            self.done()
 
999
        else:
 
1000
            raise errors.SmartProtocolError(
 
1001
                'Bad message kind byte: %r' % (message_part_kind,))
 
1002
 
 
1003
    def _state_accept_expecting_one_byte(self):
 
1004
        byte = self._extract_single_byte()
 
1005
        self.state_accept = self._state_accept_expecting_message_part
 
1006
        try:
 
1007
            self.message_handler.byte_part_received(byte)
 
1008
        except:
 
1009
            raise errors.SmartMessageHandlerError(sys.exc_info())
 
1010
 
 
1011
    def _state_accept_expecting_bytes(self):
 
1012
        # XXX: this should not buffer whole message part, but instead deliver
 
1013
        # the bytes as they arrive.
 
1014
        prefixed_bytes = self._extract_length_prefixed_bytes()
 
1015
        self.state_accept = self._state_accept_expecting_message_part
 
1016
        try:
 
1017
            self.message_handler.bytes_part_received(prefixed_bytes)
 
1018
        except:
 
1019
            raise errors.SmartMessageHandlerError(sys.exc_info())
 
1020
 
 
1021
    def _state_accept_expecting_structure(self):
 
1022
        structure = self._extract_prefixed_bencoded_data()
 
1023
        self.state_accept = self._state_accept_expecting_message_part
 
1024
        try:
 
1025
            self.message_handler.structure_part_received(structure)
 
1026
        except:
 
1027
            raise errors.SmartMessageHandlerError(sys.exc_info())
 
1028
 
 
1029
    def done(self):
 
1030
        self.unused_data = self._get_in_buffer()
 
1031
        self._set_in_buffer(None)
 
1032
        self.state_accept = self._state_accept_reading_unused
 
1033
        try:
 
1034
            self.message_handler.end_received()
 
1035
        except:
 
1036
            raise errors.SmartMessageHandlerError(sys.exc_info())
 
1037
 
 
1038
    def _state_accept_reading_unused(self):
 
1039
        self.unused_data = self._get_in_buffer()
 
1040
        self._set_in_buffer(None)
 
1041
 
 
1042
    def next_read_size(self):
 
1043
        if self.state_accept == self._state_accept_reading_unused:
 
1044
            return 0
 
1045
        elif self.decoding_failed:
 
1046
            # An exception occured while processing this message, probably from
 
1047
            # self.message_handler.  We're not sure that this state machine is
 
1048
            # in a consistent state, so just signal that we're done (i.e. give
 
1049
            # up).
 
1050
            return 0
 
1051
        else:
 
1052
            if self._number_needed_bytes is not None:
 
1053
                return self._number_needed_bytes - self._in_buffer_len
 
1054
            else:
 
1055
                raise AssertionError("don't know how many bytes are expected!")
 
1056
 
 
1057
 
 
1058
class _ProtocolThreeEncoder(object):
 
1059
 
 
1060
    response_marker = request_marker = MESSAGE_VERSION_THREE
 
1061
 
 
1062
    def __init__(self, write_func):
 
1063
        self._buf = []
 
1064
        self._real_write_func = write_func
 
1065
 
 
1066
    def _write_func(self, bytes):
 
1067
        self._buf.append(bytes)
 
1068
        if len(self._buf) > 100:
 
1069
            self.flush()
 
1070
 
 
1071
    def flush(self):
 
1072
        if self._buf:
 
1073
            self._real_write_func(''.join(self._buf))
 
1074
            del self._buf[:]
 
1075
 
 
1076
    def _serialise_offsets(self, offsets):
 
1077
        """Serialise a readv offset list."""
 
1078
        txt = []
 
1079
        for start, length in offsets:
 
1080
            txt.append('%d,%d' % (start, length))
 
1081
        return '\n'.join(txt)
 
1082
 
 
1083
    def _write_protocol_version(self):
 
1084
        self._write_func(MESSAGE_VERSION_THREE)
 
1085
 
 
1086
    def _write_prefixed_bencode(self, structure):
 
1087
        bytes = bencode(structure)
 
1088
        self._write_func(struct.pack('!L', len(bytes)))
 
1089
        self._write_func(bytes)
 
1090
 
 
1091
    def _write_headers(self, headers):
 
1092
        self._write_prefixed_bencode(headers)
 
1093
 
 
1094
    def _write_structure(self, args):
 
1095
        self._write_func('s')
 
1096
        utf8_args = []
 
1097
        for arg in args:
 
1098
            if type(arg) is unicode:
 
1099
                utf8_args.append(arg.encode('utf8'))
 
1100
            else:
 
1101
                utf8_args.append(arg)
 
1102
        self._write_prefixed_bencode(utf8_args)
 
1103
 
 
1104
    def _write_end(self):
 
1105
        self._write_func('e')
 
1106
        self.flush()
 
1107
 
 
1108
    def _write_prefixed_body(self, bytes):
 
1109
        self._write_func('b')
 
1110
        self._write_func(struct.pack('!L', len(bytes)))
 
1111
        self._write_func(bytes)
 
1112
 
 
1113
    def _write_chunked_body_start(self):
 
1114
        self._write_func('oC')
 
1115
 
 
1116
    def _write_error_status(self):
 
1117
        self._write_func('oE')
 
1118
 
 
1119
    def _write_success_status(self):
 
1120
        self._write_func('oS')
 
1121
 
 
1122
 
 
1123
class ProtocolThreeResponder(_ProtocolThreeEncoder):
 
1124
 
 
1125
    def __init__(self, write_func):
 
1126
        _ProtocolThreeEncoder.__init__(self, write_func)
 
1127
        self.response_sent = False
 
1128
        self._headers = {'Software version': bzrlib.__version__}
 
1129
 
 
1130
    def send_error(self, exception):
 
1131
        if self.response_sent:
 
1132
            raise AssertionError(
 
1133
                "send_error(%s) called, but response already sent."
 
1134
                % (exception,))
 
1135
        if isinstance(exception, errors.UnknownSmartMethod):
 
1136
            failure = request.FailedSmartServerResponse(
 
1137
                ('UnknownMethod', exception.verb))
 
1138
            self.send_response(failure)
 
1139
            return
 
1140
        self.response_sent = True
 
1141
        self._write_protocol_version()
 
1142
        self._write_headers(self._headers)
 
1143
        self._write_error_status()
 
1144
        self._write_structure(('error', str(exception)))
 
1145
        self._write_end()
 
1146
 
 
1147
    def send_response(self, response):
 
1148
        if self.response_sent:
 
1149
            raise AssertionError(
 
1150
                "send_response(%r) called, but response already sent."
 
1151
                % (response,))
 
1152
        self.response_sent = True
 
1153
        self._write_protocol_version()
 
1154
        self._write_headers(self._headers)
 
1155
        if response.is_successful():
 
1156
            self._write_success_status()
 
1157
        else:
 
1158
            self._write_error_status()
 
1159
        self._write_structure(response.args)
 
1160
        if response.body is not None:
 
1161
            self._write_prefixed_body(response.body)
 
1162
        elif response.body_stream is not None:
 
1163
            for exc_info, chunk in _iter_with_errors(response.body_stream):
 
1164
                if exc_info is not None:
 
1165
                    self._write_error_status()
 
1166
                    error_struct = request._translate_error(exc_info[1])
 
1167
                    self._write_structure(error_struct)
 
1168
                    break
 
1169
                else:
 
1170
                    if isinstance(chunk, request.FailedSmartServerResponse):
 
1171
                        self._write_error_status()
 
1172
                        self._write_structure(chunk.args)
 
1173
                        break
 
1174
                    self._write_prefixed_body(chunk)
 
1175
        self._write_end()
 
1176
 
 
1177
 
 
1178
def _iter_with_errors(iterable):
 
1179
    """Handle errors from iterable.next().
 
1180
 
 
1181
    Use like::
 
1182
 
 
1183
        for exc_info, value in _iter_with_errors(iterable):
 
1184
            ...
 
1185
 
 
1186
    This is a safer alternative to::
 
1187
 
 
1188
        try:
 
1189
            for value in iterable:
 
1190
               ...
 
1191
        except:
 
1192
            ...
 
1193
 
 
1194
    Because the latter will catch errors from the for-loop body, not just
 
1195
    iterable.next()
 
1196
 
 
1197
    If an error occurs, exc_info will be a exc_info tuple, and the generator
 
1198
    will terminate.  Otherwise exc_info will be None, and value will be the
 
1199
    value from iterable.next().  Note that KeyboardInterrupt and SystemExit
 
1200
    will not be itercepted.
 
1201
    """
 
1202
    iterator = iter(iterable)
 
1203
    while True:
 
1204
        try:
 
1205
            yield None, iterator.next()
 
1206
        except StopIteration:
 
1207
            return
 
1208
        except (KeyboardInterrupt, SystemExit):
 
1209
            raise
 
1210
        except Exception:
 
1211
            yield sys.exc_info(), None
 
1212
            return
 
1213
 
 
1214
 
 
1215
class ProtocolThreeRequester(_ProtocolThreeEncoder, Requester):
 
1216
 
 
1217
    def __init__(self, medium_request):
 
1218
        _ProtocolThreeEncoder.__init__(self, medium_request.accept_bytes)
 
1219
        self._medium_request = medium_request
 
1220
        self._headers = {}
 
1221
 
 
1222
    def set_headers(self, headers):
 
1223
        self._headers = headers.copy()
 
1224
 
 
1225
    def call(self, *args):
 
1226
        if 'hpss' in debug.debug_flags:
 
1227
            mutter('hpss call:   %s', repr(args)[1:-1])
 
1228
            base = getattr(self._medium_request._medium, 'base', None)
 
1229
            if base is not None:
 
1230
                mutter('             (to %s)', base)
 
1231
            self._request_start_time = time.time()
 
1232
        self._write_protocol_version()
 
1233
        self._write_headers(self._headers)
 
1234
        self._write_structure(args)
 
1235
        self._write_end()
 
1236
        self._medium_request.finished_writing()
 
1237
 
 
1238
    def call_with_body_bytes(self, args, body):
 
1239
        """Make a remote call of args with body bytes 'body'.
 
1240
 
 
1241
        After calling this, call read_response_tuple to find the result out.
 
1242
        """
 
1243
        if 'hpss' in debug.debug_flags:
 
1244
            mutter('hpss call w/body: %s (%r...)', repr(args)[1:-1], body[:20])
 
1245
            path = getattr(self._medium_request._medium, '_path', None)
 
1246
            if path is not None:
 
1247
                mutter('                  (to %s)', path)
 
1248
            mutter('              %d bytes', len(body))
 
1249
            self._request_start_time = time.time()
 
1250
        self._write_protocol_version()
 
1251
        self._write_headers(self._headers)
 
1252
        self._write_structure(args)
 
1253
        self._write_prefixed_body(body)
 
1254
        self._write_end()
 
1255
        self._medium_request.finished_writing()
 
1256
 
 
1257
    def call_with_body_readv_array(self, args, body):
 
1258
        """Make a remote call with a readv array.
 
1259
 
 
1260
        The body is encoded with one line per readv offset pair. The numbers in
 
1261
        each pair are separated by a comma, and no trailing \n is emitted.
 
1262
        """
 
1263
        if 'hpss' in debug.debug_flags:
 
1264
            mutter('hpss call w/readv: %s', repr(args)[1:-1])
 
1265
            path = getattr(self._medium_request._medium, '_path', None)
 
1266
            if path is not None:
 
1267
                mutter('                  (to %s)', path)
 
1268
            self._request_start_time = time.time()
 
1269
        self._write_protocol_version()
 
1270
        self._write_headers(self._headers)
 
1271
        self._write_structure(args)
 
1272
        readv_bytes = self._serialise_offsets(body)
 
1273
        if 'hpss' in debug.debug_flags:
 
1274
            mutter('              %d bytes in readv request', len(readv_bytes))
 
1275
        self._write_prefixed_body(readv_bytes)
 
1276
        self._write_end()
 
1277
        self._medium_request.finished_writing()
 
1278
 
 
1279
    def call_with_body_stream(self, args, stream):
 
1280
        if 'hpss' in debug.debug_flags:
 
1281
            mutter('hpss call w/body stream: %r', args)
 
1282
            path = getattr(self._medium_request._medium, '_path', None)
 
1283
            if path is not None:
 
1284
                mutter('                  (to %s)', path)
 
1285
            self._request_start_time = time.time()
 
1286
        self._write_protocol_version()
 
1287
        self._write_headers(self._headers)
 
1288
        self._write_structure(args)
 
1289
        # TODO: notice if the server has sent an early error reply before we
 
1290
        #       have finished sending the stream.  We would notice at the end
 
1291
        #       anyway, but if the medium can deliver it early then it's good
 
1292
        #       to short-circuit the whole request...
 
1293
        for exc_info, part in _iter_with_errors(stream):
 
1294
            if exc_info is not None:
 
1295
                # Iterating the stream failed.  Cleanly abort the request.
 
1296
                self._write_error_status()
 
1297
                # Currently the client unconditionally sends ('error',) as the
 
1298
                # error args.
 
1299
                self._write_structure(('error',))
 
1300
                self._write_end()
 
1301
                self._medium_request.finished_writing()
 
1302
                raise exc_info[0], exc_info[1], exc_info[2]
 
1303
            else:
 
1304
                self._write_prefixed_body(part)
 
1305
                self.flush()
 
1306
        self._write_end()
 
1307
        self._medium_request.finished_writing()
343
1308