librbd/crypto: fix memory leak in when DataCryptor fails

If DataCryptor fails, either in init_context or update_context,
the encryption context is not returned, which causes a memory leak.
This commit fixes this issue.

Signed-off-by: Or Ozeri <oro@il.ibm.com>

diff --git a/src/librbd/crypto/BlockCrypto.cc b/src/librbd/crypto/BlockCrypto.cc
index 6e6a9efeeba28a43e4eb26876ba1bb83cb556283..f9e2a9d21903ae3798f83cb22ca4f1cdd53bc94e 100644 (file)
--- a/src/librbd/crypto/BlockCrypto.cc
+++ b/src/librbd/crypto/BlockCrypto.cc
@@ -4,6 +4,7 @@
 #include "librbd/crypto/BlockCrypto.h"
 #include "include/byteorder.h"
 #include "include/ceph_assert.h"
+#include "include/scope_guard.h"
 
 #include <stdlib.h>
 
@@ -53,6 +54,10 @@ int BlockCrypto<T>::crypt(ceph::bufferlist* data, uint64_t image_offset,
     lderr(m_cct) << "unable to get crypt context" << dendl;
     return -EIO;
   }
+
+  auto sg = make_scope_guard([&] {
+      m_data_cryptor->return_context(ctx, mode); });
+
   auto sector_number = image_offset / 512;
   auto appender = data->get_contiguous_appender(src.length());
   unsigned char* out_buf_ptr = nullptr;
@@ -107,8 +112,6 @@ int BlockCrypto<T>::crypt(ceph::bufferlist* data, uint64_t image_offset,
     }
   }
 
-  m_data_cryptor->return_context(ctx, mode);
-
   return 0;
 }
 

diff --git a/src/test/librbd/crypto/test_mock_BlockCrypto.cc b/src/test/librbd/crypto/test_mock_BlockCrypto.cc
index 4ddc43d3bb7c0afcab099565f295671312e3c3af..07dc19437de33d2d17f0cb520e1fc5c9777aae3c 100644 (file)
--- a/src/test/librbd/crypto/test_mock_BlockCrypto.cc
+++ b/src/test/librbd/crypto/test_mock_BlockCrypto.cc
@@ -11,7 +11,9 @@ template class librbd::crypto::BlockCrypto<
 
 using ::testing::ExpectationSet;
 using ::testing::internal::ExpectationBase;
+using ::testing::Invoke;
 using ::testing::Return;
+using ::testing::WithArg;
 using ::testing::_;
 
 namespace librbd {
@@ -54,6 +56,15 @@ struct TestMockCryptoBlockCrypto : public TestFixture {
                       new MockCryptoContext())));
     }
 
+    void expect_return_context(CipherMode mode) {
+      _set_last_expectation(
+              EXPECT_CALL(*cryptor, return_context(_, mode))
+              .After(*expectation_set).WillOnce(WithArg<0>(
+                      Invoke([](MockCryptoContext* ctx) {
+                        delete ctx;
+                      }))));
+    }
+
     void expect_init_context(const std::string& iv) {
       _set_last_expectation(
               EXPECT_CALL(*cryptor, init_context(_, CompareArrayToString(iv),
@@ -95,7 +106,7 @@ TEST_F(TestMockCryptoBlockCrypto, Encrypt) {
   expect_update_context(std::string(2048, '1') + std::string(2048, '2'), 4096);
   expect_init_context(std::string("\x38\x12\0\0\0\0\0\0\0\0\0\0\0\0\0\0", 16));
   expect_update_context(std::string(2048, '2') + std::string(2048, '3'), 4096);
-  EXPECT_CALL(*cryptor, return_context(_, CipherMode::CIPHER_MODE_ENC));
+  expect_return_context(CipherMode::CIPHER_MODE_ENC);
 
   ASSERT_EQ(0, bc->encrypt(&data, image_offset));
 
@@ -127,6 +138,7 @@ TEST_F(TestMockCryptoBlockCrypto, InitContextError) {
   data.append(std::string(4096, '1'));
   expect_get_context(CipherMode::CIPHER_MODE_ENC);
   EXPECT_CALL(*cryptor, init_context(_, _, _)).WillOnce(Return(-123));
+  expect_return_context(CipherMode::CIPHER_MODE_ENC);
   ASSERT_EQ(-123, bc->encrypt(&data, 0));
 }
 
@@ -136,6 +148,7 @@ TEST_F(TestMockCryptoBlockCrypto, UpdateContextError) {
   expect_get_context(CipherMode::CIPHER_MODE_ENC);
   EXPECT_CALL(*cryptor, init_context(_, _, _));
   EXPECT_CALL(*cryptor, update_context(_, _, _, _)).WillOnce(Return(-123));
+  expect_return_context(CipherMode::CIPHER_MODE_ENC);
   ASSERT_EQ(-123, bc->encrypt(&data, 0));
 }