def test_incorrect_sized_chunks(self):
assembler = util.DataAssembler(3)
with self.assertRaises(ValueError):
assembler.add(b'\x00' * (assembler.alignment + 1), 0)
assembler.add(b'\xff\xfe\xfd', 3)
with self.assertRaises(ValueError):
assembler.add(b'\x00' * (assembler.alignment + 1), 0)
with self.assertRaises(util.UnexpectedChunkLengthError):
assembler.add(b'\x00' * (assembler.alignment - 1), 0)
with self.assertRaises(ValueError):
assembler.add(b'\x00' * (assembler.alignment + 1), 6)
self.assertEqual(b'\x00\x00\x00\xff\xfe\xfd',
assembler.getbytes(incomplete=True))
# Test with predefined length
assembler = util.DataAssembler(3, length=5)
with self.assertRaises(ValueError):
assembler.add(b'\x00\x00\x00\x00', 0)
with self.assertRaises(ValueError):
assembler.add(b'\x00\x00\x00\x00', 3)
with self.assertRaises(util.UnexpectedChunkLengthError):
assembler.add(b'\x00', 3)
with self.assertRaises(util.UnexpectedChunkLengthError):
assembler.add(b'\x00\x00\x00', 3)
assembler.add(b'\xba\xbe', 3)
self.assertEqual(b'\x00\x00\x00\xba\xbe',
assembler.getbytes(incomplete=True))
assembler = util.DataAssembler(3, length=6)
with self.assertRaises(util.UnexpectedChunkLengthError):
assembler.add(b'\x00', 3)
with self.assertRaises(util.UnexpectedChunkLengthError):
assembler.add(b'\x00\x00', 3)
assembler.add(b'\xfe\xba\xbe', 3)
self.assertEqual(b'\x00\x00\x00\xfe\xba\xbe',
assembler.getbytes(incomplete=True))
def _test_unsized_add(self, length, alignment):
data, chunks = self._generate_data(length, alignment)
# Test all possible permutations
for sequence in itertools.permutations(chunks):
assembler = util.DataAssembler(alignment)
for chunk in sequence:
assembler.add(chunk.data, chunk.offset)
self.assertEqual(data, assembler.getbytes())
# Same test with add_chunk()
assembler = util.DataAssembler(alignment)
for chunk in sequence:
assembler.add_chunk(chunk)
self.assertEqual(data, assembler.getbytes())
def test_out_of_bound_offsets(self):
assembler = util.DataAssembler(3)
with self.assertRaises(ValueError):
assembler.add(b'\x00\x00\x00', -1)
assembler = util.DataAssembler(3, length=3)
with self.assertRaises(ValueError):
assembler.add(b'\x00\x00\x00', -1)
with self.assertRaises(util.ChunkPastEndError):
assembler.add(b'\x00\x00\x00', 3)
assembler = util.DataAssembler(3, length=4)
with self.assertRaises(ValueError):
assembler.add(b'\x00\x00\x00', -1)
with self.assertRaises(util.ChunkPastEndError):
assembler.add(b'\x00\x00\x00', 6)
def test_length_deduction(self):
assembler = util.DataAssembler(3)
self.assertIsNone(assembler.length)
assembler.add(b'\xff\xfe', 0)
self.assertEqual(2, assembler.length)
self.assertEqual(b'\xff\xfe', assembler.getbytes())
assembler = util.DataAssembler(3)
self.assertIsNone(assembler.length)
assembler.add(b'\xff\xfe\xfd', 0)
self.assertIsNone(assembler.length)
assembler.add(b'\x01\x02', 6)
self.assertEqual(8, assembler.length)
assembler.add(b'\xfc\xfb\x00', 3)
self.assertEqual(b'\xff\xfe\xfd\xfc\xfb\x00\x01\x02',
assembler.getbytes())
def _update_request_data(self, data, offset):
if self.request_data is None:
self.request_data = util.DataAssembler(30, length=self.request_length)
try:
self.request_data.add(data, offset)
except util.UnexpectedChunkError:
self.collision = True
return
if self.request_data.length is not None:
self._update_request_length(self.request_data.length)
def test_getbytes_incomplete(self):
assembler = util.DataAssembler(2, length=5)
with self.assertRaises(util.IncompleteDataError):
assembler.getbytes()
self.assertEqual(b'\x00\x00\x00\x00\x00',
assembler.getbytes(incomplete=True))
assembler.add(b'\x01\x02', 2)
with self.assertRaises(util.IncompleteDataError):
assembler.getbytes()
self.assertEqual(b'\x00\x00\x01\x02\x00',
assembler.getbytes(incomplete=True))
def test_maligned_offsets(self):
assembler = util.DataAssembler(3)
with self.assertRaises(ValueError):
assembler.add(b'\x00\x00\x00', 1)
with self.assertRaises(ValueError):
assembler.add(b'\x00\x00\x00', 2)
assembler.add(b'\xca\xfe\xba', 0)
with self.assertRaises(ValueError):
assembler.add(b'\x00\x00\x00', 4)
with self.assertRaises(ValueError):
assembler.add(b'\x00\x00\x00', 5)
assembler.add(b'\xbe\x00\xff', 3)
self.assertEqual(b'\xca\xfe\xba\xbe\x00\xff', assembler.getbytes())
def test_init_bad_args(self):
with self.assertRaises(TypeError):
util.DataAssembler(None) # alignment must be int
with self.assertRaises(TypeError):
util.DataAssembler(1.1) # alignment must be int
with self.assertRaises(ValueError):
util.DataAssembler(-1) # alignment must be positive
with self.assertRaises(ValueError):
util.DataAssembler(0) # alignment must be positive
with self.assertRaises(TypeError):
util.DataAssembler(1, length=1.1) # size must be int or None
with self.assertRaises(TypeError):
util.DataAssembler(1, length='asdf') # size must be int or None
with self.assertRaises(ValueError):
util.DataAssembler(1, length=-1) # size must be non-negative
def _update_response_data(self, segment_length, segment_offset, chunk):
if self.response_data is None:
self.response_data = util.DataAssembler(3, length=self.response_length)
# Response data is queried by split segments. chunk is a piece of data
# within one of the segments. In order to reassemble the entire data, we
# must first construct each segment from chunks, then assemble segments into
# final data. Fortunately, the known implementation uses segment_length=48,
# which is a multiple of 3; therefore a single DataAssembler can be used.
absolute_offset = segment_offset + chunk.offset
try:
self.response_data.add(chunk.data, absolute_offset)
except util.UnexpectedChunkError:
self.collision = True
return
if segment_length < 48:
try:
self.response_data.length = segment_offset + segment_length
except ValueError:
self.collision = True
if self.response_data.length is not None:
self._update_response_length(self.response_data.length)
def test_collisions(self):
assembler = util.DataAssembler(3)
assembler.add(b'\x00\x01\x00', 3)
with self.assertRaises(util.ChunkCollisionError):
assembler.add(b'\x00\x02\x00', 3)
assembler.add(b'\x00\x01\x00', 3) # same content
assembler.add(b'\x00\x03\x00', 0)
with self.assertRaises(util.ChunkCollisionError):
assembler.add(b'\x00\x04\x00', 0)
assembler.add(b'\x00\x03\x00', 0) # same content
assembler.add(b'\x00\x05', 6)
with self.assertRaises(util.ChunkCollisionError):
assembler.add(b'\x00\x06', 6)
with self.assertRaises(util.ChunkCollisionError):
assembler.add(b'\x00', 6)
with self.assertRaises(util.ChunkCollisionError):
assembler.add(b'\x00\x00\x00', 6)
# Collision with too long chunks should be handled as ValueError.
# Note that the previous add has determined the length the whole data.
with self.assertRaises(ValueError):
assembler.add(b'\x00\x00\x00\x00', 6)
assembler.add(b'\x00\x05', 6) # same content
self.assertEqual(b'\x00\x03\x00\x00\x01\x00\x00\x05',
assembler.getbytes())