/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

Merge up bzr.dev.

Show diffs side-by-side

added added

removed removed

Lines of Context:
 
1
# Copyright (C) 2006, 2007 Canonical Ltd
 
2
#
 
3
# This program is free software; you can redistribute it and/or modify
 
4
# it under the terms of the GNU General Public License as published by
 
5
# the Free Software Foundation; either version 2 of the License, or
 
6
# (at your option) any later version.
 
7
#
 
8
# This program is distributed in the hope that it will be useful,
 
9
# but WITHOUT ANY WARRANTY; without even the implied warranty of
 
10
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 
11
# GNU General Public License for more details.
 
12
#
 
13
# You should have received a copy of the GNU General Public License
 
14
# along with this program; if not, write to the Free Software
 
15
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 
16
 
 
17
"""Wire-level encoding and decoding of requests and responses for the smart
 
18
client and server.
 
19
"""
 
20
 
 
21
import collections
 
22
from cStringIO import StringIO
 
23
import time
 
24
 
 
25
from bzrlib import debug
 
26
from bzrlib import errors
 
27
from bzrlib.smart import request
 
28
from bzrlib.trace import log_exception_quietly, mutter
 
29
 
 
30
 
 
31
# Protocol version strings.  These are sent as prefixes of bzr requests and
 
32
# responses to identify the protocol version being used. (There are no version
 
33
# one strings because that version doesn't send any).
 
34
REQUEST_VERSION_TWO = 'bzr request 2\n'
 
35
RESPONSE_VERSION_TWO = 'bzr response 2\n'
 
36
 
 
37
 
 
38
def _recv_tuple(from_file):
 
39
    req_line = from_file.readline()
 
40
    return _decode_tuple(req_line)
 
41
 
 
42
 
 
43
def _decode_tuple(req_line):
 
44
    if req_line is None or req_line == '':
 
45
        return None
 
46
    if req_line[-1] != '\n':
 
47
        raise errors.SmartProtocolError("request %r not terminated" % req_line)
 
48
    return tuple(req_line[:-1].split('\x01'))
 
49
 
 
50
 
 
51
def _encode_tuple(args):
 
52
    """Encode the tuple args to a bytestream."""
 
53
    return '\x01'.join(args) + '\n'
 
54
 
 
55
 
 
56
class SmartProtocolBase(object):
 
57
    """Methods common to client and server"""
 
58
 
 
59
    # TODO: this only actually accomodates a single block; possibly should
 
60
    # support multiple chunks?
 
61
    def _encode_bulk_data(self, body):
 
62
        """Encode body as a bulk data chunk."""
 
63
        return ''.join(('%d\n' % len(body), body, 'done\n'))
 
64
 
 
65
    def _serialise_offsets(self, offsets):
 
66
        """Serialise a readv offset list."""
 
67
        txt = []
 
68
        for start, length in offsets:
 
69
            txt.append('%d,%d' % (start, length))
 
70
        return '\n'.join(txt)
 
71
        
 
72
 
 
73
class SmartServerRequestProtocolOne(SmartProtocolBase):
 
74
    """Server-side encoding and decoding logic for smart version 1."""
 
75
    
 
76
    def __init__(self, backing_transport, write_func, root_client_path='/'):
 
77
        self._backing_transport = backing_transport
 
78
        self._root_client_path = root_client_path
 
79
        self.excess_buffer = ''
 
80
        self._finished = False
 
81
        self.in_buffer = ''
 
82
        self.has_dispatched = False
 
83
        self.request = None
 
84
        self._body_decoder = None
 
85
        self._write_func = write_func
 
86
 
 
87
    def accept_bytes(self, bytes):
 
88
        """Take bytes, and advance the internal state machine appropriately.
 
89
        
 
90
        :param bytes: must be a byte string
 
91
        """
 
92
        if not isinstance(bytes, str):
 
93
            raise ValueError(bytes)
 
94
        self.in_buffer += bytes
 
95
        if not self.has_dispatched:
 
96
            if '\n' not in self.in_buffer:
 
97
                # no command line yet
 
98
                return
 
99
            self.has_dispatched = True
 
100
            try:
 
101
                first_line, self.in_buffer = self.in_buffer.split('\n', 1)
 
102
                first_line += '\n'
 
103
                req_args = _decode_tuple(first_line)
 
104
                self.request = request.SmartServerRequestHandler(
 
105
                    self._backing_transport, commands=request.request_handlers,
 
106
                    root_client_path=self._root_client_path)
 
107
                self.request.dispatch_command(req_args[0], req_args[1:])
 
108
                if self.request.finished_reading:
 
109
                    # trivial request
 
110
                    self.excess_buffer = self.in_buffer
 
111
                    self.in_buffer = ''
 
112
                    self._send_response(self.request.response)
 
113
            except KeyboardInterrupt:
 
114
                raise
 
115
            except Exception, exception:
 
116
                # everything else: pass to client, flush, and quit
 
117
                log_exception_quietly()
 
118
                self._send_response(request.FailedSmartServerResponse(
 
119
                    ('error', str(exception))))
 
120
                return
 
121
 
 
122
        if self.has_dispatched:
 
123
            if self._finished:
 
124
                # nothing to do.XXX: this routine should be a single state 
 
125
                # machine too.
 
126
                self.excess_buffer += self.in_buffer
 
127
                self.in_buffer = ''
 
128
                return
 
129
            if self._body_decoder is None:
 
130
                self._body_decoder = LengthPrefixedBodyDecoder()
 
131
            self._body_decoder.accept_bytes(self.in_buffer)
 
132
            self.in_buffer = self._body_decoder.unused_data
 
133
            body_data = self._body_decoder.read_pending_data()
 
134
            self.request.accept_body(body_data)
 
135
            if self._body_decoder.finished_reading:
 
136
                self.request.end_of_body()
 
137
                if not self.request.finished_reading:
 
138
                    raise AssertionError("no more body, request not finished")
 
139
            if self.request.response is not None:
 
140
                self._send_response(self.request.response)
 
141
                self.excess_buffer = self.in_buffer
 
142
                self.in_buffer = ''
 
143
            else:
 
144
                if self.request.finished_reading:
 
145
                    raise AssertionError(
 
146
                        "no response and we have finished reading.")
 
147
 
 
148
    def _send_response(self, response):
 
149
        """Send a smart server response down the output stream."""
 
150
        if self._finished:
 
151
            raise AssertionError('response already sent')
 
152
        args = response.args
 
153
        body = response.body
 
154
        self._finished = True
 
155
        self._write_protocol_version()
 
156
        self._write_success_or_failure_prefix(response)
 
157
        self._write_func(_encode_tuple(args))
 
158
        if body is not None:
 
159
            if not isinstance(body, str):
 
160
                raise ValueError(body)
 
161
            bytes = self._encode_bulk_data(body)
 
162
            self._write_func(bytes)
 
163
 
 
164
    def _write_protocol_version(self):
 
165
        """Write any prefixes this protocol requires.
 
166
        
 
167
        Version one doesn't send protocol versions.
 
168
        """
 
169
 
 
170
    def _write_success_or_failure_prefix(self, response):
 
171
        """Write the protocol specific success/failure prefix.
 
172
 
 
173
        For SmartServerRequestProtocolOne this is omitted but we
 
174
        call is_successful to ensure that the response is valid.
 
175
        """
 
176
        response.is_successful()
 
177
 
 
178
    def next_read_size(self):
 
179
        if self._finished:
 
180
            return 0
 
181
        if self._body_decoder is None:
 
182
            return 1
 
183
        else:
 
184
            return self._body_decoder.next_read_size()
 
185
 
 
186
 
 
187
class SmartServerRequestProtocolTwo(SmartServerRequestProtocolOne):
 
188
    r"""Version two of the server side of the smart protocol.
 
189
   
 
190
    This prefixes responses with the value of RESPONSE_VERSION_TWO.
 
191
    """
 
192
 
 
193
    def _write_success_or_failure_prefix(self, response):
 
194
        """Write the protocol specific success/failure prefix."""
 
195
        if response.is_successful():
 
196
            self._write_func('success\n')
 
197
        else:
 
198
            self._write_func('failed\n')
 
199
 
 
200
    def _write_protocol_version(self):
 
201
        r"""Write any prefixes this protocol requires.
 
202
        
 
203
        Version two sends the value of RESPONSE_VERSION_TWO.
 
204
        """
 
205
        self._write_func(RESPONSE_VERSION_TWO)
 
206
 
 
207
    def _send_response(self, response):
 
208
        """Send a smart server response down the output stream."""
 
209
        if (self._finished):
 
210
            raise AssertionError('response already sent')
 
211
        self._finished = True
 
212
        self._write_protocol_version()
 
213
        self._write_success_or_failure_prefix(response)
 
214
        self._write_func(_encode_tuple(response.args))
 
215
        if response.body is not None:
 
216
            if not isinstance(response.body, str):
 
217
                raise AssertionError('body must be a str')
 
218
            if not (response.body_stream is None):
 
219
                raise AssertionError(
 
220
                    'body_stream and body cannot both be set')
 
221
            bytes = self._encode_bulk_data(response.body)
 
222
            self._write_func(bytes)
 
223
        elif response.body_stream is not None:
 
224
            _send_stream(response.body_stream, self._write_func)
 
225
 
 
226
 
 
227
def _send_stream(stream, write_func):
 
228
    write_func('chunked\n')
 
229
    _send_chunks(stream, write_func)
 
230
    write_func('END\n')
 
231
 
 
232
 
 
233
def _send_chunks(stream, write_func):
 
234
    for chunk in stream:
 
235
        if isinstance(chunk, str):
 
236
            bytes = "%x\n%s" % (len(chunk), chunk)
 
237
            write_func(bytes)
 
238
        elif isinstance(chunk, request.FailedSmartServerResponse):
 
239
            write_func('ERR\n')
 
240
            _send_chunks(chunk.args, write_func)
 
241
            return
 
242
        else:
 
243
            raise errors.BzrError(
 
244
                'Chunks must be str or FailedSmartServerResponse, got %r'
 
245
                % chunk)
 
246
 
 
247
 
 
248
class _StatefulDecoder(object):
 
249
 
 
250
    def __init__(self):
 
251
        self.finished_reading = False
 
252
        self.unused_data = ''
 
253
        self.bytes_left = None
 
254
 
 
255
    def accept_bytes(self, bytes):
 
256
        """Decode as much of bytes as possible.
 
257
 
 
258
        If 'bytes' contains too much data it will be appended to
 
259
        self.unused_data.
 
260
 
 
261
        finished_reading will be set when no more data is required.  Further
 
262
        data will be appended to self.unused_data.
 
263
        """
 
264
        # accept_bytes is allowed to change the state
 
265
        current_state = self.state_accept
 
266
        self.state_accept(bytes)
 
267
        while current_state != self.state_accept:
 
268
            current_state = self.state_accept
 
269
            self.state_accept('')
 
270
 
 
271
 
 
272
class ChunkedBodyDecoder(_StatefulDecoder):
 
273
    """Decoder for chunked body data.
 
274
 
 
275
    This is very similar the HTTP's chunked encoding.  See the description of
 
276
    streamed body data in `doc/developers/network-protocol.txt` for details.
 
277
    """
 
278
 
 
279
    def __init__(self):
 
280
        _StatefulDecoder.__init__(self)
 
281
        self.state_accept = self._state_accept_expecting_header
 
282
        self._in_buffer = ''
 
283
        self.chunk_in_progress = None
 
284
        self.chunks = collections.deque()
 
285
        self.error = False
 
286
        self.error_in_progress = None
 
287
    
 
288
    def next_read_size(self):
 
289
        # Note: the shortest possible chunk is 2 bytes: '0\n', and the
 
290
        # end-of-body marker is 4 bytes: 'END\n'.
 
291
        if self.state_accept == self._state_accept_reading_chunk:
 
292
            # We're expecting more chunk content.  So we're expecting at least
 
293
            # the rest of this chunk plus an END chunk.
 
294
            return self.bytes_left + 4
 
295
        elif self.state_accept == self._state_accept_expecting_length:
 
296
            if self._in_buffer == '':
 
297
                # We're expecting a chunk length.  There's at least two bytes
 
298
                # left: a digit plus '\n'.
 
299
                return 2
 
300
            else:
 
301
                # We're in the middle of reading a chunk length.  So there's at
 
302
                # least one byte left, the '\n' that terminates the length.
 
303
                return 1
 
304
        elif self.state_accept == self._state_accept_reading_unused:
 
305
            return 1
 
306
        elif self.state_accept == self._state_accept_expecting_header:
 
307
            return max(0, len('chunked\n') - len(self._in_buffer))
 
308
        else:
 
309
            raise AssertionError("Impossible state: %r" % (self.state_accept,))
 
310
 
 
311
    def read_next_chunk(self):
 
312
        try:
 
313
            return self.chunks.popleft()
 
314
        except IndexError:
 
315
            return None
 
316
 
 
317
    def _extract_line(self):
 
318
        pos = self._in_buffer.find('\n')
 
319
        if pos == -1:
 
320
            # We haven't read a complete length prefix yet, so there's nothing
 
321
            # to do.
 
322
            return None
 
323
        line = self._in_buffer[:pos]
 
324
        # Trim the prefix (including '\n' delimiter) from the _in_buffer.
 
325
        self._in_buffer = self._in_buffer[pos+1:]
 
326
        return line
 
327
 
 
328
    def _finished(self):
 
329
        self.unused_data = self._in_buffer
 
330
        self._in_buffer = None
 
331
        self.state_accept = self._state_accept_reading_unused
 
332
        if self.error:
 
333
            error_args = tuple(self.error_in_progress)
 
334
            self.chunks.append(request.FailedSmartServerResponse(error_args))
 
335
            self.error_in_progress = None
 
336
        self.finished_reading = True
 
337
 
 
338
    def _state_accept_expecting_header(self, bytes):
 
339
        self._in_buffer += bytes
 
340
        prefix = self._extract_line()
 
341
        if prefix is None:
 
342
            # We haven't read a complete length prefix yet, so there's nothing
 
343
            # to do.
 
344
            return
 
345
        elif prefix == 'chunked':
 
346
            self.state_accept = self._state_accept_expecting_length
 
347
        else:
 
348
            raise errors.SmartProtocolError(
 
349
                'Bad chunked body header: "%s"' % (prefix,))
 
350
 
 
351
    def _state_accept_expecting_length(self, bytes):
 
352
        self._in_buffer += bytes
 
353
        prefix = self._extract_line()
 
354
        if prefix is None:
 
355
            # We haven't read a complete length prefix yet, so there's nothing
 
356
            # to do.
 
357
            return
 
358
        elif prefix == 'ERR':
 
359
            self.error = True
 
360
            self.error_in_progress = []
 
361
            self._state_accept_expecting_length('')
 
362
            return
 
363
        elif prefix == 'END':
 
364
            # We've read the end-of-body marker.
 
365
            # Any further bytes are unused data, including the bytes left in
 
366
            # the _in_buffer.
 
367
            self._finished()
 
368
            return
 
369
        else:
 
370
            self.bytes_left = int(prefix, 16)
 
371
            self.chunk_in_progress = ''
 
372
            self.state_accept = self._state_accept_reading_chunk
 
373
 
 
374
    def _state_accept_reading_chunk(self, bytes):
 
375
        self._in_buffer += bytes
 
376
        in_buffer_len = len(self._in_buffer)
 
377
        self.chunk_in_progress += self._in_buffer[:self.bytes_left]
 
378
        self._in_buffer = self._in_buffer[self.bytes_left:]
 
379
        self.bytes_left -= in_buffer_len
 
380
        if self.bytes_left <= 0:
 
381
            # Finished with chunk
 
382
            self.bytes_left = None
 
383
            if self.error:
 
384
                self.error_in_progress.append(self.chunk_in_progress)
 
385
            else:
 
386
                self.chunks.append(self.chunk_in_progress)
 
387
            self.chunk_in_progress = None
 
388
            self.state_accept = self._state_accept_expecting_length
 
389
        
 
390
    def _state_accept_reading_unused(self, bytes):
 
391
        self.unused_data += bytes
 
392
 
 
393
 
 
394
class LengthPrefixedBodyDecoder(_StatefulDecoder):
 
395
    """Decodes the length-prefixed bulk data."""
 
396
    
 
397
    def __init__(self):
 
398
        _StatefulDecoder.__init__(self)
 
399
        self.state_accept = self._state_accept_expecting_length
 
400
        self.state_read = self._state_read_no_data
 
401
        self._in_buffer = ''
 
402
        self._trailer_buffer = ''
 
403
    
 
404
    def next_read_size(self):
 
405
        if self.bytes_left is not None:
 
406
            # Ideally we want to read all the remainder of the body and the
 
407
            # trailer in one go.
 
408
            return self.bytes_left + 5
 
409
        elif self.state_accept == self._state_accept_reading_trailer:
 
410
            # Just the trailer left
 
411
            return 5 - len(self._trailer_buffer)
 
412
        elif self.state_accept == self._state_accept_expecting_length:
 
413
            # There's still at least 6 bytes left ('\n' to end the length, plus
 
414
            # 'done\n').
 
415
            return 6
 
416
        else:
 
417
            # Reading excess data.  Either way, 1 byte at a time is fine.
 
418
            return 1
 
419
        
 
420
    def read_pending_data(self):
 
421
        """Return any pending data that has been decoded."""
 
422
        return self.state_read()
 
423
 
 
424
    def _state_accept_expecting_length(self, bytes):
 
425
        self._in_buffer += bytes
 
426
        pos = self._in_buffer.find('\n')
 
427
        if pos == -1:
 
428
            return
 
429
        self.bytes_left = int(self._in_buffer[:pos])
 
430
        self._in_buffer = self._in_buffer[pos+1:]
 
431
        self.bytes_left -= len(self._in_buffer)
 
432
        self.state_accept = self._state_accept_reading_body
 
433
        self.state_read = self._state_read_in_buffer
 
434
 
 
435
    def _state_accept_reading_body(self, bytes):
 
436
        self._in_buffer += bytes
 
437
        self.bytes_left -= len(bytes)
 
438
        if self.bytes_left <= 0:
 
439
            # Finished with body
 
440
            if self.bytes_left != 0:
 
441
                self._trailer_buffer = self._in_buffer[self.bytes_left:]
 
442
                self._in_buffer = self._in_buffer[:self.bytes_left]
 
443
            self.bytes_left = None
 
444
            self.state_accept = self._state_accept_reading_trailer
 
445
        
 
446
    def _state_accept_reading_trailer(self, bytes):
 
447
        self._trailer_buffer += bytes
 
448
        # TODO: what if the trailer does not match "done\n"?  Should this raise
 
449
        # a ProtocolViolation exception?
 
450
        if self._trailer_buffer.startswith('done\n'):
 
451
            self.unused_data = self._trailer_buffer[len('done\n'):]
 
452
            self.state_accept = self._state_accept_reading_unused
 
453
            self.finished_reading = True
 
454
    
 
455
    def _state_accept_reading_unused(self, bytes):
 
456
        self.unused_data += bytes
 
457
 
 
458
    def _state_read_no_data(self):
 
459
        return ''
 
460
 
 
461
    def _state_read_in_buffer(self):
 
462
        result = self._in_buffer
 
463
        self._in_buffer = ''
 
464
        return result
 
465
 
 
466
 
 
467
class SmartClientRequestProtocolOne(SmartProtocolBase):
 
468
    """The client-side protocol for smart version 1."""
 
469
 
 
470
    def __init__(self, request):
 
471
        """Construct a SmartClientRequestProtocolOne.
 
472
 
 
473
        :param request: A SmartClientMediumRequest to serialise onto and
 
474
            deserialise from.
 
475
        """
 
476
        self._request = request
 
477
        self._body_buffer = None
 
478
        self._request_start_time = None
 
479
        self._last_verb = None
 
480
 
 
481
    def call(self, *args):
 
482
        if 'hpss' in debug.debug_flags:
 
483
            mutter('hpss call:   %s', repr(args)[1:-1])
 
484
            if getattr(self._request._medium, 'base', None) is not None:
 
485
                mutter('             (to %s)', self._request._medium.base)
 
486
            self._request_start_time = time.time()
 
487
        self._write_args(args)
 
488
        self._request.finished_writing()
 
489
        self._last_verb = args[0]
 
490
 
 
491
    def call_with_body_bytes(self, args, body):
 
492
        """Make a remote call of args with body bytes 'body'.
 
493
 
 
494
        After calling this, call read_response_tuple to find the result out.
 
495
        """
 
496
        if 'hpss' in debug.debug_flags:
 
497
            mutter('hpss call w/body: %s (%r...)', repr(args)[1:-1], body[:20])
 
498
            if getattr(self._request._medium, '_path', None) is not None:
 
499
                mutter('                  (to %s)', self._request._medium._path)
 
500
            mutter('              %d bytes', len(body))
 
501
            self._request_start_time = time.time()
 
502
            if 'hpssdetail' in debug.debug_flags:
 
503
                mutter('hpss body content: %s', body)
 
504
        self._write_args(args)
 
505
        bytes = self._encode_bulk_data(body)
 
506
        self._request.accept_bytes(bytes)
 
507
        self._request.finished_writing()
 
508
        self._last_verb = args[0]
 
509
 
 
510
    def call_with_body_readv_array(self, args, body):
 
511
        """Make a remote call with a readv array.
 
512
 
 
513
        The body is encoded with one line per readv offset pair. The numbers in
 
514
        each pair are separated by a comma, and no trailing \n is emitted.
 
515
        """
 
516
        if 'hpss' in debug.debug_flags:
 
517
            mutter('hpss call w/readv: %s', repr(args)[1:-1])
 
518
            if getattr(self._request._medium, '_path', None) is not None:
 
519
                mutter('                  (to %s)', self._request._medium._path)
 
520
            self._request_start_time = time.time()
 
521
        self._write_args(args)
 
522
        readv_bytes = self._serialise_offsets(body)
 
523
        bytes = self._encode_bulk_data(readv_bytes)
 
524
        self._request.accept_bytes(bytes)
 
525
        self._request.finished_writing()
 
526
        if 'hpss' in debug.debug_flags:
 
527
            mutter('              %d bytes in readv request', len(readv_bytes))
 
528
        self._last_verb = args[0]
 
529
 
 
530
    def cancel_read_body(self):
 
531
        """After expecting a body, a response code may indicate one otherwise.
 
532
 
 
533
        This method lets the domain client inform the protocol that no body
 
534
        will be transmitted. This is a terminal method: after calling it the
 
535
        protocol is not able to be used further.
 
536
        """
 
537
        self._request.finished_reading()
 
538
 
 
539
    def read_response_tuple(self, expect_body=False):
 
540
        """Read a response tuple from the wire.
 
541
 
 
542
        This should only be called once.
 
543
        """
 
544
        result = self._recv_tuple()
 
545
        if 'hpss' in debug.debug_flags:
 
546
            if self._request_start_time is not None:
 
547
                mutter('   result:   %6.3fs  %s',
 
548
                       time.time() - self._request_start_time,
 
549
                       repr(result)[1:-1])
 
550
                self._request_start_time = None
 
551
            else:
 
552
                mutter('   result:   %s', repr(result)[1:-1])
 
553
        self._response_is_unknown_method(result)
 
554
        if not expect_body:
 
555
            self._request.finished_reading()
 
556
        return result
 
557
 
 
558
    def _response_is_unknown_method(self, result_tuple):
 
559
        """Raise UnexpectedSmartServerResponse if the response is an 'unknonwn
 
560
        method' response to the request.
 
561
        
 
562
        :param response: The response from a smart client call_expecting_body
 
563
            call.
 
564
        :param verb: The verb used in that call.
 
565
        :raises: UnexpectedSmartServerResponse
 
566
        """
 
567
        if (result_tuple == ('error', "Generic bzr smart protocol error: "
 
568
                "bad request '%s'" % self._last_verb) or
 
569
              result_tuple == ('error', "Generic bzr smart protocol error: "
 
570
                "bad request u'%s'" % self._last_verb)):
 
571
            # The response will have no body, so we've finished reading.
 
572
            self._request.finished_reading()
 
573
            raise errors.UnknownSmartMethod(self._last_verb)
 
574
        
 
575
    def read_body_bytes(self, count=-1):
 
576
        """Read bytes from the body, decoding into a byte stream.
 
577
        
 
578
        We read all bytes at once to ensure we've checked the trailer for 
 
579
        errors, and then feed the buffer back as read_body_bytes is called.
 
580
        """
 
581
        if self._body_buffer is not None:
 
582
            return self._body_buffer.read(count)
 
583
        _body_decoder = LengthPrefixedBodyDecoder()
 
584
 
 
585
        # Read no more than 64k at a time so that we don't risk error 10055 (no
 
586
        # buffer space available) on Windows.
 
587
        max_read = 64 * 1024
 
588
        while not _body_decoder.finished_reading:
 
589
            bytes_wanted = min(_body_decoder.next_read_size(), max_read)
 
590
            bytes = self._request.read_bytes(bytes_wanted)
 
591
            _body_decoder.accept_bytes(bytes)
 
592
        self._request.finished_reading()
 
593
        self._body_buffer = StringIO(_body_decoder.read_pending_data())
 
594
        # XXX: TODO check the trailer result.
 
595
        if 'hpss' in debug.debug_flags:
 
596
            mutter('              %d body bytes read',
 
597
                   len(self._body_buffer.getvalue()))
 
598
        return self._body_buffer.read(count)
 
599
 
 
600
    def _recv_tuple(self):
 
601
        """Receive a tuple from the medium request."""
 
602
        return _decode_tuple(self._recv_line())
 
603
 
 
604
    def _recv_line(self):
 
605
        """Read an entire line from the medium request."""
 
606
        line = ''
 
607
        while not line or line[-1] != '\n':
 
608
            # TODO: this is inefficient - but tuples are short.
 
609
            new_char = self._request.read_bytes(1)
 
610
            if new_char == '':
 
611
                # end of file encountered reading from server
 
612
                raise errors.ConnectionReset(
 
613
                    "please check connectivity and permissions",
 
614
                    "(and try -Dhpss if further diagnosis is required)")
 
615
            line += new_char
 
616
        return line
 
617
 
 
618
    def query_version(self):
 
619
        """Return protocol version number of the server."""
 
620
        self.call('hello')
 
621
        resp = self.read_response_tuple()
 
622
        if resp == ('ok', '1'):
 
623
            return 1
 
624
        elif resp == ('ok', '2'):
 
625
            return 2
 
626
        else:
 
627
            raise errors.SmartProtocolError("bad response %r" % (resp,))
 
628
 
 
629
    def _write_args(self, args):
 
630
        self._write_protocol_version()
 
631
        bytes = _encode_tuple(args)
 
632
        self._request.accept_bytes(bytes)
 
633
 
 
634
    def _write_protocol_version(self):
 
635
        """Write any prefixes this protocol requires.
 
636
        
 
637
        Version one doesn't send protocol versions.
 
638
        """
 
639
 
 
640
 
 
641
class SmartClientRequestProtocolTwo(SmartClientRequestProtocolOne):
 
642
    """Version two of the client side of the smart protocol.
 
643
    
 
644
    This prefixes the request with the value of REQUEST_VERSION_TWO.
 
645
    """
 
646
 
 
647
    def read_response_tuple(self, expect_body=False):
 
648
        """Read a response tuple from the wire.
 
649
 
 
650
        This should only be called once.
 
651
        """
 
652
        version = self._request.read_line()
 
653
        if version != RESPONSE_VERSION_TWO:
 
654
            raise errors.SmartProtocolError('bad protocol marker %r' % version)
 
655
        response_status = self._recv_line()
 
656
        if response_status not in ('success\n', 'failed\n'):
 
657
            raise errors.SmartProtocolError(
 
658
                'bad protocol status %r' % response_status)
 
659
        self.response_status = response_status == 'success\n'
 
660
        return SmartClientRequestProtocolOne.read_response_tuple(self, expect_body)
 
661
 
 
662
    def _write_protocol_version(self):
 
663
        """Write any prefixes this protocol requires.
 
664
        
 
665
        Version two sends the value of REQUEST_VERSION_TWO.
 
666
        """
 
667
        self._request.accept_bytes(REQUEST_VERSION_TWO)
 
668
 
 
669
    def read_streamed_body(self):
 
670
        """Read bytes from the body, decoding into a byte stream.
 
671
        """
 
672
        # Read no more than 64k at a time so that we don't risk error 10055 (no
 
673
        # buffer space available) on Windows.
 
674
        max_read = 64 * 1024
 
675
        _body_decoder = ChunkedBodyDecoder()
 
676
        while not _body_decoder.finished_reading:
 
677
            bytes_wanted = min(_body_decoder.next_read_size(), max_read)
 
678
            bytes = self._request.read_bytes(bytes_wanted)
 
679
            _body_decoder.accept_bytes(bytes)
 
680
            for body_bytes in iter(_body_decoder.read_next_chunk, None):
 
681
                if 'hpss' in debug.debug_flags:
 
682
                    mutter('              %d byte chunk read',
 
683
                           len(body_bytes))
 
684
                yield body_bytes
 
685
        self._request.finished_reading()
 
686