Implement IBufferClient::close on BufferClient

The close() function will mark a BufferClient as closed (set mClosed =
true), remove it from service's mClientList, and remove all of its token
in service's mTokenMap.

Add test cases to make sure it works in all situations.

Test: BufferHubBuffer_test (passed)
Bug: 119623209
Change-Id: Ic1d17ced97b67ef5432c9d341469d8e6105e2717
diff --git a/libs/ui/tests/BufferHubBuffer_test.cpp b/libs/ui/tests/BufferHubBuffer_test.cpp
index 553ee0b..5f9fa67 100644
--- a/libs/ui/tests/BufferHubBuffer_test.cpp
+++ b/libs/ui/tests/BufferHubBuffer_test.cpp
@@ -127,7 +127,7 @@
     return;
 }
 
-TEST_F(BufferHubBufferTest, AllocateBuffer) {
+TEST_F(BufferHubBufferTest, AllocateAndFreeBuffer) {
     // TODO(b/116681016): directly test on BufferHubBuffer instead of the service.
     sp<IBufferHub> bufferHub = IBufferHub::getService();
     ASSERT_NE(nullptr, bufferHub.get());
@@ -138,11 +138,51 @@
     HardwareBufferDescription desc;
     memcpy(&desc, &aDesc, sizeof(HardwareBufferDescription));
 
-    IBufferHub::allocateBuffer_cb callback = [](const auto& client, const auto& status) {
-        EXPECT_EQ(status, BufferHubStatus::NO_ERROR);
-        EXPECT_NE(nullptr, client.get());
+    sp<IBufferClient> client;
+    BufferHubStatus ret;
+    IBufferHub::allocateBuffer_cb callback = [&](const auto& outClient, const auto& outStatus) {
+        client = outClient;
+        ret = outStatus;
     };
     EXPECT_TRUE(bufferHub->allocateBuffer(desc, kUserMetadataSize, callback).isOk());
+    EXPECT_EQ(ret, BufferHubStatus::NO_ERROR);
+    ASSERT_NE(nullptr, client.get());
+
+    EXPECT_EQ(BufferHubStatus::NO_ERROR, client->close());
+    EXPECT_EQ(BufferHubStatus::CLIENT_CLOSED, client->close());
+}
+
+TEST_F(BufferHubBufferTest, DuplicateFreedBuffer) {
+    // TODO(b/116681016): directly test on BufferHubBuffer instead of the service.
+    sp<IBufferHub> bufferHub = IBufferHub::getService();
+    ASSERT_NE(nullptr, bufferHub.get());
+
+    // Stride is an output, rfu0 and rfu1 are reserved data slot for future use.
+    AHardwareBuffer_Desc aDesc = {kWidth, kHeight,        kLayerCount,  kFormat,
+                                  kUsage, /*stride=*/0UL, /*rfu0=*/0UL, /*rfu1=*/0ULL};
+    HardwareBufferDescription desc;
+    memcpy(&desc, &aDesc, sizeof(HardwareBufferDescription));
+
+    sp<IBufferClient> client;
+    BufferHubStatus ret;
+    IBufferHub::allocateBuffer_cb callback = [&](const auto& outClient, const auto& outStatus) {
+        client = outClient;
+        ret = outStatus;
+    };
+    EXPECT_TRUE(bufferHub->allocateBuffer(desc, kUserMetadataSize, callback).isOk());
+    EXPECT_EQ(ret, BufferHubStatus::NO_ERROR);
+    ASSERT_NE(nullptr, client.get());
+
+    EXPECT_EQ(BufferHubStatus::NO_ERROR, client->close());
+
+    hidl_handle token;
+    IBufferClient::duplicate_cb dup_cb = [&](const auto& outToken, const auto& status) {
+        token = outToken;
+        ret = status;
+    };
+    EXPECT_TRUE(client->duplicate(dup_cb).isOk());
+    EXPECT_EQ(ret, BufferHubStatus::CLIENT_CLOSED);
+    EXPECT_EQ(token.getNativeHandle(), nullptr);
 }
 
 TEST_F(BufferHubBufferTest, DuplicateAndImportBuffer) {
@@ -230,5 +270,50 @@
     native_handle_delete(tokenHandle);
 }
 
+TEST_F(BufferHubBufferTest, ImportFreedBuffer) {
+    // TODO(b/116681016): directly test on BufferHubBuffer instead of the service.
+    sp<IBufferHub> bufferhub = IBufferHub::getService();
+    ASSERT_NE(nullptr, bufferhub.get());
+
+    // Stride is an output, rfu0 and rfu1 are reserved data slot for future use.
+    AHardwareBuffer_Desc aDesc = {kWidth, kHeight,        kLayerCount,  kFormat,
+                                  kUsage, /*stride=*/0UL, /*rfu0=*/0UL, /*rfu1=*/0ULL};
+    HardwareBufferDescription desc;
+    memcpy(&desc, &aDesc, sizeof(HardwareBufferDescription));
+
+    sp<IBufferClient> client;
+    BufferHubStatus ret;
+    IBufferHub::allocateBuffer_cb alloc_cb = [&](const auto& outClient, const auto& status) {
+        client = outClient;
+        ret = status;
+    };
+    ASSERT_TRUE(bufferhub->allocateBuffer(desc, kUserMetadataSize, alloc_cb).isOk());
+    EXPECT_EQ(ret, BufferHubStatus::NO_ERROR);
+    ASSERT_NE(nullptr, client.get());
+
+    hidl_handle token;
+    IBufferClient::duplicate_cb dup_cb = [&](const auto& outToken, const auto& status) {
+        token = outToken;
+        ret = status;
+    };
+    ASSERT_TRUE(client->duplicate(dup_cb).isOk());
+    EXPECT_EQ(ret, BufferHubStatus::NO_ERROR);
+    ASSERT_NE(token.getNativeHandle(), nullptr);
+    EXPECT_EQ(token->numInts, 1);
+    EXPECT_EQ(token->numFds, 0);
+
+    // Close the client. Now the token should be invalid.
+    client->close();
+
+    sp<IBufferClient> client2;
+    IBufferHub::importBuffer_cb import_cb = [&](const auto& outClient, const auto& status) {
+        client2 = outClient;
+        ret = status;
+    };
+    EXPECT_TRUE(bufferhub->importBuffer(token, import_cb).isOk());
+    EXPECT_EQ(ret, BufferHubStatus::INVALID_TOKEN);
+    EXPECT_EQ(nullptr, client2.get());
+}
+
 } // namespace
 } // namespace android
diff --git a/services/bufferhub/BufferClient.cpp b/services/bufferhub/BufferClient.cpp
index 37fd75f..7459517 100644
--- a/services/bufferhub/BufferClient.cpp
+++ b/services/bufferhub/BufferClient.cpp
@@ -39,7 +39,29 @@
     return new BufferClient(service, node);
 }
 
+BufferClient::~BufferClient() {
+    close();
+}
+
+Return<BufferHubStatus> BufferClient::close() {
+    std::lock_guard<std::mutex> lock(mClosedMutex);
+    if (mClosed) {
+        return BufferHubStatus::CLIENT_CLOSED;
+    }
+
+    getService()->onClientClosed(this);
+    mBufferNode.reset();
+    mClosed = true;
+    return BufferHubStatus::NO_ERROR;
+}
+
 Return<void> BufferClient::duplicate(duplicate_cb _hidl_cb) {
+    std::lock_guard<std::mutex> lock(mClosedMutex);
+    if (mClosed) {
+        _hidl_cb(/*token=*/hidl_handle(), /*status=*/BufferHubStatus::CLIENT_CLOSED);
+        return Void();
+    }
+
     if (!mBufferNode) {
         // Should never happen
         ALOGE("%s: node is missing.", __FUNCTION__);
@@ -47,15 +69,19 @@
         return Void();
     }
 
+    const hidl_handle token = getService()->registerToken(this);
+    _hidl_cb(/*token=*/token, /*status=*/BufferHubStatus::NO_ERROR);
+    return Void();
+}
+
+sp<BufferHubService> BufferClient::getService() {
     sp<BufferHubService> service = mService.promote();
     if (service == nullptr) {
         // Should never happen. Kill the process.
         LOG_FATAL("%s: service died.", __FUNCTION__);
     }
 
-    const hidl_handle token = service->registerToken(this);
-    _hidl_cb(/*token=*/token, /*status=*/BufferHubStatus::NO_ERROR);
-    return Void();
+    return service;
 }
 
 } // namespace implementation
diff --git a/services/bufferhub/BufferHubService.cpp b/services/bufferhub/BufferHubService.cpp
index 6f97f0d..409f2ea 100644
--- a/services/bufferhub/BufferHubService.cpp
+++ b/services/bufferhub/BufferHubService.cpp
@@ -111,6 +111,30 @@
     return returnToken;
 }
 
+void BufferHubService::onClientClosed(const BufferClient* client) {
+    removeTokenByClient(client);
+
+    std::lock_guard<std::mutex> lock(mClientListMutex);
+    auto iter = std::find(mClientList.begin(), mClientList.end(), client);
+    if (iter != mClientList.end()) {
+        mClientList.erase(iter);
+    }
+}
+
+void BufferHubService::removeTokenByClient(const BufferClient* client) {
+    std::lock_guard<std::mutex> lock(mTokenMapMutex);
+    auto iter = mTokenMap.begin();
+    while (iter != mTokenMap.end()) {
+        if (iter->second == client) {
+            auto oldIter = iter;
+            ++iter;
+            mTokenMap.erase(oldIter);
+        } else {
+            ++iter;
+        }
+    }
+}
+
 } // namespace implementation
 } // namespace V1_0
 } // namespace bufferhub
diff --git a/services/bufferhub/include/bufferhub/BufferClient.h b/services/bufferhub/include/bufferhub/BufferClient.h
index 769ec86..7f5d3a6 100644
--- a/services/bufferhub/include/bufferhub/BufferClient.h
+++ b/services/bufferhub/include/bufferhub/BufferClient.h
@@ -44,14 +44,22 @@
     // Creates a BufferClient from an existing BufferClient. Will share the same BufferNode.
     explicit BufferClient(const BufferClient& other)
           : mService(other.mService), mBufferNode(other.mBufferNode) {}
+    ~BufferClient();
 
+    Return<BufferHubStatus> close() override;
     Return<void> duplicate(duplicate_cb _hidl_cb) override;
 
 private:
     BufferClient(wp<BufferHubService> service, const std::shared_ptr<BufferNode>& node)
           : mService(service), mBufferNode(node) {}
 
+    sp<BufferHubService> getService();
+
     wp<BufferHubService> mService;
+
+    std::mutex mClosedMutex;
+    bool mClosed GUARDED_BY(mClosedMutex) = false;
+
     std::shared_ptr<BufferNode> mBufferNode;
 };
 
diff --git a/services/bufferhub/include/bufferhub/BufferHubService.h b/services/bufferhub/include/bufferhub/BufferHubService.h
index 6535659..1c328b3 100644
--- a/services/bufferhub/include/bufferhub/BufferHubService.h
+++ b/services/bufferhub/include/bufferhub/BufferHubService.h
@@ -48,7 +48,12 @@
     // Internal help function for IBufferClient::duplicate.
     hidl_handle registerToken(const wp<BufferClient>& client);
 
+    void onClientClosed(const BufferClient* client);
+
 private:
+    // Helper function to remove all the token belongs to a specific client.
+    void removeTokenByClient(const BufferClient* client);
+
     // List of active BufferClient for bookkeeping.
     std::mutex mClientListMutex;
     std::vector<wp<BufferClient>> mClientList GUARDED_BY(mClientListMutex);