Skip to content

Commit

Permalink
WriteEncryptionHeader build ok
Browse files Browse the repository at this point in the history
  • Loading branch information
@acelyc111
acelyc111 committed Jul 25, 2023
1 parent ad0c0e7 commit 5d15816
Show file tree
Hide file tree
Showing 3 changed files with 116 additions and 82 deletions.
138 changes: 90 additions & 48 deletions encryption/encryption.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,9 @@

#include "encryption/encryption.h"

#include <openssl/err.h>
#include <openssl/opensslv.h>
#include <openssl/rand.h>

#include <algorithm>
#include <limits>
Expand All @@ -20,6 +22,8 @@ namespace ROCKSDB_NAMESPACE {
namespace encryption {

namespace {
const char* const kEncryptionHeaderMagic = "pegsenc";

uint64_t GetBigEndian64(const unsigned char* buf) {
if (port::kLittleEndian) {
return (static_cast<uint64_t>(buf[0]) << 56) +
Expand Down Expand Up @@ -238,44 +242,11 @@ Status NewAESCTRCipherStream(EncryptionMethod method, const std::string& file_ke
return Status::OK();
}

Status AESEncryptionProvider::AddCipher(const std::string& /*descriptor*/, const char* cipher,
size_t len, bool /*for_write*/) {
if (cipher_) {
return Status::NotSupported("Cannot add keys to AESEncryptionProvider");
} else if (strcmp(encryption::AESBlockCipher::kClassName(), cipher) == 0) {
const EVP_CIPHER* evp_cipher = GetEVPCipher(method_);
if (evp_cipher == nullptr) {
return Status::InvalidArgument("Unsupported encryption method: " +
std::to_string(static_cast<int>(method_)));
}
if (server_key_.size() != KeySize(method_)) {
return Status::InvalidArgument("Encryption key size mismatch. " +
std::to_string(server_key_.size()) + "(actual) vs. " +
std::to_string(KeySize(method_)) + "(expected).");
}
if (server_key_iv_.size() != AES_BLOCK_SIZE) {
return Status::InvalidArgument(
"server key iv size not equal to block cipher block size: " + std::to_string(server_key_iv_.size()) +
"(actual) vs. " + std::to_string(AES_BLOCK_SIZE) + "(expected).");
}
uint64_t iv_high =
GetBigEndian64(reinterpret_cast<const unsigned char*>(server_key_iv_.data()));
uint64_t iv_low = GetBigEndian64(
reinterpret_cast<const unsigned char*>(server_key_iv_.data() + sizeof(uint64_t)));

cipher_.reset(new encryption::AESBlockCipher(evp_cipher, server_key_, iv_high, iv_low));
return Status::OK();
} else {
return BlockCipher::CreateFromString(ConfigOptions(), std::string(cipher),
&cipher_);
}
}

Status AESEncryptionProvider::ReadEncryptionHeader(Slice prefix, FileEncryptionInfo* file_info) {
if (!prefix.starts_with("pegsenc")) {
Status AESEncryptionProvider::ReadEncryptionHeader(Slice prefix, FileEncryptionInfo* file_info) const {
if (!prefix.starts_with(kEncryptionHeaderMagic)) {
return Status::Corruption("Invalid encryption header");
}
EncryptionMethod method = EncryptionMethod(prefix[7]);
auto method = EncryptionMethod(prefix[7]);
size_t key_size = KeySize(method);
if (key_size == 0) {
return Status::Corruption("Unknown encryption algorithm");
Expand All @@ -301,14 +272,93 @@ Status AESEncryptionProvider::ReadEncryptionHeader(Slice prefix, FileEncryptionI
return Status::OK();
}

std::string GetOpenSSLErrors() {
std::ostringstream serr;
uint32_t l;
int line, flags;
const char *file, *data;
bool is_first = true;
while ((l = ERR_get_error_line_data(&file, &line, &data, &flags)) != 0) {
if (is_first) {
is_first = false;
} else {
serr << " ";
}

char buf[256];
ERR_error_string_n(l, buf, sizeof(buf));
serr << buf << ":" << file << ":" << line;
if (flags & ERR_TXT_STRING) {
serr << ":" << data;
}
}
return serr.str();
}

#define OPENSSL_RET_NOT_OK(call, msg) \
if ((call) <= 0) { \
return Status::IOError((msg), GetOpenSSLErrors()); \
}

Status GenerateHeader(EncryptionMethod method,
char* file_key) {
size_t key_size = KeySize(method);
OPENSSL_RET_NOT_OK(RAND_bytes(reinterpret_cast<unsigned char*>(file_key), key_size / 8),
"Failed to generate random key");
return Status::OK();
}

Status AESEncryptionProvider::WriteEncryptionHeader(
char* header_buf) const {
auto method = EncryptionMethod(method_);
size_t key_size = KeySize(method_);
if (key_size == 0) {
return Status::Corruption("Unknown encryption algorithm");
}

AlignedBuffer file_key_buff;
file_key_buff.Alignment(16);
file_key_buff.AllocateNewBuffer(64);

auto append_size = file_key_buff.Append(kEncryptionHeaderMagic, 7);
assert(7 == append_size);
char method_buff[1];
method_buff[0] = static_cast<typename std::underlying_type<EncryptionMethod>::type>(method_);
append_size = file_key_buff.Append(method_buff, 1);
assert(1 == append_size);
char file_key[32] = {0};
Status s = GenerateHeader(method, file_key);
if (!s.ok()) {
return s;
}
s = ROCKSDB_NAMESPACE::encryption::Cipher(method_, server_key_, 0, 0, 0, file_key, 32, true /*is_encrypt*/);
if (!s.ok()) {
return s;
}
append_size = file_key_buff.Append(file_key, 32);
assert(32 == append_size);
file_key_buff.PadToAlignmentWith(0);
assert(64 == file_key_buff.CurrentSize());
// TODO(yingchun): write IV
// TODO(yingchun): reduce memcpy.
memcpy(header_buf, file_key_buff.BufferStart(), file_key_buff.CurrentSize());
return Status::OK();
}

Status AESEncryptionProvider::CreateNewPrefix(
const std::string& /*fname*/, char* prefix, size_t prefix_length) const {
assert(64 == prefix_length);
auto s = WriteEncryptionHeader(prefix);
if (!s.ok()) {
return s;
}
return Status::OK();
}

Status AESEncryptionProvider::CreateCipherStream(
const std::string& fname, const EnvOptions& /*options*/, Slice& prefix,
std::unique_ptr<BlockAccessCipherStream>* result) {
assert(result != nullptr);
if (!cipher_) {
return Status::NotSupported("AddCipher before CreateCipherStream.");
}

FileEncryptionInfo file_info;
Status s = ReadEncryptionHeader(prefix, &file_info);
if (!s.ok()) {
Expand Down Expand Up @@ -441,14 +491,6 @@ static std::unordered_map<std::string, OptionTypeInfo>
OptionVerificationType::kNormal, OptionTypeFlags::kNone}},
};

AESBlockCipher::AESBlockCipher(const EVP_CIPHER* cipher, const std::string& server_key,
uint64_t iv_high, uint64_t iv_low)
: cipher_(cipher), server_key_(server_key),
iv_high_(iv_high), iv_low_(iv_low) {
// RegisterOptions("AESBlockCipherOptions", &blockSize_,
// &aes_block_cipher_type_info);
}

Env* NewKeyManagedEncryptedEnv(
Env* base_env, const std::string& server_key,
const std::string& server_key_iv, EncryptionMethod method) {
Expand Down
42 changes: 10 additions & 32 deletions encryption/encryption.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -101,28 +101,30 @@ class AESEncryptionProvider : public EncryptionProvider {
: server_key_(std::move(server_key)), server_key_iv_(std::move(server_key_iv)), method_(method) {}
virtual ~AESEncryptionProvider() = default;

const char* Name() const override { return "AESEncryptionProvider"; }
static const char* kClassName() { return "AESEncryptionProvider"; }
const char* Name() const override { return kClassName(); }

// TODO(yingchun): the file header size, including padding. Should be a multiple of the page size.
size_t GetPrefixLength() const override { return 0; }
size_t GetPrefixLength() const override { return 64; }

// TODO(yingchun): the FileEncryptionInfo should be stored in the file header.
Status CreateNewPrefix(const std::string& /*fname*/, char* /*prefix*/,
size_t /*prefix_length*/) const override {
return Status::OK();
}
Status CreateNewPrefix(const std::string& /*fname*/, char* prefix,
size_t prefix_length) const override;

// TODO(yingchun): add the Server Key?
Status AddCipher(const std::string& /*descriptor*/, const char* /*cipher*/,
size_t /*len*/, bool /*for_write*/) override;
size_t /*len*/, bool /*for_write*/) override {
return Status::NotSupported();
}

// TODO(yingchun): generate a File Key for the file according to the internal Server Key?
Status CreateCipherStream(
const std::string& fname, const EnvOptions& options, Slice& prefix,
std::unique_ptr<BlockAccessCipherStream>* result) override;

private:
Status ReadEncryptionHeader(Slice prefix, FileEncryptionInfo* file_info);
Status WriteEncryptionHeader(char* header_buf) const;
Status ReadEncryptionHeader(Slice prefix, FileEncryptionInfo* file_info) const;

// TODO(yingchun): remove key_manager_
const std::string server_key_;
Expand All @@ -131,30 +133,6 @@ class AESEncryptionProvider : public EncryptionProvider {
std::shared_ptr<BlockCipher> cipher_;
};

// Implements a BlockCipher using AES.
class AESBlockCipher : public BlockCipher {
public:
explicit AESBlockCipher(const EVP_CIPHER* cipher, const std::string& server_key,
uint64_t iv_high, uint64_t iv_low);

static const char* kClassName() { return "AESBlockCipher"; }
const char* Name() const override { return kClassName(); }

size_t BlockSize() override { return AES_BLOCK_SIZE; }
Status Encrypt(char* data) override {
return Status::NotSupported();
}
Status Decrypt(char* data) override {
return Status::NotSupported();
}

private:
const EVP_CIPHER* cipher_;
const std::string server_key_;
uint64_t iv_high_;
uint64_t iv_low_;
};

} // namespace encryption
} // namespace ROCKSDB_NAMESPACE

Expand Down
18 changes: 16 additions & 2 deletions env/env_encryption.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1147,6 +1147,22 @@ static void RegisterEncryptionBuiltins() {
return guard->get();
});

// Match "AES" or "AES://test"
lib->AddFactory<EncryptionProvider>(
ObjectLibrary::PatternEntry(encryption::AESEncryptionProvider::kClassName(), true)
.AddSuffix("://test"),
[](const std::string& uri, std::unique_ptr<EncryptionProvider>* guard,
std::string* /*errmsg*/) {
if (EndsWith(uri, "://test")) {
guard->reset(new encryption::AESEncryptionProvider("test_serverkey", "test_serverkey_iv",
encryption::EncryptionMethod::kAES128_CTR));
} else {
// TODO(yingchun): deal with the default encryption provider
// guard->reset(new AESEncryptionProvider());
}
return guard->get();
});

lib->AddFactory<EncryptionProvider>(
"1://test", [](const std::string& /*uri*/,
std::unique_ptr<EncryptionProvider>* guard,
Expand All @@ -1173,8 +1189,6 @@ static void RegisterEncryptionBuiltins() {

return guard->get();
});

// TODO(yingchun): add AESBlockCipher
});
}
} // namespace
Expand Down

0 comments on commit 5d15816

Please sign in to comment.