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/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/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,