Merge pull request #6 from nvidianz/fix-double-free

Fixed memory leaks in tests

src/processing/plugins/mock_processor.cc

@@ -57,7 +57,10 @@ void* MockProcessor::HandleGHPairs(std::size_t *size, void *buffer, std::size_t
     }
   }
 
-  return buffer;
+  auto result = malloc(buf_size);
+  memcpy(result, buffer, buf_size);
+
+  return result;
 }
 
 void *MockProcessor::ProcessAggregation(std::size_t *size, std::map<int, std::vector<int>> nodes) {

src/processing/plugins/mock_processor.h

@@ -21,15 +21,27 @@ class MockProcessor: public processing::Processor {
     std::vector<int> slots_;
 
  public:
+
+    ~MockProcessor() {
+        if (gh_pairs_) {
+            gh_pairs_->clear();
+            delete gh_pairs_;
+        }
+    }
+
     void Initialize(bool active, std::map<std::string, std::string> params) override {
         this->active_ = active;
         this->params_ = &params;
     }
 
     void Shutdown() override {
-        this->gh_pairs_ = nullptr;
-        this->cuts_.clear();
-        this->slots_.clear();
+        if (gh_pairs_) {
+            gh_pairs_->clear();
+            delete gh_pairs_;
+        }
+        gh_pairs_ = nullptr;
+        cuts_.clear();
+        slots_.clear();
     }
 
     void FreeBuffer(void *buffer) override {

src/processing/processor.h

@@ -18,6 +18,12 @@ const char kLoadFunc[] = "LoadProcessor";
 /*! \brief An processor interface to handle tasks that require external library through plugins */
 class Processor {
  public:
+    /*!
+     * \brief Virtual destructor
+     *
+     */
+    virtual ~Processor() = default;
+
     /*!
      * \brief Initialize the processor
      *

src/processing/processor_loader.cc

@@ -78,9 +78,13 @@ namespace processing {
 
     void ProcessorLoader::unload() {
 #if defined(_WIN32)
-        FreeLibrary(handle_);
+        if (handle_) {
+            FreeLibrary(handle_);
+        }
 #else
-        dlclose(handle_);
+        if (handle_) {
+            dlclose(handle_);
+        }
 #endif
     }
 }  // namespace processing

tests/cpp/processing/test_processor.cc

@@ -13,10 +13,12 @@ class ProcessorTest : public testing::Test {
         auto loader = processing::ProcessorLoader();
         processor_ = loader.load(processing::kMockProcessor);
         processor_->Initialize(true, {});
+        loader.unload();
     }
 
     void TearDown() override {
         processor_->Shutdown();
+        delete processor_;
         processor_ = nullptr;
     }
 
@@ -62,13 +64,13 @@ TEST_F(ProcessorTest, TestGHEncoding) {
     ASSERT_EQ(0, memcmp(buffer, new_buffer, buf_size));
 
     // Clean up
-    free(buffer);
-    free(new_buffer);
+    processor_->FreeBuffer(buffer);
+    processor_->FreeBuffer(new_buffer);
 }
 
 TEST_F(ProcessorTest, TestAggregation) {
     size_t buf_size;
-    processor_->ProcessGHPairs(&buf_size, gh_pairs_);  // Pass the GH pairs to the plugin
+    auto gh_buffer = processor_->ProcessGHPairs(&buf_size, gh_pairs_);  // Pass the GH pairs to the plugin
 
     processor_->InitAggregationContext(cuts_, slots_);
     auto buffer = processor_->ProcessAggregation(&buf_size, nodes_);
@@ -87,4 +89,7 @@ TEST_F(ProcessorTest, TestAggregation) {
     for (size_t i = 0; i < histos.size(); ++i) {
         EXPECT_NEAR(expected_result[i], histos[i], kError) << "Histogram differs at index " << i;
     }
+
+    processor_->FreeBuffer(buffer);
+    processor_->FreeBuffer(gh_buffer);
 }