support encryption for data store in disk (#247)

librocksdb_sys/crocksdb/c.cc

@@ -15,6 +15,7 @@
 #include "rocksdb/convenience.h"
 #include "rocksdb/db.h"
 #include "rocksdb/env.h"
+#include "rocksdb/env_encryption.h"
 #include "rocksdb/filter_policy.h"
 #include "rocksdb/iterator.h"
 #include "rocksdb/ldb_tool.h"
@@ -31,12 +32,11 @@
 #include "rocksdb/table_properties.h"
 #include "rocksdb/universal_compaction.h"
 #include "rocksdb/utilities/backupable_db.h"
+#include "rocksdb/utilities/db_ttl.h"
 #include "rocksdb/utilities/debug.h"
 #include "rocksdb/utilities/options_util.h"
-#include "rocksdb/utilities/db_ttl.h"
 #include "rocksdb/write_batch.h"
 
-
 #include "db/column_family.h"
 #include "table/sst_file_writer_collectors.h"
 #include "table/table_reader.h"
@@ -71,6 +71,10 @@ using rocksdb::DBWithTTL;
 using rocksdb::DBOptions;
 using rocksdb::Env;
 using rocksdb::EnvOptions;
+using rocksdb::EncryptionProvider;
+using rocksdb::BlockCipher;
+using rocksdb::CTREncryptionProvider;
+using rocksdb::NewEncryptedEnv;
 using rocksdb::ExternalFileIngestionInfo;
 using rocksdb::EventListener;
 using rocksdb::InfoLogLevel;
@@ -460,6 +464,8 @@ struct crocksdb_mergeoperator_t : public MergeOperator {
 struct crocksdb_env_t {
   Env* rep;
   bool is_default;
+  EncryptionProvider* encryption_provoider;
+  BlockCipher* block_cipher;
 };
 
 struct crocksdb_slicetransform_t : public SliceTransform {
@@ -3229,20 +3235,64 @@ void crocksdb_cache_set_capacity(crocksdb_cache_t* cache, size_t capacity) {
   cache->rep->SetCapacity(capacity);
 }
 
-crocksdb_env_t* crocksdb_create_default_env() {
+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->is_default = true;
   return result;
 }
 
-crocksdb_env_t* crocksdb_create_mem_env() {
+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->is_default = false;
   return result;
 }
 
+struct CTRBlockCipher : public BlockCipher {
+  CTRBlockCipher(size_t block_size, const std::string& cipertext)
+      : block_size_(block_size), cipertext_(cipertext) {
+    assert(block_size == cipertext.size());
+  }
+
+  virtual size_t BlockSize() { return block_size_; }
+
+  virtual Status Encrypt(char* data) {
+    const char* ciper_ptr = cipertext_.c_str();
+    for (size_t i = 0; i < block_size_; i++) {
+      data[i] = data[i] ^ ciper_ptr[i];
+    }
+
+    return Status::OK();
+  }
+
+  virtual Status Decrypt(char* data) {
+    Encrypt(data);
+    return Status::OK();
+  }
+
+ protected:
+  std::string cipertext_;
+  size_t block_size_;
+};
+
+crocksdb_env_t* crocksdb_default_ctr_encrypted_env_create(
+    const char* ciphertext, size_t ciphertext_len) {
+  auto result = new crocksdb_env_t;
+  result->block_cipher = new CTRBlockCipher(
+      ciphertext_len, std::string(ciphertext, ciphertext_len));
+  result->encryption_provoider =
+      new CTREncryptionProvider(*result->block_cipher);
+  result->rep = NewEncryptedEnv(Env::Default(), result->encryption_provoider);
+  result->is_default = true;
+
+  return result;
+}
+
 void crocksdb_env_set_background_threads(crocksdb_env_t* env, int n) {
   env->rep->SetBackgroundThreads(n);
 }
@@ -3265,6 +3315,8 @@ 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;
   delete env;
 }
 

librocksdb_sys/crocksdb/crocksdb/c.h

@@ -1305,8 +1305,11 @@ extern C_ROCKSDB_LIBRARY_API void crocksdb_cache_set_capacity(
 
 /* Env */
 
-extern C_ROCKSDB_LIBRARY_API crocksdb_env_t* crocksdb_create_default_env();
-extern C_ROCKSDB_LIBRARY_API crocksdb_env_t* crocksdb_create_mem_env();
+extern C_ROCKSDB_LIBRARY_API crocksdb_env_t* crocksdb_default_env_create();
+extern C_ROCKSDB_LIBRARY_API crocksdb_env_t* crocksdb_mem_env_create();
+extern C_ROCKSDB_LIBRARY_API crocksdb_env_t*
+crocksdb_default_ctr_encrypted_env_create(const char* ciphertext,
+                                          size_t ciphertext_len);
 extern C_ROCKSDB_LIBRARY_API void crocksdb_env_set_background_threads(
     crocksdb_env_t* env, int n);
 extern C_ROCKSDB_LIBRARY_API void

librocksdb_sys/src/lib.rs

@@ -1125,8 +1125,12 @@ extern "C" {
     pub fn crocksdb_compactionfilter_destroy(filter: *mut DBCompactionFilter);
 
     // Env
-    pub fn crocksdb_create_default_env() -> *mut DBEnv;
-    pub fn crocksdb_create_mem_env() -> *mut DBEnv;
+    pub fn crocksdb_default_env_create() -> *mut DBEnv;
+    pub fn crocksdb_mem_env_create() -> *mut DBEnv;
+    pub fn crocksdb_default_ctr_encrypted_env_create(
+        ciphertext: *const c_char,
+        ciphertext_len: size_t,
+    ) -> *mut DBEnv;
     pub fn crocksdb_env_file_exists(env: *mut DBEnv, path: *const c_char, err: *mut *mut c_char);
     pub fn crocksdb_env_delete_file(env: *mut DBEnv, path: *const c_char, err: *mut *mut c_char);
     pub fn crocksdb_env_destroy(env: *mut DBEnv);

src/lib.rs

@@ -64,3 +64,4 @@ mod table_filter;
 mod table_properties;
 mod table_properties_collector;
 mod table_properties_collector_factory;
+mod util;

src/rocksdb.rs

@@ -28,12 +28,14 @@ use std::collections::BTreeMap;
 use std::ffi::{CStr, CString};
 use std::fmt::{self, Debug, Formatter};
 use std::io;
+use std::mem;
 use std::ops::Deref;
 use std::path::{Path, PathBuf};
 use std::str::from_utf8;
 use std::sync::Arc;
 use std::{fs, ptr, slice};
 use table_properties::TablePropertiesCollection;
+use util::is_power_of_two;
 
 pub struct CFHandle {
     inner: *mut DBCFHandle,
@@ -2090,7 +2092,7 @@ impl Default for Env {
     fn default() -> Env {
         unsafe {
             Env {
-                inner: crocksdb_ffi::crocksdb_create_default_env(),
+                inner: crocksdb_ffi::crocksdb_default_env_create(),
             }
         }
     }
@@ -2100,11 +2102,31 @@ impl Env {
     pub fn new_mem() -> Env {
         unsafe {
             Env {
-                inner: crocksdb_ffi::crocksdb_create_mem_env(),
+                inner: crocksdb_ffi::crocksdb_mem_env_create(),
             }
         }
     }
 
+    // Create a ctr encrypted env with the default env and a given ciper text.
+    // The length of ciper text must be 2^n, and must be less or equal to 2048.
+    // The recommanded block size are 1024, 512 and 256.
+    pub fn new_default_ctr_encrypted_env(ciphertext: &[u8]) -> Result<Env, String> {
+        let len = ciphertext.len();
+        if len > 2048 || !is_power_of_two(len) {
+            return Err(
+                "ciphertext length must be less or equal to 2048, and must be power of 2"
+                    .to_owned(),
+            );
+        }
+        let env = unsafe {
+            crocksdb_ffi::crocksdb_default_ctr_encrypted_env_create(
+                mem::transmute(&ciphertext[0]),
+                len,
+            )
+        };
+        Ok(Env { inner: env })
+    }
+
     pub fn new_sequential_file(
         &self,
         path: &str,

src/util.rs

+// Copyright 2018 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+pub fn is_power_of_two(mut n: usize) -> bool {
+    if n == 0 {
+        return false;
+    }
+    while n % 2 == 0 {
+        n = n / 2;
+    }
+    n == 1
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_is_power_of_two() {
+        assert_eq!(is_power_of_two(0), false);
+        assert_eq!(is_power_of_two(1), true);
+        assert_eq!(is_power_of_two(2), true);
+        assert_eq!(is_power_of_two(4), true);
+        assert_eq!(is_power_of_two(8), true);
+        assert_eq!(is_power_of_two(5), false);
+    }
+}

tests/cases/mod.rs

@@ -3,6 +3,7 @@ mod test_compact_range;
 mod test_compaction_filter;
 mod test_delete_files_in_range;
 mod test_delete_range;
+mod test_encryption;
 mod test_event_listener;
 mod test_ingest_external_file;
 mod test_iterator;

tests/cases/test_encryption.rs

+// Copyright 2018 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use rocksdb::Env;
+use rocksdb::{DBOptions, Writable, DB};
+use std::sync::Arc;
+use tempdir::TempDir;
+
+#[test]
+fn test_ctr_encrypted_env() {
+    let test_cipher_texts: &[&[u8]] = &[
+        &[16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1],
+        &[8, 7, 6, 5, 4, 3, 2, 1],
+    ];
+    for ciphertext in test_cipher_texts {
+        test_ctr_encrypted_env_impl(Arc::new(
+            Env::new_default_ctr_encrypted_env(ciphertext).unwrap(),
+        ));
+    }
+}
+
+fn test_ctr_encrypted_env_impl(encrypted_env: Arc<Env>) {
+    let path = TempDir::new("_rust_rocksdb_cryption_env").expect("");
+    let path_str = path.path().to_str().unwrap();
+
+    let mut opts = DBOptions::new();
+    opts.create_if_missing(true);
+    opts.set_env(encrypted_env.clone());
+    let db = DB::open(opts, path_str).unwrap();
+
+    let samples = vec![
+        (b"key1".to_vec(), b"value1".to_vec()),
+        (b"key2".to_vec(), b"value2".to_vec()),
+        (b"key3".to_vec(), b"value3".to_vec()),
+        (b"key4".to_vec(), b"value4".to_vec()),
+    ];
+    for &(ref k, ref v) in &samples {
+        db.put(k, v).unwrap();
+
+        // check value
+        assert_eq!(v.as_slice(), &*db.get(k).unwrap().unwrap());
+    }
+
+    // flush to sst file
+    db.flush(true).unwrap();
+
+    // check value in db
+    for &(ref k, ref v) in &samples {
+        assert_eq!(v.as_slice(), &*db.get(k).unwrap().unwrap());
+    }
+
+    // close db and open again.
+    drop(db);
+    let mut opts = DBOptions::new();
+    opts.create_if_missing(true);
+    opts.set_env(encrypted_env);
+    let db = DB::open(opts, path_str).unwrap();
+
+    // check value in db again
+    for &(ref k, ref v) in &samples {
+        assert_eq!(v.as_slice(), &*db.get(k).unwrap().unwrap());
+    }
+}