~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/smart/protocol.py

  • Committer: Martin Pool
  • Date: 2005-07-21 21:32:13 UTC
  • Revision ID: mbp@sourcefrog.net-20050721213213-c6ac0e8b06eaad0f
- bzr update-hashes shows some stats on what it did

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