add proper error handling for checksum mismatch

src/sagas/flash-firmware.test.ts

@@ -184,7 +184,7 @@ describe('flashFirmware', () => {
 
             // hub indicates success
 
-            saga.put(programResponse(0, totalFirmwareSize));
+            saga.put(programResponse(0xffffff42, totalFirmwareSize));
 
             action = await saga.take();
             expect(action).toEqual(didProgress(1));
@@ -1113,7 +1113,7 @@ describe('flashFirmware', () => {
 
             // hub indicates incorrect size
 
-            saga.put(programResponse(0, totalFirmwareSize - 1));
+            saga.put(programResponse(0xffffff33, totalFirmwareSize - 1));
 
             // should get a hub error
 
@@ -1132,6 +1132,156 @@ describe('flashFirmware', () => {
 
             await saga.end();
         });
+
+        test('checksum mismatch after flashing', async () => {
+            const metadata: FirmwareMetadata = {
+                'metadata-version': '1.0.0',
+                'device-id': HubType.MoveHub,
+                'checksum-type': 'sum',
+                'firmware-version': '1.2.3',
+                'max-firmware-size': 1024,
+                'mpy-abi-version': 5,
+                'mpy-cross-options': ['-mno-unicode'],
+                'user-mpy-offset': 100,
+            };
+
+            const zip = new JSZip();
+            zip.file('firmware-base.bin', new Uint8Array(64));
+            zip.file('firmware.metadata.json', JSON.stringify(metadata));
+            zip.file('main.py', 'print("test")');
+            zip.file('ReadMe_OSS.txt', 'test');
+
+            jest.spyOn(window, 'fetch').mockResolvedValueOnce(
+                new Response(await zip.generateAsync({ type: 'blob' })),
+            );
+
+            const saga = new AsyncSaga(
+                flashFirmware,
+                {
+                    bootloader: { connection: BootloaderConnectionState.Disconnected },
+                    settings: { flashCurrentProgram: false },
+                },
+                {
+                    nextMessageId: createCountFunc(),
+                },
+            );
+
+            // saga is triggered by this action
+
+            saga.put(flashFirmwareAction());
+
+            // first step is to connect to the hub bootloader
+
+            let action = await saga.take();
+            expect(action).toEqual(connect());
+
+            saga.updateState({
+                bootloader: { connection: BootloaderConnectionState.Connected },
+            });
+            saga.put(didConnect());
+
+            // then find out what kind of hub it is
+
+            action = await saga.take();
+            expect(action).toEqual(infoRequest(0));
+
+            saga.put(didRequest(0));
+            saga.put(infoResponse(0x01000000, 0x08005000, 0x081f800, HubType.MoveHub));
+
+            // then compile main.py to .mpy
+
+            action = await saga.take();
+            expect(action).toMatchInlineSnapshot(`
+                Object {
+                  "options": Array [
+                    "-mno-unicode",
+                  ],
+                  "script": "print(\\"test\\")",
+                  "type": "mpy.action.compile",
+                }
+            `);
+
+            const mpySize = 20;
+            const mpyBinaryData = new Uint8Array(mpySize);
+            saga.put(didCompile(mpyBinaryData));
+
+            // then start flashing the firmware
+
+            // should get didStart action just before starting to erase
+            action = await saga.take();
+            expect(action).toEqual(didStart());
+
+            // erase first
+
+            action = await saga.take();
+            expect(action).toEqual(eraseRequest(1));
+
+            saga.put(didRequest(1));
+            saga.put(eraseResponse(Result.OK));
+
+            // then write the new firmware
+
+            const totalFirmwareSize = metadata['user-mpy-offset'] + mpySize + 8;
+            action = await saga.take();
+            expect(action).toEqual(initRequest(2, totalFirmwareSize));
+
+            saga.put(didRequest(2));
+            saga.put(initResponse(Result.OK));
+
+            const dummyPayload = new ArrayBuffer(0);
+            let id = 2;
+            for (let count = 1, offset = 0; ; count++, offset += 14) {
+                action = await saga.take();
+                expect(action).toEqual(
+                    programRequest(++id, 0x08005000 + offset, dummyPayload),
+                );
+                expect(
+                    (action as BootloaderProgramRequestAction).payload.byteLength,
+                ).toBe(Math.min(14, totalFirmwareSize - offset));
+
+                saga.put(didRequest(id));
+
+                action = await saga.take();
+                expect(action).toEqual(didProgress(offset / totalFirmwareSize));
+
+                // Have to be careful that a checksum request is not sent after
+                // last payload is sent, otherwise the hub gets confused.
+
+                if (offset + 14 >= totalFirmwareSize) {
+                    expect(count).toBe(10);
+                    break;
+                }
+
+                if (count % 10 === 0) {
+                    action = await saga.take();
+                    expect(action).toEqual(checksumRequest(++id));
+
+                    saga.put(didRequest(id));
+                    saga.put(checksumResponse(0));
+                }
+            }
+
+            // hub indicates incorrect checksum
+
+            saga.put(programResponse(0xffffffff, totalFirmwareSize));
+
+            // should get a hub error
+
+            action = await saga.take();
+            expect(action).toEqual(
+                didFailToFinish(
+                    FailToFinishReasonType.HubError,
+                    HubError.ChecksumMismatch,
+                ),
+            );
+
+            // should request to disconnect after failure
+
+            action = await saga.take();
+            expect(action).toEqual(disconnect());
+
+            await saga.end();
+        });
     });
 
     describe('user supplied firmware.zip', () => {
@@ -1265,7 +1415,7 @@ describe('flashFirmware', () => {
 
             // hub indicates success
 
-            saga.put(programResponse(0, totalFirmwareSize));
+            saga.put(programResponse(0xffffff97, totalFirmwareSize));
 
             action = await saga.take();
             expect(action).toEqual(didProgress(1));
@@ -1794,7 +1944,7 @@ describe('flashFirmware', () => {
 
         // hub indicates success
 
-        saga.put(programResponse(0, totalFirmwareSize));
+        saga.put(programResponse(0xffffff33, totalFirmwareSize));
 
         action = await saga.take();
         expect(action).toEqual(didProgress(1));

src/sagas/flash-firmware.ts

@@ -150,7 +150,7 @@ function* firmwareIterator(data: DataView, maxSize: number): Generator<number> {
 function* loadFirmware(
     data: ArrayBuffer,
     program: string | undefined,
-): SagaGenerator<{ firmware: Uint8Array; deviceId: HubType }> {
+): SagaGenerator<{ firmware: Uint8Array; deviceId: HubType; checksum: number }> {
     const [reader, readerErr] = yield* call(() => maybe(FirmwareReader.load(data)));
 
     if (readerErr) {
@@ -224,13 +224,13 @@ function* loadFirmware(
         yield* disconnectAndCancel();
     }
 
-    firmwareView.setUint32(
-        checksumOffset,
-        sumComplement32(firmwareIterator(firmwareView, metadata['max-firmware-size'])),
-        true,
+    const checksum = sumComplement32(
+        firmwareIterator(firmwareView, metadata['max-firmware-size']),
     );
 
-    return { firmware, deviceId: metadata['device-id'] };
+    firmwareView.setUint32(checksumOffset, checksum, true);
+
+    return { firmware, deviceId: metadata['device-id'], checksum };
 }
 
 /**
@@ -258,6 +258,7 @@ function* flashFirmware(action: FlashFirmwareFlashAction): Generator {
     try {
         let firmware: Uint8Array | undefined = undefined;
         let deviceId: HubType | undefined = undefined;
+        let checksum: number | undefined = undefined;
 
         let program: string | undefined = undefined;
 
@@ -278,7 +279,10 @@ function* flashFirmware(action: FlashFirmwareFlashAction): Generator {
         }
 
         if (action.data !== undefined) {
-            ({ firmware, deviceId } = yield* loadFirmware(action.data, program));
+            ({ firmware, deviceId, checksum } = yield* loadFirmware(
+                action.data,
+                program,
+            ));
         }
 
         yield* put(connect());
@@ -327,7 +331,7 @@ function* flashFirmware(action: FlashFirmwareFlashAction): Generator {
             }
 
             const data = yield* call(() => response.arrayBuffer());
-            ({ firmware, deviceId } = yield* loadFirmware(data, program));
+            ({ firmware, deviceId, checksum } = yield* loadFirmware(data, program));
 
             if (deviceId !== undefined && info.hubType !== deviceId) {
                 yield* put(didFailToFinish(FailToFinishReasonType.DeviceMismatch));
@@ -409,6 +413,7 @@ function* flashFirmware(action: FlashFirmwareFlashAction): Generator {
             BootloaderResponseActionType.Program,
             5000,
         );
+
         if (flash.count !== firmware.length) {
             yield* put(
                 didFailToFinish(
@@ -419,6 +424,16 @@ function* flashFirmware(action: FlashFirmwareFlashAction): Generator {
             yield* disconnectAndCancel();
         }
 
+        if (~flash.checksum !== checksum) {
+            yield* put(
+                didFailToFinish(
+                    FailToFinishReasonType.HubError,
+                    HubError.ChecksumMismatch,
+                ),
+            );
+            yield* disconnectAndCancel();
+        }
+
         yield* put(didProgress(1));
 
         // this will cause the remote device to disconnect and reboot