/brz/remove-bazaar

To get this branch, use:
bzr branch http://gegoxaren.bato24.eu/bzr/brz/remove-bazaar

« back to all changes in this revision

Viewing changes to bzrlib/smart/protocol.py

  • Committer: Robert Collins
  • Date: 2007-04-19 02:27:44 UTC
  • mto: This revision was merged to the branch mainline in revision 2426.
  • Revision ID: robertc@robertcollins.net-20070419022744-pfdqz42kp1wizh43
``make docs`` now creates a man page at ``man1/bzr.1`` fixing bug 107388.
(Robert Collins)

Show diffs side-by-side

added added

removed removed

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