1  // Copyright (C) 2024, Mark Qvist
  2  
  3  // This program is free software: you can redistribute it and/or modify
  4  // it under the terms of the GNU General Public License as published by
  5  // the Free Software Foundation, either version 3 of the License, or
  6  // (at your option) any later version.
  7  
  8  // This program is distributed in the hope that it will be useful,
  9  // but WITHOUT ANY WARRANTY; without even the implied warranty of
 10  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 11  // GNU General Public License for more details.
 12  
 13  // You should have received a copy of the GNU General Public License
 14  // along with this program.  If not, see <https://www.gnu.org/licenses/>.
 15  
 16  #include <Ed25519.h>
 17  
 18  #if MCU_VARIANT == MCU_ESP32
 19  #include "mbedtls/md.h"
 20  #include "esp_ota_ops.h"
 21  #include "esp_flash_partitions.h"
 22  #include "esp_partition.h"
 23  
 24  #elif MCU_VARIANT == MCU_NRF52
 25  #include "Adafruit_nRFCrypto.h"
 26  
 27  // size of chunk to retrieve from flash sector
 28  #define CHUNK_SIZE 128
 29  
 30  #define END_SECTION_SIZE 256
 31  
 32  #if defined(NRF52840_XXAA)
 33  // https://learn.adafruit.com/introducing-the-adafruit-nrf52840-feather/hathach-memory-map
 34  // each section follows along from one another, in this order
 35  // this is always at the start of the memory map
 36  #define APPLICATION_START 0x26000
 37  
 38  #define USER_DATA_START 0xED000
 39  
 40  #define IMG_SIZE_START 0xFF008
 41  #endif
 42  
 43  #endif
 44  
 45  // Forward declaration from Utilities.h
 46  void eeprom_update(int mapped_addr, uint8_t byte);
 47  uint8_t eeprom_read(uint32_t addr);
 48  void hard_reset(void);
 49  
 50  #if !HAS_EEPROM && MCU_VARIANT == MCU_NRF52
 51    void eeprom_flush();
 52  #endif
 53  
 54  const uint8_t dev_keys [] PROGMEM = {
 55     0x0f, 0x15, 0x86, 0x74, 0xa0, 0x7d, 0xf2, 0xde, 0x32, 0x11, 0x29, 0xc1, 0x0d, 0xda, 0xcc, 0xc3,
 56     0xe1, 0x9b, 0xac, 0xf2, 0x27, 0x06, 0xee, 0x89, 0x1f, 0x7a, 0xfc, 0xc3, 0x6a, 0xf5, 0x38, 0x08
 57  };
 58  
 59  #define DEV_SIG_LEN 64
 60  uint8_t dev_sig[DEV_SIG_LEN];
 61  
 62  #define DEV_KEY_LEN 32
 63  uint8_t dev_k_prv[DEV_KEY_LEN];
 64  uint8_t dev_k_pub[DEV_KEY_LEN];
 65  
 66  #define DEV_HASH_LEN 32
 67  uint8_t dev_hash[DEV_HASH_LEN];
 68  uint8_t dev_partition_table_hash[DEV_HASH_LEN];
 69  uint8_t dev_bootloader_hash[DEV_HASH_LEN];
 70  uint8_t dev_firmware_hash[DEV_HASH_LEN];
 71  uint8_t dev_firmware_hash_target[DEV_HASH_LEN];
 72  
 73  #define EEPROM_SIG_LEN 128
 74  uint8_t dev_eeprom_signature[EEPROM_SIG_LEN];
 75  
 76  bool dev_signature_validated = false;
 77  bool fw_signature_validated = true;
 78  
 79  #define DEV_SIG_OFFSET EEPROM_SIZE-EEPROM_RESERVED-DEV_SIG_LEN
 80  #define dev_sig_addr(a) (a+DEV_SIG_OFFSET)
 81  
 82  #define DEV_FWHASH_OFFSET EEPROM_SIZE-EEPROM_RESERVED-DEV_SIG_LEN-DEV_HASH_LEN
 83  #define dev_fwhash_addr(a) (a+DEV_FWHASH_OFFSET)
 84  
 85  bool device_signatures_ok() {
 86    return dev_signature_validated && fw_signature_validated;
 87  }
 88  
 89  void device_validate_signature() {
 90    int n_keys = sizeof(dev_keys)/DEV_KEY_LEN;
 91    bool valid_signature_found = false;
 92    for (int i = 0; i < n_keys; i++) {
 93      memcpy(dev_k_pub, dev_keys+DEV_KEY_LEN*i, DEV_KEY_LEN);
 94      if (Ed25519::verify(dev_sig, dev_k_pub, dev_hash, DEV_HASH_LEN)) {
 95          valid_signature_found = true;
 96      }
 97    }
 98  
 99    if (valid_signature_found) {
100      dev_signature_validated = true;
101    } else {
102      dev_signature_validated = false;
103    }
104  }
105  
106  void device_save_signature() {
107    device_validate_signature();
108    if (dev_signature_validated) {
109      for (uint8_t i = 0; i < DEV_SIG_LEN; i++) {
110        eeprom_update(dev_sig_addr(i), dev_sig[i]);
111      }
112    }
113  }
114  
115  void device_load_signature() {
116    for (uint8_t i = 0; i < DEV_SIG_LEN; i++) {
117      #if HAS_EEPROM
118          dev_sig[i] = EEPROM.read(dev_sig_addr(i));
119      #elif MCU_VARIANT == MCU_NRF52
120          dev_sig[i] = eeprom_read(dev_sig_addr(i));
121      #endif
122    }
123  }
124  
125  void device_load_firmware_hash() {
126    for (uint8_t i = 0; i < DEV_HASH_LEN; i++) {
127      #if HAS_EEPROM
128          dev_firmware_hash_target[i] = EEPROM.read(dev_fwhash_addr(i));
129      #elif MCU_VARIANT == MCU_NRF52
130          dev_firmware_hash_target[i] = eeprom_read(dev_fwhash_addr(i));
131      #endif
132    }
133  }
134  
135  void device_save_firmware_hash() {
136    for (uint8_t i = 0; i < DEV_HASH_LEN; i++) {
137      eeprom_update(dev_fwhash_addr(i), dev_firmware_hash_target[i]);
138    }
139    #if !HAS_EEPROM && MCU_VARIANT == MCU_NRF52
140      eeprom_flush();
141    #endif
142    if (!fw_signature_validated) hard_reset();
143  }
144  
145  #if MCU_VARIANT == MCU_NRF52
146  uint32_t retrieve_application_size() {
147      uint8_t bytes[4];
148      memcpy(bytes, (const void*)IMG_SIZE_START, 4);
149      uint32_t fw_len = bytes[0] | bytes[1] << 8 | bytes[2] << 16 | bytes[3] << 24;
150      return fw_len;
151  }
152  
153  void calculate_region_hash(unsigned long long start, unsigned long long end, uint8_t* return_hash) {
154      // this function calculates the hash digest of a region of memory,
155      // currently it is only designed to work for the application region
156      uint8_t chunk[CHUNK_SIZE] = {0};
157  
158      // to store potential last chunk of program
159      uint8_t chunk_next[CHUNK_SIZE] = {0};
160      nRFCrypto_Hash hash;
161  
162      hash.begin(CRYS_HASH_SHA256_mode);
163  
164      uint8_t size;
165  
166      while (start < end ) {
167          const void* src = (const void*)start;
168          if (start + CHUNK_SIZE >= end) {
169              size = end - start;
170          }
171          else {
172              size = CHUNK_SIZE;
173          }
174  
175          memcpy(chunk, src, CHUNK_SIZE);
176  
177          hash.update(chunk, size);
178  
179          start += CHUNK_SIZE;
180      }
181      hash.end(return_hash);
182  }
183  #endif
184  
185  void device_validate_partitions() {
186    device_load_firmware_hash();
187    #if MCU_VARIANT == MCU_ESP32
188    esp_partition_t partition;
189    partition.address   = ESP_PARTITION_TABLE_OFFSET;
190    partition.size      = ESP_PARTITION_TABLE_MAX_LEN;
191    partition.type      = ESP_PARTITION_TYPE_DATA;
192    esp_partition_get_sha256(&partition, dev_partition_table_hash);
193    partition.address   = ESP_BOOTLOADER_OFFSET;
194    partition.size      = ESP_PARTITION_TABLE_OFFSET;
195    partition.type      = ESP_PARTITION_TYPE_APP;
196    esp_partition_get_sha256(&partition, dev_bootloader_hash);
197    esp_partition_get_sha256(esp_ota_get_running_partition(), dev_firmware_hash);
198    #elif MCU_VARIANT == MCU_NRF52
199    // todo, add bootloader, partition table, or softdevice?
200    calculate_region_hash(APPLICATION_START, APPLICATION_START+retrieve_application_size(), dev_firmware_hash);
201    #endif
202    #if VALIDATE_FIRMWARE
203      for (uint8_t i = 0; i < DEV_HASH_LEN; i++) {
204        if (dev_firmware_hash_target[i] != dev_firmware_hash[i]) {
205          fw_signature_validated = false;
206          break;
207        }
208      }
209    #endif
210  }
211  
212  bool device_firmware_ok() {
213    return fw_signature_validated;
214  }
215  
216  #if MCU_VARIANT == MCU_ESP32 || MCU_VARIANT == MCU_NRF52
217  bool device_init() {
218    if (bt_ready) {
219      #if MCU_VARIANT == MCU_ESP32
220      for (uint8_t i=0; i<EEPROM_SIG_LEN; i++){dev_eeprom_signature[i]=EEPROM.read(eeprom_addr(ADDR_SIGNATURE+i));}
221      mbedtls_md_context_t ctx;
222      mbedtls_md_type_t md_type = MBEDTLS_MD_SHA256;     
223      mbedtls_md_init(&ctx);
224      mbedtls_md_setup(&ctx, mbedtls_md_info_from_type(md_type), 0);
225      mbedtls_md_starts(&ctx);
226      #if HAS_BLUETOOTH == true || HAS_BLE == true
227        mbedtls_md_update(&ctx, dev_bt_mac, BT_DEV_ADDR_LEN);
228      #else
229        // TODO: Get from BLE stack instead
230        // mbedtls_md_update(&ctx, dev_bt_mac, BT_DEV_ADDR_LEN);
231      #endif
232      mbedtls_md_update(&ctx, dev_eeprom_signature, EEPROM_SIG_LEN);
233      mbedtls_md_finish(&ctx, dev_hash);
234      mbedtls_md_free(&ctx);
235      #elif MCU_VARIANT == MCU_NRF52
236      for (uint8_t i=0; i<EEPROM_SIG_LEN; i++){dev_eeprom_signature[i]=eeprom_read(eeprom_addr(ADDR_SIGNATURE+i));}
237      nRFCrypto.begin();
238  
239      nRFCrypto_Hash hash;
240  
241      hash.begin(CRYS_HASH_SHA256_mode);
242  
243      #if HAS_BLUETOOTH == true || HAS_BLE == true
244        hash.update(dev_bt_mac, BT_DEV_ADDR_LEN);
245      #else
246        // TODO: Get from BLE stack instead
247        // hash.update(dev_bt_mac, BT_DEV_ADDR_LEN);
248      #endif
249      hash.update(dev_eeprom_signature, EEPROM_SIG_LEN);
250  
251      hash.end(dev_hash);
252      #endif
253      device_load_signature();
254      device_validate_signature();
255  
256      device_validate_partitions();
257  
258      #if MCU_VARIANT == MCU_NRF52
259      nRFCrypto.end();
260      #endif
261      device_init_done = true;
262      return device_init_done && fw_signature_validated;
263    } else {
264      return false;
265    }
266  }
267  #endif