def testSingleEmptyChunk(self):
buf = io.BytesIO()
chunked.Write(buf, b"")
buf.seek(0, io.SEEK_SET)
self.assertEqual(chunked.Read(buf), b"")
self.assertIsNone(chunked.Read(buf))
def testMultiple(self):
chunks = [b"foo", b"bar", b"baz"]
buf = io.BytesIO(b"".join(map(chunked.Encode, chunks)))
self.assertEqual(chunked.Read(buf), b"foo")
self.assertEqual(chunked.Read(buf), b"bar")
self.assertEqual(chunked.Read(buf), b"baz")
self.assertIsNone(chunked.Read(buf))
def testMultipleChunks(self):
buf = io.BytesIO()
chunked.Write(buf, b"foo")
chunked.Write(buf, b"bar")
chunked.Write(buf, b"baz")
buf.seek(0, io.SEEK_SET)
self.assertEqual(chunked.Read(buf), b"foo")
self.assertEqual(chunked.Read(buf), b"bar")
self.assertEqual(chunked.Read(buf), b"baz")
self.assertIsNone(chunked.Read(buf))
def testContentTooShort(self):
buf = io.BytesIO()
chunked.Write(buf, b"foobarbaz")
buf = io.BytesIO(buf.getvalue()[:-2])
with self.assertRaises(ValueError):
chunked.Read(buf)
def testMalformedInputWithoutMaxChunkSizeSet(self):
buf1 = io.BytesIO()
chunked.Write(buf1, b"foobarbaz")
buf2 = io.BytesIO(buf1.getvalue()[:-2])
with self.assertRaises(chunked.ChunkTruncatedError):
chunked.Read(buf2)
def testMalformedInputUnseekable(self):
class UnseekableBytesIO(io.BytesIO):
def seekable(self) -> bool:
return False
buf = io.BytesIO()
chunked.Write(buf, b"foobarbaz")
buf = UnseekableBytesIO(buf.getvalue()[:-2])
with self.assertRaisesRegex(ValueError, "Malformed input"):
chunked.Read(buf)
def testRawGzchunkedEmpty(self):
client_id = db_test_utils.InitializeClient(data_store.REL_DB)
flow_id = _WriteTimeline(client_id, [])
args = api_timeline.ApiGetCollectedTimelineArgs()
args.client_id = client_id
args.flow_id = flow_id
args.format = api_timeline.ApiGetCollectedTimelineArgs.Format.RAW_GZCHUNKED
content = b"".join(self.handler.Handle(args).GenerateContent())
buf = io.BytesIO(content)
self.assertIsNone(chunked.Read(buf))
def Deserialize(stream):
"""Deserializes a stream a chunks into a stream of data.
Args:
stream: A stream of serialized chunks (in the gzchunked format).
Yields:
A stream of deserialized data.
"""
for chunk in stream:
buf = io.BytesIO(chunk)
with gzip.GzipFile(fileobj=buf, mode="rb") as filedesc:
while True:
data = chunked.Read(filedesc) # pytype: disable=wrong-arg-types
if data is None:
break
yield data
def Deserialize(stream: Iterator[bytes]) -> Iterator[bytes]:
"""Deserializes a stream a chunks into a stream of data.
Args:
stream: A stream of serialized chunks (in the gzchunked format).
Yields:
A stream of deserialized data.
"""
for chunk in stream:
buf = io.BytesIO(chunk)
with gzip.GzipFile(fileobj=buf, mode="rb") as filedesc:
filedesc.seek(0, os.SEEK_END)
fd_size = filedesc.tell()
filedesc.seek(0, os.SEEK_SET)
while True:
data = chunked.Read(filedesc, max_chunk_size=fd_size) # pytype: disable=wrong-arg-types
if data is None:
break
fd_size -= len(data)
yield data
def testEmpty(self):
chunks = [b""]
buf = io.BytesIO(b"".join(map(chunked.Encode, chunks)))
self.assertEqual(chunked.Read(buf), b"")
self.assertIsNone(chunked.Read(buf))
def testIncorrectSizeTag(self):
buf = io.BytesIO(b"\x00\xff\xee")
with self.assertRaises(ValueError):
chunked.Read(buf)
def testEmptyBuffer(self):
buf = io.BytesIO()
self.assertIsNone(chunked.Read(buf))
def testMalformedInputSeekable(self):
buf = io.BytesIO(b"\xff" * 1024)
with self.assertRaisesRegex(ValueError, "Malformed input"):
chunked.Read(buf)
def testMalformedInpuWithMaxChunkSizeSet(self):
buf = io.BytesIO(b"\xff" * 1024)
with self.assertRaises(chunked.ChunkSizeTooBigError):
chunked.Read(buf, max_chunk_size=1024)
