InfiniTime/src/components/ble/DfuService.h

#pragma once

#include <cstdint>
#include <array>

#define min // workaround: nimble's min/max macros conflict with libstdc++
#define max
#include <host/ble_gap.h>
#undef max
#undef min

namespace Pinetime {
  namespace System {
    class SystemTask;
  }

  namespace Drivers {
    class SpiNorFlash;
  }

  namespace Controllers {
    class Ble;
    class Settings;
    class NotificationManager;

    class DfuService {
    public:
      DfuService(Pinetime::System::SystemTask& systemTask,
                 Pinetime::Controllers::Ble& bleController,
                 Pinetime::Drivers::SpiNorFlash& spiNorFlash);
      void Init();
      int OnServiceData(uint16_t connectionHandle, uint16_t attributeHandle, ble_gatt_access_ctxt* context);
      void OnTimeout();
      void Reset();

      class NotificationManager {
      public:
        NotificationManager();
        bool AsyncSend(uint16_t connection, uint16_t charactHandle, uint8_t* data, size_t size);
        void Send(uint16_t connection, uint16_t characteristicHandle, const uint8_t* data, const size_t s);

      private:
        TimerHandle_t timer;
        uint16_t connectionHandle = 0;
        uint16_t characteristicHandle = 0;
        size_t size = 0;
        uint8_t buffer[10];

      public:
        void OnNotificationTimer();
        void Reset();
      };

      class DfuImage {
      public:
        DfuImage(Pinetime::Drivers::SpiNorFlash& spiNorFlash) : spiNorFlash {spiNorFlash} {
        }

        void Init(size_t chunkSize, size_t totalSize, uint16_t expectedCrc);
        void Erase();
        void Append(uint8_t* data, size_t size);
        bool Validate();
        bool IsComplete();

      private:
        Pinetime::Drivers::SpiNorFlash& spiNorFlash;
        static constexpr size_t bufferSize = 200;
        bool ready = false;
        size_t chunkSize = 0;
        size_t totalSize = 0;
        size_t maxSize = 475136;
        size_t bufferWriteIndex = 0;
        size_t totalWriteIndex = 0;
        static constexpr size_t writeOffset = 0x40000;
        uint8_t tempBuffer[bufferSize];
        uint16_t expectedCrc = 0;

        void WriteMagicNumber();
        uint16_t ComputeCrc(uint8_t const* p_data, uint32_t size, uint16_t const* p_crc);
      };

      static constexpr ble_uuid128_t serviceUuid {
        .u {.type = BLE_UUID_TYPE_128},
        .value = {0x23, 0xD1, 0xBC, 0xEA, 0x5F, 0x78, 0x23, 0x15, 0xDE, 0xEF, 0x12, 0x12, 0x30, 0x15, 0x00, 0x00}};

    private:
      Pinetime::System::SystemTask& systemTask;
      Pinetime::Controllers::Ble& bleController;
      DfuImage dfuImage;
      NotificationManager notificationManager;

      static constexpr const char denyAlert[] = "InfiniTime\0Firmware update attempted, but disabled in settings.";
      static constexpr const uint8_t denyAlertLength = sizeof(denyAlert); // for this to work denyAlert MUST be array

      static constexpr uint16_t dfuServiceId {0x1530};
      static constexpr uint16_t packetCharacteristicId {0x1532};
      static constexpr uint16_t controlPointCharacteristicId {0x1531};
      static constexpr uint16_t revisionCharacteristicId {0x1534};

      uint16_t revision {0x0008};

      static constexpr ble_uuid128_t packetCharacteristicUuid {
        .u {.type = BLE_UUID_TYPE_128},
        .value = {0x23, 0xD1, 0xBC, 0xEA, 0x5F, 0x78, 0x23, 0x15, 0xDE, 0xEF, 0x12, 0x12, 0x32, 0x15, 0x00, 0x00}};

      static constexpr ble_uuid128_t controlPointCharacteristicUuid {
        .u {.type = BLE_UUID_TYPE_128},
        .value = {0x23, 0xD1, 0xBC, 0xEA, 0x5F, 0x78, 0x23, 0x15, 0xDE, 0xEF, 0x12, 0x12, 0x31, 0x15, 0x00, 0x00}};

      static constexpr ble_uuid128_t revisionCharacteristicUuid {
        .u {.type = BLE_UUID_TYPE_128},
        .value = {0x23, 0xD1, 0xBC, 0xEA, 0x5F, 0x78, 0x23, 0x15, 0xDE, 0xEF, 0x12, 0x12, 0x34, 0x15, 0x00, 0x00}};

      struct ble_gatt_chr_def characteristicDefinition[4];
      struct ble_gatt_svc_def serviceDefinition[2];
      uint16_t packetCharacteristicHandle;
      uint16_t controlPointCharacteristicHandle;
      uint16_t revisionCharacteristicHandle;

      enum class States : uint8_t { Idle, Init, Start, Data, Validate, Validated };
      States state = States::Idle;

      enum class ImageTypes : uint8_t {
        NoImage = 0x00,
        SoftDevice = 0x01,
        Bootloader = 0x02,
        SoftDeviceAndBootloader = 0x03,
        Application = 0x04
      };

      enum class Opcodes : uint8_t {
        StartDFU = 0x01,
        InitDFUParameters = 0x02,
        ReceiveFirmwareImage = 0x03,
        ValidateFirmware = 0x04,
        ActivateImageAndReset = 0x05,
        PacketReceiptNotificationRequest = 0x08,
        Response = 0x10,
        PacketReceiptNotification = 0x11
      };

      enum class ErrorCodes {
        NoError = 0x01,
        InvalidState = 0x02,
        NotSupported = 0x03,
        DataSizeExceedsLimits = 0x04,
        CrcError = 0x05,
        OperationFailed = 0x06
      };

      uint8_t nbPacketsToNotify = 0;
      uint32_t nbPacketReceived = 0;
      uint32_t bytesReceived = 0;

      uint32_t softdeviceSize = 0;
      uint32_t bootloaderSize = 0;
      uint32_t applicationSize = 0;
      uint16_t expectedCrc = 0;

      int SendDfuRevision(os_mbuf* om) const;
      int WritePacketHandler(uint16_t connectionHandle, os_mbuf* om);
      int ControlPointHandler(uint16_t connectionHandle, os_mbuf* om);

      TimerHandle_t timeoutTimer;
    };
  }
}