feat: add ability to abort running OTA

Author: Nerivec

src/controller/helpers/ota.ts

@@ -445,7 +445,7 @@ export class OtaSession {
         );
     }
 
-    public async run(): Promise<TZclFrame<"genOta", "upgradeEndRequest">> {
+    public async run(abortSignal: AbortSignal): Promise<TZclFrame<"genOta", "upgradeEndRequest">> {
         // can take a long time, use max (int32 - 1), ~24 days
         const upgradeEndRequest = this.waitForOtaCommand<"upgradeEndRequest">(this.endpoint.ID, UPGRADE_END_REQUEST_ID, undefined, 2147483647);
 
@@ -470,15 +470,20 @@ export class OtaSession {
                             request.header.transactionSequenceNumber,
                             pageOffset,
                             pagePayload.pageSize,
+                            abortSignal.aborted,
                         );
+
+                        abortSignal.throwIfAborted();
                     }
                 } else {
                     await this.sendImageBlockResponse(
                         request.payload as TClusterPayload<"genOta", "imageBlockRequest">,
                         request.header.transactionSequenceNumber,
                         0,
                         0,
+                        abortSignal.aborted,
                     );
+                    abortSignal.throwIfAborted();
                 }
                 /* v8 ignore start */
             }
@@ -494,9 +499,12 @@ export class OtaSession {
             upgradeEndRequest.cancel();
 
             const err = error as Error;
-            err.message = `Device ${this.ieeeAddr} did not start/finish firmware download after being notified. (${err.message})`;
 
-            throw err;
+            if (err.name === "AbortError") {
+                throw new Error(`OTA for device ${this.ieeeAddr} was aborted`);
+            }
+
+            throw new Error(`Device ${this.ieeeAddr} did not start/finish firmware download after being notified. (${err.message})`);
         }
     }
 
@@ -533,6 +541,7 @@ export class OtaSession {
         requestTsn: number,
         pageOffset: number,
         pageSize: number,
+        abort: boolean,
     ): Promise<number> {
         // throttle if needed
         let callNow = performance.now();
@@ -548,6 +557,16 @@ export class OtaSession {
 
         this.#lastBlockResponseTime = callNow;
 
+        if (abort) {
+            try {
+                await this.endpoint.commandResponse("genOta", "imageBlockResponse", {status: Zcl.Status.ABORT}, undefined, requestTsn);
+            } catch (error) {
+                logger.debug(() => `Abort image block response failed for ${this.ieeeAddr}: ${(error as Error).message}`, NS);
+            }
+
+            return 0;
+        }
+
         try {
             const blockPayload = buildImageBlockPayload(this.image, requestPayload, pageOffset, pageSize, this.dataSettings.baseSize);
 

src/controller/model/device.ts

@@ -78,6 +78,7 @@ export class Device extends Entity<ControllerEventMap> {
     private _gpSecurityKey?: number[];
     #scheduledOta: OtaSource | undefined;
     #otaInProgress = false;
+    #otaAbortController: AbortController | undefined;
 
     // Getters/setters
     get ieeeAddr(): string {
@@ -1716,9 +1717,12 @@ export class Device extends Entity<ControllerEventMap> {
         let endResult: TZclFrame<"genOta", "upgradeEndRequest">;
 
         try {
-            endResult = await session.run();
+            this.#otaAbortController = new AbortController();
+
+            endResult = await session.run(this.#otaAbortController.signal);
         } finally {
             this.#otaInProgress = false;
+            this.#otaAbortController = undefined;
         }
 
         logger.debug(() => `Received upgrade end request for ${this.ieeeAddr}: ${JSON.stringify(endResult.payload)}`, NS);
@@ -1808,6 +1812,13 @@ export class Device extends Entity<ControllerEventMap> {
         }
     }
 
+    /**
+     * Abort running OTA if any. Send `ABORT` with next block response to device.
+     */
+    abortOta(): void {
+        this.#otaAbortController?.abort();
+    }
+
     scheduleOta(source: OtaSource): void {
         assert(
             this.endpoints.some((e) => e.supportsOutputCluster("genOta")),

test/device-ota.test.ts

@@ -25,31 +25,33 @@ const IMAGE_PAGE_REQUEST_ID = Zcl.Clusters.genOta.commands.imagePageRequest.ID;
 
 type OtaDeviceBehavior = {
     baseSize: number;
-    /** Mimick the device sending image page requests instead of image block requests */
+    /** Mimic the device sending image page requests instead of image block requests */
     usePageRequests?: boolean;
     pageSize?: number;
-    /** Mimick the device sending non-value (0xff) as `maximumDataSize` */
+    /** Mimic the device sending non-value (0xff) as `maximumDataSize` */
     useNonValueDataSize?: boolean;
-    /** Mimick the device stopping image block/page requests at specified block (i.e. stalling) */
+    /** Mimic the device stopping image block/page requests at specified block (i.e. stalling) */
     stopAfterBlocks?: number;
-    /** Mimick a received default response after block 1. Will repeat last block. */
+    /** Mimic a received default response after block 1. Will repeat last block. */
     triggerDefaultResponse?: Zcl.Status;
-    /** Mimick the device sending out-of-order offset for block/page request, block 2 swapped with block 3 */
+    /** Trigger abort after block 1. */
+    abort?: boolean;
+    /** Mimic the device sending out-of-order offset for block/page request, block 2 swapped with block 3 */
     shuffleOffsets?: boolean;
     /**
      * TODO: implement this
-     * Mimick the device sending block/page request with an offset that is lower or higher than expected flow of "previous offset+data size" at block 2:
+     * Mimic the device sending block/page request with an offset that is lower or higher than expected flow of "previous offset+data size" at block 2:
      * - normal flow would be something like: block1=[offset=0, dataSize=50], block2=[offset=50, dataSize=50], block3=[offset=100, dataSize=50]
      * - with this block2 has this applied to offset: block1=[offset=0, dataSize=50], block2=[offset=(dataSize-misalignedOffset), dataSize=50], block3=[offset=(dataSize*2-misalignedOffset), dataSize=50]
      */
     misalignedOffset?: number;
-    /** Mimick failing block 2 response (mimick device sending new image block/page request for same offset) */
-    failBlockResponse?: boolean;
-    /** Mimick the device sending or not of `upgradeEndRequest` (at end of block/page requests, or as specified by other behaviors) */
+    /** Mimic failing given block response (mimic device sending new image block/page request for same offset OR failure to abort) */
+    failBlockResponse?: number;
+    /** Mimic the device sending or not of `upgradeEndRequest` (at end of block/page requests, or as specified by other behaviors) */
     sendUpgradeEnd?: boolean;
-    /** Mimick the device sending that specific status in `upgradeEndRequest` */
+    /** Mimic the device sending that specific status in `upgradeEndRequest` */
     upgradeEndStatus?: Zcl.Status;
-    /** Mimick the device sending `upgradeEndRequest` after that specific block/page request */
+    /** Mimic the device sending `upgradeEndRequest` after that specific block/page request */
     upgradeEndAfterBlocks?: number;
 };
 
@@ -68,6 +70,7 @@ const createEndpointStub = () => {
 };
 
 const createOtaDeviceWaitFor = (
+    device: Device,
     endpoint: Endpoint,
     image: OtaImage,
     current: TClusterCommandPayload<"genOta", "queryNextImageRequest">,
@@ -113,10 +116,10 @@ const createOtaDeviceWaitFor = (
         });
     };
 
-    if (settings.failBlockResponse) {
+    if (settings.failBlockResponse !== undefined) {
         endpoint.commandResponse = vi.fn((_clusterKey, commandKey, _payload, _options, _transactionSequenceNumber) => {
             if (commandKey === "imageBlockResponse") {
-                if (blocksServed === 1) {
+                if (blocksServed === settings.failBlockResponse) {
                     repeatLastBlock = true;
                     return Promise.reject(new Error("block-fail"));
                 }
@@ -220,6 +223,10 @@ const createOtaDeviceWaitFor = (
             maybeScheduleUpgradeEnd();
         }
 
+        if (blocksServed === 2 && settings.abort) {
+            device.abortOta();
+        }
+
         return {
             promise: Promise.resolve({
                 clusterID: OTA_CLUSTER_ID,
@@ -567,19 +574,18 @@ const createDevice = ({
     database.write = () => {};
     Entity.injectDatabase(database);
 
+    const device = Device.create("Router", "0x1", 0x1001, 1, manufacturerName, "Mains", modelID, InterviewState.Successful, undefined);
     const endpoint = createEndpointStub();
     const currentPayload: TClusterCommandPayload<"genOta", "queryNextImageRequest"> = requestPayload ?? {
         fieldControl: 0,
         manufacturerCode: image.header.manufacturerCode,
         imageType: image.header.imageType,
         fileVersion: image.header.fileVersion + (source?.downgrade ? 1 : -1),
     };
-    const waitFor = createOtaDeviceWaitFor(endpoint, image, currentPayload, {baseSize: dataSettings.baseSize, ...behavior});
+    const waitFor = createOtaDeviceWaitFor(device, endpoint, image, currentPayload, {baseSize: dataSettings.baseSize, ...behavior});
     const adapter = {waitFor, hasZdoMessageOverhead: false} as unknown as Adapter;
     Entity.injectAdapter(adapter);
 
-    const device = Device.create("Router", "0x1", 0x1001, 1, manufacturerName, "Mains", modelID, InterviewState.Successful, undefined);
-
     if (autoAnnounce) {
         const originalOnce = device.once.bind(device);
 
@@ -1200,6 +1206,61 @@ describe("Device OTA", () => {
             expect(device.otaInProgress).toStrictEqual(false);
         });
 
+        it("aborts an in-progress update", async () => {
+            const fileName = OTA_FILES[0];
+            const [image] = await loadImage(fileName);
+            firmwareBuffer = image.raw;
+            const requestPayload: TClusterCommandPayload<"genOta", "queryNextImageRequest"> = {
+                fieldControl: 0,
+                manufacturerCode: image.header.manufacturerCode,
+                imageType: image.header.imageType,
+                fileVersion: image.header.fileVersion - 1,
+            };
+            const baseSize = 55;
+
+            const {device, endpoint, run} = createDevice({
+                image,
+                source: {},
+                requestPayload,
+                dataSettings: {requestTimeout: 1000, responseDelay: 0, baseSize},
+                behavior: {baseSize, sendUpgradeEnd: true, abort: true},
+            });
+
+            await expect(run()).rejects.toThrow(/OTA.*was aborted/);
+            const calls = getResponses(endpoint, "imageBlockResponse");
+            expect(getResponses(endpoint, "imageBlockResponse").length).toStrictEqual(2);
+            expect(calls[1][2]).toStrictEqual({status: Zcl.Status.ABORT});
+            expect(device.otaInProgress).toStrictEqual(false);
+        });
+
+        it("fails to abort an in-progress update", async () => {
+            // this is same as success for ZH, only the device may not have aborted itself
+            const fileName = OTA_FILES[0];
+            const [image] = await loadImage(fileName);
+            firmwareBuffer = image.raw;
+            const requestPayload: TClusterCommandPayload<"genOta", "queryNextImageRequest"> = {
+                fieldControl: 0,
+                manufacturerCode: image.header.manufacturerCode,
+                imageType: image.header.imageType,
+                fileVersion: image.header.fileVersion - 1,
+            };
+            const baseSize = 55;
+
+            const {device, endpoint, run} = createDevice({
+                image,
+                source: {},
+                requestPayload,
+                dataSettings: {requestTimeout: 1000, responseDelay: 0, baseSize},
+                behavior: {baseSize, sendUpgradeEnd: true, abort: true, failBlockResponse: 2},
+            });
+
+            await expect(run()).rejects.toThrow(/OTA.*was aborted/);
+            const calls = getResponses(endpoint, "imageBlockResponse");
+            expect(getResponses(endpoint, "imageBlockResponse").length).toStrictEqual(2);
+            expect(calls[1][2]).toStrictEqual({status: Zcl.Status.ABORT});
+            expect(device.otaInProgress).toStrictEqual(false);
+        });
+
         it("considers an upgrade successful even if no device announce", async () => {
             const fileName = OTA_FILES[0];
             const [image] = await loadImage(fileName);
@@ -2202,7 +2263,7 @@ describe("Device OTA", () => {
                 image,
                 source: {},
                 dataSettings,
-                behavior: {baseSize, sendUpgradeEnd: true, failBlockResponse: true},
+                behavior: {baseSize, sendUpgradeEnd: true, failBlockResponse: 1},
             });
 
             const [from, to] = await run();