← Back to branch summary

~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/smart/protocol.py

  • Committer: Jelmer Vernooij
  • Date: 2010-12-20 11:57:14 UTC
  • mto: This revision was merged to the branch mainline in revision 5577.
  • Revision ID: jelmer@samba.org-20101220115714-2ru3hfappjweeg7q
Don't use no-plugins.

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.
20
20
 
21
21
import collections
22
22
from cStringIO import StringIO
 
23
import struct
 
24
import sys
 
25
import thread
 
26
import threading
23
27
import time
24
28
 
25
 
from bzrlib import debug
26
 
from bzrlib import errors
27
 
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
28
36
from bzrlib.trace import log_exception_quietly, mutter
 
37
from bzrlib.bencode import bdecode_as_tuple, bencode
29
38
 
30
39
 
31
40
# Protocol version strings.  These are sent as prefixes of bzr requests and
34
43
REQUEST_VERSION_TWO = 'bzr request 2\n'
35
44
RESPONSE_VERSION_TWO = 'bzr response 2\n'
36
45
 
 
46
MESSAGE_VERSION_THREE = 'bzr message 3 (bzr 1.6)\n'
 
47
RESPONSE_VERSION_THREE = REQUEST_VERSION_THREE = MESSAGE_VERSION_THREE
 
48
 
37
49
 
38
50
def _recv_tuple(from_file):
39
51
    req_line = from_file.readline()
41
53
 
42
54
 
43
55
def _decode_tuple(req_line):
44
 
    if req_line == None or req_line == '':
 
56
    if req_line is None or req_line == '':
45
57
        return None
46
58
    if req_line[-1] != '\n':
47
59
        raise errors.SmartProtocolError("request %r not terminated" % req_line)
50
62
 
51
63
def _encode_tuple(args):
52
64
    """Encode the tuple args to a bytestream."""
53
 
    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)
54
106
 
55
107
 
56
108
class SmartProtocolBase(object):
68
120
        for start, length in offsets:
69
121
            txt.append('%d,%d' % (start, length))
70
122
        return '\n'.join(txt)
71
 
        
 
123
 
72
124
 
73
125
class SmartServerRequestProtocolOne(SmartProtocolBase):
74
126
    """Server-side encoding and decoding logic for smart version 1."""
75
 
    
76
 
    def __init__(self, backing_transport, write_func):
 
127
 
 
128
    def __init__(self, backing_transport, write_func, root_client_path='/',
 
129
            jail_root=None):
77
130
        self._backing_transport = backing_transport
78
 
        self.excess_buffer = ''
 
131
        self._root_client_path = root_client_path
 
132
        self._jail_root = jail_root
 
133
        self.unused_data = ''
79
134
        self._finished = False
80
135
        self.in_buffer = ''
81
 
        self.has_dispatched = False
 
136
        self._has_dispatched = False
82
137
        self.request = None
83
138
        self._body_decoder = None
84
139
        self._write_func = write_func
85
140
 
86
141
    def accept_bytes(self, bytes):
87
142
        """Take bytes, and advance the internal state machine appropriately.
88
 
        
 
143
 
89
144
        :param bytes: must be a byte string
90
145
        """
91
 
        assert isinstance(bytes, str)
 
146
        if not isinstance(bytes, str):
 
147
            raise ValueError(bytes)
92
148
        self.in_buffer += bytes
93
 
        if not self.has_dispatched:
 
149
        if not self._has_dispatched:
94
150
            if '\n' not in self.in_buffer:
95
151
                # no command line yet
96
152
                return
97
 
            self.has_dispatched = True
 
153
            self._has_dispatched = True
98
154
            try:
99
155
                first_line, self.in_buffer = self.in_buffer.split('\n', 1)
100
156
                first_line += '\n'
101
157
                req_args = _decode_tuple(first_line)
102
158
                self.request = request.SmartServerRequestHandler(
103
 
                    self._backing_transport, commands=request.request_handlers)
104
 
                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)
105
163
                if self.request.finished_reading:
106
164
                    # trivial request
107
 
                    self.excess_buffer = self.in_buffer
 
165
                    self.unused_data = self.in_buffer
108
166
                    self.in_buffer = ''
109
167
                    self._send_response(self.request.response)
110
168
            except KeyboardInterrupt:
111
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
112
177
            except Exception, exception:
113
178
                # everything else: pass to client, flush, and quit
114
179
                log_exception_quietly()
116
181
                    ('error', str(exception))))
117
182
                return
118
183
 
119
 
        if self.has_dispatched:
 
184
        if self._has_dispatched:
120
185
            if self._finished:
121
 
                # nothing to do.XXX: this routine should be a single state 
 
186
                # nothing to do.XXX: this routine should be a single state
122
187
                # machine too.
123
 
                self.excess_buffer += self.in_buffer
 
188
                self.unused_data += self.in_buffer
124
189
                self.in_buffer = ''
125
190
                return
126
191
            if self._body_decoder is None:
131
196
            self.request.accept_body(body_data)
132
197
            if self._body_decoder.finished_reading:
133
198
                self.request.end_of_body()
134
 
                assert self.request.finished_reading, \
135
 
                    "no more body, request not finished"
 
199
                if not self.request.finished_reading:
 
200
                    raise AssertionError("no more body, request not finished")
136
201
            if self.request.response is not None:
137
202
                self._send_response(self.request.response)
138
 
                self.excess_buffer = self.in_buffer
 
203
                self.unused_data = self.in_buffer
139
204
                self.in_buffer = ''
140
205
            else:
141
 
                assert not self.request.finished_reading, \
142
 
                    "no response and we have finished reading."
 
206
                if self.request.finished_reading:
 
207
                    raise AssertionError(
 
208
                        "no response and we have finished reading.")
143
209
 
144
210
    def _send_response(self, response):
145
211
        """Send a smart server response down the output stream."""
146
 
        assert not self._finished, 'response already sent'
 
212
        if self._finished:
 
213
            raise AssertionError('response already sent')
147
214
        args = response.args
148
215
        body = response.body
149
216
        self._finished = True
151
218
        self._write_success_or_failure_prefix(response)
152
219
        self._write_func(_encode_tuple(args))
153
220
        if body is not None:
154
 
            assert isinstance(body, str), 'body must be a str'
 
221
            if not isinstance(body, str):
 
222
                raise ValueError(body)
155
223
            bytes = self._encode_bulk_data(body)
156
224
            self._write_func(bytes)
157
225
 
158
226
    def _write_protocol_version(self):
159
227
        """Write any prefixes this protocol requires.
160
 
        
 
228
 
161
229
        Version one doesn't send protocol versions.
162
230
        """
163
231
 
180
248
 
181
249
class SmartServerRequestProtocolTwo(SmartServerRequestProtocolOne):
182
250
    r"""Version two of the server side of the smart protocol.
183
 
   
 
251
 
184
252
    This prefixes responses with the value of RESPONSE_VERSION_TWO.
185
253
    """
186
254
 
 
255
    response_marker = RESPONSE_VERSION_TWO
 
256
    request_marker = REQUEST_VERSION_TWO
 
257
 
187
258
    def _write_success_or_failure_prefix(self, response):
188
259
        """Write the protocol specific success/failure prefix."""
189
260
        if response.is_successful():
193
264
 
194
265
    def _write_protocol_version(self):
195
266
        r"""Write any prefixes this protocol requires.
196
 
        
 
267
 
197
268
        Version two sends the value of RESPONSE_VERSION_TWO.
198
269
        """
199
 
        self._write_func(RESPONSE_VERSION_TWO)
 
270
        self._write_func(self.response_marker)
200
271
 
201
272
    def _send_response(self, response):
202
273
        """Send a smart server response down the output stream."""
203
 
        assert not self._finished, 'response already sent'
 
274
        if (self._finished):
 
275
            raise AssertionError('response already sent')
204
276
        self._finished = True
205
277
        self._write_protocol_version()
206
278
        self._write_success_or_failure_prefix(response)
207
279
        self._write_func(_encode_tuple(response.args))
208
280
        if response.body is not None:
209
 
            assert isinstance(response.body, str), 'body must be a str'
210
 
            assert response.body_stream is None, (
211
 
                'body_stream and body cannot both be set')
 
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')
212
286
            bytes = self._encode_bulk_data(response.body)
213
287
            self._write_func(bytes)
214
288
        elif response.body_stream is not None:
236
310
                % chunk)
237
311
 
238
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
 
239
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
    """
240
337
 
241
338
    def __init__(self):
242
339
        self.finished_reading = False
 
340
        self._in_buffer_list = []
 
341
        self._in_buffer_len = 0
243
342
        self.unused_data = ''
244
343
        self.bytes_left = None
 
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
245
383
 
246
384
    def accept_bytes(self, bytes):
247
385
        """Decode as much of bytes as possible.
253
391
        data will be appended to self.unused_data.
254
392
        """
255
393
        # accept_bytes is allowed to change the state
256
 
        current_state = self.state_accept
257
 
        self.state_accept(bytes)
258
 
        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.
259
401
            current_state = self.state_accept
260
 
            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
261
413
 
262
414
 
263
415
class ChunkedBodyDecoder(_StatefulDecoder):
270
422
    def __init__(self):
271
423
        _StatefulDecoder.__init__(self)
272
424
        self.state_accept = self._state_accept_expecting_header
273
 
        self._in_buffer = ''
274
425
        self.chunk_in_progress = None
275
426
        self.chunks = collections.deque()
276
427
        self.error = False
277
428
        self.error_in_progress = None
278
 
    
 
429
 
279
430
    def next_read_size(self):
280
431
        # Note: the shortest possible chunk is 2 bytes: '0\n', and the
281
432
        # end-of-body marker is 4 bytes: 'END\n'.
284
435
            # the rest of this chunk plus an END chunk.
285
436
            return self.bytes_left + 4
286
437
        elif self.state_accept == self._state_accept_expecting_length:
287
 
            if self._in_buffer == '':
 
438
            if self._in_buffer_len == 0:
288
439
                # We're expecting a chunk length.  There's at least two bytes
289
440
                # left: a digit plus '\n'.
290
441
                return 2
295
446
        elif self.state_accept == self._state_accept_reading_unused:
296
447
            return 1
297
448
        elif self.state_accept == self._state_accept_expecting_header:
298
 
            return max(0, len('chunked\n') - len(self._in_buffer))
 
449
            return max(0, len('chunked\n') - self._in_buffer_len)
299
450
        else:
300
451
            raise AssertionError("Impossible state: %r" % (self.state_accept,))
301
452
 
306
457
            return None
307
458
 
308
459
    def _extract_line(self):
309
 
        pos = self._in_buffer.find('\n')
 
460
        in_buf = self._get_in_buffer()
 
461
        pos = in_buf.find('\n')
310
462
        if pos == -1:
311
 
            # We haven't read a complete length prefix yet, so there's nothing
312
 
            # to do.
313
 
            return None
314
 
        line = self._in_buffer[:pos]
 
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]
315
467
        # Trim the prefix (including '\n' delimiter) from the _in_buffer.
316
 
        self._in_buffer = self._in_buffer[pos+1:]
 
468
        self._set_in_buffer(in_buf[pos+1:])
317
469
        return line
318
470
 
319
471
    def _finished(self):
320
 
        self.unused_data = self._in_buffer
321
 
        self._in_buffer = None
 
472
        self.unused_data = self._get_in_buffer()
 
473
        self._in_buffer_list = []
 
474
        self._in_buffer_len = 0
322
475
        self.state_accept = self._state_accept_reading_unused
323
476
        if self.error:
324
477
            error_args = tuple(self.error_in_progress)
326
479
            self.error_in_progress = None
327
480
        self.finished_reading = True
328
481
 
329
 
    def _state_accept_expecting_header(self, bytes):
330
 
        self._in_buffer += bytes
 
482
    def _state_accept_expecting_header(self):
331
483
        prefix = self._extract_line()
332
 
        if prefix is None:
333
 
            # We haven't read a complete length prefix yet, so there's nothing
334
 
            # to do.
335
 
            return
336
 
        elif prefix == 'chunked':
 
484
        if prefix == 'chunked':
337
485
            self.state_accept = self._state_accept_expecting_length
338
486
        else:
339
487
            raise errors.SmartProtocolError(
340
488
                'Bad chunked body header: "%s"' % (prefix,))
341
489
 
342
 
    def _state_accept_expecting_length(self, bytes):
343
 
        self._in_buffer += bytes
 
490
    def _state_accept_expecting_length(self):
344
491
        prefix = self._extract_line()
345
 
        if prefix is None:
346
 
            # We haven't read a complete length prefix yet, so there's nothing
347
 
            # to do.
348
 
            return
349
 
        elif prefix == 'ERR':
 
492
        if prefix == 'ERR':
350
493
            self.error = True
351
494
            self.error_in_progress = []
352
 
            self._state_accept_expecting_length('')
 
495
            self._state_accept_expecting_length()
353
496
            return
354
497
        elif prefix == 'END':
355
498
            # We've read the end-of-body marker.
362
505
            self.chunk_in_progress = ''
363
506
            self.state_accept = self._state_accept_reading_chunk
364
507
 
365
 
    def _state_accept_reading_chunk(self, bytes):
366
 
        self._in_buffer += bytes
367
 
        in_buffer_len = len(self._in_buffer)
368
 
        self.chunk_in_progress += self._in_buffer[:self.bytes_left]
369
 
        self._in_buffer = self._in_buffer[self.bytes_left:]
 
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:])
370
513
        self.bytes_left -= in_buffer_len
371
514
        if self.bytes_left <= 0:
372
515
            # Finished with chunk
377
520
                self.chunks.append(self.chunk_in_progress)
378
521
            self.chunk_in_progress = None
379
522
            self.state_accept = self._state_accept_expecting_length
380
 
        
381
 
    def _state_accept_reading_unused(self, bytes):
382
 
        self.unused_data += bytes
 
523
 
 
524
    def _state_accept_reading_unused(self):
 
525
        self.unused_data += self._get_in_buffer()
 
526
        self._in_buffer_list = []
383
527
 
384
528
 
385
529
class LengthPrefixedBodyDecoder(_StatefulDecoder):
386
530
    """Decodes the length-prefixed bulk data."""
387
 
    
 
531
 
388
532
    def __init__(self):
389
533
        _StatefulDecoder.__init__(self)
390
534
        self.state_accept = self._state_accept_expecting_length
391
535
        self.state_read = self._state_read_no_data
392
 
        self._in_buffer = ''
 
536
        self._body = ''
393
537
        self._trailer_buffer = ''
394
 
    
 
538
 
395
539
    def next_read_size(self):
396
540
        if self.bytes_left is not None:
397
541
            # Ideally we want to read all the remainder of the body and the
407
551
        else:
408
552
            # Reading excess data.  Either way, 1 byte at a time is fine.
409
553
            return 1
410
 
        
 
554
 
411
555
    def read_pending_data(self):
412
556
        """Return any pending data that has been decoded."""
413
557
        return self.state_read()
414
558
 
415
 
    def _state_accept_expecting_length(self, bytes):
416
 
        self._in_buffer += bytes
417
 
        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')
418
562
        if pos == -1:
419
563
            return
420
 
        self.bytes_left = int(self._in_buffer[:pos])
421
 
        self._in_buffer = self._in_buffer[pos+1:]
422
 
        self.bytes_left -= len(self._in_buffer)
 
564
        self.bytes_left = int(in_buf[:pos])
 
565
        self._set_in_buffer(in_buf[pos+1:])
423
566
        self.state_accept = self._state_accept_reading_body
424
 
        self.state_read = self._state_read_in_buffer
 
567
        self.state_read = self._state_read_body_buffer
425
568
 
426
 
    def _state_accept_reading_body(self, bytes):
427
 
        self._in_buffer += bytes
428
 
        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)
429
574
        if self.bytes_left <= 0:
430
575
            # Finished with body
431
576
            if self.bytes_left != 0:
432
 
                self._trailer_buffer = self._in_buffer[self.bytes_left:]
433
 
                self._in_buffer = self._in_buffer[:self.bytes_left]
 
577
                self._trailer_buffer = self._body[self.bytes_left:]
 
578
                self._body = self._body[:self.bytes_left]
434
579
            self.bytes_left = None
435
580
            self.state_accept = self._state_accept_reading_trailer
436
 
        
437
 
    def _state_accept_reading_trailer(self, bytes):
438
 
        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)
439
585
        # TODO: what if the trailer does not match "done\n"?  Should this raise
440
586
        # a ProtocolViolation exception?
441
587
        if self._trailer_buffer.startswith('done\n'):
442
588
            self.unused_data = self._trailer_buffer[len('done\n'):]
443
589
            self.state_accept = self._state_accept_reading_unused
444
590
            self.finished_reading = True
445
 
    
446
 
    def _state_accept_reading_unused(self, bytes):
447
 
        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)
448
595
 
449
596
    def _state_read_no_data(self):
450
597
        return ''
451
598
 
452
 
    def _state_read_in_buffer(self):
453
 
        result = self._in_buffer
454
 
        self._in_buffer = ''
 
599
    def _state_read_body_buffer(self):
 
600
        result = self._body
 
601
        self._body = ''
455
602
        return result
456
603
 
457
604
 
458
 
class SmartClientRequestProtocolOne(SmartProtocolBase):
 
605
class SmartClientRequestProtocolOne(SmartProtocolBase, Requester,
 
606
                                    message.ResponseHandler):
459
607
    """The client-side protocol for smart version 1."""
460
608
 
461
609
    def __init__(self, request):
467
615
        self._request = request
468
616
        self._body_buffer = None
469
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)
470
623
 
471
624
    def call(self, *args):
472
625
        if 'hpss' in debug.debug_flags:
473
626
            mutter('hpss call:   %s', repr(args)[1:-1])
474
 
            self._request_start_time = time.time()
 
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()
475
630
        self._write_args(args)
476
631
        self._request.finished_writing()
 
632
        self._last_verb = args[0]
477
633
 
478
634
    def call_with_body_bytes(self, args, body):
479
635
        """Make a remote call of args with body bytes 'body'.
482
638
        """
483
639
        if 'hpss' in debug.debug_flags:
484
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)
485
643
            mutter('              %d bytes', len(body))
486
 
            self._request_start_time = time.time()
 
644
            self._request_start_time = osutils.timer_func()
 
645
            if 'hpssdetail' in debug.debug_flags:
 
646
                mutter('hpss body content: %s', body)
487
647
        self._write_args(args)
488
648
        bytes = self._encode_bulk_data(body)
489
649
        self._request.accept_bytes(bytes)
490
650
        self._request.finished_writing()
 
651
        self._last_verb = args[0]
491
652
 
492
653
    def call_with_body_readv_array(self, args, body):
493
654
        """Make a remote call with a readv array.
497
658
        """
498
659
        if 'hpss' in debug.debug_flags:
499
660
            mutter('hpss call w/readv: %s', repr(args)[1:-1])
500
 
            self._request_start_time = time.time()
 
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()
501
664
        self._write_args(args)
502
665
        readv_bytes = self._serialise_offsets(body)
503
666
        bytes = self._encode_bulk_data(readv_bytes)
505
668
        self._request.finished_writing()
506
669
        if 'hpss' in debug.debug_flags:
507
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])
508
680
 
509
681
    def cancel_read_body(self):
510
682
        """After expecting a body, a response code may indicate one otherwise.
515
687
        """
516
688
        self._request.finished_reading()
517
689
 
518
 
    def read_response_tuple(self, expect_body=False):
519
 
        """Read a response tuple from the wire.
520
 
 
521
 
        This should only be called once.
522
 
        """
 
690
    def _read_response_tuple(self):
523
691
        result = self._recv_tuple()
524
692
        if 'hpss' in debug.debug_flags:
525
693
            if self._request_start_time is not None:
526
694
                mutter('   result:   %6.3fs  %s',
527
 
                       time.time() - self._request_start_time,
 
695
                       osutils.timer_func() - self._request_start_time,
528
696
                       repr(result)[1:-1])
529
697
                self._request_start_time = None
530
698
            else:
531
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)
532
710
        if not expect_body:
533
711
            self._request.finished_reading()
534
712
        return result
535
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
 
536
760
    def read_body_bytes(self, count=-1):
537
761
        """Read bytes from the body, decoding into a byte stream.
538
 
        
539
 
        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
540
764
        errors, and then feed the buffer back as read_body_bytes is called.
541
765
        """
542
766
        if self._body_buffer is not None:
544
768
        _body_decoder = LengthPrefixedBodyDecoder()
545
769
 
546
770
        while not _body_decoder.finished_reading:
547
 
            bytes_wanted = _body_decoder.next_read_size()
548
 
            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.")
549
776
            _body_decoder.accept_bytes(bytes)
550
777
        self._request.finished_reading()
551
778
        self._body_buffer = StringIO(_body_decoder.read_pending_data())
557
784
 
558
785
    def _recv_tuple(self):
559
786
        """Receive a tuple from the medium request."""
560
 
        return _decode_tuple(self._recv_line())
561
 
 
562
 
    def _recv_line(self):
563
 
        """Read an entire line from the medium request."""
564
 
        line = ''
565
 
        while not line or line[-1] != '\n':
566
 
            # TODO: this is inefficient - but tuples are short.
567
 
            new_char = self._request.read_bytes(1)
568
 
            if new_char == '':
569
 
                # end of file encountered reading from server
570
 
                raise errors.ConnectionReset(
571
 
                    "please check connectivity and permissions",
572
 
                    "(and try -Dhpss if further diagnosis is required)")
573
 
            line += new_char
574
 
        return line
 
787
        return _decode_tuple(self._request.read_line())
575
788
 
576
789
    def query_version(self):
577
790
        """Return protocol version number of the server."""
591
804
 
592
805
    def _write_protocol_version(self):
593
806
        """Write any prefixes this protocol requires.
594
 
        
 
807
 
595
808
        Version one doesn't send protocol versions.
596
809
        """
597
810
 
598
811
 
599
812
class SmartClientRequestProtocolTwo(SmartClientRequestProtocolOne):
600
813
    """Version two of the client side of the smart protocol.
601
 
    
 
814
 
602
815
    This prefixes the request with the value of REQUEST_VERSION_TWO.
603
816
    """
604
817
 
 
818
    response_marker = RESPONSE_VERSION_TWO
 
819
    request_marker = REQUEST_VERSION_TWO
 
820
 
605
821
    def read_response_tuple(self, expect_body=False):
606
822
        """Read a response tuple from the wire.
607
823
 
608
824
        This should only be called once.
609
825
        """
610
826
        version = self._request.read_line()
611
 
        if version != RESPONSE_VERSION_TWO:
612
 
            raise errors.SmartProtocolError('bad protocol marker %r' % version)
613
 
        response_status = self._recv_line()
614
 
        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:
615
843
            raise errors.SmartProtocolError(
616
844
                'bad protocol status %r' % response_status)
617
 
        self.response_status = response_status == 'success\n'
618
 
        return SmartClientRequestProtocolOne.read_response_tuple(self, expect_body)
619
845
 
620
846
    def _write_protocol_version(self):
621
847
        """Write any prefixes this protocol requires.
622
 
        
 
848
 
623
849
        Version two sends the value of REQUEST_VERSION_TWO.
624
850
        """
625
 
        self._request.accept_bytes(REQUEST_VERSION_TWO)
 
851
        self._request.accept_bytes(self.request_marker)
626
852
 
627
853
    def read_streamed_body(self):
628
854
        """Read bytes from the body, decoding into a byte stream.
629
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.
630
858
        _body_decoder = ChunkedBodyDecoder()
631
859
        while not _body_decoder.finished_reading:
632
 
            bytes_wanted = _body_decoder.next_read_size()
633
 
            bytes = self._request.read_bytes(bytes_wanted)
 
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.")
634
865
            _body_decoder.accept_bytes(bytes)
635
866
            for body_bytes in iter(_body_decoder.read_next_chunk, None):
636
 
                if 'hpss' in debug.debug_flags:
 
867
                if 'hpss' in debug.debug_flags and type(body_bytes) is str:
637
868
                    mutter('              %d byte chunk read',
638
869
                           len(body_bytes))
639
870
                yield body_bytes
640
871
        self._request.finished_reading()
641
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: It is probably more appropriate to use sum(map(len, _buf))
 
1085
        #       for total number of bytes to write, rather than buffer based on
 
1086
        #       the number of write() calls
 
1087
        # TODO: Another possibility would be to turn this into an async model.
 
1088
        #       Where we let another thread know that we have some bytes if
 
1089
        #       they want it, but we don't actually block for it
 
1090
        #       Note that osutils.send_all always sends 64kB chunks anyway, so
 
1091
        #       we might just push out smaller bits at a time?
 
1092
        self._buf.append(bytes)
 
1093
        self._buf_len += len(bytes)
 
1094
        if self._buf_len > self.BUFFER_SIZE:
 
1095
            self.flush()
 
1096
 
 
1097
    def flush(self):
 
1098
        if self._buf:
 
1099
            self._real_write_func(''.join(self._buf))
 
1100
            del self._buf[:]
 
1101
            self._buf_len = 0
 
1102
 
 
1103
    def _serialise_offsets(self, offsets):
 
1104
        """Serialise a readv offset list."""
 
1105
        txt = []
 
1106
        for start, length in offsets:
 
1107
            txt.append('%d,%d' % (start, length))
 
1108
        return '\n'.join(txt)
 
1109
 
 
1110
    def _write_protocol_version(self):
 
1111
        self._write_func(MESSAGE_VERSION_THREE)
 
1112
 
 
1113
    def _write_prefixed_bencode(self, structure):
 
1114
        bytes = bencode(structure)
 
1115
        self._write_func(struct.pack('!L', len(bytes)))
 
1116
        self._write_func(bytes)
 
1117
 
 
1118
    def _write_headers(self, headers):
 
1119
        self._write_prefixed_bencode(headers)
 
1120
 
 
1121
    def _write_structure(self, args):
 
1122
        self._write_func('s')
 
1123
        utf8_args = []
 
1124
        for arg in args:
 
1125
            if type(arg) is unicode:
 
1126
                utf8_args.append(arg.encode('utf8'))
 
1127
            else:
 
1128
                utf8_args.append(arg)
 
1129
        self._write_prefixed_bencode(utf8_args)
 
1130
 
 
1131
    def _write_end(self):
 
1132
        self._write_func('e')
 
1133
        self.flush()
 
1134
 
 
1135
    def _write_prefixed_body(self, bytes):
 
1136
        self._write_func('b')
 
1137
        self._write_func(struct.pack('!L', len(bytes)))
 
1138
        self._write_func(bytes)
 
1139
 
 
1140
    def _write_chunked_body_start(self):
 
1141
        self._write_func('oC')
 
1142
 
 
1143
    def _write_error_status(self):
 
1144
        self._write_func('oE')
 
1145
 
 
1146
    def _write_success_status(self):
 
1147
        self._write_func('oS')
 
1148
 
 
1149
 
 
1150
class ProtocolThreeResponder(_ProtocolThreeEncoder):
 
1151
 
 
1152
    def __init__(self, write_func):
 
1153
        _ProtocolThreeEncoder.__init__(self, write_func)
 
1154
        self.response_sent = False
 
1155
        self._headers = {'Software version': bzrlib.__version__}
 
1156
        if 'hpss' in debug.debug_flags:
 
1157
            self._thread_id = thread.get_ident()
 
1158
            self._response_start_time = None
 
1159
 
 
1160
    def _trace(self, action, message, extra_bytes=None, include_time=False):
 
1161
        if self._response_start_time is None:
 
1162
            self._response_start_time = osutils.timer_func()
 
1163
        if include_time:
 
1164
            t = '%5.3fs ' % (time.clock() - self._response_start_time)
 
1165
        else:
 
1166
            t = ''
 
1167
        if extra_bytes is None:
 
1168
            extra = ''
 
1169
        else:
 
1170
            extra = ' ' + repr(extra_bytes[:40])
 
1171
            if len(extra) > 33:
 
1172
                extra = extra[:29] + extra[-1] + '...'
 
1173
        mutter('%12s: [%s] %s%s%s'
 
1174
               % (action, self._thread_id, t, message, extra))
 
1175
 
 
1176
    def send_error(self, exception):
 
1177
        if self.response_sent:
 
1178
            raise AssertionError(
 
1179
                "send_error(%s) called, but response already sent."
 
1180
                % (exception,))
 
1181
        if isinstance(exception, errors.UnknownSmartMethod):
 
1182
            failure = request.FailedSmartServerResponse(
 
1183
                ('UnknownMethod', exception.verb))
 
1184
            self.send_response(failure)
 
1185
            return
 
1186
        if 'hpss' in debug.debug_flags:
 
1187
            self._trace('error', str(exception))
 
1188
        self.response_sent = True
 
1189
        self._write_protocol_version()
 
1190
        self._write_headers(self._headers)
 
1191
        self._write_error_status()
 
1192
        self._write_structure(('error', str(exception)))
 
1193
        self._write_end()
 
1194
 
 
1195
    def send_response(self, response):
 
1196
        if self.response_sent:
 
1197
            raise AssertionError(
 
1198
                "send_response(%r) called, but response already sent."
 
1199
                % (response,))
 
1200
        self.response_sent = True
 
1201
        self._write_protocol_version()
 
1202
        self._write_headers(self._headers)
 
1203
        if response.is_successful():
 
1204
            self._write_success_status()
 
1205
        else:
 
1206
            self._write_error_status()
 
1207
        if 'hpss' in debug.debug_flags:
 
1208
            self._trace('response', repr(response.args))
 
1209
        self._write_structure(response.args)
 
1210
        if response.body is not None:
 
1211
            self._write_prefixed_body(response.body)
 
1212
            if 'hpss' in debug.debug_flags:
 
1213
                self._trace('body', '%d bytes' % (len(response.body),),
 
1214
                            response.body, include_time=True)
 
1215
        elif response.body_stream is not None:
 
1216
            count = num_bytes = 0
 
1217
            first_chunk = None
 
1218
            for exc_info, chunk in _iter_with_errors(response.body_stream):
 
1219
                count += 1
 
1220
                if exc_info is not None:
 
1221
                    self._write_error_status()
 
1222
                    error_struct = request._translate_error(exc_info[1])
 
1223
                    self._write_structure(error_struct)
 
1224
                    break
 
1225
                else:
 
1226
                    if isinstance(chunk, request.FailedSmartServerResponse):
 
1227
                        self._write_error_status()
 
1228
                        self._write_structure(chunk.args)
 
1229
                        break
 
1230
                    num_bytes += len(chunk)
 
1231
                    if first_chunk is None:
 
1232
                        first_chunk = chunk
 
1233
                    self._write_prefixed_body(chunk)
 
1234
                    self.flush()
 
1235
                    if 'hpssdetail' in debug.debug_flags:
 
1236
                        # Not worth timing separately, as _write_func is
 
1237
                        # actually buffered
 
1238
                        self._trace('body chunk',
 
1239
                                    '%d bytes' % (len(chunk),),
 
1240
                                    chunk, suppress_time=True)
 
1241
            if 'hpss' in debug.debug_flags:
 
1242
                self._trace('body stream',
 
1243
                            '%d bytes %d chunks' % (num_bytes, count),
 
1244
                            first_chunk)
 
1245
        self._write_end()
 
1246
        if 'hpss' in debug.debug_flags:
 
1247
            self._trace('response end', '', include_time=True)
 
1248
 
 
1249
 
 
1250
def _iter_with_errors(iterable):
 
1251
    """Handle errors from iterable.next().
 
1252
 
 
1253
    Use like::
 
1254
 
 
1255
        for exc_info, value in _iter_with_errors(iterable):
 
1256
            ...
 
1257
 
 
1258
    This is a safer alternative to::
 
1259
 
 
1260
        try:
 
1261
            for value in iterable:
 
1262
               ...
 
1263
        except:
 
1264
            ...
 
1265
 
 
1266
    Because the latter will catch errors from the for-loop body, not just
 
1267
    iterable.next()
 
1268
 
 
1269
    If an error occurs, exc_info will be a exc_info tuple, and the generator
 
1270
    will terminate.  Otherwise exc_info will be None, and value will be the
 
1271
    value from iterable.next().  Note that KeyboardInterrupt and SystemExit
 
1272
    will not be itercepted.
 
1273
    """
 
1274
    iterator = iter(iterable)
 
1275
    while True:
 
1276
        try:
 
1277
            yield None, iterator.next()
 
1278
        except StopIteration:
 
1279
            return
 
1280
        except (KeyboardInterrupt, SystemExit):
 
1281
            raise
 
1282
        except Exception:
 
1283
            mutter('_iter_with_errors caught error')
 
1284
            log_exception_quietly()
 
1285
            yield sys.exc_info(), None
 
1286
            return
 
1287
 
 
1288
 
 
1289
class ProtocolThreeRequester(_ProtocolThreeEncoder, Requester):
 
1290
 
 
1291
    def __init__(self, medium_request):
 
1292
        _ProtocolThreeEncoder.__init__(self, medium_request.accept_bytes)
 
1293
        self._medium_request = medium_request
 
1294
        self._headers = {}
 
1295
 
 
1296
    def set_headers(self, headers):
 
1297
        self._headers = headers.copy()
 
1298
 
 
1299
    def call(self, *args):
 
1300
        if 'hpss' in debug.debug_flags:
 
1301
            mutter('hpss call:   %s', repr(args)[1:-1])
 
1302
            base = getattr(self._medium_request._medium, 'base', None)
 
1303
            if base is not None:
 
1304
                mutter('             (to %s)', base)
 
1305
            self._request_start_time = osutils.timer_func()
 
1306
        self._write_protocol_version()
 
1307
        self._write_headers(self._headers)
 
1308
        self._write_structure(args)
 
1309
        self._write_end()
 
1310
        self._medium_request.finished_writing()
 
1311
 
 
1312
    def call_with_body_bytes(self, args, body):
 
1313
        """Make a remote call of args with body bytes 'body'.
 
1314
 
 
1315
        After calling this, call read_response_tuple to find the result out.
 
1316
        """
 
1317
        if 'hpss' in debug.debug_flags:
 
1318
            mutter('hpss call w/body: %s (%r...)', repr(args)[1:-1], body[:20])
 
1319
            path = getattr(self._medium_request._medium, '_path', None)
 
1320
            if path is not None:
 
1321
                mutter('                  (to %s)', path)
 
1322
            mutter('              %d bytes', len(body))
 
1323
            self._request_start_time = osutils.timer_func()
 
1324
        self._write_protocol_version()
 
1325
        self._write_headers(self._headers)
 
1326
        self._write_structure(args)
 
1327
        self._write_prefixed_body(body)
 
1328
        self._write_end()
 
1329
        self._medium_request.finished_writing()
 
1330
 
 
1331
    def call_with_body_readv_array(self, args, body):
 
1332
        """Make a remote call with a readv array.
 
1333
 
 
1334
        The body is encoded with one line per readv offset pair. The numbers in
 
1335
        each pair are separated by a comma, and no trailing \n is emitted.
 
1336
        """
 
1337
        if 'hpss' in debug.debug_flags:
 
1338
            mutter('hpss call w/readv: %s', repr(args)[1:-1])
 
1339
            path = getattr(self._medium_request._medium, '_path', None)
 
1340
            if path is not None:
 
1341
                mutter('                  (to %s)', path)
 
1342
            self._request_start_time = osutils.timer_func()
 
1343
        self._write_protocol_version()
 
1344
        self._write_headers(self._headers)
 
1345
        self._write_structure(args)
 
1346
        readv_bytes = self._serialise_offsets(body)
 
1347
        if 'hpss' in debug.debug_flags:
 
1348
            mutter('              %d bytes in readv request', len(readv_bytes))
 
1349
        self._write_prefixed_body(readv_bytes)
 
1350
        self._write_end()
 
1351
        self._medium_request.finished_writing()
 
1352
 
 
1353
    def call_with_body_stream(self, args, stream):
 
1354
        if 'hpss' in debug.debug_flags:
 
1355
            mutter('hpss call w/body stream: %r', args)
 
1356
            path = getattr(self._medium_request._medium, '_path', None)
 
1357
            if path is not None:
 
1358
                mutter('                  (to %s)', path)
 
1359
            self._request_start_time = osutils.timer_func()
 
1360
        self._write_protocol_version()
 
1361
        self._write_headers(self._headers)
 
1362
        self._write_structure(args)
 
1363
        # TODO: notice if the server has sent an early error reply before we
 
1364
        #       have finished sending the stream.  We would notice at the end
 
1365
        #       anyway, but if the medium can deliver it early then it's good
 
1366
        #       to short-circuit the whole request...
 
1367
        for exc_info, part in _iter_with_errors(stream):
 
1368
            if exc_info is not None:
 
1369
                # Iterating the stream failed.  Cleanly abort the request.
 
1370
                self._write_error_status()
 
1371
                # Currently the client unconditionally sends ('error',) as the
 
1372
                # error args.
 
1373
                self._write_structure(('error',))
 
1374
                self._write_end()
 
1375
                self._medium_request.finished_writing()
 
1376
                raise exc_info[0], exc_info[1], exc_info[2]
 
1377
            else:
 
1378
                self._write_prefixed_body(part)
 
1379
                self.flush()
 
1380
        self._write_end()
 
1381
        self._medium_request.finished_writing()
 
1382