diff --git a/engineio/payload.py b/engineio/payload.py index fbf9cbd2..51825839 100644 --- a/engineio/payload.py +++ b/engineio/payload.py @@ -14,31 +14,19 @@ def __init__(self, packets=None, encoded_payload=None): if encoded_payload is not None: self.decode(encoded_payload) - def encode(self, b64=False, jsonp_index=None): + def encode(self, jsonp_index=None): """Encode the payload for transmission.""" - encoded_payload = b'' + encoded_payload = '' for pkt in self.packets: - encoded_packet = pkt.encode(b64=b64) - packet_len = len(encoded_packet) - if b64: - encoded_payload += str(packet_len).encode('utf-8') + b':' + \ - encoded_packet - else: - binary_len = b'' - while packet_len != 0: - binary_len = six.int2byte(packet_len % 10) + binary_len - packet_len = int(packet_len / 10) - if not pkt.binary: - encoded_payload += b'\0' - else: - encoded_payload += b'\1' - encoded_payload += binary_len + b'\xff' + encoded_packet + if encoded_payload: + encoded_payload += '\x1e' + encoded_payload += pkt.encode(b64=True) if jsonp_index is not None: - encoded_payload = b'___eio[' + \ - str(jsonp_index).encode() + \ - b']("' + \ - encoded_payload.replace(b'"', b'\\"') + \ - b'");' + encoded_payload = '___eio[' + \ + str(jsonp_index) + \ + ']("' + \ + encoded_payload.replace('"', '\\"') + \ + '");' return encoded_payload def decode(self, encoded_payload): @@ -49,33 +37,12 @@ def decode(self, encoded_payload): return # JSONP POST payload starts with 'd=' - if encoded_payload.startswith(b'd='): + if encoded_payload.startswith('d='): encoded_payload = urllib.parse.parse_qs( - encoded_payload)[b'd'][0] + encoded_payload)['d'][0] - i = 0 - if six.byte2int(encoded_payload[0:1]) <= 1: - # binary encoding - while i < len(encoded_payload): - if len(self.packets) >= self.max_decode_packets: - raise ValueError('Too many packets in payload') - packet_len = 0 - i += 1 - while six.byte2int(encoded_payload[i:i + 1]) != 255: - packet_len = packet_len * 10 + six.byte2int( - encoded_payload[i:i + 1]) - i += 1 - self.packets.append(packet.Packet( - encoded_packet=encoded_payload[i + 1:i + 1 + packet_len])) - i += packet_len + 1 - else: - # assume text encoding - encoded_payload = encoded_payload.decode('utf-8') - while i < len(encoded_payload): - if len(self.packets) >= self.max_decode_packets: - raise ValueError('Too many packets in payload') - j = encoded_payload.find(':', i) - packet_len = int(encoded_payload[i:j]) - pkt = encoded_payload[j + 1:j + 1 + packet_len] - self.packets.append(packet.Packet(encoded_packet=pkt)) - i = j + 1 + packet_len + encoded_packets = encoded_payload.split('\x1e') + if len(encoded_packets) > self.max_decode_packets: + raise ValueError('Too many packets in payload') + self.packets = [packet.Packet(encoded_packet=encoded_packet) + for encoded_packet in encoded_packets] diff --git a/tests/common/test_payload.py b/tests/common/test_payload.py index 262d43f0..f801f971 100644 --- a/tests/common/test_payload.py +++ b/tests/common/test_payload.py @@ -11,87 +11,59 @@ class TestPayload(unittest.TestCase): def test_encode_empty_payload(self): p = payload.Payload() assert p.packets == [] - assert p.encode() == b'' + assert p.encode() == '' def test_decode_empty_payload(self): - p = payload.Payload(encoded_payload=b'') - assert p.encode() == b'' + p = payload.Payload(encoded_payload='') + assert p.encode() == '' - def test_encode_payload_xhr2(self): + def test_encode_payload_text(self): pkt = packet.Packet(packet.MESSAGE, data=six.text_type('abc')) p = payload.Payload([pkt]) assert p.packets == [pkt] - assert p.encode() == b'\x00\x04\xff4abc' + assert p.encode() == '4abc' - def test_decode_payload_xhr2(self): - p = payload.Payload(encoded_payload=b'\x00\x04\xff4abc') - assert p.encode() == b'\x00\x04\xff4abc' - - def test_encode_payload_xhr_text(self): + def test_encode_payload_text_multiple(self): pkt = packet.Packet(packet.MESSAGE, data=six.text_type('abc')) - p = payload.Payload([pkt]) - assert p.packets == [pkt] - assert p.encode(b64=True) == b'4:4abc' + pkt2 = packet.Packet(packet.MESSAGE, data=six.text_type('def')) + p = payload.Payload([pkt, pkt2]) + assert p.packets == [pkt, pkt2] + assert p.encode() == '4abc\x1e4def' - def test_decode_payload_xhr_text(self): - p = payload.Payload(encoded_payload=b'4:4abc') - assert p.encode() == b'\x00\x04\xff4abc' - - def test_encode_payload_xhr_binary(self): - pkt = packet.Packet(packet.MESSAGE, data=b'\x00\x01\x02', binary=True) + def test_encode_payload_binary(self): + pkt = packet.Packet(packet.MESSAGE, data=b'\x00\x01\x02') p = payload.Payload([pkt]) assert p.packets == [pkt] - assert p.encode(b64=True) == b'6:b4AAEC' + assert p.encode() == 'bAAEC' + + def test_encode_payload_binary_multiple(self): + pkt = packet.Packet(packet.MESSAGE, data=b'\x00\x01\x02') + pkt2 = packet.Packet(packet.MESSAGE, data=b'\x03\x04\x05\x06') + p = payload.Payload([pkt, pkt2]) + assert p.packets == [pkt, pkt2] + assert p.encode() == 'bAAEC\x1ebAwQFBg==' - def test_decode_payload_xhr_binary(self): - p = payload.Payload(encoded_payload=b'6:b4AAEC') - assert p.encode() == b'\x01\x04\xff\x04\x00\x01\x02' + def test_encode_payload_text_binary_multiple(self): + pkt = packet.Packet(packet.MESSAGE, data=six.text_type('abc')) + pkt2 = packet.Packet(packet.MESSAGE, data=b'\x03\x04\x05\x06') + p = payload.Payload([pkt, pkt2, pkt2, pkt]) + assert p.packets == [pkt, pkt2, pkt2, pkt] + assert p.encode() == '4abc\x1ebAwQFBg==\x1ebAwQFBg==\x1e4abc' def test_encode_jsonp_payload(self): pkt = packet.Packet(packet.MESSAGE, data=six.text_type('abc')) p = payload.Payload([pkt]) assert p.packets == [pkt] - assert p.encode(jsonp_index=233) == b'___eio[233]("\x00\x04\xff4abc");' - assert p.encode(jsonp_index=233, b64=True) == b'___eio[233]("4:4abc");' + assert p.encode(jsonp_index=233) == '___eio[233]("4abc");' def test_decode_jsonp_payload(self): - p = payload.Payload(encoded_payload=b'd=4:4abc') - assert p.encode() == b'\x00\x04\xff4abc' + p = payload.Payload(encoded_payload='d=4abc') + assert p.encode() == '4abc' def test_decode_invalid_payload(self): with pytest.raises(ValueError): - payload.Payload(encoded_payload=b'bad payload') - - # performance improvements in the payload parser assume packets in a - # payload are either all binary or all text, so the following test does - # no work anymore. - # - # def test_decode_multi_payload(self): - # p = payload.Payload(encoded_payload=b'4:4abc\x00\x04\xff4def') - # self.assertEqual(len(p.packets), 2) - # self.assertEqual(p.packets[0].data, 'abc') - # self.assertEqual(p.packets[1].data, 'def') - - def test_decode_multi_binary_payload(self): - p = payload.Payload( - encoded_payload=b'\x00\x04\xff4abc\x00\x04\xff4def' - ) - assert len(p.packets) == 2 - assert p.packets[0].data == 'abc' - assert p.packets[1].data == 'def' - - def test_decode_multi_text_payload(self): - p = payload.Payload(encoded_payload=b'4:4abc4:4def') - assert len(p.packets) == 2 - assert p.packets[0].data == 'abc' - assert p.packets[1].data == 'def' - - def test_decode_multi_binary_payload_with_too_many_packets(self): - with pytest.raises(ValueError): - payload.Payload( - encoded_payload=b'\x00\x04\xff4abc\x00\x04\xff4def' * 9 - ) + payload.Payload(encoded_payload='bad payload') - def test_decode_multi_text_payload_with_too_many_packets(self): + def test_decode_multi_payload_with_too_many_packets(self): with pytest.raises(ValueError): - payload.Payload(encoded_payload=b'4:4abc4:4def' * 9) + payload.Payload(encoded_payload='4abc\x1e4def\x1e' * 9 + '6')