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
17
17
"""Wire-level encoding and decoding of requests and responses for the smart
22
22
from cStringIO import StringIO
28
from bzrlib import debug
29
from bzrlib import errors
35
30
from bzrlib.smart import message, request
36
31
from bzrlib.trace import log_exception_quietly, mutter
37
from bzrlib.bencode import bdecode_as_tuple, bencode
32
from bzrlib.util.bencode import bdecode, bencode
40
35
# Protocol version strings. These are sent as prefixes of bzr requests and
63
58
def _encode_tuple(args):
64
59
"""Encode the tuple args to a bytestream."""
65
joined = '\x01'.join(args) + '\n'
66
if type(joined) is unicode:
67
# XXX: We should fix things so this never happens! -AJB, 20100304
68
mutter('response args contain unicode, should be only bytes: %r',
70
joined = joined.encode('ascii')
60
return '\x01'.join(args) + '\n'
74
63
class Requester(object):
120
109
for start, length in offsets:
121
110
txt.append('%d,%d' % (start, length))
122
111
return '\n'.join(txt)
125
114
class SmartServerRequestProtocolOne(SmartProtocolBase):
126
115
"""Server-side encoding and decoding logic for smart version 1."""
128
def __init__(self, backing_transport, write_func, root_client_path='/',
117
def __init__(self, backing_transport, write_func, root_client_path='/'):
130
118
self._backing_transport = backing_transport
131
119
self._root_client_path = root_client_path
132
self._jail_root = jail_root
133
120
self.unused_data = ''
134
121
self._finished = False
135
122
self.in_buffer = ''
157
144
req_args = _decode_tuple(first_line)
158
145
self.request = request.SmartServerRequestHandler(
159
146
self._backing_transport, commands=request.request_handlers,
160
root_client_path=self._root_client_path,
161
jail_root=self._jail_root)
162
self.request.args_received(req_args)
147
root_client_path=self._root_client_path)
148
self.request.dispatch_command(req_args[0], req_args[1:])
163
149
if self.request.finished_reading:
164
150
# trivial request
165
151
self.unused_data = self.in_buffer
338
324
def __init__(self):
339
325
self.finished_reading = False
340
self._in_buffer_list = []
341
self._in_buffer_len = 0
342
327
self.unused_data = ''
343
328
self.bytes_left = None
344
329
self._number_needed_bytes = None
346
def _get_in_buffer(self):
347
if len(self._in_buffer_list) == 1:
348
return self._in_buffer_list[0]
349
in_buffer = ''.join(self._in_buffer_list)
350
if len(in_buffer) != self._in_buffer_len:
351
raise AssertionError(
352
"Length of buffer did not match expected value: %s != %s"
353
% self._in_buffer_len, len(in_buffer))
354
self._in_buffer_list = [in_buffer]
357
def _get_in_bytes(self, count):
358
"""Grab X bytes from the input_buffer.
360
Callers should have already checked that self._in_buffer_len is >
361
count. Note, this does not consume the bytes from the buffer. The
362
caller will still need to call _get_in_buffer() and then
363
_set_in_buffer() if they actually need to consume the bytes.
365
# check if we can yield the bytes from just the first entry in our list
366
if len(self._in_buffer_list) == 0:
367
raise AssertionError('Callers must be sure we have buffered bytes'
368
' before calling _get_in_bytes')
369
if len(self._in_buffer_list[0]) > count:
370
return self._in_buffer_list[0][:count]
371
# We can't yield it from the first buffer, so collapse all buffers, and
373
in_buf = self._get_in_buffer()
374
return in_buf[:count]
376
def _set_in_buffer(self, new_buf):
377
if new_buf is not None:
378
self._in_buffer_list = [new_buf]
379
self._in_buffer_len = len(new_buf)
381
self._in_buffer_list = []
382
self._in_buffer_len = 0
384
331
def accept_bytes(self, bytes):
385
332
"""Decode as much of bytes as possible.
391
338
data will be appended to self.unused_data.
393
340
# accept_bytes is allowed to change the state
341
current_state = self.state_accept
394
342
self._number_needed_bytes = None
395
# lsprof puts a very large amount of time on this specific call for
397
self._in_buffer_list.append(bytes)
398
self._in_buffer_len += len(bytes)
343
self._in_buffer += bytes
400
345
# Run the function for the current state.
401
current_state = self.state_accept
402
346
self.state_accept()
403
347
while current_state != self.state_accept:
404
348
# The current state has changed. Run the function for the new
459
403
def _extract_line(self):
460
in_buf = self._get_in_buffer()
461
pos = in_buf.find('\n')
404
pos = self._in_buffer.find('\n')
463
406
# We haven't read a complete line yet, so request more bytes before
465
408
raise _NeedMoreBytes(1)
409
line = self._in_buffer[:pos]
467
410
# Trim the prefix (including '\n' delimiter) from the _in_buffer.
468
self._set_in_buffer(in_buf[pos+1:])
411
self._in_buffer = self._in_buffer[pos+1:]
471
414
def _finished(self):
472
self.unused_data = self._get_in_buffer()
473
self._in_buffer_list = []
474
self._in_buffer_len = 0
415
self.unused_data = self._in_buffer
475
417
self.state_accept = self._state_accept_reading_unused
477
419
error_args = tuple(self.error_in_progress)
506
448
self.state_accept = self._state_accept_reading_chunk
508
450
def _state_accept_reading_chunk(self):
509
in_buf = self._get_in_buffer()
510
in_buffer_len = len(in_buf)
511
self.chunk_in_progress += in_buf[:self.bytes_left]
512
self._set_in_buffer(in_buf[self.bytes_left:])
451
in_buffer_len = len(self._in_buffer)
452
self.chunk_in_progress += self._in_buffer[:self.bytes_left]
453
self._in_buffer = self._in_buffer[self.bytes_left:]
513
454
self.bytes_left -= in_buffer_len
514
455
if self.bytes_left <= 0:
515
456
# Finished with chunk
520
461
self.chunks.append(self.chunk_in_progress)
521
462
self.chunk_in_progress = None
522
463
self.state_accept = self._state_accept_expecting_length
524
465
def _state_accept_reading_unused(self):
525
self.unused_data += self._get_in_buffer()
526
self._in_buffer_list = []
466
self.unused_data += self._in_buffer
529
470
class LengthPrefixedBodyDecoder(_StatefulDecoder):
530
471
"""Decodes the length-prefixed bulk data."""
532
473
def __init__(self):
533
474
_StatefulDecoder.__init__(self)
534
475
self.state_accept = self._state_accept_expecting_length
535
476
self.state_read = self._state_read_no_data
537
478
self._trailer_buffer = ''
539
480
def next_read_size(self):
540
481
if self.bytes_left is not None:
541
482
# Ideally we want to read all the remainder of the body and the
552
493
# Reading excess data. Either way, 1 byte at a time is fine.
555
496
def read_pending_data(self):
556
497
"""Return any pending data that has been decoded."""
557
498
return self.state_read()
559
500
def _state_accept_expecting_length(self):
560
in_buf = self._get_in_buffer()
561
pos = in_buf.find('\n')
501
pos = self._in_buffer.find('\n')
564
self.bytes_left = int(in_buf[:pos])
565
self._set_in_buffer(in_buf[pos+1:])
504
self.bytes_left = int(self._in_buffer[:pos])
505
self._in_buffer = self._in_buffer[pos+1:]
566
506
self.state_accept = self._state_accept_reading_body
567
507
self.state_read = self._state_read_body_buffer
569
509
def _state_accept_reading_body(self):
570
in_buf = self._get_in_buffer()
572
self.bytes_left -= len(in_buf)
573
self._set_in_buffer(None)
510
self._body += self._in_buffer
511
self.bytes_left -= len(self._in_buffer)
574
513
if self.bytes_left <= 0:
575
514
# Finished with body
576
515
if self.bytes_left != 0:
578
517
self._body = self._body[:self.bytes_left]
579
518
self.bytes_left = None
580
519
self.state_accept = self._state_accept_reading_trailer
582
521
def _state_accept_reading_trailer(self):
583
self._trailer_buffer += self._get_in_buffer()
584
self._set_in_buffer(None)
522
self._trailer_buffer += self._in_buffer
585
524
# TODO: what if the trailer does not match "done\n"? Should this raise
586
525
# a ProtocolViolation exception?
587
526
if self._trailer_buffer.startswith('done\n'):
588
527
self.unused_data = self._trailer_buffer[len('done\n'):]
589
528
self.state_accept = self._state_accept_reading_unused
590
529
self.finished_reading = True
592
531
def _state_accept_reading_unused(self):
593
self.unused_data += self._get_in_buffer()
594
self._set_in_buffer(None)
532
self.unused_data += self._in_buffer
596
535
def _state_read_no_data(self):
626
565
mutter('hpss call: %s', repr(args)[1:-1])
627
566
if getattr(self._request._medium, 'base', None) is not None:
628
567
mutter(' (to %s)', self._request._medium.base)
629
self._request_start_time = osutils.timer_func()
568
self._request_start_time = time.time()
630
569
self._write_args(args)
631
570
self._request.finished_writing()
632
571
self._last_verb = args[0]
641
580
if getattr(self._request._medium, '_path', None) is not None:
642
581
mutter(' (to %s)', self._request._medium._path)
643
582
mutter(' %d bytes', len(body))
644
self._request_start_time = osutils.timer_func()
583
self._request_start_time = time.time()
645
584
if 'hpssdetail' in debug.debug_flags:
646
585
mutter('hpss body content: %s', body)
647
586
self._write_args(args)
660
599
mutter('hpss call w/readv: %s', repr(args)[1:-1])
661
600
if getattr(self._request._medium, '_path', None) is not None:
662
601
mutter(' (to %s)', self._request._medium._path)
663
self._request_start_time = osutils.timer_func()
602
self._request_start_time = time.time()
664
603
self._write_args(args)
665
604
readv_bytes = self._serialise_offsets(body)
666
605
bytes = self._encode_bulk_data(readv_bytes)
670
609
mutter(' %d bytes in readv request', len(readv_bytes))
671
610
self._last_verb = args[0]
673
def call_with_body_stream(self, args, stream):
674
# Protocols v1 and v2 don't support body streams. So it's safe to
675
# assume that a v1/v2 server doesn't support whatever method we're
676
# trying to call with a body stream.
677
self._request.finished_writing()
678
self._request.finished_reading()
679
raise errors.UnknownSmartMethod(args[0])
681
612
def cancel_read_body(self):
682
613
"""After expecting a body, a response code may indicate one otherwise.
756
687
# The response will have no body, so we've finished reading.
757
688
self._request.finished_reading()
758
689
raise errors.UnknownSmartMethod(self._last_verb)
760
691
def read_body_bytes(self, count=-1):
761
692
"""Read bytes from the body, decoding into a byte stream.
763
We read all bytes at once to ensure we've checked the trailer for
694
We read all bytes at once to ensure we've checked the trailer for
764
695
errors, and then feed the buffer back as read_body_bytes is called.
766
697
if self._body_buffer is not None:
874
805
def build_server_protocol_three(backing_transport, write_func,
875
root_client_path, jail_root=None):
876
807
request_handler = request.SmartServerRequestHandler(
877
808
backing_transport, commands=request.request_handlers,
878
root_client_path=root_client_path, jail_root=jail_root)
809
root_client_path=root_client_path)
879
810
responder = ProtocolThreeResponder(write_func)
880
811
message_handler = message.ConventionalRequestHandler(request_handler, responder)
881
812
return ProtocolThreeDecoder(message_handler)
911
842
# We do *not* set self.decoding_failed here. The message handler
912
843
# has raised an error, but the decoder is still able to parse bytes
913
844
# and determine when this message ends.
914
if not isinstance(exception.exc_value, errors.UnknownSmartMethod):
915
log_exception_quietly()
845
log_exception_quietly()
916
846
self.message_handler.protocol_error(exception.exc_value)
917
847
# The state machine is ready to continue decoding, but the
918
848
# exception has interrupted the loop that runs the state machine.
935
865
self.message_handler.protocol_error(exception)
937
867
def _extract_length_prefixed_bytes(self):
938
if self._in_buffer_len < 4:
868
if len(self._in_buffer) < 4:
939
869
# A length prefix by itself is 4 bytes, and we don't even have that
941
871
raise _NeedMoreBytes(4)
942
(length,) = struct.unpack('!L', self._get_in_bytes(4))
872
(length,) = struct.unpack('!L', self._in_buffer[:4])
943
873
end_of_bytes = 4 + length
944
if self._in_buffer_len < end_of_bytes:
874
if len(self._in_buffer) < end_of_bytes:
945
875
# We haven't yet read as many bytes as the length-prefix says there
947
877
raise _NeedMoreBytes(end_of_bytes)
948
878
# Extract the bytes from the buffer.
949
in_buf = self._get_in_buffer()
950
bytes = in_buf[4:end_of_bytes]
951
self._set_in_buffer(in_buf[end_of_bytes:])
879
bytes = self._in_buffer[4:end_of_bytes]
880
self._in_buffer = self._in_buffer[end_of_bytes:]
954
883
def _extract_prefixed_bencoded_data(self):
955
884
prefixed_bytes = self._extract_length_prefixed_bytes()
957
decoded = bdecode_as_tuple(prefixed_bytes)
886
decoded = bdecode(prefixed_bytes)
958
887
except ValueError:
959
888
raise errors.SmartProtocolError(
960
889
'Bytes %r not bencoded' % (prefixed_bytes,))
963
892
def _extract_single_byte(self):
964
if self._in_buffer_len == 0:
893
if self._in_buffer == '':
965
894
# The buffer is empty
966
895
raise _NeedMoreBytes(1)
967
in_buf = self._get_in_buffer()
969
self._set_in_buffer(in_buf[1:])
896
one_byte = self._in_buffer[0]
897
self._in_buffer = self._in_buffer[1:]
972
900
def _state_accept_expecting_protocol_version(self):
973
needed_bytes = len(MESSAGE_VERSION_THREE) - self._in_buffer_len
974
in_buf = self._get_in_buffer()
901
needed_bytes = len(MESSAGE_VERSION_THREE) - len(self._in_buffer)
975
902
if needed_bytes > 0:
976
903
# We don't have enough bytes to check if the protocol version
977
904
# marker is right. But we can check if it is already wrong by
981
908
# len(MESSAGE_VERSION_THREE) bytes. So if the bytes we have so far
982
909
# are wrong then we should just raise immediately rather than
984
if not MESSAGE_VERSION_THREE.startswith(in_buf):
911
if not MESSAGE_VERSION_THREE.startswith(self._in_buffer):
985
912
# We have enough bytes to know the protocol version is wrong
986
raise errors.UnexpectedProtocolVersionMarker(in_buf)
913
raise errors.UnexpectedProtocolVersionMarker(self._in_buffer)
987
914
raise _NeedMoreBytes(len(MESSAGE_VERSION_THREE))
988
if not in_buf.startswith(MESSAGE_VERSION_THREE):
989
raise errors.UnexpectedProtocolVersionMarker(in_buf)
990
self._set_in_buffer(in_buf[len(MESSAGE_VERSION_THREE):])
915
if not self._in_buffer.startswith(MESSAGE_VERSION_THREE):
916
raise errors.UnexpectedProtocolVersionMarker(self._in_buffer)
917
self._in_buffer = self._in_buffer[len(MESSAGE_VERSION_THREE):]
991
918
self.state_accept = self._state_accept_expecting_headers
993
920
def _state_accept_expecting_headers(self):
1073
1000
class _ProtocolThreeEncoder(object):
1075
1002
response_marker = request_marker = MESSAGE_VERSION_THREE
1076
BUFFER_SIZE = 1024*1024 # 1 MiB buffer before flushing
1078
1004
def __init__(self, write_func):
1081
1006
self._real_write_func = write_func
1083
1008
def _write_func(self, bytes):
1084
# TODO: It is probably more appropriate to use sum(map(len, _buf))
1085
# for total number of bytes to write, rather than buffer based on
1086
# the number of write() calls
1087
# TODO: Another possibility would be to turn this into an async model.
1088
# Where we let another thread know that we have some bytes if
1089
# they want it, but we don't actually block for it
1090
# Note that osutils.send_all always sends 64kB chunks anyway, so
1091
# we might just push out smaller bits at a time?
1092
self._buf.append(bytes)
1093
self._buf_len += len(bytes)
1094
if self._buf_len > self.BUFFER_SIZE:
1097
1011
def flush(self):
1099
self._real_write_func(''.join(self._buf))
1013
self._real_write_func(self._buf)
1103
1016
def _serialise_offsets(self, offsets):
1104
1017
"""Serialise a readv offset list."""
1153
1063
_ProtocolThreeEncoder.__init__(self, write_func)
1154
1064
self.response_sent = False
1155
1065
self._headers = {'Software version': bzrlib.__version__}
1156
if 'hpss' in debug.debug_flags:
1157
self._thread_id = thread.get_ident()
1158
self._response_start_time = None
1160
def _trace(self, action, message, extra_bytes=None, include_time=False):
1161
if self._response_start_time is None:
1162
self._response_start_time = osutils.timer_func()
1164
t = '%5.3fs ' % (time.clock() - self._response_start_time)
1167
if extra_bytes is None:
1170
extra = ' ' + repr(extra_bytes[:40])
1172
extra = extra[:29] + extra[-1] + '...'
1173
mutter('%12s: [%s] %s%s%s'
1174
% (action, self._thread_id, t, message, extra))
1176
1067
def send_error(self, exception):
1177
1068
if self.response_sent:
1204
1093
self._write_success_status()
1206
1095
self._write_error_status()
1207
if 'hpss' in debug.debug_flags:
1208
self._trace('response', repr(response.args))
1209
1096
self._write_structure(response.args)
1210
1097
if response.body is not None:
1211
1098
self._write_prefixed_body(response.body)
1212
if 'hpss' in debug.debug_flags:
1213
self._trace('body', '%d bytes' % (len(response.body),),
1214
response.body, include_time=True)
1215
1099
elif response.body_stream is not None:
1216
count = num_bytes = 0
1218
for exc_info, chunk in _iter_with_errors(response.body_stream):
1220
if exc_info is not None:
1221
self._write_error_status()
1222
error_struct = request._translate_error(exc_info[1])
1223
self._write_structure(error_struct)
1226
if isinstance(chunk, request.FailedSmartServerResponse):
1227
self._write_error_status()
1228
self._write_structure(chunk.args)
1230
num_bytes += len(chunk)
1231
if first_chunk is None:
1233
self._write_prefixed_body(chunk)
1234
if 'hpssdetail' in debug.debug_flags:
1235
# Not worth timing separately, as _write_func is
1237
self._trace('body chunk',
1238
'%d bytes' % (len(chunk),),
1239
chunk, suppress_time=True)
1240
if 'hpss' in debug.debug_flags:
1241
self._trace('body stream',
1242
'%d bytes %d chunks' % (num_bytes, count),
1100
for chunk in response.body_stream:
1101
self._write_prefixed_body(chunk)
1244
1103
self._write_end()
1245
if 'hpss' in debug.debug_flags:
1246
self._trace('response end', '', include_time=True)
1249
def _iter_with_errors(iterable):
1250
"""Handle errors from iterable.next().
1254
for exc_info, value in _iter_with_errors(iterable):
1257
This is a safer alternative to::
1260
for value in iterable:
1265
Because the latter will catch errors from the for-loop body, not just
1268
If an error occurs, exc_info will be a exc_info tuple, and the generator
1269
will terminate. Otherwise exc_info will be None, and value will be the
1270
value from iterable.next(). Note that KeyboardInterrupt and SystemExit
1271
will not be itercepted.
1273
iterator = iter(iterable)
1276
yield None, iterator.next()
1277
except StopIteration:
1279
except (KeyboardInterrupt, SystemExit):
1282
mutter('_iter_with_errors caught error')
1283
log_exception_quietly()
1284
yield sys.exc_info(), None
1288
1106
class ProtocolThreeRequester(_ProtocolThreeEncoder, Requester):
1295
1113
def set_headers(self, headers):
1296
1114
self._headers = headers.copy()
1298
1116
def call(self, *args):
1299
1117
if 'hpss' in debug.debug_flags:
1300
1118
mutter('hpss call: %s', repr(args)[1:-1])
1301
1119
base = getattr(self._medium_request._medium, 'base', None)
1302
1120
if base is not None:
1303
1121
mutter(' (to %s)', base)
1304
self._request_start_time = osutils.timer_func()
1122
self._request_start_time = time.time()
1305
1123
self._write_protocol_version()
1306
1124
self._write_headers(self._headers)
1307
1125
self._write_structure(args)
1349
1167
self._write_end()
1350
1168
self._medium_request.finished_writing()
1352
def call_with_body_stream(self, args, stream):
1353
if 'hpss' in debug.debug_flags:
1354
mutter('hpss call w/body stream: %r', args)
1355
path = getattr(self._medium_request._medium, '_path', None)
1356
if path is not None:
1357
mutter(' (to %s)', path)
1358
self._request_start_time = osutils.timer_func()
1359
self._write_protocol_version()
1360
self._write_headers(self._headers)
1361
self._write_structure(args)
1362
# TODO: notice if the server has sent an early error reply before we
1363
# have finished sending the stream. We would notice at the end
1364
# anyway, but if the medium can deliver it early then it's good
1365
# to short-circuit the whole request...
1366
for exc_info, part in _iter_with_errors(stream):
1367
if exc_info is not None:
1368
# Iterating the stream failed. Cleanly abort the request.
1369
self._write_error_status()
1370
# Currently the client unconditionally sends ('error',) as the
1372
self._write_structure(('error',))
1374
self._medium_request.finished_writing()
1375
raise exc_info[0], exc_info[1], exc_info[2]
1377
self._write_prefixed_body(part)
1380
self._medium_request.finished_writing()
added
removed
bzrlib/smart/protocol.py
