Add KeyManagedEncryptedEnv API (#441)

Add `KeyManagedEncryptedEnv` API and related `EncryptionKeyManager` trait.
Signed-off-by: Yi Wu <yiwu@pingcap.com>

.gitmodules

 [submodule "rocksdb"]
 	path = librocksdb_sys/rocksdb
 	url = https://github.com/tikv/rocksdb.git
-	branch = 6.4.tikv
 [submodule "titan"]
 	path = librocksdb_sys/libtitan_sys/titan
 	url = https://github.com/tikv/titan.git

.travis.yml

@@ -18,6 +18,9 @@ matrix:
   - os: linux
     arch: arm64
     rust: stable
+  - os: linux
+    rust: nightly
+    env: FEATURES="encryption,jemalloc,portable,sse"
 
 env:
   global:
@@ -34,5 +37,5 @@ script:
   # copile rocksdb may cost more than 10 minutes, see https://docs.travis-ci.com/user/common-build-problems/#build-times-out-because-no-output-was-received
   - travis_wait make clippy
   - cargo fmt --all -- --check
-  - cargo build
-  - cargo test --all
+  - cargo build --features=$FEATURES
+  - cargo test --all --features=$FEATURES

Cargo.toml

@@ -16,6 +16,7 @@ exclude = [
 
 [features]
 default = []
+encryption = ["librocksdb_sys/encryption"]
 jemalloc = ["librocksdb_sys/jemalloc"]
 portable = ["librocksdb_sys/portable"]
 sse = ["librocksdb_sys/sse"]

librocksdb_sys/Cargo.toml

@@ -10,6 +10,7 @@ bzip2-sys = "0.1.8+1.0.8"
 libc = "0.2.11"
 libtitan_sys = { path = "libtitan_sys" }
 libz-sys = { version = "1.0.25", features = ["static"] }
+openssl-sys = { version = "0.9.54", optional = true, features = ["vendored"] }
 zstd-sys = "1.4.15+zstd.1.4.4"
 
 [dev-dependencies]
@@ -17,6 +18,7 @@ tempfile = "3.1"
 
 [features]
 default = []
+encryption = ["openssl-sys"]
 jemalloc = ["jemalloc-sys"]
 # portable doesn't require static link, though it's meaningless
 # when not using with static-link right now in this crate.
@@ -29,7 +31,7 @@ cmake = "0.1"
 bindgen = "0.51"
 
 [dependencies.jemalloc-sys]
-version = "0.3.2"
+version = "0.1.7"
 optional = true
 features = ["unprefixed_malloc_on_supported_platforms"]
 

librocksdb_sys/build.rs

@@ -125,6 +125,10 @@ fn link_cpp(build: &mut Build) {
 fn build_rocksdb() -> Build {
     let target = env::var("TARGET").expect("TARGET was not set");
     let mut cfg = Config::new("rocksdb");
+    if cfg!(feature = "encryption") {
+        cfg.register_dep("OPENSSL").define("WITH_OPENSSL", "ON");
+        println!("cargo:rustc-link-lib=static=crypto");
+    }
     if cfg!(feature = "jemalloc") && NO_JEMALLOC_TARGETS.iter().all(|i| !target.contains(i)) {
         cfg.register_dep("JEMALLOC").define("WITH_JEMALLOC", "ON");
         println!("cargo:rustc-link-lib=static=jemalloc");
@@ -192,6 +196,9 @@ fn build_rocksdb() -> Build {
     // Adding rocksdb specific compile macros.
     // TODO: should make sure crocksdb compile options is the same as rocksdb and titan.
     build.define("ROCKSDB_SUPPORT_THREAD_LOCAL", None);
+    if cfg!(feature = "encryption") {
+        build.define("OPENSSL", None);
+    }
 
     println!("cargo:rustc-link-lib=static=rocksdb");
     println!("cargo:rustc-link-lib=static=titan");

librocksdb_sys/crocksdb/c.cc

@@ -14,6 +14,7 @@
 #include "rocksdb/comparator.h"
 #include "rocksdb/convenience.h"
 #include "rocksdb/db.h"
+#include "rocksdb/encryption.h"
 #include "rocksdb/env.h"
 #include "rocksdb/env_encryption.h"
 #include "rocksdb/filter_policy.h"
@@ -178,6 +179,13 @@ using rocksdb::titandb::TitanBlobRunMode;
 
 using rocksdb::MemoryAllocator;
 
+#ifdef OPENSSL
+using rocksdb::encryption::EncryptionMethod;
+using rocksdb::encryption::FileEncryptionInfo;
+using rocksdb::encryption::KeyManager;
+using rocksdb::encryption::NewKeyManagedEncryptedEnv;
+#endif
+
 using std::shared_ptr;
 
 extern "C" {
@@ -500,7 +508,7 @@ struct crocksdb_mergeoperator_t : public MergeOperator {
 struct crocksdb_env_t {
   Env* rep;
   bool is_default;
-  EncryptionProvider* encryption_provoider;
+  EncryptionProvider* encryption_provider;
   BlockCipher* block_cipher;
 };
 
@@ -544,6 +552,16 @@ struct crocksdb_universal_compaction_options_t {
   rocksdb::CompactionOptionsUniversal *rep;
 };
 
+#ifdef OPENSSL
+struct crocksdb_file_encryption_info_t {
+  FileEncryptionInfo* rep;
+};
+
+struct crocksdb_encryption_key_manager_t {
+  std::shared_ptr<KeyManager> rep;
+};
+#endif
+
 static bool SaveError(char** errptr, const Status& s) {
   assert(errptr != nullptr);
   if (s.ok()) {
@@ -3510,7 +3528,7 @@ crocksdb_env_t* crocksdb_default_env_create() {
   crocksdb_env_t* result = new crocksdb_env_t;
   result->rep = Env::Default();
   result->block_cipher = nullptr;
-  result->encryption_provoider = nullptr;
+  result->encryption_provider = nullptr;
   result->is_default = true;
   return result;
 }
@@ -3519,7 +3537,7 @@ crocksdb_env_t* crocksdb_mem_env_create() {
   crocksdb_env_t* result = new crocksdb_env_t;
   result->rep = rocksdb::NewMemEnv(Env::Default());
   result->block_cipher = nullptr;
-  result->encryption_provoider = nullptr;
+  result->encryption_provider = nullptr;
   result->is_default = false;
   return result;
 }
@@ -3557,9 +3575,9 @@ crocksdb_ctr_encrypted_env_create(crocksdb_env_t* base_env,
   auto result = new crocksdb_env_t;
   result->block_cipher = new CTRBlockCipher(
       ciphertext_len, std::string(ciphertext, ciphertext_len));
-  result->encryption_provoider =
+  result->encryption_provider =
       new CTREncryptionProvider(*result->block_cipher);
-  result->rep = NewEncryptedEnv(base_env->rep, result->encryption_provoider);
+  result->rep = NewEncryptedEnv(base_env->rep, result->encryption_provider);
   result->is_default = false;
 
   return result;
@@ -3588,7 +3606,7 @@ void crocksdb_env_delete_file(crocksdb_env_t* env, const char* path, char** errp
 void crocksdb_env_destroy(crocksdb_env_t* env) {
   if (!env->is_default) delete env->rep;
   if (env->block_cipher) delete env->block_cipher;
-  if (env->encryption_provoider) delete env->encryption_provoider;
+  if (env->encryption_provider) delete env->encryption_provider;
   delete env;
 }
 
@@ -3630,6 +3648,267 @@ void crocksdb_sequential_file_destroy(crocksdb_sequential_file_t* file) {
   delete file;
 }
 
+#ifdef OPENSSL
+crocksdb_file_encryption_info_t* crocksdb_file_encryption_info_create() {
+  crocksdb_file_encryption_info_t* file_info = new crocksdb_file_encryption_info_t;
+  file_info->rep = new FileEncryptionInfo;
+  return file_info;
+}
+
+void crocksdb_file_encryption_info_destroy(
+    crocksdb_file_encryption_info_t* file_info) {
+  delete file_info->rep;
+  delete file_info;
+}
+
+crocksdb_encryption_method_t crocksdb_file_encryption_info_method(
+    crocksdb_file_encryption_info_t* file_info) {
+  assert(file_info != nullptr);
+  assert(file_info->rep != nullptr);
+  switch (file_info->rep->method) {
+    case EncryptionMethod::kUnknown:
+      return crocksdb_encryption_method_t::kUnknown;
+    case EncryptionMethod::kPlaintext:
+      return crocksdb_encryption_method_t::kPlaintext;
+    case EncryptionMethod::kAES128_CTR:
+      return crocksdb_encryption_method_t::kAES128_CTR;
+    case EncryptionMethod::kAES192_CTR:
+      return crocksdb_encryption_method_t::kAES192_CTR;
+    case EncryptionMethod::kAES256_CTR:
+      return crocksdb_encryption_method_t::kAES256_CTR;
+    default:
+      assert(false);
+  }
+}
+
+const char* crocksdb_file_encryption_info_key(
+    crocksdb_file_encryption_info_t* file_info, size_t* keylen) {
+  assert(file_info != nullptr);
+  assert(file_info->rep != nullptr);
+  assert(keylen != nullptr);
+  *keylen = file_info->rep->key.size();
+  return file_info->rep->key.c_str();
+}
+
+const char* crocksdb_file_encryption_info_iv(
+    crocksdb_file_encryption_info_t* file_info, size_t* ivlen) {
+  assert(file_info != nullptr);
+  assert(file_info->rep != nullptr);
+  assert(ivlen != nullptr);
+  *ivlen = file_info->rep->iv.size();
+  return file_info->rep->iv.c_str();
+}
+
+void crocksdb_file_encryption_info_set_method(
+    crocksdb_file_encryption_info_t* file_info, crocksdb_encryption_method_t method) {
+  assert(file_info != nullptr);
+  switch (method) {
+    case kUnknown:
+      file_info->rep->method = EncryptionMethod::kUnknown;
+      break;
+    case kPlaintext:
+      file_info->rep->method = EncryptionMethod::kPlaintext;
+      break;
+    case kAES128_CTR:
+      file_info->rep->method = EncryptionMethod::kAES128_CTR;
+      break;
+    case kAES192_CTR:
+      file_info->rep->method = EncryptionMethod::kAES192_CTR;
+      break;
+    case kAES256_CTR:
+      file_info->rep->method = EncryptionMethod::kAES256_CTR;
+      break;
+    default:
+      assert(false);
+  };
+}
+
+void crocksdb_file_encryption_info_set_key(
+    crocksdb_file_encryption_info_t* file_info, const char* key, size_t keylen) {
+  assert(file_info != nullptr);
+  file_info->rep->key = std::string(key, keylen);
+}
+
+void crocksdb_file_encryption_info_set_iv(
+    crocksdb_file_encryption_info_t* file_info, const char* iv, size_t ivlen) {
+  assert(file_info != nullptr);
+  file_info->rep->iv = std::string(iv, ivlen);
+}
+
+struct crocksdb_encryption_key_manager_impl_t : public KeyManager {
+  void* state;
+  void (*destructor)(void*);
+  crocksdb_encryption_key_manager_get_file_cb get_file;
+  crocksdb_encryption_key_manager_new_file_cb new_file;
+  crocksdb_encryption_key_manager_delete_file_cb delete_file;
+  crocksdb_encryption_key_manager_link_file_cb link_file;
+  crocksdb_encryption_key_manager_rename_file_cb rename_file;
+
+  virtual ~crocksdb_encryption_key_manager_impl_t() {
+    destructor(state);
+  }
+
+  Status GetFile(
+      const std::string& fname, FileEncryptionInfo* file_info) override {
+    crocksdb_file_encryption_info_t info;
+    info.rep = file_info;
+    const char* ret = get_file(state, fname.c_str(), &info);
+    Status s;
+    if (ret != nullptr) {
+      s = Status::Corruption(std::string(ret));
+      delete ret;
+    }
+    return s;
+  }
+  
+  Status NewFile(
+      const std::string& fname, FileEncryptionInfo* file_info) override {
+    crocksdb_file_encryption_info_t info;
+    info.rep = file_info;
+    const char* ret = new_file(state, fname.c_str(), &info);
+    Status s;
+    if (ret != nullptr) {
+      s = Status::Corruption(std::string(ret));
+      delete ret;
+    }
+    return s;
+  }
+  
+  Status DeleteFile(const std::string& fname) override {
+    const char* ret = delete_file(state, fname.c_str());
+    Status s;
+    if (ret != nullptr) {
+      s = Status::Corruption(std::string(ret));
+      delete ret;
+    }
+    return s;
+  }
+
+  Status LinkFile(
+      const std::string& src_fname, const std::string& dst_fname) override {
+    const char* ret = link_file(state, src_fname.c_str(), dst_fname.c_str());
+    Status s;
+    if (ret != nullptr) {
+      s = Status::Corruption(std::string(ret));
+      delete ret;
+    }
+    return s;
+  }
+
+  Status RenameFile(
+      const std::string& src_fname, const std::string& dst_fname) override {
+    const char* ret = rename_file(state, src_fname.c_str(), dst_fname.c_str());
+    Status s;
+    if (ret != nullptr) {
+      s = Status::Corruption(std::string(ret));
+      delete ret;
+    }
+    return s;
+  }
+};
+
+crocksdb_encryption_key_manager_t* crocksdb_encryption_key_manager_create(
+    void* state, void (*destructor)(void*),
+    crocksdb_encryption_key_manager_get_file_cb get_file,
+    crocksdb_encryption_key_manager_new_file_cb new_file,
+    crocksdb_encryption_key_manager_delete_file_cb delete_file,
+    crocksdb_encryption_key_manager_link_file_cb link_file,
+    crocksdb_encryption_key_manager_rename_file_cb rename_file) {
+  std::shared_ptr<crocksdb_encryption_key_manager_impl_t> key_manager_impl =
+      std::make_shared<crocksdb_encryption_key_manager_impl_t>();
+  key_manager_impl->state = state;
+  key_manager_impl->destructor = destructor;
+  key_manager_impl->get_file = get_file;
+  key_manager_impl->new_file = new_file;
+  key_manager_impl->delete_file = delete_file;
+  key_manager_impl->link_file = link_file;
+  key_manager_impl->rename_file = rename_file;
+  crocksdb_encryption_key_manager_t* key_manager = new crocksdb_encryption_key_manager_t;
+  key_manager->rep = key_manager_impl;
+  return key_manager;
+}
+
+void crocksdb_encryption_key_manager_destroy(crocksdb_encryption_key_manager_t* key_manager) {
+  delete key_manager;
+}
+
+
+const char* crocksdb_encryption_key_manager_get_file(
+    crocksdb_encryption_key_manager_t* key_manager, const char* fname,
+    crocksdb_file_encryption_info_t* file_info) {
+  assert(key_manager != nullptr && key_manager->rep != nullptr);
+  assert(fname != nullptr);
+  assert(file_info != nullptr && file_info->rep != nullptr);
+  Status s = key_manager->rep->GetFile(fname, file_info->rep);
+  if (!s.ok()) {
+    return strdup(s.ToString().c_str());
+  }
+  return nullptr;
+}
+
+const char* crocksdb_encryption_key_manager_new_file(
+    crocksdb_encryption_key_manager_t* key_manager, const char* fname,
+    crocksdb_file_encryption_info_t* file_info) {
+  assert(key_manager != nullptr && key_manager->rep != nullptr);
+  assert(fname != nullptr);
+  assert(file_info != nullptr && file_info->rep != nullptr);
+  Status s = key_manager->rep->NewFile(fname, file_info->rep);
+  if (!s.ok()) {
+    return strdup(s.ToString().c_str());
+  }
+  return nullptr;
+}
+
+const char* crocksdb_encryption_key_manager_delete_file(
+    crocksdb_encryption_key_manager_t* key_manager, const char* fname) {
+  assert(key_manager != nullptr && key_manager->rep != nullptr);
+  assert(fname != nullptr);
+  Status s = key_manager->rep->DeleteFile(fname);
+  if (!s.ok()) {
+    return strdup(s.ToString().c_str());
+  }
+  return nullptr;
+}
+
+const char* crocksdb_encryption_key_manager_link_file(
+    crocksdb_encryption_key_manager_t* key_manager, const char* src_fname,
+    const char* dst_fname) {
+  assert(key_manager != nullptr && key_manager->rep != nullptr);
+  assert(src_fname != nullptr);
+  assert(dst_fname != nullptr);
+  Status s = key_manager->rep->LinkFile(src_fname, dst_fname);
+  if (!s.ok()) {
+    return strdup(s.ToString().c_str());
+  }
+  return nullptr;
+}
+
+const char* crocksdb_encryption_key_manager_rename_file(
+    crocksdb_encryption_key_manager_t* key_manager, const char* src_fname,
+    const char* dst_fname) {
+  assert(key_manager != nullptr && key_manager->rep != nullptr);
+  assert(src_fname != nullptr);
+  assert(dst_fname != nullptr);
+  Status s = key_manager->rep->RenameFile(src_fname, dst_fname);
+  if (!s.ok()) {
+    return strdup(s.ToString().c_str());
+  }
+  return nullptr;
+}
+
+crocksdb_env_t* crocksdb_key_managed_encrypted_env_create(
+    crocksdb_env_t* base_env, crocksdb_encryption_key_manager_t* key_manager) {
+  assert(base_env != nullptr);
+  assert(key_manager != nullptr);
+  crocksdb_env_t* result = new crocksdb_env_t;
+  result->rep = NewKeyManagedEncryptedEnv(base_env->rep, key_manager->rep);
+  result->block_cipher = nullptr;
+  result->encryption_provider = nullptr;
+  result->is_default = false;
+  return result;
+}
+#endif
+
 crocksdb_sstfilereader_t* crocksdb_sstfilereader_create(
     const crocksdb_options_t* io_options) {
   auto reader = new crocksdb_sstfilereader_t;

librocksdb_sys/crocksdb/crocksdb/c.h

@@ -181,6 +181,19 @@ typedef enum crocksdb_backgrounderrorreason_t {
   kMemTable = 4,
 } crocksdb_backgrounderrorreason_t;
 
+#ifdef OPENSSL
+typedef enum crocksdb_encryption_method_t {
+  kUnknown = 0,
+  kPlaintext = 1,
+  kAES128_CTR = 2,
+  kAES192_CTR = 3,
+  kAES256_CTR = 4,
+} crocksdb_encryption_method_t;
+
+typedef struct crocksdb_file_encryption_info_t crocksdb_file_encryption_info_t;
+typedef struct crocksdb_encryption_key_manager_t crocksdb_encryption_key_manager_t;
+#endif
+
 /* DB operations */
 
 extern C_ROCKSDB_LIBRARY_API crocksdb_t* crocksdb_open(
@@ -1439,6 +1452,72 @@ extern C_ROCKSDB_LIBRARY_API void crocksdb_sequential_file_skip(
 extern C_ROCKSDB_LIBRARY_API void crocksdb_sequential_file_destroy(
     crocksdb_sequential_file_t*);
 
+/* KeyManagedEncryptedEnv */
+
+#ifdef OPENSSL
+extern C_ROCKSDB_LIBRARY_API crocksdb_file_encryption_info_t*
+crocksdb_file_encryption_info_create();
+extern C_ROCKSDB_LIBRARY_API void crocksdb_file_encryption_info_destroy(
+    crocksdb_file_encryption_info_t* file_info);
+extern C_ROCKSDB_LIBRARY_API crocksdb_encryption_method_t
+crocksdb_file_encryption_info_method(crocksdb_file_encryption_info_t* file_info);
+extern C_ROCKSDB_LIBRARY_API const char* crocksdb_file_encryption_info_key(
+    crocksdb_file_encryption_info_t* file_info, size_t* keylen);
+extern C_ROCKSDB_LIBRARY_API const char* crocksdb_file_encryption_info_iv(
+    crocksdb_file_encryption_info_t* file_info, size_t* ivlen);
+extern C_ROCKSDB_LIBRARY_API void crocksdb_file_encryption_info_set_method(
+    crocksdb_file_encryption_info_t* file_info, crocksdb_encryption_method_t method);
+extern C_ROCKSDB_LIBRARY_API void crocksdb_file_encryption_info_set_key(
+    crocksdb_file_encryption_info_t* file_info, const char* key, size_t keylen);
+extern C_ROCKSDB_LIBRARY_API void crocksdb_file_encryption_info_set_iv(
+    crocksdb_file_encryption_info_t* file_info, const char* iv, size_t ivlen);
+
+typedef const char* (*crocksdb_encryption_key_manager_get_file_cb)(
+    void* state, const char* fname, crocksdb_file_encryption_info_t* file_info);
+typedef const char* (*crocksdb_encryption_key_manager_new_file_cb)(
+    void* state, const char* fname, crocksdb_file_encryption_info_t* file_info);
+typedef const char* (*crocksdb_encryption_key_manager_delete_file_cb)(
+    void* state, const char* fname);
+typedef const char* (*crocksdb_encryption_key_manager_link_file_cb)(
+    void* state, const char* src_fname, const char* dst_fname);
+typedef const char* (*crocksdb_encryption_key_manager_rename_file_cb)(
+    void* state, const char* src_fname, const char* dst_fname);
+
+extern C_ROCKSDB_LIBRARY_API crocksdb_encryption_key_manager_t*
+crocksdb_encryption_key_manager_create(
+    void* state, void (*destructor)(void*),
+    crocksdb_encryption_key_manager_get_file_cb get_file,
+    crocksdb_encryption_key_manager_new_file_cb new_file,
+    crocksdb_encryption_key_manager_delete_file_cb delete_file,
+    crocksdb_encryption_key_manager_link_file_cb link_file,
+    crocksdb_encryption_key_manager_rename_file_cb rename_file);
+extern C_ROCKSDB_LIBRARY_API void crocksdb_encryption_key_manager_destroy(
+    crocksdb_encryption_key_manager_t*);
+extern C_ROCKSDB_LIBRARY_API const char*
+crocksdb_encryption_key_manager_get_file(
+    crocksdb_encryption_key_manager_t* key_manager, const char* fname,
+    crocksdb_file_encryption_info_t* file_info);
+extern C_ROCKSDB_LIBRARY_API const char*
+crocksdb_encryption_key_manager_new_file(
+    crocksdb_encryption_key_manager_t* key_manager, const char* fname,
+    crocksdb_file_encryption_info_t* file_info);
+extern C_ROCKSDB_LIBRARY_API const char*
+crocksdb_encryption_key_manager_delete_file(
+    crocksdb_encryption_key_manager_t* key_manager, const char* fname);
+extern C_ROCKSDB_LIBRARY_API const char*
+crocksdb_encryption_key_manager_link_file(
+    crocksdb_encryption_key_manager_t* key_manager, const char* src_fname,
+    const char* dst_fname);
+extern C_ROCKSDB_LIBRARY_API const char*
+crocksdb_encryption_key_manager_rename_file(
+    crocksdb_encryption_key_manager_t* key_manager, const char* src_fname,
+    const char* dst_fname);
+
+extern C_ROCKSDB_LIBRARY_API crocksdb_env_t*
+crocksdb_key_managed_encrypted_env_create(
+    crocksdb_env_t*, crocksdb_encryption_key_manager_t*);
+#endif
+
 /* SstFile */
 
 extern C_ROCKSDB_LIBRARY_API crocksdb_sstfilereader_t*

librocksdb_sys/crocksdb/format-diff.sh

@@ -41,7 +41,7 @@
 # access directly without any path.
 if [ -z $CLANG_FORMAT_DIFF ]
 then
-CLANG_FORMAT_DIFF="clang-format-diff.py"
+CLANG_FORMAT_DIFF="clang-format-diff"
 fi
 
 # Check clang-format-diff.py
@@ -99,10 +99,10 @@ LAST_MASTER=`git merge-base master HEAD`
 if [ -z "$uncommitted_code" ]
 then
   # Check the format of last commit
-  diffs=$(git diff -U0 $LAST_MASTER^ | $CLANG_FORMAT_DIFF -p 1)
+  diffs=$(git diff -U0 $LAST_MASTER^ | $CLANG_FORMAT_DIFF -p1)
 else
   # Check the format of uncommitted lines,
-  diffs=$(git diff -U0 HEAD | $CLANG_FORMAT_DIFF -p 1)
+  diffs=$(git diff -U0 HEAD | $CLANG_FORMAT_DIFF -p1)
 fi
 
 if [ -z "$diffs" ]
@@ -141,9 +141,9 @@ fi
 # Do in-place format adjustment.
 if [ -z "$uncommitted_code" ]
 then
-  git diff -U0 $LAST_MASTER^ | $CLANG_FORMAT_DIFF -i -p 1
+  git diff -U0 $LAST_MASTER^ | $CLANG_FORMAT_DIFF -i -p1
 else
-  git diff -U0 HEAD^ | $CLANG_FORMAT_DIFF -i -p 1
+  git diff -U0 HEAD^ | $CLANG_FORMAT_DIFF -i -p1
 fi
 echo "Files reformatted!"
 

librocksdb_sys/src/lib.rs

@@ -154,6 +154,12 @@ pub struct DBWriteStallInfo(c_void);
 pub struct DBStatusPtr(c_void);
 #[repr(C)]
 pub struct DBMapProperty(c_void);
+#[cfg(feature = "encryption")]
+#[repr(C)]
+pub struct DBFileEncryptionInfo(c_void);
+#[cfg(feature = "encryption")]
+#[repr(C)]
+pub struct DBEncryptionKeyManagerInstance(c_void);
 
 #[derive(Copy, Clone, Debug, Eq, PartialEq)]
 #[repr(C)]
@@ -385,6 +391,24 @@ pub enum DBBackgroundErrorReason {
     MemTable = 4,
 }
 
+#[cfg(feature = "encryption")]
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+#[repr(C)]
+pub enum DBEncryptionMethod {
+    Unknown = 0,
+    Plaintext = 1,
+    Aes128Ctr = 2,
+    Aes192Ctr = 3,
+    Aes256Ctr = 4,
+}
+
+#[cfg(feature = "encryption")]
+impl fmt::Display for DBEncryptionMethod {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{:?}", self)
+    }
+}
+
 /// # Safety
 ///
 /// ptr must point to a valid CStr value
@@ -1433,6 +1457,101 @@ extern "C" {
     );
     pub fn crocksdb_ingestexternalfileoptions_destroy(opt: *mut IngestExternalFileOptions);
 
+    // KeyManagedEncryptedEnv
+    #[cfg(feature = "encryption")]
+    pub fn crocksdb_file_encryption_info_create() -> *mut DBFileEncryptionInfo;
+    #[cfg(feature = "encryption")]
+    pub fn crocksdb_file_encryption_info_destroy(file_info: *mut DBFileEncryptionInfo);
+    #[cfg(feature = "encryption")]
+    pub fn crocksdb_file_encryption_info_method(
+        file_info: *mut DBFileEncryptionInfo,
+    ) -> DBEncryptionMethod;
+    #[cfg(feature = "encryption")]
+    pub fn crocksdb_file_encryption_info_key(
+        file_info: *mut DBFileEncryptionInfo,
+        key_len: *mut size_t,
+    ) -> *const c_char;
+    #[cfg(feature = "encryption")]
+    pub fn crocksdb_file_encryption_info_iv(
+        file_info: *mut DBFileEncryptionInfo,
+        iv_len: *mut size_t,
+    ) -> *const c_char;
+    #[cfg(feature = "encryption")]
+    pub fn crocksdb_file_encryption_info_set_method(
+        file_info: *mut DBFileEncryptionInfo,
+        method: DBEncryptionMethod,
+    );
+    #[cfg(feature = "encryption")]
+    pub fn crocksdb_file_encryption_info_set_key(
+        file_info: *mut DBFileEncryptionInfo,
+        key: *const c_char,
+        key_len: size_t,
+    );
+    #[cfg(feature = "encryption")]
+    pub fn crocksdb_file_encryption_info_set_iv(
+        file_info: *mut DBFileEncryptionInfo,
+        iv: *const c_char,
+        iv_len: size_t,
+    );
+
+    #[cfg(feature = "encryption")]
+    pub fn crocksdb_encryption_key_manager_create(
+        state: *mut c_void,
+        destructor: extern "C" fn(*mut c_void),
+        get_file: extern "C" fn(
+            *mut c_void,
+            *const c_char,
+            *mut DBFileEncryptionInfo,
+        ) -> *const c_char,
+        new_file: extern "C" fn(
+            *mut c_void,
+            *const c_char,
+            *mut DBFileEncryptionInfo,
+        ) -> *const c_char,
+        delete_file: extern "C" fn(*mut c_void, *const c_char) -> *const c_char,
+        link_file: extern "C" fn(*mut c_void, *const c_char, *const c_char) -> *const c_char,
+        rename_file: extern "C" fn(*mut c_void, *const c_char, *const c_char) -> *const c_char,
+    ) -> *mut DBEncryptionKeyManagerInstance;
+    #[cfg(feature = "encryption")]
+    pub fn crocksdb_encryption_key_manager_destroy(
+        key_manager: *mut DBEncryptionKeyManagerInstance,
+    );
+    #[cfg(feature = "encryption")]
+    pub fn crocksdb_encryption_key_manager_get_file(
+        key_manager: *mut DBEncryptionKeyManagerInstance,
+        fname: *const c_char,
+        file_info: *mut DBFileEncryptionInfo,
+    ) -> *const c_char;
+    #[cfg(feature = "encryption")]
+    pub fn crocksdb_encryption_key_manager_new_file(
+        key_manager: *mut DBEncryptionKeyManagerInstance,
+        fname: *const c_char,
+        file_info: *mut DBFileEncryptionInfo,
+    ) -> *const c_char;
+    #[cfg(feature = "encryption")]
+    pub fn crocksdb_encryption_key_manager_delete_file(
+        key_manager: *mut DBEncryptionKeyManagerInstance,
+        fname: *const c_char,
+    ) -> *const c_char;
+    #[cfg(feature = "encryption")]
+    pub fn crocksdb_encryption_key_manager_link_file(
+        key_manager: *mut DBEncryptionKeyManagerInstance,
+        src_fname: *const c_char,
+        dst_fname: *const c_char,
+    ) -> *const c_char;
+    #[cfg(feature = "encryption")]
+    pub fn crocksdb_encryption_key_manager_rename_file(
+        key_manager: *mut DBEncryptionKeyManagerInstance,
+        src_fname: *const c_char,
+        dst_fname: *const c_char,
+    ) -> *const c_char;
+
+    #[cfg(feature = "encryption")]
+    pub fn crocksdb_key_managed_encrypted_env_create(
+        base_env: *mut DBEnv,
+        key_manager: *mut DBEncryptionKeyManagerInstance,
+    ) -> *mut DBEnv;
+
     // SstFileReader
     pub fn crocksdb_sstfilereader_create(io_options: *const Options) -> *mut SstFileReader;
 

src/encryption.rs

+// Copyright 2020 TiKV Project Authors. Licensed under Apache-2.0.
+
+pub use crocksdb_ffi::{
+    self, DBEncryptionKeyManagerInstance, DBEncryptionMethod, DBFileEncryptionInfo,
+};
+
+use libc::{c_char, c_void, size_t};
+use std::ffi::{CStr, CString};
+use std::fmt::{self, Debug, Formatter};
+use std::io::Result;
+use std::ptr;
+use std::sync::Arc;
+
+#[derive(Clone, PartialEq, Eq)]
+pub struct FileEncryptionInfo {
+    pub method: DBEncryptionMethod,
+    pub key: Vec<u8>,
+    pub iv: Vec<u8>,
+}
+
+impl Default for FileEncryptionInfo {
+    fn default() -> Self {
+        FileEncryptionInfo {
+            method: DBEncryptionMethod::Unknown,
+            key: vec![],
+            iv: vec![],
+        }
+    }
+}
+
+impl Debug for FileEncryptionInfo {
+    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+        write!(
+            f,
+            "FileEncryptionInfo [method={}, key=...<{} bytes>, iv=...<{} bytes>]",
+            self.method,
+            self.key.len(),
+            self.iv.len()
+        )
+    }
+}
+
+impl FileEncryptionInfo {
+    pub unsafe fn copy_to(&self, file_info: *mut DBFileEncryptionInfo) {
+        crocksdb_ffi::crocksdb_file_encryption_info_set_method(file_info, self.method);
+        crocksdb_ffi::crocksdb_file_encryption_info_set_key(
+            file_info,
+            self.key.as_ptr() as *const c_char,
+            self.key.len() as size_t,
+        );
+        crocksdb_ffi::crocksdb_file_encryption_info_set_iv(
+            file_info,
+            self.iv.as_ptr() as *const c_char,
+            self.iv.len() as size_t,
+        );
+    }
+}
+
+pub trait EncryptionKeyManager: Sync + Send {
+    fn get_file(&self, fname: &str) -> Result<FileEncryptionInfo>;
+    fn new_file(&self, fname: &str) -> Result<FileEncryptionInfo>;
+    fn delete_file(&self, fname: &str) -> Result<()>;
+    fn link_file(&self, src_fname: &str, dst_fname: &str) -> Result<()>;
+    fn rename_file(&self, src_fname: &str, dst_fname: &str) -> Result<()>;
+}
+
+// Copy rust-owned error message to C-owned string. Caller is responsible to delete the result.
+fn copy_error<T: Into<Vec<u8>>>(err: T) -> *const c_char {
+    let cstr = CString::new(err).unwrap();
+    unsafe { libc::strdup(cstr.as_ptr()) }
+}
+
+extern "C" fn encryption_key_manager_destructor(ctx: *mut c_void) {
+    unsafe {
+        // Recover from raw pointer and implicitly drop.
+        Box::from_raw(ctx as *mut Arc<dyn EncryptionKeyManager>);
+    }
+}
+
+extern "C" fn encryption_key_manager_get_file(
+    ctx: *mut c_void,
+    fname: *const c_char,
+    file_info: *mut DBFileEncryptionInfo,
+) -> *const c_char {
+    let key_manager = unsafe { &*(ctx as *mut Arc<dyn EncryptionKeyManager>) };
+    let fname = match unsafe { CStr::from_ptr(fname).to_str() } {
+        Ok(ret) => ret,
+        Err(err) => {
+            return copy_error(format!(
+                "Encryption key manager encounter non-utf8 file name: {}",
+                err
+            ));
+        }
+    };
+    match key_manager.get_file(fname) {
+        Ok(ret) => {
+            unsafe {
+                ret.copy_to(file_info);
+            }
+            ptr::null()
+        }
+        Err(err) => copy_error(format!("Encryption key manager get file failure: {}", err)),
+    }
+}
+
+extern "C" fn encryption_key_manager_new_file(
+    ctx: *mut c_void,
+    fname: *const c_char,
+    file_info: *mut DBFileEncryptionInfo,
+) -> *const c_char {
+    let key_manager = unsafe { &*(ctx as *mut Arc<dyn EncryptionKeyManager>) };
+    let fname = match unsafe { CStr::from_ptr(fname).to_str() } {
+        Ok(ret) => ret,
+        Err(err) => {
+            return copy_error(format!(
+                "Encryption key manager encounter non-utf8 file name: {}",
+                err
+            ));
+        }
+    };
+    match key_manager.new_file(fname) {
+        Ok(ret) => {
+            unsafe {
+                ret.copy_to(file_info);
+            }
+            ptr::null()
+        }
+        Err(err) => copy_error(format!("Encryption key manager new file failure: {}", err)),
+    }
+}
+
+extern "C" fn encryption_key_manager_delete_file(
+    ctx: *mut c_void,
+    fname: *const c_char,
+) -> *const c_char {
+    let key_manager = unsafe { &*(ctx as *mut Arc<dyn EncryptionKeyManager>) };
+    let fname = match unsafe { CStr::from_ptr(fname).to_str() } {
+        Ok(ret) => ret,
+        Err(err) => {
+            return copy_error(format!(
+                "Encryption key manager encounter non-utf8 file name: {}",
+                err
+            ));
+        }
+    };
+    match key_manager.delete_file(fname) {
+        Ok(()) => ptr::null(),
+        Err(err) => copy_error(format!(
+            "Encryption key manager delete file failure: {}",
+            err
+        )),
+    }
+}
+
+extern "C" fn encryption_key_manager_link_file(
+    ctx: *mut c_void,
+    src_fname: *const c_char,
+    dst_fname: *const c_char,
+) -> *const c_char {
+    let key_manager = unsafe { &*(ctx as *mut Arc<dyn EncryptionKeyManager>) };
+    let src_fname = match unsafe { CStr::from_ptr(src_fname).to_str() } {
+        Ok(ret) => ret,
+        Err(err) => {
+            return copy_error(format!(
+                "Encryption key manager encounter non-utf8 file name: {}",
+                err
+            ));
+        }
+    };
+    let dst_fname = match unsafe { CStr::from_ptr(dst_fname).to_str() } {
+        Ok(ret) => ret,
+        Err(err) => {
+            return copy_error(format!(
+                "Encryption key manager encounter non-utf8 file name: {}",
+                err
+            ));
+        }
+    };
+    match key_manager.link_file(src_fname, dst_fname) {
+        Ok(()) => ptr::null(),
+        Err(err) => copy_error(format!(
+            "Encryption key manager delete file failure: {}",
+            err
+        )),
+    }
+}
+
+extern "C" fn encryption_key_manager_rename_file(
+    ctx: *mut c_void,
+    src_fname: *const c_char,
+    dst_fname: *const c_char,
+) -> *const c_char {
+    let key_manager = unsafe { &*(ctx as *mut Arc<dyn EncryptionKeyManager>) };
+    let src_fname = match unsafe { CStr::from_ptr(src_fname).to_str() } {
+        Ok(ret) => ret,
+        Err(err) => {
+            return copy_error(format!(
+                "Encryption key manager encounter non-utf8 file name: {}",
+                err
+            ));
+        }
+    };
+    let dst_fname = match unsafe { CStr::from_ptr(dst_fname).to_str() } {
+        Ok(ret) => ret,
+        Err(err) => {
+            return copy_error(format!(
+                "Encryption key manager encounter non-utf8 file name: {}",
+                err
+            ));
+        }
+    };
+    match key_manager.rename_file(src_fname, dst_fname) {
+        Ok(()) => ptr::null(),
+        Err(err) => copy_error(format!(
+            "Encryption key manager delete file failure: {}",
+            err
+        )),
+    }
+}
+
+pub struct DBEncryptionKeyManager {
+    pub inner: *mut DBEncryptionKeyManagerInstance,
+}
+
+unsafe impl Send for DBEncryptionKeyManager {}
+unsafe impl Sync for DBEncryptionKeyManager {}
+
+impl DBEncryptionKeyManager {
+    pub fn new(key_manager: Arc<dyn EncryptionKeyManager>) -> DBEncryptionKeyManager {
+        // Size of Arc<dyn T>::into_raw is of 128-bits, which couldn't be used as C-style pointer.
+        // Bixing it to make a 64-bits pointer.
+        let ctx = Box::into_raw(Box::new(key_manager)) as *mut c_void;
+        let instance = unsafe {
+            crocksdb_ffi::crocksdb_encryption_key_manager_create(
+                ctx,
+                encryption_key_manager_destructor,
+                encryption_key_manager_get_file,
+                encryption_key_manager_new_file,
+                encryption_key_manager_delete_file,
+                encryption_key_manager_link_file,
+                encryption_key_manager_rename_file,
+            )
+        };
+        DBEncryptionKeyManager { inner: instance }
+    }
+}
+
+impl Drop for DBEncryptionKeyManager {
+    fn drop(&mut self) {
+        unsafe {
+            crocksdb_ffi::crocksdb_encryption_key_manager_destroy(self.inner);
+        }
+    }
+}
+
+// The implementation of EncryptionKeyManager is used to test calling the methods through FFI.
+#[cfg(test)]
+impl EncryptionKeyManager for DBEncryptionKeyManager {
+    fn get_file(&self, fname: &str) -> Result<FileEncryptionInfo> {
+        use std::io::{Error, ErrorKind};
+        use std::mem;
+        use std::slice;
+        let ret: Result<FileEncryptionInfo>;
+        unsafe {
+            let file_info = crocksdb_ffi::crocksdb_file_encryption_info_create();
+            let err = crocksdb_ffi::crocksdb_encryption_key_manager_get_file(
+                self.inner,
+                CString::new(fname).unwrap().as_ptr(),
+                file_info,
+            );
+            if err == ptr::null() {
+                let mut key_len: size_t = 0;
+                let mut iv_len: size_t = 0;
+                let key: *const u8 = mem::transmute(
+                    crocksdb_ffi::crocksdb_file_encryption_info_key(file_info, &mut key_len),
+                );
+                let iv: *const u8 = mem::transmute(crocksdb_ffi::crocksdb_file_encryption_info_iv(
+                    file_info,
+                    &mut iv_len,
+                ));
+                ret = Ok(FileEncryptionInfo {
+                    method: crocksdb_ffi::crocksdb_file_encryption_info_method(file_info),
+                    key: slice::from_raw_parts(key, key_len).to_vec(),
+                    iv: slice::from_raw_parts(iv, iv_len).to_vec(),
+                });
+            } else {
+                ret = Err(Error::new(
+                    ErrorKind::Other,
+                    format!("{}", CStr::from_ptr(err).to_str().unwrap()),
+                ));
+                libc::free(err as _);
+            }
+            crocksdb_ffi::crocksdb_file_encryption_info_destroy(file_info);
+        }
+        ret
+    }
+
+    fn new_file(&self, fname: &str) -> Result<FileEncryptionInfo> {
+        use std::io::{Error, ErrorKind};
+        use std::mem;
+        use std::slice;
+        let ret: Result<FileEncryptionInfo>;
+        unsafe {
+            let file_info = crocksdb_ffi::crocksdb_file_encryption_info_create();
+            let err = crocksdb_ffi::crocksdb_encryption_key_manager_new_file(
+                self.inner,
+                CString::new(fname).unwrap().as_ptr(),
+                file_info,
+            );
+            if err == ptr::null() {
+                let mut key_len: size_t = 0;
+                let mut iv_len: size_t = 0;
+                let key: *const u8 = mem::transmute(
+                    crocksdb_ffi::crocksdb_file_encryption_info_key(file_info, &mut key_len),
+                );
+                let iv: *const u8 = mem::transmute(crocksdb_ffi::crocksdb_file_encryption_info_iv(
+                    file_info,
+                    &mut iv_len,
+                ));
+                ret = Ok(FileEncryptionInfo {
+                    method: crocksdb_ffi::crocksdb_file_encryption_info_method(file_info),
+                    key: slice::from_raw_parts(key, key_len).to_vec(),
+                    iv: slice::from_raw_parts(iv, iv_len).to_vec(),
+                });
+            } else {
+                ret = Err(Error::new(
+                    ErrorKind::Other,
+                    format!("{}", CStr::from_ptr(err).to_str().unwrap()),
+                ));
+                libc::free(err as _);
+            }
+            crocksdb_ffi::crocksdb_file_encryption_info_destroy(file_info);
+        }
+        ret
+    }
+
+    fn delete_file(&self, fname: &str) -> Result<()> {
+        use std::io::{Error, ErrorKind};
+        let ret: Result<()>;
+        unsafe {
+            let err = crocksdb_ffi::crocksdb_encryption_key_manager_delete_file(
+                self.inner,
+                CString::new(fname).unwrap().as_ptr(),
+            );
+            if err == ptr::null() {
+                ret = Ok(());
+            } else {
+                ret = Err(Error::new(
+                    ErrorKind::Other,
+                    format!("{}", CStr::from_ptr(err).to_str().unwrap()),
+                ));
+                libc::free(err as _);
+            }
+        }
+        ret
+    }
+
+    fn link_file(&self, src_fname: &str, dst_fname: &str) -> Result<()> {
+        use std::io::{Error, ErrorKind};
+        let ret: Result<()>;
+        unsafe {
+            let err = crocksdb_ffi::crocksdb_encryption_key_manager_link_file(
+                self.inner,
+                CString::new(src_fname).unwrap().as_ptr(),
+                CString::new(dst_fname).unwrap().as_ptr(),
+            );
+            if err == ptr::null() {
+                ret = Ok(());
+            } else {
+                ret = Err(Error::new(
+                    ErrorKind::Other,
+                    format!("{}", CStr::from_ptr(err).to_str().unwrap()),
+                ));
+                libc::free(err as _);
+            }
+        }
+        ret
+    }
+
+    fn rename_file(&self, src_fname: &str, dst_fname: &str) -> Result<()> {
+        use std::io::{Error, ErrorKind};
+        let ret: Result<()>;
+        unsafe {
+            let err = crocksdb_ffi::crocksdb_encryption_key_manager_rename_file(
+                self.inner,
+                CString::new(src_fname).unwrap().as_ptr(),
+                CString::new(dst_fname).unwrap().as_ptr(),
+            );
+            if err == ptr::null() {
+                ret = Ok(());
+            } else {
+                ret = Err(Error::new(
+                    ErrorKind::Other,
+                    format!("{}", CStr::from_ptr(err).to_str().unwrap()),
+                ));
+                libc::free(err as _);
+            }
+        }
+        ret
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use std::io::{Error, ErrorKind};
+    use std::sync::atomic::{AtomicUsize, Ordering};
+    use std::sync::{Arc, Mutex};
+
+    use super::*;
+
+    struct TestDrop {
+        called: Arc<AtomicUsize>,
+    }
+
+    impl Drop for TestDrop {
+        fn drop(&mut self) {
+            self.called.fetch_add(1, Ordering::SeqCst);
+        }
+    }
+
+    struct TestEncryptionKeyManager {
+        pub get_file_called: AtomicUsize,
+        pub new_file_called: AtomicUsize,
+        pub delete_file_called: AtomicUsize,
+        pub link_file_called: AtomicUsize,
+        pub rename_file_called: AtomicUsize,
+        pub fname: Mutex<String>,
+        pub dst_fname: Mutex<String>,
+        pub return_value: Option<FileEncryptionInfo>,
+        pub drop: Option<TestDrop>,
+    }
+
+    impl Default for TestEncryptionKeyManager {
+        fn default() -> Self {
+            TestEncryptionKeyManager {
+                get_file_called: AtomicUsize::new(0),
+                new_file_called: AtomicUsize::new(0),
+                delete_file_called: AtomicUsize::new(0),
+                link_file_called: AtomicUsize::new(0),
+                rename_file_called: AtomicUsize::new(0),
+                fname: Mutex::new("".to_string()),
+                dst_fname: Mutex::new("".to_string()),
+                return_value: None,
+                drop: None,
+            }
+        }
+    }
+
+    impl EncryptionKeyManager for Mutex<TestEncryptionKeyManager> {
+        fn get_file(&self, fname: &str) -> Result<FileEncryptionInfo> {
+            let key_manager = self.lock().unwrap();
+            key_manager.get_file_called.fetch_add(1, Ordering::SeqCst);
+            key_manager.fname.lock().unwrap().insert_str(0, fname);
+            match &key_manager.return_value {
+                Some(file_info) => Ok(file_info.clone()),
+                None => Err(Error::new(ErrorKind::Other, "")),
+            }
+        }
+
+        fn new_file(&self, fname: &str) -> Result<FileEncryptionInfo> {
+            let key_manager = self.lock().unwrap();
+            key_manager.new_file_called.fetch_add(1, Ordering::SeqCst);
+            key_manager.fname.lock().unwrap().insert_str(0, fname);
+            match &key_manager.return_value {
+                Some(file_info) => Ok(file_info.clone()),
+                None => Err(Error::new(ErrorKind::Other, "")),
+            }
+        }
+
+        fn delete_file(&self, fname: &str) -> Result<()> {
+            let key_manager = self.lock().unwrap();
+            key_manager
+                .delete_file_called
+                .fetch_add(1, Ordering::SeqCst);
+            key_manager.fname.lock().unwrap().insert_str(0, fname);
+            match &key_manager.return_value {
+                Some(_) => Ok(()),
+                None => Err(Error::new(ErrorKind::Other, "")),
+            }
+        }
+
+        fn link_file(&self, src_fname: &str, dst_fname: &str) -> Result<()> {
+            let key_manager = self.lock().unwrap();
+            key_manager.link_file_called.fetch_add(1, Ordering::SeqCst);
+            key_manager.fname.lock().unwrap().insert_str(0, src_fname);
+            key_manager
+                .dst_fname
+                .lock()
+                .unwrap()
+                .insert_str(0, dst_fname);
+            match &key_manager.return_value {
+                Some(_) => Ok(()),
+                None => Err(Error::new(ErrorKind::Other, "")),
+            }
+        }
+
+        fn rename_file(&self, src_fname: &str, dst_fname: &str) -> Result<()> {
+            let key_manager = self.lock().unwrap();
+            key_manager
+                .rename_file_called
+                .fetch_add(1, Ordering::SeqCst);
+            key_manager.fname.lock().unwrap().insert_str(0, src_fname);
+            key_manager
+                .dst_fname
+                .lock()
+                .unwrap()
+                .insert_str(0, dst_fname);
+            match &key_manager.return_value {
+                Some(_) => Ok(()),
+                None => Err(Error::new(ErrorKind::Other, "")),
+            }
+        }
+    }
+
+    #[test]
+    fn create_and_destroy() {
+        let drop_called = Arc::new(AtomicUsize::new(0));
+        let key_manager = Arc::new(Mutex::new(TestEncryptionKeyManager {
+            drop: Some(TestDrop {
+                called: drop_called.clone(),
+            }),
+            ..Default::default()
+        }));
+        let db_key_manager = DBEncryptionKeyManager::new(key_manager.clone());
+        drop(key_manager);
+        assert_eq!(0, drop_called.load(Ordering::SeqCst));
+        drop(db_key_manager);
+        assert_eq!(1, drop_called.load(Ordering::SeqCst));
+    }
+
+    #[test]
+    fn get_file() {
+        let key_manager = Arc::new(Mutex::new(TestEncryptionKeyManager {
+            return_value: Some(FileEncryptionInfo {
+                method: DBEncryptionMethod::Aes128Ctr,
+                key: b"test_key_get_file".to_vec(),
+                iv: b"test_iv_get_file".to_vec(),
+            }),
+            ..Default::default()
+        }));
+        let db_key_manager = DBEncryptionKeyManager::new(key_manager.clone());
+        let file_info = db_key_manager.get_file("get_file_path").unwrap();
+        assert_eq!(DBEncryptionMethod::Aes128Ctr, file_info.method);
+        assert_eq!(b"test_key_get_file", file_info.key.as_slice());
+        assert_eq!(b"test_iv_get_file", file_info.iv.as_slice());
+        let record = key_manager.lock().unwrap();
+        assert_eq!(1, record.get_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.new_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.delete_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.link_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.rename_file_called.load(Ordering::SeqCst));
+        assert_eq!("get_file_path", record.fname.lock().unwrap().as_str());
+    }
+
+    #[test]
+    fn get_file_error() {
+        let key_manager = Arc::new(Mutex::new(TestEncryptionKeyManager::default()));
+        let db_key_manager = DBEncryptionKeyManager::new(key_manager.clone());
+        assert!(db_key_manager.get_file("get_file_path").is_err());
+        let record = key_manager.lock().unwrap();
+        assert_eq!(1, record.get_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.new_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.delete_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.link_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.rename_file_called.load(Ordering::SeqCst));
+        assert_eq!("get_file_path", record.fname.lock().unwrap().as_str());
+    }
+
+    #[test]
+    fn new_file() {
+        let key_manager = Arc::new(Mutex::new(TestEncryptionKeyManager {
+            return_value: Some(FileEncryptionInfo {
+                method: DBEncryptionMethod::Aes256Ctr,
+                key: b"test_key_new_file".to_vec(),
+                iv: b"test_iv_new_file".to_vec(),
+            }),
+            ..Default::default()
+        }));
+        let db_key_manager = DBEncryptionKeyManager::new(key_manager.clone());
+        let file_info = db_key_manager.new_file("new_file_path").unwrap();
+        assert_eq!(DBEncryptionMethod::Aes256Ctr, file_info.method);
+        assert_eq!(b"test_key_new_file", file_info.key.as_slice());
+        assert_eq!(b"test_iv_new_file", file_info.iv.as_slice());
+        let record = key_manager.lock().unwrap();
+        assert_eq!(0, record.get_file_called.load(Ordering::SeqCst));
+        assert_eq!(1, record.new_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.delete_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.link_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.rename_file_called.load(Ordering::SeqCst));
+        assert_eq!("new_file_path", record.fname.lock().unwrap().as_str());
+    }
+
+    #[test]
+    fn new_file_error() {
+        let key_manager = Arc::new(Mutex::new(TestEncryptionKeyManager::default()));
+        let db_key_manager = DBEncryptionKeyManager::new(key_manager.clone());
+        assert!(db_key_manager.new_file("new_file_path").is_err());
+        let record = key_manager.lock().unwrap();
+        assert_eq!(0, record.get_file_called.load(Ordering::SeqCst));
+        assert_eq!(1, record.new_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.delete_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.link_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.rename_file_called.load(Ordering::SeqCst));
+        assert_eq!("new_file_path", record.fname.lock().unwrap().as_str());
+    }
+
+    #[test]
+    fn delete_file() {
+        let key_manager = Arc::new(Mutex::new(TestEncryptionKeyManager {
+            return_value: Some(FileEncryptionInfo::default()),
+            ..Default::default()
+        }));
+        let db_key_manager = DBEncryptionKeyManager::new(key_manager.clone());
+        assert!(db_key_manager.delete_file("delete_file_path").is_ok());
+        let record = key_manager.lock().unwrap();
+        assert_eq!(0, record.get_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.new_file_called.load(Ordering::SeqCst));
+        assert_eq!(1, record.delete_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.link_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.rename_file_called.load(Ordering::SeqCst));
+        assert_eq!("delete_file_path", record.fname.lock().unwrap().as_str());
+    }
+
+    #[test]
+    fn delete_file_error() {
+        let key_manager = Arc::new(Mutex::new(TestEncryptionKeyManager::default()));
+        let db_key_manager = DBEncryptionKeyManager::new(key_manager.clone());
+        assert!(db_key_manager.delete_file("delete_file_path").is_err());
+        let record = key_manager.lock().unwrap();
+        assert_eq!(0, record.get_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.new_file_called.load(Ordering::SeqCst));
+        assert_eq!(1, record.delete_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.link_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.rename_file_called.load(Ordering::SeqCst));
+        assert_eq!("delete_file_path", record.fname.lock().unwrap().as_str());
+    }
+
+    #[test]
+    fn link_file() {
+        let key_manager = Arc::new(Mutex::new(TestEncryptionKeyManager {
+            return_value: Some(FileEncryptionInfo::default()),
+            ..Default::default()
+        }));
+        let db_key_manager = DBEncryptionKeyManager::new(key_manager.clone());
+        assert!(db_key_manager
+            .link_file("src_link_file_path", "dst_link_file_path")
+            .is_ok());
+        let record = key_manager.lock().unwrap();
+        assert_eq!(0, record.get_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.new_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.delete_file_called.load(Ordering::SeqCst));
+        assert_eq!(1, record.link_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.rename_file_called.load(Ordering::SeqCst));
+        assert_eq!("src_link_file_path", record.fname.lock().unwrap().as_str());
+        assert_eq!(
+            "dst_link_file_path",
+            record.dst_fname.lock().unwrap().as_str()
+        );
+    }
+
+    #[test]
+    fn link_file_error() {
+        let key_manager = Arc::new(Mutex::new(TestEncryptionKeyManager::default()));
+        let db_key_manager = DBEncryptionKeyManager::new(key_manager.clone());
+        assert!(db_key_manager
+            .link_file("src_link_file_path", "dst_link_file_path")
+            .is_err());
+        let record = key_manager.lock().unwrap();
+        assert_eq!(0, record.get_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.new_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.delete_file_called.load(Ordering::SeqCst));
+        assert_eq!(1, record.link_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.rename_file_called.load(Ordering::SeqCst));
+        assert_eq!("src_link_file_path", record.fname.lock().unwrap().as_str());
+        assert_eq!(
+            "dst_link_file_path",
+            record.dst_fname.lock().unwrap().as_str()
+        );
+    }
+
+    #[test]
+    fn rename_file() {
+        let key_manager = Arc::new(Mutex::new(TestEncryptionKeyManager {
+            return_value: Some(FileEncryptionInfo::default()),
+            ..Default::default()
+        }));
+        let db_key_manager = DBEncryptionKeyManager::new(key_manager.clone());
+        assert!(db_key_manager
+            .rename_file("src_rename_file_path", "dst_rename_file_path")
+            .is_ok());
+        let record = key_manager.lock().unwrap();
+        assert_eq!(0, record.get_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.new_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.delete_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.link_file_called.load(Ordering::SeqCst));
+        assert_eq!(1, record.rename_file_called.load(Ordering::SeqCst));
+        assert_eq!(
+            "src_rename_file_path",
+            record.fname.lock().unwrap().as_str()
+        );
+        assert_eq!(
+            "dst_rename_file_path",
+            record.dst_fname.lock().unwrap().as_str()
+        );
+    }
+
+    #[test]
+    fn rename_file_error() {
+        let key_manager = Arc::new(Mutex::new(TestEncryptionKeyManager::default()));
+        let db_key_manager = DBEncryptionKeyManager::new(key_manager.clone());
+        assert!(db_key_manager
+            .rename_file("src_rename_file_path", "dst_rename_file_path")
+            .is_err());
+        let record = key_manager.lock().unwrap();
+        assert_eq!(0, record.get_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.new_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.delete_file_called.load(Ordering::SeqCst));
+        assert_eq!(0, record.link_file_called.load(Ordering::SeqCst));
+        assert_eq!(1, record.rename_file_called.load(Ordering::SeqCst));
+        assert_eq!(
+            "src_rename_file_path",
+            record.fname.lock().unwrap().as_str()
+        );
+        assert_eq!(
+            "dst_rename_file_path",
+            record.dst_fname.lock().unwrap().as_str()
+        );
+    }
+}

src/lib.rs

@@ -29,6 +29,8 @@ pub use compaction_filter::{
     CompactionFilter, CompactionFilterContext, CompactionFilterFactory,
     CompactionFilterFactoryHandle, CompactionFilterHandle, DBCompactionFilter,
 };
+#[cfg(feature = "encryption")]
+pub use encryption::{DBEncryptionMethod, EncryptionKeyManager, FileEncryptionInfo};
 pub use event_listener::{
     CompactionJobInfo, EventListener, FlushJobInfo, IngestionInfo, WriteStallInfo,
 };
@@ -67,6 +69,8 @@ pub use rocksdb::Kv;
 
 mod compaction_filter;
 pub mod comparator;
+#[cfg(feature = "encryption")]
+mod encryption;
 mod event_listener;
 pub mod merge_operator;
 mod metadata;

src/rocksdb.rs

@@ -39,6 +39,8 @@ use std::str::from_utf8;
 use std::sync::Arc;
 use std::{fs, ptr, slice};
 
+#[cfg(feature = "encryption")]
+use encryption::{DBEncryptionKeyManager, EncryptionKeyManager};
 use table_properties::{TableProperties, TablePropertiesCollection};
 use table_properties_rc::TablePropertiesCollection as RcTablePropertiesCollection;
 use titan::TitanDBOptions;
@@ -2525,6 +2527,25 @@ impl Env {
         Env::new_ctr_encrypted_env(Arc::new(Env::default()), ciphertext)
     }
 
+    // Create an encrypted env that accepts an external key manager.
+    #[cfg(feature = "encryption")]
+    pub fn new_key_managed_encrypted_env(
+        base_env: Arc<Env>,
+        key_manager: Arc<dyn EncryptionKeyManager>,
+    ) -> Result<Env, String> {
+        let db_key_manager = DBEncryptionKeyManager::new(key_manager);
+        let env = unsafe {
+            crocksdb_ffi::crocksdb_key_managed_encrypted_env_create(
+                base_env.inner,
+                db_key_manager.inner,
+            )
+        };
+        Ok(Env {
+            inner: env,
+            base: Some(base_env),
+        })
+    }
+
     pub fn new_sequential_file(
         &self,
         path: &str,