← Back to branch summary

~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/smart/protocol.py

  • Committer: Martin Packman
  • Date: 2011-12-23 19:38:22 UTC
  • mto: This revision was merged to the branch mainline in revision 6405.
  • Revision ID: martin.packman@canonical.com-20111223193822-hesheea4o8aqwexv
Accept and document passing the medium rather than transport for smart connections

Show diffs side-by-side

added added

removed removed

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