Finish admitted OTA blocks before timed stops

tools/lora_ota/ota_block_drain.py

@@ -0,0 +1,157 @@
+"""Finish admitted OTA blocks at a soft deadline without opening another block.
+
+The verified baseline must come from an owned idle-to-pull session while DATA
+is closed and the source cache is empty. Aggregate receiver counts cannot
+prove a drain when an unrelated active fetch already has untracked RAM slots.
+"""
+from __future__ import annotations
+
+import math
+import struct
+import time
+
+
+class BlockDrainError(ValueError):
+    pass
+
+
+class BlockDrainSeeder:
+    def __init__(self, blob, *, factory=None, clock=time.monotonic, on_change=None):
+        if factory is None:
+            from ota_wizard_repeater import SeederData
+            factory = SeederData
+        self.delegate = factory(blob)
+        self.info = self.delegate.info
+        self.clock, self.on_change = clock, on_change
+        self.state = "closed"
+        self.baseline = self.deadline = self.required_blocks = None
+        self.admitted = set()
+        self.blocked_requests = self.payload_bytes = 0
+        self.last_verified = None
+
+    def snapshot(self):
+        return dict(state=self.state, baseline_blocks=self.baseline,
+                    admitted_blocks=sorted(self.admitted), required_blocks=self.required_blocks,
+                    last_verified_blocks=self.last_verified, blocked_data_requests=self.blocked_requests,
+                    host_payload_bytes=self.payload_bytes)
+
+    def changed(self):
+        if self.on_change is not None:
+            # Called before a newly admitted DATA response is returned, so a
+            # caller can persist its admission journal before USB transmits it.
+            try: self.on_change(self.snapshot())
+            except BaseException:
+                # A retry must not bypass a failed durability barrier merely
+                # because the index was added to the in-memory set already.
+                self.state = "failed"
+                raise
+
+    def arm(self, verified_blocks, *, fresh_resume, deadline=None):
+        total = math.ceil(self.info.payload_size / self.info.block_size)
+        if (self.state != "closed" or self.baseline is not None or fresh_resume is not True
+                or type(verified_blocks) is not int or not 0 <= verified_blocks <= total):
+            raise BlockDrainError("A closed, owned resume needs a verified baseline")
+        self.baseline = self.last_verified = verified_blocks
+        self.deadline = deadline
+        self.state = "open"
+        self.changed()
+        self.check_deadline()
+
+    def freeze(self):
+        if self.state == "open":
+            self.state = "draining"
+            self.required_blocks = self.baseline + len(self.admitted)
+            self.changed()
+        return self.required_blocks
+
+    def restore_owned(self, snapshot, *, deadline=None):
+        """Restore a journal whose receiver/package/owner scope was validated."""
+        total = math.ceil(self.info.payload_size / self.info.block_size)
+        baseline = snapshot.get("baseline_blocks")
+        blocks = snapshot.get("admitted_blocks")
+        count = snapshot.get("last_verified_blocks")
+        state = snapshot.get("state")
+        if (self.baseline is not None or type(baseline) is not int or not 0 <= baseline <= total
+                or not isinstance(blocks, list) or any(type(b) is not int or not 0 <= b < total for b in blocks)
+                or len(set(blocks)) != len(blocks) or baseline + len(blocks) > total
+                or type(count) is not int or not baseline <= count <= baseline + len(blocks)
+                or state not in ("open", "draining", "closed")):
+            raise BlockDrainError("Invalid owned block admission journal")
+        self.baseline, self.last_verified = baseline, count
+        self.admitted = set(blocks)
+        self.deadline = deadline
+        self.state = "open" if state == "open" else "draining"
+        if self.state == "draining": self.required_blocks = baseline + len(blocks)
+        self.changed()
+        self.check_deadline()
+
+    def check_deadline(self):
+        if self.state == "open" and self.deadline is not None and self.clock() >= self.deadline:
+            self.freeze()
+
+    def close(self):
+        self.state = "closed"
+        self.changed()
+
+    def observe_verified(self, count):
+        """Only a receiver's validated, matching-MID FETCH count belongs here."""
+        if self.baseline is None or type(count) is not int:
+            raise BlockDrainError("No verified baseline")
+        maximum = self.baseline + len(self.admitted)
+        if count < self.last_verified or count > maximum:
+            raise BlockDrainError("Receiver progress includes untracked or lost blocks")
+        self.last_verified = count
+        return self.state == "draining" and count == self.required_blocks
+
+    def blocked(self):
+        self.blocked_requests += 1
+        return 1, b""
+
+    def can_serve(self, block):
+        self.check_deadline()
+        return self.state == "open" or (self.state == "draining" and block in self.admitted)
+
+    def admit(self, block):
+        # Recheck after compression/read work. An event loop can be awaiting
+        # a slow CLI command at cutoff; DATA itself enforces the deadline.
+        if not self.can_serve(block): return False
+        if block not in self.admitted:
+            total = math.ceil(self.info.payload_size / self.info.block_size)
+            if self.baseline + len(self.admitted) >= total:
+                raise BlockDrainError("More DATA indices than missing receiver blocks")
+            self.admitted.add(block)
+            self.changed()
+        return True
+
+    def request(self, op, args):
+        info = self.info
+        if op == 3 and len(args) == 7:
+            index, offset, length = struct.unpack("<BIH", args)
+            if index == 0 and 0 < length <= 192 and offset + length <= len(self.delegate.blob):
+                left = max(offset, info.payload_offset)
+                right = min(offset + length, info.payload_offset + info.payload_size)
+                if left < right:
+                    first = (left - info.payload_offset) // info.block_size
+                    last = (right - 1 - info.payload_offset) // info.block_size
+                    if first != last or not self.can_serve(first): return self.blocked()
+                    status, data = self.delegate.request(op, args)
+                    if status == 0:
+                        if not self.admit(first): return self.blocked()
+                        self.payload_bytes += right - left
+                    return status, data
+        if op == 9 and len(args) == 7:
+            index, block, offset, length = struct.unpack("<BHHH", args)
+            total = math.ceil(info.payload_size / info.block_size)
+            if index == 0 and block < total and length <= 190 and (length or not offset):
+                # Compressed-size probes carry no DATA and must neither open
+                # a block nor prevent a subsequent raw fallback.
+                if not length: return self.delegate.request(op, args)
+                if not self.can_serve(block): return self.blocked()
+                status, data = self.delegate.request(op, args)
+                if status == 0:
+                    if not self.admit(block): return self.blocked()
+                    self.payload_bytes += len(data) - 2
+                return status, data
+        # Manifest and Merkle proof metadata remain readable throughout the
+        # drain. The delegate implements only read-only seeder operations.
+        return self.delegate.request(op, args)

tools/lora_ota/ota_wizard_repeater.py

@@ -146,11 +146,13 @@ class RepeaterTransport:
 
     def __init__(self, port, baud=115200, *, command_timeout=12,
                  remote_timeout=40, poll_seconds=.25, serial_factory=None,
-                 open_settle=.75, known_repeater: str | None = None, acl_only=False):
+                 open_settle=.75, known_repeater: str | None = None, acl_only=False,
+                 seeder_factory=None):
         self.port, self.baud = port, baud
         self.command_timeout, self.remote_timeout = command_timeout, remote_timeout
         self.poll_seconds, self.open_settle = poll_seconds, open_settle
         self._serial_factory = serial_factory
+        self._seeder_factory = seeder_factory or SeederData
         self._serial = self._reader_task = None
         self._decoder = StreamDecoder()
         self._lines = asyncio.Queue(maxsize=64)
@@ -510,7 +512,7 @@ class RepeaterTransport:
     async def serve(self, blob, stop_event):
         if self._serve_task is not None:
             raise RepeaterError("The USB repeater is already serving an update")
-        self._seeder = SeederData(blob)
+        self._seeder = self._seeder_factory(blob)
         self._serve_task, self._serve_stop = asyncio.current_task(), stop_event
         stopped = failed = None
         try:

tools/lora_ota/test_ota_wizard_block_drain.py

@@ -0,0 +1,117 @@
+"""Offline soft-deadline DATA gating, including raw/DEFLATE and resume scope."""
+import struct
+import unittest
+
+from ota_block_drain import BlockDrainError, BlockDrainSeeder
+from test_lora_ota import firmware, mota_blob
+
+
+def package():
+    return mota_blob(firmware(b"A" * 2048 + b"B" * 2048 + b"C" * 2048, 0x01120001),
+                     version=0x01120001, block_size_log2=11)
+
+
+def raw(gate, block, offset=0, length=16):
+    return gate.request(3, struct.pack("<BIH", 0,
+        gate.info.payload_offset + block * gate.info.block_size + offset, length))
+
+
+class BlockDrainTests(unittest.TestCase):
+    def gate(self):
+        now = [0]
+        gate = BlockDrainSeeder(package(), clock=lambda: now[0])
+        return gate, now
+
+    def test_closed_baseline_blocks_data_but_not_manifest_or_proofs(self):
+        gate, _ = self.gate()
+        self.assertEqual(raw(gate, 0), (1, b""))
+        self.assertEqual(gate.request(3, struct.pack("<BIH", 0, 0, 8))[0], 0)
+        end = gate.info.payload_offset + gate.info.payload_size
+        if end < len(gate.delegate.blob):
+            self.assertEqual(gate.request(3, struct.pack("<BIH", 0, end, 1))[0], 0)
+        with self.assertRaises(BlockDrainError): gate.arm(0, fresh_resume=False)
+        self.assertEqual(gate.admitted, set())
+
+    def test_deadline_itself_stops_new_indices_but_keeps_started_block_retries(self):
+        gate, now = self.gate()
+        gate.arm(0, fresh_resume=True, deadline=10)
+        self.assertEqual(raw(gate, 0)[0], 0)
+        now[0] = 10
+        self.assertEqual(raw(gate, 1), (1, b""))
+        self.assertEqual(gate.state, "draining")
+        self.assertEqual(gate.required_blocks, 1)
+        self.assertEqual(raw(gate, 0, 100, 50)[0], 0)
+        self.assertFalse(gate.observe_verified(0))
+        self.assertTrue(gate.observe_verified(1))
+        gate.close()
+        self.assertEqual(raw(gate, 0), (1, b""))
+
+    def test_noncontiguous_two_block_flight_requires_both_verified(self):
+        gate, _ = self.gate()
+        gate.arm(2, fresh_resume=True)
+        self.assertEqual(raw(gate, 3)[0], 0)
+        self.assertEqual(raw(gate, 0)[0], 0)
+        self.assertEqual(gate.freeze(), 4)
+        self.assertFalse(gate.observe_verified(3))
+        self.assertTrue(gate.observe_verified(4))
+        self.assertEqual(gate.admitted, {0, 3})
+
+    def test_deflate_size_probe_and_invalid_ranges_do_not_admit(self):
+        gate, _ = self.gate()
+        gate.arm(0, fresh_resume=True)
+        for block in (0, 1, 2):
+            self.assertEqual(gate.request(9, struct.pack("<BHHH", 0, block, 0, 0))[0], 0)
+        self.assertEqual(gate.admitted, set())
+        self.assertEqual(gate.request(9, struct.pack("<BHHH", 0, 0, 60000, 10)), (1, b""))
+        self.assertEqual(gate.admitted, set())
+        self.assertEqual(gate.request(9, struct.pack("<BHHH", 0, 1, 0, 4))[0], 0)
+        gate.freeze()
+        self.assertEqual(raw(gate, 0), (1, b""))
+        self.assertEqual(gate.request(9, struct.pack("<BHHH", 0, 1, 4, 4))[0], 0)
+
+    def test_cross_block_read_is_rejected_atomically(self):
+        gate, _ = self.gate()
+        gate.arm(0, fresh_resume=True)
+        self.assertEqual(raw(gate, 0, 2040, 20), (1, b""))
+        self.assertEqual(gate.admitted, set())
+
+    def test_admission_is_journaled_before_data_is_returned(self):
+        events = []
+        gate = BlockDrainSeeder(package(), on_change=lambda state: events.append(state))
+        gate.arm(0, fresh_resume=True)
+        self.assertEqual(raw(gate, 2)[0], 0)
+        self.assertEqual(events[-1]["admitted_blocks"], [2])
+        gate.on_change = lambda state: (_ for _ in ()).throw(OSError("journal failed"))
+        with self.assertRaisesRegex(OSError, "journal failed"): raw(gate, 0)
+        self.assertEqual(gate.state, "failed")
+        self.assertEqual(raw(gate, 0), (1, b""))
+
+    def test_owned_restart_keeps_all_indices_and_existing_drain(self):
+        gate, _ = self.gate()
+        gate.arm(1, fresh_resume=True)
+        raw(gate, 0)
+        raw(gate, 2)
+        gate.observe_verified(2)
+        gate.freeze()
+        restored, _ = self.gate()
+        restored.restore_owned(gate.snapshot(), deadline=100)
+        self.assertEqual(restored.state, "draining")
+        self.assertEqual(raw(restored, 3), (1, b""))
+        self.assertEqual(raw(restored, 0)[0], 0)
+        self.assertTrue(restored.observe_verified(3))
+
+    def test_unknown_progress_or_corrupt_journal_cannot_claim_a_drain(self):
+        gate, _ = self.gate()
+        gate.arm(1, fresh_resume=True)
+        raw(gate, 0)
+        gate.freeze()
+        for count in (0, 3, None):
+            with self.subTest(count=count), self.assertRaises(BlockDrainError): gate.observe_verified(count)
+        for patch in ({"admitted_blocks": [0, 0]}, {"admitted_blocks": [999]},
+                      {"baseline_blocks": -1}, {"last_verified_blocks": 4}):
+            restored, _ = self.gate()
+            snapshot = dict(gate.snapshot(), **patch)
+            with self.assertRaises(BlockDrainError): restored.restore_owned(snapshot)
+
+
+if __name__ == "__main__": unittest.main()

tools/lora_ota/test_ota_wizard_repeater.py

@@ -603,6 +603,31 @@ class RepeaterTransportTest(unittest.IsolatedAsyncioTestCase):
         self.assertTrue(self.serial.folder_off)
         self.assertIsNone(self.link._seeder)
 
+    async def test_seeder_factory_installs_data_policy_before_folder_attachment(self):
+        await self.link.close()
+        self.serial = FakeSerial()
+        policies = []
+        def factory(blob):
+            from ota_block_drain import BlockDrainSeeder
+            gate = BlockDrainSeeder(blob)
+            policies.append(gate)
+            return gate
+        self.link = RepeaterTransport('fake', serial_factory=lambda **kwargs: self.serial,
+            open_settle=0, command_timeout=.5, seeder_factory=factory)
+        await self.link.open()
+        stopped = asyncio.Event()
+        serving = asyncio.create_task(self.link.serve(package(), stopped))
+        for _ in range(100):
+            if policies and not self.serial.attaching and len(self.serial.binary) >= 3: break
+            await asyncio.sleep(.01)
+        self.assertEqual(len(policies), 1)
+        self.assertIs(self.link._seeder, policies[0])
+        gate = policies[0]
+        offset = gate.info.payload_offset
+        self.assertEqual(gate.request(3, struct.pack('<BIH', 0, offset, 16)), (1, b''))
+        stopped.set()
+        await asyncio.wait_for(serving, 1)
+
     async def test_stale_job_and_unknown_results_do_not_resend_mutations(self):
         await self.link.remote_command(TARGET, "password", "ver")
         self.serial.stale_job = True