Fix AlignedBuffer's usage in Encryption Env (#5396)

Summary:
The usage of `AlignedBuffer` in env_encryption.cc writes and reads to/from the AlignedBuffer's internal buffer directly without going through AlignedBuffer's APIs (like `Append` and `Read`), causing encapsulation to break in some cases. The writes are especially problematic as after the data is written to the buffer (directly using either memmove or memcpy), the size of the buffer is not updated ... causing the AlignedBuffer to lose track of the encapsulated buffer's current size.
Fixed this by updating the buffer size after every write.

Todo for later:
Add an overloaded method to AlignedBuffer to support a memmove in addition to a memcopy. Encryption env does a memmove, and hence I couldn't switch to using `AlignedBuffer.Append()`.
Pull Request resolved: https://github.com/facebook/rocksdb/pull/5396

Test Plan: `make check`

Differential Revision: D15764756

Pulled By: sagar0

fbshipit-source-id: 2e24b52bd3b4b5056c5c1da157f91ddf89370183

env/env_encryption.cc

@@ -204,12 +204,15 @@ class EncryptedWritableFile : public WritableFileWrapper {
       // Encrypt in cloned buffer
       buf.Alignment(GetRequiredBufferAlignment());
       buf.AllocateNewBuffer(data.size());
+      // TODO (sagar0): Modify AlignedBuffer.Append to allow doing a memmove
+      // so that the next two lines can be replaced with buf.Append().
       memmove(buf.BufferStart(), data.data(), data.size());
-      status = stream_->Encrypt(offset, buf.BufferStart(), data.size());
+      buf.Size(data.size());
+      status = stream_->Encrypt(offset, buf.BufferStart(), buf.CurrentSize());
       if (!status.ok()) {
         return status;
       }
-      dataToAppend = Slice(buf.BufferStart(), data.size());
+      dataToAppend = Slice(buf.BufferStart(), buf.CurrentSize());
     }
     status = file_->Append(dataToAppend);
     if (!status.ok()) {
@@ -228,11 +231,12 @@ class EncryptedWritableFile : public WritableFileWrapper {
       buf.Alignment(GetRequiredBufferAlignment());
       buf.AllocateNewBuffer(data.size());
       memmove(buf.BufferStart(), data.data(), data.size());
-      status = stream_->Encrypt(offset, buf.BufferStart(), data.size());
+      buf.Size(data.size());
+      status = stream_->Encrypt(offset, buf.BufferStart(), buf.CurrentSize());
       if (!status.ok()) {
         return status;
       }
-      dataToAppend = Slice(buf.BufferStart(), data.size());
+      dataToAppend = Slice(buf.BufferStart(), buf.CurrentSize());
     }
     status = file_->PositionedAppend(dataToAppend, offset);
     if (!status.ok()) {
@@ -332,11 +336,12 @@ class EncryptedRandomRWFile : public RandomRWFile {
       buf.Alignment(GetRequiredBufferAlignment());
       buf.AllocateNewBuffer(data.size());
       memmove(buf.BufferStart(), data.data(), data.size());
-      status = stream_->Encrypt(offset, buf.BufferStart(), data.size());
+      buf.Size(data.size());
+      status = stream_->Encrypt(offset, buf.BufferStart(), buf.CurrentSize());
       if (!status.ok()) {
         return status;
       }
-      dataToWrite = Slice(buf.BufferStart(), data.size());
+      dataToWrite = Slice(buf.BufferStart(), buf.CurrentSize());
     }
     status = file_->Write(offset, dataToWrite);
     return status;
@@ -400,6 +405,7 @@ class EncryptedEnv : public EnvWrapper {
       if (!status.ok()) {
         return status;
       }
+      prefixBuf.Size(prefixLength);
     }
     // Create cipher stream
     std::unique_ptr<BlockAccessCipherStream> stream;
@@ -437,6 +443,7 @@ class EncryptedEnv : public EnvWrapper {
       if (!status.ok()) {
         return status;
       }
+      prefixBuf.Size(prefixLength);
     }
     // Create cipher stream
     std::unique_ptr<BlockAccessCipherStream> stream;
@@ -471,7 +478,8 @@ class EncryptedEnv : public EnvWrapper {
       prefixBuf.Alignment(underlying->GetRequiredBufferAlignment());
       prefixBuf.AllocateNewBuffer(prefixLength);
       provider_->CreateNewPrefix(fname, prefixBuf.BufferStart(), prefixLength);
-      prefixSlice = Slice(prefixBuf.BufferStart(), prefixLength);
+      prefixBuf.Size(prefixLength);
+      prefixSlice = Slice(prefixBuf.BufferStart(), prefixBuf.CurrentSize());
       // Write prefix
       status = underlying->Append(prefixSlice);
       if (!status.ok()) {
@@ -517,7 +525,8 @@ class EncryptedEnv : public EnvWrapper {
       prefixBuf.Alignment(underlying->GetRequiredBufferAlignment());
       prefixBuf.AllocateNewBuffer(prefixLength);
       provider_->CreateNewPrefix(fname, prefixBuf.BufferStart(), prefixLength);
-      prefixSlice = Slice(prefixBuf.BufferStart(), prefixLength);
+      prefixBuf.Size(prefixLength);
+      prefixSlice = Slice(prefixBuf.BufferStart(), prefixBuf.CurrentSize());
       // Write prefix
       status = underlying->Append(prefixSlice);
       if (!status.ok()) {
@@ -558,7 +567,8 @@ class EncryptedEnv : public EnvWrapper {
       prefixBuf.Alignment(underlying->GetRequiredBufferAlignment());
       prefixBuf.AllocateNewBuffer(prefixLength);
       provider_->CreateNewPrefix(fname, prefixBuf.BufferStart(), prefixLength);
-      prefixSlice = Slice(prefixBuf.BufferStart(), prefixLength);
+      prefixBuf.Size(prefixLength);
+      prefixSlice = Slice(prefixBuf.BufferStart(), prefixBuf.CurrentSize());
       // Write prefix
       status = underlying->Append(prefixSlice);
       if (!status.ok()) {
@@ -609,10 +619,12 @@ class EncryptedEnv : public EnvWrapper {
         if (!status.ok()) {
           return status;
         }
+        prefixBuf.Size(prefixLength);
       } else {
         // File is new, initialize & write prefix
         provider_->CreateNewPrefix(fname, prefixBuf.BufferStart(), prefixLength);
-        prefixSlice = Slice(prefixBuf.BufferStart(), prefixLength);
+        prefixBuf.Size(prefixLength);
+        prefixSlice = Slice(prefixBuf.BufferStart(), prefixBuf.CurrentSize());
         // Write prefix
         status = underlying->Write(0, prefixSlice);
         if (!status.ok()) {
@@ -630,7 +642,7 @@ class EncryptedEnv : public EnvWrapper {
     return Status::OK();
   }
 
-    // Store in *result the attributes of the children of the specified directory.
+  // Store in *result the attributes of the children of the specified directory.
   // In case the implementation lists the directory prior to iterating the files
   // and files are concurrently deleted, the deleted files will be omitted from
   // result.
@@ -670,7 +682,6 @@ class EncryptedEnv : public EnvWrapper {
   EncryptionProvider *provider_;
 };
 
-
 // Returns an Env that encrypts data when stored on disk and decrypts data when
 // read from disk.
 Env* NewEncryptedEnv(Env* base_env, EncryptionProvider* provider) {