From 78eea601c753e9d08053b4cd823f47e5e2250b08 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?L=C3=A9a=20Saviot?= <lea.saviot@numworks.com>
Date: Tue, 24 Sep 2019 18:37:20 +0200
Subject: [PATCH] [ion/device] Prevent the user from leaving DFU during
 software download

We prevent the use of the Back key to leave DFU after a flash erase
operation. This way, we prevent the user from interrupting a software
download. After every software download, the calculator resets, which
unlocks the "exit on pressing back"
---
 ion/src/device/shared/usb/calculator.cpp    |  2 +-
 ion/src/device/shared/usb/calculator.h      |  1 +
 ion/src/device/shared/usb/dfu_interface.cpp |  1 +
 ion/src/device/shared/usb/dfu_interface.h   | 11 ++++++++++-
 4 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/ion/src/device/shared/usb/calculator.cpp b/ion/src/device/shared/usb/calculator.cpp
index 84fed03c5..73d6bf579 100644
--- a/ion/src/device/shared/usb/calculator.cpp
+++ b/ion/src/device/shared/usb/calculator.cpp
@@ -20,7 +20,7 @@ void Calculator::PollAndReset(bool exitWithKeyboard) {
 
   Ion::Device::Keyboard::activateRow(exitKeyRow);
 
-  while (!(exitWithKeyboard && Ion::Device::Keyboard::columnIsActive(exitKeyColumn)) &&
+  while (!(exitWithKeyboard && !c.isErasingAndWriting() && Ion::Device::Keyboard::columnIsActive(exitKeyColumn)) &&
       Ion::USB::isPlugged() &&
       !c.isSoftDisconnected()) {
     c.poll();
diff --git a/ion/src/device/shared/usb/calculator.h b/ion/src/device/shared/usb/calculator.h
index d35ba2a24..2532356e3 100644
--- a/ion/src/device/shared/usb/calculator.h
+++ b/ion/src/device/shared/usb/calculator.h
@@ -115,6 +115,7 @@ public:
   {
   }
   uint32_t addressPointer() const { return m_dfuInterface.addressPointer(); }
+  bool isErasingAndWriting() const { return m_dfuInterface.isErasingAndWriting(); }
 protected:
   virtual Descriptor * descriptor(uint8_t type, uint8_t index) override;
   virtual void setActiveConfiguration(uint8_t configurationIndex) override {
diff --git a/ion/src/device/shared/usb/dfu_interface.cpp b/ion/src/device/shared/usb/dfu_interface.cpp
index 5316f7b46..2b98274b4 100644
--- a/ion/src/device/shared/usb/dfu_interface.cpp
+++ b/ion/src/device/shared/usb/dfu_interface.cpp
@@ -209,6 +209,7 @@ void DFUInterface::eraseMemoryIfNeeded() {
     return;
   }
 
+  willErase();
   if (m_erasePage == Flash::TotalNumberOfSectors()) {
     Flash::MassErase();
   } else {
diff --git a/ion/src/device/shared/usb/dfu_interface.h b/ion/src/device/shared/usb/dfu_interface.h
index d36694dba..0ae25d874 100644
--- a/ion/src/device/shared/usb/dfu_interface.h
+++ b/ion/src/device/shared/usb/dfu_interface.h
@@ -27,12 +27,14 @@ public:
     m_largeBuffer{0},
     m_largeBufferLength(0),
     m_writeAddress(0),
-    m_bInterfaceAlternateSetting(bInterfaceAlternateSetting)
+    m_bInterfaceAlternateSetting(bInterfaceAlternateSetting),
+    m_isErasingAndWriting(false)
   {
   }
   uint32_t addressPointer() const { return m_addressPointer; }
   void wholeDataReceivedCallback(SetupPacket * request, uint8_t * transferBuffer, uint16_t * transferBufferLength) override;
   void wholeDataSentCallback(SetupPacket * request, uint8_t * transferBuffer, uint16_t * transferBufferLength) override;
+  bool isErasingAndWriting() const { return m_isErasingAndWriting; }
 
 protected:
   void setActiveInterfaceAlternative(uint8_t interfaceAlternativeIndex) override {
@@ -151,6 +153,12 @@ private:
   bool dfuAbort(uint16_t * transferBufferLength);
   // Leave DFU
   void leaveDFUAndReset();
+  /* Erase and Write state. After starting the erase of flash memory, the user
+   * can no longer leave DFU mode by pressing the Back key of the keyboard. This
+   * way, we prevent the user from interrupting a software download. After every
+   * software download, the calculator resets, which unlocks the "exit on
+   * pressing back". */
+  void willErase() { m_isErasingAndWriting = true; }
 
   Device * m_device;
   Status m_status;
@@ -162,6 +170,7 @@ private:
   uint16_t m_largeBufferLength;
   uint32_t m_writeAddress;
   uint8_t m_bInterfaceAlternateSetting;
+  bool m_isErasingAndWriting;
 };
 
 }