From d2ecde5c54b32165d22641dcb3f96b2c0cc3cc4c Mon Sep 17 00:00:00 2001
From: Xusong Wang <xusongw@google.com>
Date: Mon, 23 Aug 2021 11:36:19 -0700
Subject: [PATCH] Fix use-after-free crash in VtsHalNeuralnetworksTargetTest.

Prior to this CL, the AHardwareBuffer in TestBlobAHWB is released in the
destructor, but later used (unlock) during the destruction of the
mMapping member. This CL fixed this issue by managing the lifetime of
AHardwareBuffer with SharedMemory.

Bug: 197199690
Test: VtsHalNeuralnetworksTargetTest
Change-Id: I00748aaaa1a3a3d9b3b62bedb77a655ddb6e210f
---
 neuralnetworks/aidl/vts/functional/Utils.cpp | 19 ++++++-------------
 neuralnetworks/aidl/vts/functional/Utils.h   |  3 +--
 2 files changed, 7 insertions(+), 15 deletions(-)

diff --git a/neuralnetworks/aidl/vts/functional/Utils.cpp b/neuralnetworks/aidl/vts/functional/Utils.cpp
index 9af362ea8a..325a436f79 100644
--- a/neuralnetworks/aidl/vts/functional/Utils.cpp
+++ b/neuralnetworks/aidl/vts/functional/Utils.cpp
@@ -153,26 +153,19 @@ void TestBlobAHWB::initialize(uint32_t size) {
             .stride = size,
     };
 
-    ASSERT_EQ(AHardwareBuffer_allocate(&desc, &mAhwb), 0);
-    ASSERT_NE(mAhwb, nullptr);
+    AHardwareBuffer* ahwb = nullptr;
+    ASSERT_EQ(AHardwareBuffer_allocate(&desc, &ahwb), 0);
+    ASSERT_NE(ahwb, nullptr);
 
-    const auto sharedMemory =
-            nn::createSharedMemoryFromAHWB(mAhwb, /*takeOwnership=*/false).value();
-    mMapping = nn::map(sharedMemory).value();
+    mMemory = nn::createSharedMemoryFromAHWB(ahwb, /*takeOwnership=*/true).value();
+    mMapping = nn::map(mMemory).value();
     mPtr = static_cast<uint8_t*>(std::get<void*>(mMapping.pointer));
     CHECK_NE(mPtr, nullptr);
-    mAidlMemory = utils::convert(sharedMemory).value();
+    mAidlMemory = utils::convert(mMemory).value();
 
     mIsValid = true;
 }
 
-TestBlobAHWB::~TestBlobAHWB() {
-    if (mAhwb) {
-        AHardwareBuffer_unlock(mAhwb, nullptr);
-        AHardwareBuffer_release(mAhwb);
-    }
-}
-
 std::string gtestCompliantName(std::string name) {
     // gtest test names must only contain alphanumeric characters
     std::replace_if(
diff --git a/neuralnetworks/aidl/vts/functional/Utils.h b/neuralnetworks/aidl/vts/functional/Utils.h
index 9dd73592dd..ca81418417 100644
--- a/neuralnetworks/aidl/vts/functional/Utils.h
+++ b/neuralnetworks/aidl/vts/functional/Utils.h
@@ -102,11 +102,10 @@ class TestBlobAHWB : public TestMemoryBase {
     // The constructor calls initialize, which constructs the memory resources. This is a
     // workaround that gtest macros cannot be used directly in a constructor.
     TestBlobAHWB(uint32_t size) { initialize(size); }
-    ~TestBlobAHWB();
 
   private:
     void initialize(uint32_t size);
-    AHardwareBuffer* mAhwb = nullptr;
+    nn::SharedMemory mMemory;
     nn::Mapping mMapping;
 };