Nigel
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!31077 fix crypto bug 

Merge pull request !31077 from yeyunpeng2020/crypto
This commit is contained in:
i-robot 2022-03-10 11:05:57 +00:00 committed by Gitee
parent c331203a2e 9b5c662262
commit c7129ecb6f
No known key found for this signature in database
GPG Key ID: 173E9B9CA92EEF8F
2 changed files with 33 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

44
mindspore/core/utils/crypto.cc
View File

@ -55,7 +55,7 @@ bool IsCipherFile(const std::string &file_path) {
fid.read(int_buf.data(), static_cast<int64_t>(sizeof(int32_t)));
fid.close();
auto flag = ByteToInt(reinterpret_cast<Byte *>(int_buf.data()), int_buf.size());
return static_cast<unsigned int>(flag) == MAGIC_NUM;
return static_cast<unsigned int>(flag) == GCM_MAGIC_NUM || static_cast<unsigned int>(flag) == CBC_MAGIC_NUM;
}
bool IsCipherFile(const Byte *model_data) {
@ -63,7 +63,7 @@ bool IsCipherFile(const Byte *model_data) {
std::vector<Byte> int_buf;
int_buf.assign(model_data, model_data + sizeof(int32_t));
auto flag = ByteToInt(int_buf.data(), int_buf.size());
return static_cast<unsigned int>(flag) == MAGIC_NUM;
return static_cast<unsigned int>(flag) == GCM_MAGIC_NUM || static_cast<unsigned int>(flag) == CBC_MAGIC_NUM;
}
#ifndef ENABLE_OPENSSL
std::unique_ptr<Byte[]> Encrypt(size_t *, const Byte *, size_t, const Byte *, size_t, const std::string &) {
@ -363,7 +363,10 @@ std::unique_ptr<Byte[]> Encrypt(size_t *encrypt_len, const Byte *plain_data, siz
size_t key_len, const std::string &enc_mode) {
MS_EXCEPTION_IF_NULL(plain_data);
MS_EXCEPTION_IF_NULL(key);
if (enc_mode != "AES-GCM" && enc_mode != "AES-CBC") {
MS_LOG(ERROR) << "mode only support AES-GCM|AES-CBC.";
return nullptr;
}
size_t block_enc_buf_len = MAX_BLOCK_SIZE + RESERVED_BYTE_PER_BLOCK;
size_t encrypt_buf_len = plain_len + ((plain_len + MAX_BLOCK_SIZE) / MAX_BLOCK_SIZE) * RESERVED_BYTE_PER_BLOCK;
std::vector<Byte> int_buf(sizeof(int32_t));
@ -383,8 +386,11 @@ std::unique_ptr<Byte[]> Encrypt(size_t *encrypt_len, const Byte *plain_data, siz
MS_LOG(ERROR) << "Failed to encrypt data, please check if enc_key or enc_mode is valid.";
return nullptr;
}
IntToByte(&int_buf, static_cast<int32_t>(MAGIC_NUM));
if (enc_mode == "AES-GCM") {
IntToByte(&int_buf, static_cast<int32_t>(GCM_MAGIC_NUM));
} else {
IntToByte(&int_buf, static_cast<int32_t>(CBC_MAGIC_NUM));
}
size_t capacity = std::min(encrypt_buf_len - *encrypt_len, SECUREC_MEM_MAX_LEN); // avoid dest size over 2gb
auto ret = memcpy_s(encrypt_data.get() + *encrypt_len, capacity, int_buf.data(), sizeof(int32_t));
if (ret != 0) {
@ -415,7 +421,10 @@ std::unique_ptr<Byte[]> Encrypt(size_t *encrypt_len, const Byte *plain_data, siz
std::unique_ptr<Byte[]> Decrypt(size_t *decrypt_len, const std::string &encrypt_data_path, const Byte *key,
size_t key_len, const std::string &dec_mode) {
MS_EXCEPTION_IF_NULL(key);
if (dec_mode != "AES-GCM" && dec_mode != "AES-CBC") {
MS_LOG(ERROR) << "mode only support AES-GCM|AES-CBC.";
return nullptr;
}
std::ifstream fid(encrypt_data_path, std::ios::in | std::ios::binary);
if (!fid) {
MS_LOG(ERROR) << "Open file '" << encrypt_data_path << "' failed, please check the correct of the file.";
@ -435,9 +444,12 @@ std::unique_ptr<Byte[]> Decrypt(size_t *decrypt_len, const std::string &encrypt_
*decrypt_len = 0;
while (static_cast<size_t>(fid.tellg()) < file_size) {
fid.read(int_buf.data(), static_cast<int32_t>(sizeof(int32_t)));
auto cipher_flag = ByteToInt(reinterpret_cast<Byte *>(int_buf.data()), int_buf.size());
if (static_cast<unsigned int>(cipher_flag) != MAGIC_NUM) {
MS_LOG(ERROR) << "File \"" << encrypt_data_path << "\" is not an encrypted file and cannot be decrypted";
auto cipher_flag = static_cast<unsigned int>(ByteToInt(reinterpret_cast<Byte *>(int_buf.data()), int_buf.size()));
if (dec_mode == "AES-GCM" && cipher_flag != GCM_MAGIC_NUM) {
MS_LOG(ERROR) << "File \"" << encrypt_data_path << "\" is not an encrypted AES-GCM file and cannot be decrypted";
return nullptr;
} else if (dec_mode == "AES-CBC" && cipher_flag != CBC_MAGIC_NUM) {
MS_LOG(ERROR) << "File \"" << encrypt_data_path << "\" is not an encrypted AES-CBC file and cannot be decrypted";
return nullptr;
}
@ -473,7 +485,10 @@ std::unique_ptr<Byte[]> Decrypt(size_t *decrypt_len, const Byte *model_data, siz
size_t key_len, const std::string &dec_mode) {
MS_EXCEPTION_IF_NULL(model_data);
MS_EXCEPTION_IF_NULL(key);
if (dec_mode != "AES-GCM" && dec_mode != "AES-CBC") {
MS_LOG(ERROR) << "mode only support AES-GCM|AES-CBC.";
return nullptr;
}
std::vector<char> block_buf;
std::vector<char> int_buf(sizeof(int32_t));
std::vector<Byte> decrypt_block_buf(MAX_BLOCK_SIZE);
@ -489,9 +504,12 @@ std::unique_ptr<Byte[]> Decrypt(size_t *decrypt_len, const Byte *model_data, siz
}
int_buf.assign(model_data + offset, model_data + offset + sizeof(int32_t));
offset += int_buf.size();
auto cipher_flag = ByteToInt(reinterpret_cast<Byte *>(int_buf.data()), int_buf.size());
if (static_cast<unsigned int>(cipher_flag) != MAGIC_NUM) {
MS_LOG(ERROR) << "model_data is not encrypted and therefore cannot be decrypted.";
auto cipher_flag = static_cast<unsigned int>(ByteToInt(reinterpret_cast<Byte *>(int_buf.data()), int_buf.size()));
if (dec_mode == "AES-GCM" && cipher_flag != GCM_MAGIC_NUM) {
MS_LOG(ERROR) << "model_data is not encrypted AES-GCM and therefore cannot be decrypted.";
return nullptr;
} else if (dec_mode == "AES-CBC" && cipher_flag != CBC_MAGIC_NUM) {
MS_LOG(ERROR) << "model_data is not encrypted AES-CBC and therefore cannot be decrypted.";
return nullptr;
}
unsigned char tag[Byte16];

3
mindspore/core/utils/crypto.h
View File

@ -25,7 +25,8 @@ typedef unsigned char Byte;
namespace mindspore {
constexpr size_t MAX_BLOCK_SIZE = 512 * 1024 * 1024; // Maximum ciphertext segment, units is Byte
constexpr size_t RESERVED_BYTE_PER_BLOCK = 50; // Reserved byte per block to save addition info
constexpr unsigned int MAGIC_NUM = 0x7F3A5ED8; // Magic number
constexpr unsigned int GCM_MAGIC_NUM = 0x7F3A5ED8; // Magic number
constexpr unsigned int CBC_MAGIC_NUM = 0x7F3A5ED9; // Magic number
constexpr size_t Byte16 = 16;
MS_CORE_API std::unique_ptr<Byte[]> Encrypt(size_t *encrypt_len, const Byte *plain_data, size_t plain_len,