2 files changed, 34 insertions, 8 deletions

diff --git a/src/Components/Ble/DfuService.cpp b/src/Components/Ble/DfuService.cpp
index e7b4c5f4..de525321 100644
--- a/src/Components/Ble/DfuService.cpp
+++ b/src/Components/Ble/DfuService.cpp
@@ -107,7 +107,7 @@ int DfuService::WritePacketHandler(uint16_t connectionHandle, os_mbuf *om) {
       bleController.FirmwareUpdateTotalBytes(applicationSize);
       NRF_LOG_INFO("[DFU] -> Start data received : SD size : %d, BT size : %d, app size : %d", softdeviceSize, bootloaderSize, applicationSize);
 
-      for(int erased = 0; erased < applicationSize; erased += 0x1000) {
+      for(int erased = 0; erased < maxImageSize; erased += 0x1000) {
 #if 1
         spiNorFlash.SectorErase(writeOffset + erased);
 
@@ -159,7 +159,7 @@ int DfuService::WritePacketHandler(uint16_t connectionHandle, os_mbuf *om) {
 
       bytesReceived += om->om_len;
       bleController.FirmwareUpdateCurrentBytes(bytesReceived);
-      NRF_LOG_INFO("[DFU] -> Bytes received : %d in %d packets", bytesReceived, nbPacketReceived);
+      //NRF_LOG_INFO("[DFU] -> Bytes received : %d in %d packets", bytesReceived, nbPacketReceived);
 
 
 
@@ -170,6 +170,16 @@ int DfuService::WritePacketHandler(uint16_t connectionHandle, os_mbuf *om) {
         SendNotification(connectionHandle, data, 5);
       }
       if(bytesReceived == applicationSize) {
+          if((nbPacketReceived % nbPacketsToNotify) != 0) {
+              auto remaningPacket = nbPacketReceived % nbPacketsToNotify;
+              uint32_t spiOffset = writeOffset + ((nbPacketReceived-remaningPacket)*20);
+
+              spiNorFlash.Write(writeOffset + ((nbPacketReceived-remaningPacket)*20), tempBuffer, remaningPacket * 20);
+          }
+          if(applicationSize < maxImageSize) {
+              WriteMagicNumber();
+          }
+
         uint8_t data[3]{static_cast<uint8_t>(Opcodes::Response),
                         static_cast<uint8_t>(Opcodes::ReceiveFirmwareImage),
                         static_cast<uint8_t>(ErrorCodes::NoError)};
@@ -304,14 +314,27 @@ void DfuService::Validate() {
     while(currentOffset < applicationSize) {
         uint32_t readSize = (applicationSize - currentOffset) > chunkSize ? chunkSize : (applicationSize - currentOffset);
 
-        spiNorFlash.Read(writeOffset + currentOffset, tempBuffer, chunkSize);
+        spiNorFlash.Read(writeOffset + currentOffset, tempBuffer, readSize);
         if(first) {
-            crc = ComputeCrc(tempBuffer, chunkSize, NULL);
+            crc = ComputeCrc(tempBuffer, readSize, NULL);
             first = false;
         }
         else
-            crc = ComputeCrc(tempBuffer, chunkSize, &crc);
+            crc = ComputeCrc(tempBuffer, readSize, &crc);
         currentOffset += readSize;
     }
+
     NRF_LOG_INFO("CRC : %u", crc);
 }
+
+void DfuService::WriteMagicNumber() {
+    uint32_t magic[4] = {
+            0xf395c277,
+            0x7fefd260,
+            0x0f505235,
+            0x8079b62c,
+    };
+
+    uint32_t offset = writeOffset + (maxImageSize - (4 * sizeof(uint32_t)));
+    spiNorFlash.Write(offset, reinterpret_cast<uint8_t *>(magic), 4 * sizeof(uint32_t));
+}
diff --git a/src/Components/Ble/DfuService.h b/src/Components/Ble/DfuService.h
index e22e09f0..c8351ede 100644
--- a/src/Components/Ble/DfuService.h
+++ b/src/Components/Ble/DfuService.h
@@ -18,6 +18,7 @@ namespace Pinetime {
         DfuService(Pinetime::System::SystemTask& systemTask, Pinetime::Controllers::Ble& bleController,
                    Pinetime::Drivers::SpiNorFlash& spiNorFlash);
         void Init();
+        void Validate();
 
         int OnServiceData(uint16_t connectionHandle, uint16_t attributeHandle, ble_gatt_access_ctxt *context);
       private:
@@ -89,11 +90,12 @@ namespace Pinetime {
         uint8_t nbPacketsToNotify = 0;
         uint32_t nbPacketReceived = 0;
         uint32_t bytesReceived = 0;
-        uint32_t writeOffset = 0; //0x40000;
+        uint32_t writeOffset = 0x40000;
 
         uint32_t softdeviceSize = 0;
         uint32_t bootloaderSize = 0;
-        uint32_t applicationSize = 115200;
+        uint32_t applicationSize = 0;
+        static constexpr uint32_t maxImageSize = 475136;
 
         int SendDfuRevision(os_mbuf *om) const;
         void SendNotification(uint16_t connectionHandle, const uint8_t *data, const size_t size);
@@ -103,9 +105,10 @@ namespace Pinetime {
         uint8_t tempBuffer[200];
         uint16_t ComputeCrc(uint8_t const * p_data, uint32_t size, uint16_t const * p_crc);
 
-        void Validate();
         bool firstCrc = true;
         uint16_t tempCrc = 0;
+
+        void WriteMagicNumber();
     };
   }
 }
\ No newline at end of file