From 2b2b3b8dfc33c637b3f6c1de4eead2f994204eca Mon Sep 17 00:00:00 2001
From: Yingchun Lai <laiyingchun@apache.org>
Date: Wed, 9 Aug 2023 11:04:21 +0800
Subject: [PATCH] fix renaming encrypted directory (#13)

https://github.com/apache/incubator-pegasus/issues/1575

Cherry-pick from
https://github.com/tikv/rocksdb/commit/14f36f8ccb7973a2a675fc54b02d130021563b97
(without compaction related code)

* fix renaming encrypted directory

Signed-off-by: tabokie <xy.tao@outlook.com>
---
 encryption/encryption.cc     | 11 +++++------
 include/rocksdb/encryption.h |  8 ++++++++
 test_util/testutil.h         | 12 ++++++++++++
 3 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/encryption/encryption.cc b/encryption/encryption.cc
index 0bf7cd41103..23d203a0c06 100644
--- a/encryption/encryption.cc
+++ b/encryption/encryption.cc
@@ -530,23 +530,22 @@ Status KeyManagedEncryptedEnv::RenameFile(const std::string& src_fname,
   }
   s = target()->RenameFile(src_fname, dst_fname);
   if (s.ok()) {
-    s = key_manager_->DeleteFile(src_fname);
+    s = key_manager_->DeleteFileExt(src_fname, dst_fname);
   } else {
     Status delete_status __attribute__((__unused__)) =
-        key_manager_->DeleteFile(dst_fname);
+        key_manager_->DeleteFileExt(dst_fname, src_fname);
     assert(delete_status.ok());
   }
   return s;
 }
 
 Status KeyManagedEncryptedEnv::DeleteDir(const std::string& dname) {
-  Status s = target()->DeleteDir(dname);
+  // We don't guarantee atomicity. Delete keys first.
+  Status s = key_manager_->DeleteFile(dname);
   if (!s.ok()) {
     return s;
   }
-  // We don't use a dedicated `DeleteDir` function, because RocksDB already uses
-  // `RenameFile` for both file and directory.
-  return key_manager_->DeleteFile(dname);
+  return target()->DeleteDir(dname);
 }
 
 Env* NewKeyManagedEncryptedEnv(Env* base_env,
diff --git a/include/rocksdb/encryption.h b/include/rocksdb/encryption.h
index 8bfac6ff134..c9c54760e25 100644
--- a/include/rocksdb/encryption.h
+++ b/include/rocksdb/encryption.h
@@ -55,9 +55,17 @@ class KeyManager {
                          FileEncryptionInfo* file_info) = 0;
   virtual Status NewFile(const std::string& fname,
                          FileEncryptionInfo* file_info) = 0;
+  // Used with both file and directory.
   virtual Status DeleteFile(const std::string& fname) = 0;
   virtual Status LinkFile(const std::string& src_fname,
                           const std::string& dst_fname) = 0;
+  // Provide additional hint of physical file when the key name doesn't map to
+  // one. A typical use case of this is atomically deleting a directory by
+  // renaming it first.
+  virtual Status DeleteFileExt(const std::string& fname,
+                               const std::string& /*physical_fname*/) {
+    return DeleteFile(fname);
+  }
 };
 
 // An Env with underlying files being encrypted. It holds a reference to an
diff --git a/test_util/testutil.h b/test_util/testutil.h
index 84531a4ad2d..279830d84e0 100644
--- a/test_util/testutil.h
+++ b/test_util/testutil.h
@@ -86,6 +86,18 @@ class TestKeyManager : public encryption::KeyManager {
   Status DeleteFile(const std::string& fname) override {
     std::lock_guard<std::mutex> l(mutex);
     file_set.erase(fname);
+    if (!fname.empty()) {
+      std::string copy = fname;
+      if (copy.back() != '/') {
+        copy.push_back('/');
+      }
+      auto begin = file_set.lower_bound(copy);
+      auto end = begin;
+      while (end != file_set.end() && end->compare(0, copy.size(), copy) == 0) {
+        end++;
+      }
+      file_set.erase(begin, end);
+    }
     return Status::OK();
   }