Add the debug option to flash to the current partition

src/app/debug/page.tsx

@@ -1,6 +1,6 @@
 'use client';
 
-import React, { ReactNode, useState } from 'react';
+import React, { ReactNode, useEffect, useState } from 'react';
 import {
   Alert,
   Badge,
@@ -27,6 +27,11 @@ import HexSpan from '@/components/HexSpan';
 import HexViewer from '@/components/HexViewer';
 import { downloadData } from '@/utils/download';
 import { FirmwareInfo } from '@/utils/firmwareIdentifier';
+import FileUpload, { FileUploadHandle } from '@/components/FileUpload';
+import {
+  getOfficialFirmwareVersions,
+  getCommunityFirmwareRemoteData,
+} from '@/remote/firmwareFetcher';
 
 function OtadataDebug({ otaPartition }: { otaPartition: OtaPartition }) {
   const bootPartitionLabel = otaPartition.getCurrentBootPartitionLabel();
@@ -272,8 +277,24 @@ function FirmwareIdentificationDebug({
 }
 
 export default function Debug() {
-  const { debugActions, stepData, isRunning } = useEspOperations();
+  const { actions, debugActions, stepData, isRunning } = useEspOperations();
   const [debugOutputNode, setDebugOutputNode] = useState<ReactNode>(null);
+  const appPartitionFileInput = React.useRef<FileUploadHandle>(null);
+  const [officialFirmwareVersions, setOfficialFirmwareVersions] = useState<{
+    en: string;
+    ch: string;
+  } | null>(null);
+  const [communityFirmwareVersions, setCommunityFirmwareVersions] = useState<{
+    crossPoint: { version: string; releaseDate: string };
+  } | null>(null);
+
+  useEffect(() => {
+    getOfficialFirmwareVersions().then((versions) =>
+      setOfficialFirmwareVersions(versions),
+    );
+
+    getCommunityFirmwareRemoteData().then(setCommunityFirmwareVersions);
+  }, []);
 
   return (
     <Flex direction="column" gap="20px">
@@ -388,6 +409,86 @@ export default function Debug() {
         </Stack>
       </Stack>
       <Separator />
+      <Stack gap={3} as="section">
+        <div>
+          <Heading size="xl">Overwrite current partition (Advanced)</Heading>
+          <Stack gap={1} color="grey" textStyle="sm">
+            <p>
+              These are advanced flashing options for users who want to flash
+              firmware directly to the currently selected partition, as opposed
+              to the backup partition.
+            </p>
+            <p>
+              <b>Flash to current partition</b> will download the firmware and
+              overwrite your current running firmware. The device will reboot
+              with the new firmware on the same partition.
+            </p>
+          </Stack>
+          <Alert.Root status="warning" marginTop={3}>
+            <Alert.Indicator />
+            <Alert.Content>
+              <Alert.Title>
+                Warning: Current firmware will be overwritten
+              </Alert.Title>
+              <Alert.Description>
+                Flashing to the current partition will overwrite the the backup
+                partition unchanged. Proceed with caution.
+              </Alert.Description>
+            </Alert.Content>
+          </Alert.Root>
+        </div>
+        <Stack as="section">
+          <Stack direction="column" gap={2}>
+            <Button
+              variant="subtle"
+              flexGrow={1}
+              onClick={actions.flashEnglishFirmware}
+              disabled={isRunning || !officialFirmwareVersions}
+              loading={!officialFirmwareVersions}
+            >
+              Flash English firmware ({officialFirmwareVersions?.en ?? '...'})
+              to current
+            </Button>
+            <Button
+              variant="subtle"
+              onClick={actions.flashChineseFirmware}
+              disabled={isRunning || !officialFirmwareVersions}
+              loading={!officialFirmwareVersions}
+            >
+              Flash Chinese firmware ({officialFirmwareVersions?.ch ?? '...'})
+              to current
+            </Button>
+            <Button
+              variant="subtle"
+              onClick={actions.flashCrossPointFirmware}
+              disabled={isRunning || !communityFirmwareVersions}
+              loading={!communityFirmwareVersions}
+            >
+              Flash CrossPoint firmware (
+              {communityFirmwareVersions?.crossPoint.version}) -{' '}
+              {communityFirmwareVersions?.crossPoint.releaseDate} to current
+            </Button>
+          </Stack>
+          <Stack direction="row" gap={2}>
+            <Flex grow={1}>
+              <FileUpload ref={appPartitionFileInput} />
+            </Flex>
+            <Button
+              variant="subtle"
+              flexGrow={1}
+              onClick={() =>
+                actions.flashCustomFirmware(() =>
+                  appPartitionFileInput.current?.getFile(),
+                )
+              }
+              disabled={isRunning}
+            >
+              Flash file to current
+            </Button>
+          </Stack>
+        </Stack>
+      </Stack>
+      <Separator />
       <Card.Root variant="subtle">
         <Card.Header>
           <Heading size="lg">Steps</Heading>

src/app/page.tsx

@@ -120,42 +120,50 @@ export default function Home() {
               device using <b>Save full flash</b> above.
             </p>
             <p>
-              <b>Flash English/Chinese firmware</b> will download the firmware,
+              <b>Flash to backup partition</b> will download the firmware,
               overwrite the backup partition with the new firmware, and swap
               over to using this partition (leaving your existing firmware as
-              the new backup). This is significantly faster than a full flash
-              write and will retain all your settings. If it goes wrong, it
-              should be fine to run again.
+              the new backup). This is fast and retains all settings, with the
+              option to swap back if needed. If it goes wrong, it should be fine
+              to run again.
+            </p>
+            <p>
+              For more advanced flashing options (like flashing to the current
+              partition), see the <b>Debug</b> page.
             </p>
           </Stack>
         </div>
         <Stack as="section">
-          <Button
-            variant="subtle"
-            onClick={actions.flashEnglishFirmware}
-            disabled={isRunning || !officialFirmwareVersions}
-            loading={!officialFirmwareVersions}
-          >
-            Flash English firmware ({officialFirmwareVersions?.en ?? '...'})
-          </Button>
-          <Button
-            variant="subtle"
-            onClick={actions.flashChineseFirmware}
-            disabled={isRunning || !officialFirmwareVersions}
-            loading={!officialFirmwareVersions}
-          >
-            Flash Chinese firmware ({officialFirmwareVersions?.ch ?? '...'})
-          </Button>
-          <Button
-            variant="subtle"
-            onClick={actions.flashCrossPointFirmware}
-            disabled={isRunning || !communityFirmwareVersions}
-            loading={!communityFirmwareVersions}
-          >
-            Flash CrossPoint firmware (
-            {communityFirmwareVersions?.crossPoint.version}) -{' '}
-            {communityFirmwareVersions?.crossPoint.releaseDate}
-          </Button>
+          <Stack direction="column" gap={2}>
+            <Button
+              variant="subtle"
+              onClick={actions.flashEnglishFirmwareToBackup}
+              disabled={isRunning || !officialFirmwareVersions}
+              loading={!officialFirmwareVersions}
+            >
+              Flash English firmware ({officialFirmwareVersions?.en ?? '...'})
+              to backup
+            </Button>
+            <Button
+              variant="subtle"
+              onClick={actions.flashChineseFirmwareToBackup}
+              disabled={isRunning || !officialFirmwareVersions}
+              loading={!officialFirmwareVersions}
+            >
+              Flash Chinese firmware ({officialFirmwareVersions?.ch ?? '...'})
+              to backup
+            </Button>
+            <Button
+              variant="subtle"
+              onClick={actions.flashCrossPointFirmwareToBackup}
+              disabled={isRunning || !communityFirmwareVersions}
+              loading={!communityFirmwareVersions}
+            >
+              Flash CrossPoint firmware (
+              {communityFirmwareVersions?.crossPoint.version}) -{' '}
+              {communityFirmwareVersions?.crossPoint.releaseDate} to backup
+            </Button>
+          </Stack>
           <Stack direction="row">
             <Flex grow={1}>
               <FileUpload ref={appPartitionFileInput} />
@@ -164,13 +172,13 @@ export default function Home() {
               variant="subtle"
               flexGrow={1}
               onClick={() =>
-                actions.flashCustomFirmware(() =>
+                actions.flashCustomFirmwareToBackup(() =>
                   appPartitionFileInput.current?.getFile(),
                 )
               }
               disabled={isRunning}
             >
-              Flash firmware from file
+              Flash firmware from file to backup
             </Button>
           </Stack>
           {process.env.NODE_ENV === 'development' && (

src/esp/useEspOperations.ts

@@ -30,6 +30,54 @@ export function useEspOperations() {
 
   const flashRemoteFirmware = async (
     getFirmware: () => Promise<Uint8Array>,
+  ) => {
+    initializeSteps([
+      'Connect to device',
+      'Download firmware',
+      'Read otadata partition',
+      'Flash app partition',
+      'Reset device',
+    ]);
+
+    const espController = await runStep('Connect to device', async () => {
+      const c = await EspController.fromRequestedDevice();
+      await c.connect();
+      return c;
+    });
+
+    const firmwareFile = await runStep('Download firmware', getFirmware);
+
+    const currentPartitionLabel = await runStep(
+      'Read otadata partition',
+      async (): Promise<OtaPartitionDetails['partitionLabel']> => {
+        const partition = await espController.readOtadataPartition((_, p, t) =>
+          updateStepData('Read otadata partition', {
+            progress: { current: p, total: t },
+          }),
+        );
+
+        return partition.getCurrentBootPartitionLabel();
+      },
+    );
+
+    const flashAppPartitionStepName = `Flash app partition (${currentPartitionLabel})`;
+    updateStepData('Flash app partition', { name: flashAppPartitionStepName });
+    await runStep(flashAppPartitionStepName, () =>
+      espController.writeAppPartition(
+        currentPartitionLabel,
+        firmwareFile,
+        (_, p, t) =>
+          updateStepData(flashAppPartitionStepName, {
+            progress: { current: p, total: t },
+          }),
+      ),
+    );
+
+    await runStep('Reset device', () => espController.disconnect());
+  };
+
+  const flashRemoteFirmwareToBackup = async (
+    getFirmware: () => Promise<Uint8Array>,
   ) => {
     initializeSteps([
       'Connect to device',
@@ -89,6 +137,13 @@ export function useEspOperations() {
     await runStep('Reset device', () => espController.disconnect());
   };
 
+  const flashEnglishFirmwareToBackup = async () =>
+    flashRemoteFirmwareToBackup(() => getOfficialFirmware('en'));
+  const flashChineseFirmwareToBackup = async () =>
+    flashRemoteFirmwareToBackup(() => getOfficialFirmware('ch'));
+  const flashCrossPointFirmwareToBackup = async () =>
+    flashRemoteFirmwareToBackup(() => getCommunityFirmware('CrossPoint'));
+
   const flashEnglishFirmware = async () =>
     flashRemoteFirmware(() => getOfficialFirmware('en'));
   const flashChineseFirmware = async () =>
@@ -97,6 +152,60 @@ export function useEspOperations() {
     flashRemoteFirmware(() => getCommunityFirmware('CrossPoint'));
 
   const flashCustomFirmware = async (getFile: () => File | undefined) => {
+    initializeSteps([
+      'Read file',
+      'Connect to device',
+      'Read otadata partition',
+      'Flash app partition',
+      'Reset device',
+    ]);
+
+    const fileData = await runStep('Read file', async () => {
+      const file = getFile();
+      if (!file) {
+        throw new Error('File not found');
+      }
+      return new Uint8Array(await file.arrayBuffer());
+    });
+
+    const espController = await runStep('Connect to device', async () => {
+      const c = await EspController.fromRequestedDevice();
+      await c.connect();
+      return c;
+    });
+
+    const currentPartitionLabel = await runStep(
+      'Read otadata partition',
+      async (): Promise<OtaPartitionDetails['partitionLabel']> => {
+        const partition = await espController.readOtadataPartition((_, p, t) =>
+          updateStepData('Read otadata partition', {
+            progress: { current: p, total: t },
+          }),
+        );
+
+        return partition.getCurrentBootPartitionLabel();
+      },
+    );
+
+    const flashAppPartitionStepName = `Flash app partition (${currentPartitionLabel})`;
+    updateStepData('Flash app partition', { name: flashAppPartitionStepName });
+    await runStep(flashAppPartitionStepName, () =>
+      espController.writeAppPartition(
+        currentPartitionLabel,
+        fileData,
+        (_, p, t) =>
+          updateStepData(flashAppPartitionStepName, {
+            progress: { current: p, total: t },
+          }),
+      ),
+    );
+
+    await runStep('Reset device', () => espController.disconnect());
+  };
+
+  const flashCustomFirmwareToBackup = async (
+    getFile: () => File | undefined,
+  ) => {
     initializeSteps([
       'Read file',
       'Connect to device',
@@ -482,6 +591,16 @@ export function useEspOperations() {
       flashChineseFirmware: wrapWithRunning(flashChineseFirmware),
       flashCrossPointFirmware: wrapWithRunning(flashCrossPointFirmware),
       flashCustomFirmware: wrapWithRunning(flashCustomFirmware),
+      flashEnglishFirmwareToBackup: wrapWithRunning(
+        flashEnglishFirmwareToBackup,
+      ),
+      flashChineseFirmwareToBackup: wrapWithRunning(
+        flashChineseFirmwareToBackup,
+      ),
+      flashCrossPointFirmwareToBackup: wrapWithRunning(
+        flashCrossPointFirmwareToBackup,
+      ),
+      flashCustomFirmwareToBackup: wrapWithRunning(flashCustomFirmwareToBackup),
       saveFullFlash: wrapWithRunning(saveFullFlash),
       writeFullFlash: wrapWithRunning(writeFullFlash),
       fakeWriteFullFlash: wrapWithRunning(fakeWriteFullFlash),