Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Guard crypto parts in "Update" to save resources #10630

Merged
merged 8 commits into from
Nov 25, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 19 additions & 1 deletion libraries/Update/src/Update.h
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
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
Loading