Update Updater.cpp

Author: Jason2866

Diff for: libraries/Update/src/Updater.cpp

@@ -9,6 +9,9 @@
 #include "spi_flash_mmap.h"
 #include "esp_ota_ops.h"
 #include "esp_image_format.h"
+#ifndef UPDATE_NOCRYPT
+#include "mbedtls/aes.h"
+#endif /* UPDATE_NOCRYPT */
 
 static const char *_err2str(uint8_t _error) {
   if (_error == UPDATE_ERROR_OK) {
@@ -37,6 +40,10 @@ static const char *_err2str(uint8_t _error) {
     return ("Bad Argument");
   } else if (_error == UPDATE_ERROR_ABORT) {
     return ("Aborted");
+#ifndef UPDATE_NOCRYPT
+  } else if (_error == UPDATE_ERROR_DECRYPT) {
+    return ("Decryption error");
+#endif /* UPDATE_NOCRYPT */
   }
   return ("UNKNOWN");
 }
@@ -64,7 +71,15 @@ bool UpdateClass::_enablePartition(const esp_partition_t *partition) {
 }
 
 UpdateClass::UpdateClass()
-  : _error(0), _buffer(0), _bufferLen(0), _size(0), _progress_callback(NULL), _progress(0), _paroffset(0), _command(U_FLASH), _partition(NULL) {}
+  : _error(0),
+#ifndef UPDATE_NOCRYPT
+_cryptKey(0), _cryptBuffer(0),
+#endif /* UPDATE_NOCRYPT */
+_buffer(0), _skipBuffer(0), _bufferLen(0), _size(0), _progress_callback(NULL), _progress(0), _paroffset(0), _command(U_FLASH), _partition(NULL)
+#ifndef UPDATE_NOCRYPT
+, _cryptMode(U_AES_DECRYPT_AUTO), _cryptAddress(0), _cryptCfg(0xf)
+#endif /* UPDATE_NOCRYPT */
+{}
 
 UpdateClass &UpdateClass::onProgress(THandlerFunction_Progress fn) {
   _progress_callback = fn;
@@ -79,6 +94,9 @@ void UpdateClass::_reset() {
     delete[] _skipBuffer;
   }
 
+#ifndef UPDATE_NOCRYPT
+  _cryptBuffer = nullptr;
+#endif /* UPDATE_NOCRYPT */
   _buffer = nullptr;
   _skipBuffer = nullptr;
   _bufferLen = 0;
@@ -170,6 +188,50 @@ bool UpdateClass::begin(size_t size, int command, int ledPin, uint8_t ledOn, con
   return true;
 }
 
+#ifndef UPDATE_NOCRYPT
+bool UpdateClass::setupCrypt(const uint8_t *cryptKey, size_t cryptAddress, uint8_t cryptConfig, int cryptMode) {
+  if (setCryptKey(cryptKey)) {
+    if (setCryptMode(cryptMode)) {
+      setCryptAddress(cryptAddress);
+      setCryptConfig(cryptConfig);
+      return true;
+    }
+  }
+  return false;
+}
+
+bool UpdateClass::setCryptKey(const uint8_t *cryptKey) {
+  if (!cryptKey) {
+    if (_cryptKey) {
+      delete[] _cryptKey;
+      _cryptKey = 0;
+      log_d("AES key unset");
+    }
+    return false;  //key cleared, no key to decrypt with
+  }
+  //initialize
+  if (!_cryptKey) {
+    _cryptKey = new (std::nothrow) uint8_t[ENCRYPTED_KEY_SIZE];
+  }
+  if (!_cryptKey) {
+    log_e("new failed");
+    return false;
+  }
+  memcpy(_cryptKey, cryptKey, ENCRYPTED_KEY_SIZE);
+  return true;
+}
+
+bool UpdateClass::setCryptMode(const int cryptMode) {
+  if (cryptMode >= U_AES_DECRYPT_NONE && cryptMode <= U_AES_DECRYPT_ON) {
+    _cryptMode = cryptMode;
+  } else {
+    log_e("bad crypt mode argument %i", cryptMode);
+    return false;
+  }
+  return true;
+}
+#endif /* UPDATE_NOCRYPT */
+
 void UpdateClass::_abort(uint8_t err) {
   _reset();
   _error = err;
@@ -179,7 +241,144 @@ void UpdateClass::abort() {
   _abort(UPDATE_ERROR_ABORT);
 }
 
+#ifndef UPDATE_NOCRYPT
+void UpdateClass::_cryptKeyTweak(size_t cryptAddress, uint8_t *tweaked_key) {
+  memcpy(tweaked_key, _cryptKey, ENCRYPTED_KEY_SIZE);
+  if (_cryptCfg == 0) {
+    return;  //no tweaking needed, use crypt key as-is
+  }
+
+  const uint8_t pattern[] = {23, 23, 23, 14, 23, 23, 23, 12, 23, 23, 23, 10, 23, 23, 23, 8};
+  int pattern_idx = 0;
+  int key_idx = 0;
+  int bit_len = 0;
+  uint32_t tweak = 0;
+  cryptAddress &= 0x00ffffe0;  //bit 23-5
+  cryptAddress <<= 8;          //bit23 shifted to bit31(MSB)
+  while (pattern_idx < sizeof(pattern)) {
+    tweak = cryptAddress << (23 - pattern[pattern_idx]);  //bit shift for small patterns
+    // alternative to: tweak = rotl32(tweak,8 - bit_len);
+    tweak = (tweak << (8 - bit_len)) | (tweak >> (24 + bit_len));  //rotate to line up with end of previous tweak bits
+    bit_len += pattern[pattern_idx++] - 4;                         //add number of bits in next pattern(23-4 = 19bits = 23bit to 5bit)
+    while (bit_len > 7) {
+      tweaked_key[key_idx++] ^= tweak;  //XOR byte
+      // alternative to: tweak = rotl32(tweak, 8);
+      tweak = (tweak << 8) | (tweak >> 24);  //compiler should optimize to use rotate(fast)
+      bit_len -= 8;
+    }
+    tweaked_key[key_idx] ^= tweak;  //XOR remaining bits, will XOR zeros if no remaining bits
+  }
+  if (_cryptCfg == 0xf) {
+    return;  //return with fully tweaked key
+  }
+
+  //some of tweaked key bits need to be restore back to crypt key bits
+  const uint8_t cfg_bits[] = {67, 65, 63, 61};
+  key_idx = 0;
+  pattern_idx = 0;
+  while (key_idx < ENCRYPTED_KEY_SIZE) {
+    bit_len += cfg_bits[pattern_idx];
+    if ((_cryptCfg & (1 << pattern_idx)) == 0) {  //restore crypt key bits
+      while (bit_len > 0) {
+        if (bit_len > 7 || ((_cryptCfg & (2 << pattern_idx)) == 0)) {  //restore a crypt key byte
+          tweaked_key[key_idx] = _cryptKey[key_idx];
+        } else {  //MSBits restore crypt key bits, LSBits keep as tweaked bits
+          tweaked_key[key_idx] &= (0xff >> bit_len);
+          tweaked_key[key_idx] |= (_cryptKey[key_idx] & (~(0xff >> bit_len)));
+        }
+        key_idx++;
+        bit_len -= 8;
+      }
+    } else {  //keep tweaked key bits
+      while (bit_len > 0) {
+        if (bit_len < 8 && ((_cryptCfg & (2 << pattern_idx)) == 0)) {  //MSBits keep as tweaked bits, LSBits restore crypt key bits
+          tweaked_key[key_idx] &= (~(0xff >> bit_len));
+          tweaked_key[key_idx] |= (_cryptKey[key_idx] & (0xff >> bit_len));
+        }
+        key_idx++;
+        bit_len -= 8;
+      }
+    }
+    pattern_idx++;
+  }
+}
+
+bool UpdateClass::_decryptBuffer() {
+  if (!_cryptKey) {
+    log_w("AES key not set");
+    return false;
+  }
+  if (_bufferLen % ENCRYPTED_BLOCK_SIZE != 0) {
+    log_e("buffer size error");
+    return false;
+  }
+  if (!_cryptBuffer) {
+    _cryptBuffer = new (std::nothrow) uint8_t[ENCRYPTED_BLOCK_SIZE];
+  }
+  if (!_cryptBuffer) {
+    log_e("new failed");
+    return false;
+  }
+  uint8_t tweaked_key[ENCRYPTED_KEY_SIZE];  //tweaked crypt key
+  int done = 0;
+
+  /*
+        Mbedtls functions will be replaced with esp_aes functions when hardware acceleration is available
+
+        To Do:
+        Replace mbedtls for the cases where there's no hardware acceleration
+     */
+
+  mbedtls_aes_context ctx;  //initialize AES
+  mbedtls_aes_init(&ctx);
+  while ((_bufferLen - done) >= ENCRYPTED_BLOCK_SIZE) {
+    for (int i = 0; i < ENCRYPTED_BLOCK_SIZE; i++) {
+      _cryptBuffer[(ENCRYPTED_BLOCK_SIZE - 1) - i] = _buffer[i + done];  //reverse order 16 bytes to decrypt
+    }
+    if (((_cryptAddress + _progress + done) % ENCRYPTED_TWEAK_BLOCK_SIZE) == 0 || done == 0) {
+      _cryptKeyTweak(_cryptAddress + _progress + done, tweaked_key);  //update tweaked crypt key
+      if (mbedtls_aes_setkey_enc(&ctx, tweaked_key, 256)) {
+        return false;
+      }
+      if (mbedtls_aes_setkey_dec(&ctx, tweaked_key, 256)) {
+        return false;
+      }
+    }
+    if (mbedtls_aes_crypt_ecb(&ctx, MBEDTLS_AES_ENCRYPT, _cryptBuffer, _cryptBuffer)) {  //use MBEDTLS_AES_ENCRYPT to decrypt flash code
+      return false;
+    }
+    for (int i = 0; i < ENCRYPTED_BLOCK_SIZE; i++) {
+      _buffer[i + done] = _cryptBuffer[(ENCRYPTED_BLOCK_SIZE - 1) - i];  //reverse order 16 bytes from decrypt
+    }
+    done += ENCRYPTED_BLOCK_SIZE;
+  }
+  return true;
+}
+#endif /* UPDATE_NOCRYPT */
+
 bool UpdateClass::_writeBuffer() {
+#ifndef UPDATE_NOCRYPT
+  //first bytes of loading image, check to see if loading image needs decrypting
+  if (!_progress) {
+    _cryptMode &= U_AES_DECRYPT_MODE_MASK;
+    if ((_cryptMode == U_AES_DECRYPT_ON) || ((_command == U_FLASH) && (_cryptMode & U_AES_DECRYPT_AUTO) && (_buffer[0] != ESP_IMAGE_HEADER_MAGIC))) {
+      _cryptMode |= U_AES_IMAGE_DECRYPTING_BIT;  //set to decrypt the loading image
+      log_d("Decrypting OTA Image");
+    }
+  }
+
+  if (!_target_md5_decrypted) {
+    _md5.add(_buffer, _bufferLen);
+  }
+
+  //check if data in buffer needs decrypting
+  if (_cryptMode & U_AES_IMAGE_DECRYPTING_BIT) {
+    if (!_decryptBuffer()) {
+      _abort(UPDATE_ERROR_DECRYPT);
+      return false;
+    }
+  }
+  #endif /* UPDATE_NOCRYPT */
   //first bytes of new firmware
   uint8_t skip = 0;
   if (!_progress && _command == U_FLASH) {
@@ -229,7 +428,13 @@ bool UpdateClass::_writeBuffer() {
   if (!_progress && _command == U_FLASH) {
     _buffer[0] = ESP_IMAGE_HEADER_MAGIC;
   }
-  _md5.add(_buffer, _bufferLen);
+#ifndef UPDATE_NOCRYPT
+  if (_target_md5_decrypted) {
+#endif /* UPDATE_NOCRYPT */
+    _md5.add(_buffer, _bufferLen);
+#ifndef UPDATE_NOCRYPT
+  }
+#endif /* UPDATE_NOCRYPT */
   _progress += _bufferLen;
   _bufferLen = 0;
   if (_progress_callback) {
@@ -271,12 +476,19 @@ bool UpdateClass::_verifyEnd() {
   return false;
 }
 
-bool UpdateClass::setMD5(const char *expected_md5) {
+bool UpdateClass::setMD5(const char *expected_md5
+#ifndef UPDATE_NOCRYPT
+,bool calc_post_decryption
+#endif /* UPDATE_NOCRYPT */
+) {
   if (strlen(expected_md5) != 32) {
     return false;
   }
   _target_md5 = expected_md5;
   _target_md5.toLowerCase();
+#ifndef UPDATE_NOCRYPT
+  _target_md5_decrypted = calc_post_decryption;
+#endif /* UPDATE_NOCRYPT */
   return true;
 }
 
@@ -349,10 +561,16 @@ size_t UpdateClass::writeStream(Stream &data) {
     return 0;
   }
 
-  if (!_verifyHeader(data.peek())) {
-    _reset();
-    return 0;
+#ifndef UPDATE_NOCRYPT
+  if (_command == U_FLASH && !_cryptMode) {
+#endif /* UPDATE_NOCRYPT */
+    if (!_verifyHeader(data.peek())) {
+      _reset();
+      return 0;
+    }
+#ifndef UPDATE_NOCRYPT
   }
+#endif /* UPDATE_NOCRYPT */
 
   if (_ledPin != -1) {
     pinMode(_ledPin, OUTPUT);