~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/smart/protocol.py

  • Committer: v.ladeuil+lp at free
  • Date: 2007-05-15 17:40:32 UTC
  • mto: (2485.8.44 bzr.connection.sharing)
  • mto: This revision was merged to the branch mainline in revision 2646.
  • Revision ID: v.ladeuil+lp@free.fr-20070515174032-qzdkangpv29l9e7g
Add a test that check that init connect only once. It fails.

* __init__.py:
(test_suite): Register the new test class.

* test_init.py: 
(InstrumentedTransport): A transport that can track connections.
(TransportHooks): Transport specific hooks.
(TestInit): Iniit command behavior tests.

* ftp.py:
(FtpTransport.__init__): Mark place that need fixing regarding
transport connection sharing

* builtins.py:
(cmd_init.run): Mark places that need fixing regarding transport
connection sharing.

Show diffs side-by-side

added added

removed removed

Lines of Context:
1
 
# Copyright (C) 2006-2010 Canonical Ltd
 
1
# Copyright (C) 2006, 2007 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., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
 
15
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  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
 
import collections
 
21
 
22
22
from cStringIO import StringIO
23
 
import struct
24
 
import sys
25
 
import thread
26
 
import threading
27
 
import time
28
23
 
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
 
24
from bzrlib import errors
 
25
from bzrlib.smart import request
38
26
 
39
27
 
40
28
# Protocol version strings.  These are sent as prefixes of bzr requests and
43
31
REQUEST_VERSION_TWO = 'bzr request 2\n'
44
32
RESPONSE_VERSION_TWO = 'bzr response 2\n'
45
33
 
46
 
MESSAGE_VERSION_THREE = 'bzr message 3 (bzr 1.6)\n'
47
 
RESPONSE_VERSION_THREE = REQUEST_VERSION_THREE = MESSAGE_VERSION_THREE
48
 
 
49
34
 
50
35
def _recv_tuple(from_file):
51
36
    req_line = from_file.readline()
53
38
 
54
39
 
55
40
def _decode_tuple(req_line):
56
 
    if req_line is None or req_line == '':
 
41
    if req_line == None or req_line == '':
57
42
        return None
58
43
    if req_line[-1] != '\n':
59
44
        raise errors.SmartProtocolError("request %r not terminated" % req_line)
62
47
 
63
48
def _encode_tuple(args):
64
49
    """Encode the tuple args to a bytestream."""
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)
 
50
    return '\x01'.join(args) + '\n'
106
51
 
107
52
 
108
53
class SmartProtocolBase(object):
120
65
        for start, length in offsets:
121
66
            txt.append('%d,%d' % (start, length))
122
67
        return '\n'.join(txt)
123
 
 
 
68
        
124
69
 
125
70
class SmartServerRequestProtocolOne(SmartProtocolBase):
126
71
    """Server-side encoding and decoding logic for smart version 1."""
127
 
 
128
 
    def __init__(self, backing_transport, write_func, root_client_path='/',
129
 
            jail_root=None):
 
72
    
 
73
    def __init__(self, backing_transport, write_func):
130
74
        self._backing_transport = backing_transport
131
 
        self._root_client_path = root_client_path
132
 
        self._jail_root = jail_root
133
 
        self.unused_data = ''
 
75
        self.excess_buffer = ''
134
76
        self._finished = False
135
77
        self.in_buffer = ''
136
 
        self._has_dispatched = False
 
78
        self.has_dispatched = False
137
79
        self.request = None
138
80
        self._body_decoder = None
139
81
        self._write_func = write_func
140
82
 
141
83
    def accept_bytes(self, bytes):
142
84
        """Take bytes, and advance the internal state machine appropriately.
143
 
 
 
85
        
144
86
        :param bytes: must be a byte string
145
87
        """
146
 
        if not isinstance(bytes, str):
147
 
            raise ValueError(bytes)
 
88
        assert isinstance(bytes, str)
148
89
        self.in_buffer += bytes
149
 
        if not self._has_dispatched:
 
90
        if not self.has_dispatched:
150
91
            if '\n' not in self.in_buffer:
151
92
                # no command line yet
152
93
                return
153
 
            self._has_dispatched = True
 
94
            self.has_dispatched = True
154
95
            try:
155
96
                first_line, self.in_buffer = self.in_buffer.split('\n', 1)
156
97
                first_line += '\n'
157
98
                req_args = _decode_tuple(first_line)
158
99
                self.request = request.SmartServerRequestHandler(
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)
 
100
                    self._backing_transport, commands=request.request_handlers)
 
101
                self.request.dispatch_command(req_args[0], req_args[1:])
163
102
                if self.request.finished_reading:
164
103
                    # trivial request
165
 
                    self.unused_data = self.in_buffer
 
104
                    self.excess_buffer = self.in_buffer
166
105
                    self.in_buffer = ''
167
106
                    self._send_response(self.request.response)
168
107
            except KeyboardInterrupt:
169
108
                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
177
109
            except Exception, exception:
178
110
                # everything else: pass to client, flush, and quit
179
 
                log_exception_quietly()
180
111
                self._send_response(request.FailedSmartServerResponse(
181
112
                    ('error', str(exception))))
182
113
                return
183
114
 
184
 
        if self._has_dispatched:
 
115
        if self.has_dispatched:
185
116
            if self._finished:
186
 
                # nothing to do.XXX: this routine should be a single state
 
117
                # nothing to do.XXX: this routine should be a single state 
187
118
                # machine too.
188
 
                self.unused_data += self.in_buffer
 
119
                self.excess_buffer += self.in_buffer
189
120
                self.in_buffer = ''
190
121
                return
191
122
            if self._body_decoder is None:
196
127
            self.request.accept_body(body_data)
197
128
            if self._body_decoder.finished_reading:
198
129
                self.request.end_of_body()
199
 
                if not self.request.finished_reading:
200
 
                    raise AssertionError("no more body, request not finished")
 
130
                assert self.request.finished_reading, \
 
131
                    "no more body, request not finished"
201
132
            if self.request.response is not None:
202
133
                self._send_response(self.request.response)
203
 
                self.unused_data = self.in_buffer
 
134
                self.excess_buffer = self.in_buffer
204
135
                self.in_buffer = ''
205
136
            else:
206
 
                if self.request.finished_reading:
207
 
                    raise AssertionError(
208
 
                        "no response and we have finished reading.")
 
137
                assert not self.request.finished_reading, \
 
138
                    "no response and we have finished reading."
209
139
 
210
140
    def _send_response(self, response):
211
141
        """Send a smart server response down the output stream."""
212
 
        if self._finished:
213
 
            raise AssertionError('response already sent')
 
142
        assert not self._finished, 'response already sent'
214
143
        args = response.args
215
144
        body = response.body
216
145
        self._finished = True
218
147
        self._write_success_or_failure_prefix(response)
219
148
        self._write_func(_encode_tuple(args))
220
149
        if body is not None:
221
 
            if not isinstance(body, str):
222
 
                raise ValueError(body)
 
150
            assert isinstance(body, str), 'body must be a str'
223
151
            bytes = self._encode_bulk_data(body)
224
152
            self._write_func(bytes)
225
153
 
226
154
    def _write_protocol_version(self):
227
155
        """Write any prefixes this protocol requires.
228
 
 
 
156
        
229
157
        Version one doesn't send protocol versions.
230
158
        """
231
159
 
248
176
 
249
177
class SmartServerRequestProtocolTwo(SmartServerRequestProtocolOne):
250
178
    r"""Version two of the server side of the smart protocol.
251
 
 
 
179
   
252
180
    This prefixes responses with the value of RESPONSE_VERSION_TWO.
253
181
    """
254
182
 
255
 
    response_marker = RESPONSE_VERSION_TWO
256
 
    request_marker = REQUEST_VERSION_TWO
257
 
 
258
183
    def _write_success_or_failure_prefix(self, response):
259
184
        """Write the protocol specific success/failure prefix."""
260
185
        if response.is_successful():
264
189
 
265
190
    def _write_protocol_version(self):
266
191
        r"""Write any prefixes this protocol requires.
267
 
 
 
192
        
268
193
        Version two sends the value of RESPONSE_VERSION_TWO.
269
194
        """
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
 
 
 
195
        self._write_func(RESPONSE_VERSION_TWO)
 
196
 
 
197
 
 
198
class LengthPrefixedBodyDecoder(object):
 
199
    """Decodes the length-prefixed bulk data."""
 
200
    
338
201
    def __init__(self):
 
202
        self.bytes_left = None
339
203
        self.finished_reading = False
340
 
        self._in_buffer_list = []
341
 
        self._in_buffer_len = 0
342
204
        self.unused_data = ''
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
383
 
 
 
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
    
384
210
    def accept_bytes(self, bytes):
385
211
        """Decode as much of bytes as possible.
386
212
 
391
217
        data will be appended to self.unused_data.
392
218
        """
393
219
        # accept_bytes is allowed to change the state
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.
 
220
        current_state = self.state_accept
 
221
        self.state_accept(bytes)
 
222
        while current_state != self.state_accept:
401
223
            current_state = 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 = ''
 
224
            self.state_accept('')
538
225
 
539
226
    def next_read_size(self):
540
227
        if self.bytes_left is not None:
551
238
        else:
552
239
            # Reading excess data.  Either way, 1 byte at a time is fine.
553
240
            return 1
554
 
 
 
241
        
555
242
    def read_pending_data(self):
556
243
        """Return any pending data that has been decoded."""
557
244
        return self.state_read()
558
245
 
559
 
    def _state_accept_expecting_length(self):
560
 
        in_buf = self._get_in_buffer()
561
 
        pos = in_buf.find('\n')
 
246
    def _state_accept_expecting_length(self, bytes):
 
247
        self._in_buffer += bytes
 
248
        pos = self._in_buffer.find('\n')
562
249
        if pos == -1:
563
250
            return
564
 
        self.bytes_left = int(in_buf[:pos])
565
 
        self._set_in_buffer(in_buf[pos+1:])
 
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)
566
254
        self.state_accept = self._state_accept_reading_body
567
 
        self.state_read = self._state_read_body_buffer
 
255
        self.state_read = self._state_read_in_buffer
568
256
 
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)
 
257
    def _state_accept_reading_body(self, bytes):
 
258
        self._in_buffer += bytes
 
259
        self.bytes_left -= len(bytes)
574
260
        if self.bytes_left <= 0:
575
261
            # Finished with body
576
262
            if self.bytes_left != 0:
577
 
                self._trailer_buffer = self._body[self.bytes_left:]
578
 
                self._body = self._body[:self.bytes_left]
 
263
                self._trailer_buffer = self._in_buffer[self.bytes_left:]
 
264
                self._in_buffer = self._in_buffer[:self.bytes_left]
579
265
            self.bytes_left = None
580
266
            self.state_accept = self._state_accept_reading_trailer
581
 
 
582
 
    def _state_accept_reading_trailer(self):
583
 
        self._trailer_buffer += self._get_in_buffer()
584
 
        self._set_in_buffer(None)
 
267
        
 
268
    def _state_accept_reading_trailer(self, bytes):
 
269
        self._trailer_buffer += bytes
585
270
        # TODO: what if the trailer does not match "done\n"?  Should this raise
586
271
        # a ProtocolViolation exception?
587
272
        if self._trailer_buffer.startswith('done\n'):
588
273
            self.unused_data = self._trailer_buffer[len('done\n'):]
589
274
            self.state_accept = self._state_accept_reading_unused
590
275
            self.finished_reading = True
591
 
 
592
 
    def _state_accept_reading_unused(self):
593
 
        self.unused_data += self._get_in_buffer()
594
 
        self._set_in_buffer(None)
 
276
    
 
277
    def _state_accept_reading_unused(self, bytes):
 
278
        self.unused_data += bytes
595
279
 
596
280
    def _state_read_no_data(self):
597
281
        return ''
598
282
 
599
 
    def _state_read_body_buffer(self):
600
 
        result = self._body
601
 
        self._body = ''
 
283
    def _state_read_in_buffer(self):
 
284
        result = self._in_buffer
 
285
        self._in_buffer = ''
602
286
        return result
603
287
 
604
288
 
605
 
class SmartClientRequestProtocolOne(SmartProtocolBase, Requester,
606
 
                                    message.ResponseHandler):
 
289
class SmartClientRequestProtocolOne(SmartProtocolBase):
607
290
    """The client-side protocol for smart version 1."""
608
291
 
609
292
    def __init__(self, request):
614
297
        """
615
298
        self._request = request
616
299
        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)
623
300
 
624
301
    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()
630
302
        self._write_args(args)
631
303
        self._request.finished_writing()
632
 
        self._last_verb = args[0]
633
304
 
634
305
    def call_with_body_bytes(self, args, body):
635
306
        """Make a remote call of args with body bytes 'body'.
636
307
 
637
308
        After calling this, call read_response_tuple to find the result out.
638
309
        """
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)
647
310
        self._write_args(args)
648
311
        bytes = self._encode_bulk_data(body)
649
312
        self._request.accept_bytes(bytes)
650
313
        self._request.finished_writing()
651
 
        self._last_verb = args[0]
652
314
 
653
315
    def call_with_body_readv_array(self, args, body):
654
316
        """Make a remote call with a readv array.
655
317
 
656
318
        The body is encoded with one line per readv offset pair. The numbers in
657
 
        each pair are separated by a comma, and no trailing \\n is emitted.
 
319
        each pair are separated by a comma, and no trailing \n is emitted.
658
320
        """
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()
664
321
        self._write_args(args)
665
322
        readv_bytes = self._serialise_offsets(body)
666
323
        bytes = self._encode_bulk_data(readv_bytes)
667
324
        self._request.accept_bytes(bytes)
668
325
        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])
680
326
 
681
327
    def cancel_read_body(self):
682
328
        """After expecting a body, a response code may indicate one otherwise.
687
333
        """
688
334
        self._request.finished_reading()
689
335
 
690
 
    def _read_response_tuple(self):
 
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
        """
691
341
        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)
710
342
        if not expect_body:
711
343
            self._request.finished_reading()
712
344
        return result
713
345
 
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
 
 
760
346
    def read_body_bytes(self, count=-1):
761
347
        """Read bytes from the body, decoding into a byte stream.
762
 
 
763
 
        We read all bytes at once to ensure we've checked the trailer for
 
348
        
 
349
        We read all bytes at once to ensure we've checked the trailer for 
764
350
        errors, and then feed the buffer back as read_body_bytes is called.
765
351
        """
766
352
        if self._body_buffer is not None:
768
354
        _body_decoder = LengthPrefixedBodyDecoder()
769
355
 
770
356
        while not _body_decoder.finished_reading:
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.")
 
357
            bytes_wanted = _body_decoder.next_read_size()
 
358
            bytes = self._request.read_bytes(bytes_wanted)
776
359
            _body_decoder.accept_bytes(bytes)
777
360
        self._request.finished_reading()
778
361
        self._body_buffer = StringIO(_body_decoder.read_pending_data())
779
362
        # 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()))
783
363
        return self._body_buffer.read(count)
784
364
 
785
365
    def _recv_tuple(self):
786
366
        """Receive a tuple from the medium request."""
787
 
        return _decode_tuple(self._request.read_line())
 
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
788
378
 
789
379
    def query_version(self):
790
380
        """Return protocol version number of the server."""
804
394
 
805
395
    def _write_protocol_version(self):
806
396
        """Write any prefixes this protocol requires.
807
 
 
 
397
        
808
398
        Version one doesn't send protocol versions.
809
399
        """
810
400
 
811
401
 
812
402
class SmartClientRequestProtocolTwo(SmartClientRequestProtocolOne):
813
403
    """Version two of the client side of the smart protocol.
814
 
 
 
404
    
815
405
    This prefixes the request with the value of REQUEST_VERSION_TWO.
816
406
    """
817
407
 
818
 
    response_marker = RESPONSE_VERSION_TWO
819
 
    request_marker = REQUEST_VERSION_TWO
820
 
 
821
408
    def read_response_tuple(self, expect_body=False):
822
409
        """Read a response tuple from the wire.
823
410
 
824
411
        This should only be called once.
825
412
        """
826
413
        version = self._request.read_line()
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:
 
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'):
843
418
            raise errors.SmartProtocolError(
844
419
                'bad protocol status %r' % response_status)
 
420
        self.response_status = response_status == 'success\n'
 
421
        return SmartClientRequestProtocolOne.read_response_tuple(self, expect_body)
845
422
 
846
423
    def _write_protocol_version(self):
847
 
        """Write any prefixes this protocol requires.
848
 
 
 
424
        r"""Write any prefixes this protocol requires.
 
425
        
849
426
        Version two sends the value of REQUEST_VERSION_TWO.
850
427
        """
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: 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()
 
428
        self._request.accept_bytes(REQUEST_VERSION_TWO)
1382
429