From ba098f509b7d2d21114208f700f2de8b9e45a39a Mon Sep 17 00:00:00 2001 From: Andrei Radulescu Date: Sun, 18 Aug 2024 02:31:22 +0300 Subject: [PATCH] agnos updater: set decompress max_length (#33320) * agnos: decompress max_length * flash last chunk after eof * don't decompress more than length * cleanup --------- Co-authored-by: Adeeb Shihadeh --- system/hardware/tici/agnos.py | 25 ++++++++++++++++--------- 1 file changed, 16 insertions(+), 9 deletions(-) diff --git a/system/hardware/tici/agnos.py b/system/hardware/tici/agnos.py index 7e3536f775c932..f5261953d5c468 100755 --- a/system/hardware/tici/agnos.py +++ b/system/hardware/tici/agnos.py @@ -29,16 +29,23 @@ def __init__(self, url: str) -> None: self.sha256 = hashlib.sha256() def read(self, length: int) -> bytes: - while len(self.buf) < length: - self.req.raise_for_status() + while len(self.buf) < length and not self.eof: + if self.decompressor.needs_input: + self.req.raise_for_status() - try: - compressed = next(self.it) - except StopIteration: + try: + compressed = next(self.it) + except StopIteration: + self.eof = True + break + else: + compressed = b'' + + self.buf += self.decompressor.decompress(compressed, max_length=length) + + if self.decompressor.eof: self.eof = True break - out = self.decompressor.decompress(compressed) - self.buf += out result = self.buf[:length] self.buf = self.buf[length:] @@ -83,8 +90,8 @@ def unsparsify(f: StreamingDecompressor) -> Generator[bytes, None, None]: # noop wrapper with same API as unsparsify() for non sparse images def noop(f: StreamingDecompressor) -> Generator[bytes, None, None]: - while not f.eof: - yield f.read(1024 * 1024) + while len(chunk := f.read(1024 * 1024)) > 0: + yield chunk def get_target_slot_number() -> int: