Skip to content

Commit

Permalink
Add sm4 encryption (facebook#295)
Browse files Browse the repository at this point in the history 
* Add SM4-CTR encryption algorithm

Signed-off-by: Jarvis Zheng <[email protected]>

* Adjust block size for sm4 encryption

Signed-off-by: Jarvis Zheng <[email protected]>

* Add UT for SM4 encryption

Signed-off-by: Jarvis Zheng <[email protected]>

* Adjust macros indentation for sm4

Signed-off-by: Jarvis Zheng <[email protected]>

* Fix format for adding sm4

Signed-off-by: Jarvis Zheng <[email protected]>
  • Loading branch information
@jiayang-zheng
jiayang-zheng authored Jul 11, 2022
1 parent 4037bda commit 69899a6
Show file tree
Hide file tree
Showing 5 changed files with 80 additions and 14 deletions.
59 changes: 45 additions & 14 deletions encryption/encryption.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -76,8 +76,10 @@ Status AESCTRCipherStream::Cipher(uint64_t file_offset, char* data,
return Status::IOError("Failed to create cipher context.");
}

uint64_t block_index = file_offset / AES_BLOCK_SIZE;
uint64_t block_offset = file_offset % AES_BLOCK_SIZE;
const size_t block_size = BlockSize();

uint64_t block_index = file_offset / block_size;
uint64_t block_offset = file_offset % block_size;

// In CTR mode, OpenSSL EVP API treat the IV as a 128-bit big-endien, and
// increase it by 1 for each block.
Expand All @@ -90,7 +92,7 @@ Status AESCTRCipherStream::Cipher(uint64_t file_offset, char* data,
if (std::numeric_limits<uint64_t>::max() - block_index < initial_iv_low_) {
iv_high++;
}
unsigned char iv[AES_BLOCK_SIZE];
unsigned char iv[block_size];
PutBigEndian64(iv_high, iv);
PutBigEndian64(iv_low, iv + sizeof(uint64_t));

Expand All @@ -105,13 +107,14 @@ Status AESCTRCipherStream::Cipher(uint64_t file_offset, char* data,
// multiply of block size.
ret = EVP_CIPHER_CTX_set_padding(ctx, 0);
if (ret != 1) {
FreeCipherContext(ctx);
return Status::IOError("Failed to disable padding for cipher context.");
}

uint64_t data_offset = 0;
size_t remaining_data_size = data_size;
int output_size = 0;
unsigned char partial_block[AES_BLOCK_SIZE];
unsigned char partial_block[block_size];

// In the following we assume EVP_CipherUpdate allow in and out buffer are
// the same, to save one memcpy. This is not specified in official man page.
Expand All @@ -120,37 +123,41 @@ Status AESCTRCipherStream::Cipher(uint64_t file_offset, char* data,
// buffer to fake a full block.
if (block_offset > 0) {
size_t partial_block_size =
std::min<size_t>(AES_BLOCK_SIZE - block_offset, remaining_data_size);
std::min<size_t>(block_size - block_offset, remaining_data_size);
memcpy(partial_block + block_offset, data, partial_block_size);
ret = EVP_CipherUpdate(ctx, partial_block, &output_size, partial_block,
AES_BLOCK_SIZE);
static_cast<int>(block_size));
if (ret != 1) {
FreeCipherContext(ctx);
return Status::IOError("Crypter failed for first block, offset " +
ToString(file_offset));
}
if (output_size != AES_BLOCK_SIZE) {
if (output_size != static_cast<int>(block_size)) {
FreeCipherContext(ctx);
return Status::IOError(
"Unexpected crypter output size for first block, expected " +
ToString(AES_BLOCK_SIZE) + " vs actual " + ToString(output_size));
ToString(block_size) + " vs actual " + ToString(output_size));
}
memcpy(data, partial_block + block_offset, partial_block_size);
data_offset += partial_block_size;
remaining_data_size -= partial_block_size;
}

// Handle full blocks in the middle.
if (remaining_data_size >= AES_BLOCK_SIZE) {
if (remaining_data_size >= block_size) {
size_t actual_data_size =
remaining_data_size - remaining_data_size % AES_BLOCK_SIZE;
remaining_data_size - remaining_data_size % block_size;
unsigned char* full_blocks =
reinterpret_cast<unsigned char*>(data) + data_offset;
ret = EVP_CipherUpdate(ctx, full_blocks, &output_size, full_blocks,
static_cast<int>(actual_data_size));
if (ret != 1) {
FreeCipherContext(ctx);
return Status::IOError("Crypter failed at offset " +
ToString(file_offset + data_offset));
}
if (output_size != static_cast<int>(actual_data_size)) {
FreeCipherContext(ctx);
return Status::IOError("Unexpected crypter output size, expected " +
ToString(actual_data_size) + " vs actual " +
ToString(output_size));
Expand All @@ -162,21 +169,29 @@ Status AESCTRCipherStream::Cipher(uint64_t file_offset, char* data,
// Handle partial block at the end. The parital block is copied to buffer to
// fake a full block.
if (remaining_data_size > 0) {
assert(remaining_data_size < AES_BLOCK_SIZE);
assert(remaining_data_size < block_size);
memcpy(partial_block, data + data_offset, remaining_data_size);
ret = EVP_CipherUpdate(ctx, partial_block, &output_size, partial_block,
AES_BLOCK_SIZE);
static_cast<int>(block_size));
if (ret != 1) {
FreeCipherContext(ctx);
return Status::IOError("Crypter failed for last block, offset " +
ToString(file_offset + data_offset));
}
if (output_size != AES_BLOCK_SIZE) {
if (output_size != static_cast<int>(block_size)) {
FreeCipherContext(ctx);
return Status::IOError(
"Unexpected crypter output size for last block, expected " +
ToString(AES_BLOCK_SIZE) + " vs actual " + ToString(output_size));
ToString(block_size) + " vs actual " + ToString(output_size));
}
memcpy(data + data_offset, partial_block, remaining_data_size);
}

// Since padding is disabled, and the cipher flow always passes a multiply
// of block size data while each EVP_CipherUpdate, there is no need to call
// EVP_CipherFinal_ex to finish the last block cipher.
// Reference to the implement of EVP_CipherFinal_ex:
// https://github.com/openssl/openssl/blob/OpenSSL_1_1_1-stable/crypto/evp/evp_enc.c#L219
FreeCipherContext(ctx);
return Status::OK();
#endif
Expand All @@ -197,6 +212,16 @@ Status NewAESCTRCipherStream(EncryptionMethod method, const std::string& key,
case EncryptionMethod::kAES256_CTR:
cipher = EVP_aes_256_ctr();
break;
case EncryptionMethod::kSM4_CTR:
#if OPENSSL_VERSION_NUMBER < 0x1010100fL
return Status::InvalidArgument(
"Unsupport SM4 encryption method under OpenSSL version: " +
std::string(OPENSSL_VERSION_TEXT));
#else
// Openssl support SM4 after 1.1.1 release version.
cipher = EVP_sm4_ctr();
break;
#endif
default:
return Status::InvalidArgument("Unsupported encryption method: " +
ToString(static_cast<int>(method)));
Expand Down Expand Up @@ -265,6 +290,7 @@ Status KeyManagedEncryptedEnv::NewSequentialFile(
case EncryptionMethod::kAES128_CTR:
case EncryptionMethod::kAES192_CTR:
case EncryptionMethod::kAES256_CTR:
case EncryptionMethod::kSM4_CTR:
s = encrypted_env_->NewSequentialFile(fname, result, options);
// Hack: when upgrading from TiKV <= v5.0.0-rc, the old current
// file is encrypted but it could be replaced with a plaintext
Expand Down Expand Up @@ -300,6 +326,7 @@ Status KeyManagedEncryptedEnv::NewRandomAccessFile(
case EncryptionMethod::kAES128_CTR:
case EncryptionMethod::kAES192_CTR:
case EncryptionMethod::kAES256_CTR:
case EncryptionMethod::kSM4_CTR:
s = encrypted_env_->NewRandomAccessFile(fname, result, options);
break;
default:
Expand Down Expand Up @@ -332,6 +359,7 @@ Status KeyManagedEncryptedEnv::NewWritableFile(
case EncryptionMethod::kAES128_CTR:
case EncryptionMethod::kAES192_CTR:
case EncryptionMethod::kAES256_CTR:
case EncryptionMethod::kSM4_CTR:
s = encrypted_env_->NewWritableFile(fname, result, options);
break;
default:
Expand Down Expand Up @@ -360,6 +388,7 @@ Status KeyManagedEncryptedEnv::ReopenWritableFile(
case EncryptionMethod::kAES128_CTR:
case EncryptionMethod::kAES192_CTR:
case EncryptionMethod::kAES256_CTR:
case EncryptionMethod::kSM4_CTR:
s = encrypted_env_->ReopenWritableFile(fname, result, options);
break;
default:
Expand Down Expand Up @@ -387,6 +416,7 @@ Status KeyManagedEncryptedEnv::ReuseWritableFile(
case EncryptionMethod::kAES128_CTR:
case EncryptionMethod::kAES192_CTR:
case EncryptionMethod::kAES256_CTR:
case EncryptionMethod::kSM4_CTR:
s = encrypted_env_->ReuseWritableFile(fname, old_fname, result, options);
break;
default:
Expand Down Expand Up @@ -422,6 +452,7 @@ Status KeyManagedEncryptedEnv::NewRandomRWFile(
case EncryptionMethod::kAES128_CTR:
case EncryptionMethod::kAES192_CTR:
case EncryptionMethod::kAES256_CTR:
case EncryptionMethod::kSM4_CTR:
s = encrypted_env_->NewRandomRWFile(fname, result, options);
break;
default:
Expand Down
13 changes: 13 additions & 0 deletions encryption/encryption.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,13 @@ namespace encryption {

#endif

// TODO: OpenSSL Lib does not export SM4_BLOCK_SIZE by now.
// Need to remove SM4_BLOCK_Size once Openssl lib support the definition.
// SM4 uses 128-bit block size as AES.
// Ref:
// https://github.com/openssl/openssl/blob/OpenSSL_1_1_1-stable/include/crypto/sm4.h#L24
#define SM4_BLOCK_SIZE 16

class AESCTRCipherStream : public BlockAccessCipherStream {
public:
AESCTRCipherStream(const EVP_CIPHER* cipher, const std::string& key,
Expand All @@ -51,6 +58,12 @@ class AESCTRCipherStream : public BlockAccessCipherStream {
~AESCTRCipherStream() = default;

size_t BlockSize() override {
// Openssl support SM4 after 1.1.1 release version.
#if OPENSSL_VERSION_NUMBER >= 0x1010100fL
if (EVP_CIPHER_nid(cipher_) == NID_sm4_ctr) {
return SM4_BLOCK_SIZE;
}
#endif
return AES_BLOCK_SIZE; // 16
}

Expand Down
17 changes: 17 additions & 0 deletions encryption/encryption_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,12 @@ class EncryptionTest
case EncryptionMethod::kAES256_CTR:
cipher = EVP_aes_256_ctr();
break;
#if OPENSSL_VERSION_NUMBER >= 0x1010100fL
// Openssl support SM4 after 1.1.1 release version.
case EncryptionMethod::kSM4_CTR:
cipher = EVP_sm4_ctr();
break;
#endif
default:
assert(false);
}
Expand Down Expand Up @@ -146,12 +152,23 @@ TEST_P(EncryptionTest, EncryptionTest) {
EXPECT_TRUE(TestEncryption(16, 16 * 2, IV_OVERFLOW_FULL));
}

// Openssl support SM4 after 1.1.1 release version.
#if OPENSSL_VERSION_NUMBER < 0x1010100fL
INSTANTIATE_TEST_CASE_P(
EncryptionTestInstance, EncryptionTest,
testing::Combine(testing::Bool(),
testing::Values(EncryptionMethod::kAES128_CTR,
EncryptionMethod::kAES192_CTR,
EncryptionMethod::kAES256_CTR)));
#else
INSTANTIATE_TEST_CASE_P(
EncryptionTestInstance, EncryptionTest,
testing::Combine(testing::Bool(),
testing::Values(EncryptionMethod::kAES128_CTR,
EncryptionMethod::kAES192_CTR,
EncryptionMethod::kAES256_CTR,
EncryptionMethod::kSM4_CTR)));
#endif

} // namespace encryption
} // namespace rocksdb
Expand Down
3 changes: 3 additions & 0 deletions include/rocksdb/encryption.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ enum class EncryptionMethod : int {
kAES128_CTR = 2,
kAES192_CTR = 3,
kAES256_CTR = 4,
kSM4_CTR = 5,
};

inline size_t KeySize(EncryptionMethod method) {
Expand All @@ -30,6 +31,8 @@ inline size_t KeySize(EncryptionMethod method) {
return 24;
case EncryptionMethod::kAES256_CTR:
return 32;
case EncryptionMethod::kSM4_CTR:
return 16;
default:
return 0;
};
Expand Down
2 changes: 2 additions & 0 deletions tools/db_bench_tool.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6507,6 +6507,8 @@ int db_bench_tool(int argc, char** argv) {
method = encryption::EncryptionMethod::kAES192_CTR;
} else if (!strcasecmp(FLAGS_encryption_method.c_str(), "AES256CTR")) {
method = encryption::EncryptionMethod::kAES256_CTR;
} else if (!strcasecmp(FLAGS_encryption_method.c_str(), "SM4CTR")) {
method = encryption::EncryptionMethod::kSM4_CTR;
}
if (method == encryption::EncryptionMethod::kUnknown) {
fprintf(stderr, "Unknown encryption method %s\n",
Expand Down

0 comments on commit 69899a6

Please sign in to comment.