~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/smart/protocol.py

  • Committer: Vincent Ladeuil
  • Date: 2007-06-20 14:25:06 UTC
  • mfrom: (2540 +trunk)
  • mto: This revision was merged to the branch mainline in revision 2646.
  • Revision ID: v.ladeuil+lp@free.fr-20070620142506-txsb1v8538kpsafw
merge bzr.dev @ 2540

Show diffs side-by-side

added added

removed removed

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