From 24c18d2c4b05968bc3166d8af53851dd966c4f54 Mon Sep 17 00:00:00 2001
From: Andrei Radulescu <andi.radulescu@gmail.com>
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 <adeebshihadeh@gmail.com>
old-commit-hash: ba098f509b7d2d21114208f700f2de8b9e45a39a
---
 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 7e3536f775..f5261953d5 100755
--- a/system/hardware/tici/agnos.py
+++ b/system/hardware/tici/agnos.py
@@ -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: