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 <adeebshihadeh@gmail.com>
old-commit-hash: ba098f509b7d2d21114208f700f2de8b9e45a39a
This commit is contained in:
Andrei Radulescu
2024-08-18 02:31:22 +03:00
committed by GitHub
parent 725a6fa998
commit 24c18d2c4b
+16 -9
View File
@@ -29,16 +29,23 @@ class StreamingDecompressor:
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: