#include "custom_cipher.h"
#include "rocksdb/env/composite_env_wrapper.h"
#include "rocksdb/env_encryption.h"
#include "rocksdb/file_system.h"
#include <cstdio>
using ROCKSDB_NAMESPACE::BlockAccessCipherStream;
using ROCKSDB_NAMESPACE::CompositeEnvWrapper;
using ROCKSDB_NAMESPACE::EncryptionProvider;
using ROCKSDB_NAMESPACE::Env;
using ROCKSDB_NAMESPACE::EnvOptions;
using ROCKSDB_NAMESPACE::FileSystem;
using ROCKSDB_NAMESPACE::FileSystemWrapper;
using ROCKSDB_NAMESPACE::IODebugContext;
using ROCKSDB_NAMESPACE::IOOptions;
using ROCKSDB_NAMESPACE::IOStatus;
using ROCKSDB_NAMESPACE::NewCompositeEnv;
using ROCKSDB_NAMESPACE::Slice;
using ROCKSDB_NAMESPACE::Status;
struct CustomCipher {
void *userdata;
size_t MetadataSize, BlockSize;
bool (*EncryptBlock)(void *userdata, uint64_t blockIndex, char *data,
char *metadata);
bool (*DecryptBlock)(void *userdata, uint64_t blockIndex, char *data,
const char *metadata);
void (*Destroy)(void *userdata);
};
class CustomCipherStream : public BlockAccessCipherStream {
public:
CustomCipherStream(const std::string &fname, CustomCipher cipher)
: _cipher(cipher) {
std::string meta = fname + ".meta";
_file = fopen(meta.data(), "r+");
if (!_file) {
_file = fopen(meta.data(), "w+");
}
}
size_t BlockSize() override { return _cipher.BlockSize; }
~CustomCipherStream() override { fclose(_file); }
protected:
void AllocateScratch(std::string &s) override {
s.reserve(_cipher.MetadataSize);
}
Status EncryptBlock(uint64_t blockIndex, char *data, char *scratch) override {
fseeko(_file, blockIndex * _cipher.MetadataSize, SEEK_SET);
size_t read = fread(scratch, 1, _cipher.MetadataSize, _file);
memset(scratch + read, 0, _cipher.MetadataSize - read);
if (!_cipher.EncryptBlock(_cipher.userdata, blockIndex, data, scratch)) {
return Status::Corruption();
}
fseeko(_file, blockIndex * _cipher.MetadataSize, SEEK_SET);
fwrite(scratch, _cipher.MetadataSize, 1, _file);
return Status::OK();
}
Status DecryptBlock(uint64_t blockIndex, char *data, char *scratch) override {
fseeko(_file, blockIndex * _cipher.MetadataSize, SEEK_SET);
size_t read = fread(scratch, 1, _cipher.MetadataSize, _file);
memset(scratch + read, 0, _cipher.MetadataSize - read);
if (!_cipher.DecryptBlock(_cipher.userdata, blockIndex, data, scratch)) {
return Status::Corruption();
}
return Status::OK();
}
private:
FILE *_file;
CustomCipher _cipher;
};
class CustomEncryptionProvider : public EncryptionProvider {
public:
CustomEncryptionProvider(CustomCipher cipher) : _cipher(cipher) {}
~CustomEncryptionProvider() override { _cipher.Destroy(_cipher.userdata); }
const char *Name() const override { return "CustomEncryptionProvider"; }
size_t GetPrefixLength() const override { return 0; }
Status CreateNewPrefix(const std::string &fname, char *prefix,
size_t prefixLength) const override {
return Status::OK();
}
Status AddCipher(const std::string &descriptor, const char *cipher,
size_t len, bool for_write) override {
return Status::OK();
}
Status CreateCipherStream(
const std::string &fname, const EnvOptions &options, Slice &prefix,
std::unique_ptr<BlockAccessCipherStream> *result) override {
(*result) = std::unique_ptr<BlockAccessCipherStream>(
new CustomCipherStream(fname, _cipher));
return Status::OK();
}
private:
CustomCipher _cipher;
};
class CustomFileSystem : public FileSystemWrapper {
public:
explicit CustomFileSystem(const std::shared_ptr<FileSystem> &base)
: FileSystemWrapper(base) {}
const char *Name() const override { return "CustomFileSystem"; }
IOStatus RenameFile(const std::string &src, const std::string &dest,
const IOOptions &options, IODebugContext *dbg) override {
rename((src + ".meta").data(), (dest + ".meta").data());
return FileSystemWrapper::RenameFile(src, dest, options, dbg);
}
};
extern "C" {
struct rocksdb_env_t {
Env *rep;
bool is_default;
};
rocksdb_env_t *rocksdb_create_encrypted_env(
void *userdata, size_t metadataSize, size_t blockSize,
bool (*encryptBlock)(void *userdata, uint64_t blockIndex, char *data,
char *metadata),
bool (*decryptBlock)(void *userdata, uint64_t blockIndex, char *data,
const char *metadata),
void (*Destroy)(void *userdata)) {
CustomCipher cipher = {userdata, metadataSize, blockSize, encryptBlock,
decryptBlock};
std::shared_ptr<EncryptionProvider> provider(
new CustomEncryptionProvider(cipher));
auto my_fs = std::make_shared<CustomFileSystem>(FileSystem::Default());
std::shared_ptr<Env> my_env = NewCompositeEnv(my_fs);
rocksdb_env_t *result = new rocksdb_env_t;
result->rep =
new CompositeEnvWrapper(my_env, NewEncryptedFS(my_fs, provider));
result->is_default = true;
return result;
}
}