Skip to content

Commit

Permalink
Guard crypto parts in "Update" to save resources (#10630)
Browse files Browse the repository at this point in the history 
* guard crypt update

* guard update crypt

* Update Updater.cpp

* revert logic to disable

* change disable logic

* formatting

* formatting

* remove trailing space
  • Loading branch information
@Jason2866
Jason2866 authored Nov 25, 2024
1 parent 1730e4e commit eb1933f
Show file tree
Hide file tree
Showing 2 changed files with 59 additions and 4 deletions.
20 changes: 19 additions & 1 deletion libraries/Update/src/Update.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@ class UpdateClass {
*/
bool begin(size_t size = UPDATE_SIZE_UNKNOWN, int command = U_FLASH, int ledPin = -1, uint8_t ledOn = LOW, const char *label = NULL);

#ifndef UPDATE_NOCRYPT
/*
Setup decryption configuration
Crypt Key is 32bytes(256bits) block of data, use the same key as used to encrypt image file
Expand All @@ -71,6 +72,7 @@ class UpdateClass {
Crypt Mode, used to select if image files should be decrypted or not
*/
bool setupCrypt(const uint8_t *cryptKey = 0, size_t cryptAddress = 0, uint8_t cryptConfig = 0xf, int cryptMode = U_AES_DECRYPT_AUTO);
#endif /* UPDATE_NOCRYPT */

/*
Writes a buffer to the flash and increments the address
Expand Down Expand Up @@ -99,6 +101,7 @@ class UpdateClass {
*/
bool end(bool evenIfRemaining = false);

#ifndef UPDATE_NOCRYPT
/*
sets AES256 key(32 bytes) used for decrypting image file
*/
Expand All @@ -122,6 +125,7 @@ class UpdateClass {
void setCryptConfig(const uint8_t cryptConfig) {
_cryptCfg = cryptConfig & 0x0f;
}
#endif /* UPDATE_NOCRYPT */

/*
Aborts the running update
Expand All @@ -139,7 +143,13 @@ class UpdateClass {
sets the expected MD5 for the firmware (hexString)
If calc_post_decryption is true, the update library will calculate the MD5 after the decryption, if false the calculation occurs before the decryption
*/
bool setMD5(const char *expected_md5, bool calc_post_decryption = true);
bool setMD5(
const char *expected_md5
#ifndef UPDATE_NOCRYPT
,
bool calc_post_decryption = true
#endif /* #ifdef UPDATE_NOCRYPT */
);

/*
returns the MD5 String of the successfully ended firmware
Expand Down Expand Up @@ -236,17 +246,21 @@ class UpdateClass {
private:
void _reset();
void _abort(uint8_t err);
#ifndef UPDATE_NOCRYPT
void _cryptKeyTweak(size_t cryptAddress, uint8_t *tweaked_key);
bool _decryptBuffer();
#endif /* UPDATE_NOCRYPT */
bool _writeBuffer();
bool _verifyHeader(uint8_t data);
bool _verifyEnd();
bool _enablePartition(const esp_partition_t *partition);
bool _chkDataInBlock(const uint8_t *data, size_t len) const; // check if block contains any data or is empty

uint8_t _error;
#ifndef UPDATE_NOCRYPT
uint8_t *_cryptKey;
uint8_t *_cryptBuffer;
#endif /* UPDATE_NOCRYPT */
uint8_t *_buffer;
uint8_t *_skipBuffer;
size_t _bufferLen;
Expand All @@ -258,15 +272,19 @@ class UpdateClass {
const esp_partition_t *_partition;

String _target_md5;
#ifndef UPDATE_NOCRYPT
bool _target_md5_decrypted = true;
#endif /* UPDATE_NOCRYPT */
MD5Builder _md5;

int _ledPin;
uint8_t _ledOn;

#ifndef UPDATE_NOCRYPT
uint8_t _cryptMode;
size_t _cryptAddress;
uint8_t _cryptCfg;
#endif /* UPDATE_NOCRYPT */
};

#if !defined(NO_GLOBAL_INSTANCES) && !defined(NO_GLOBAL_UPDATE)
Expand Down
43 changes: 40 additions & 3 deletions libraries/Update/src/Updater.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +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) {
Expand Down Expand Up @@ -38,8 +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");
}
Expand Down Expand Up @@ -67,8 +71,17 @@ bool UpdateClass::_enablePartition(const esp_partition_t *partition) {
}

UpdateClass::UpdateClass()
: _error(0), _cryptKey(0), _cryptBuffer(0), _buffer(0), _skipBuffer(0), _bufferLen(0), _size(0), _progress_callback(NULL), _progress(0), _paroffset(0),
_command(U_FLASH), _partition(NULL), _cryptMode(U_AES_DECRYPT_AUTO), _cryptAddress(0), _cryptCfg(0xf) {}
: _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;
Expand All @@ -83,7 +96,9 @@ void UpdateClass::_reset() {
delete[] _skipBuffer;
}

#ifndef UPDATE_NOCRYPT
_cryptBuffer = nullptr;
#endif /* UPDATE_NOCRYPT */
_buffer = nullptr;
_skipBuffer = nullptr;
_bufferLen = 0;
Expand Down Expand Up @@ -175,6 +190,7 @@ 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)) {
Expand Down Expand Up @@ -216,6 +232,7 @@ bool UpdateClass::setCryptMode(const int cryptMode) {
}
return true;
}
#endif /* UPDATE_NOCRYPT */

void UpdateClass::_abort(uint8_t err) {
_reset();
Expand All @@ -226,6 +243,7 @@ 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) {
Expand Down Expand Up @@ -338,8 +356,10 @@ bool UpdateClass::_decryptBuffer() {
}
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;
Expand All @@ -360,6 +380,7 @@ bool UpdateClass::_writeBuffer() {
return false;
}
}
#endif /* UPDATE_NOCRYPT */
//first bytes of new firmware
uint8_t skip = 0;
if (!_progress && _command == U_FLASH) {
Expand Down Expand Up @@ -409,9 +430,13 @@ bool UpdateClass::_writeBuffer() {
if (!_progress && _command == U_FLASH) {
_buffer[0] = ESP_IMAGE_HEADER_MAGIC;
}
#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) {
Expand Down Expand Up @@ -453,13 +478,21 @@ bool UpdateClass::_verifyEnd() {
return false;
}

bool UpdateClass::setMD5(const char *expected_md5, bool calc_post_decryption) {
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;
}

Expand Down Expand Up @@ -532,12 +565,16 @@ size_t UpdateClass::writeStream(Stream &data) {
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);
Expand Down

0 comments on commit eb1933f

Please sign in to comment.