Add lookupOrCreateWeak to ObjectManager

This looks for an attached object on the binder, attempts to promote it,
and returns the object if it could be promoted.
Otherwise, it will create a new object and attach it to the binder.

Test: atest binderUnitTest aidl_integration_test
Bug: 220141324
Change-Id: I2c8baa211300fa3ada50f2b2c93862f8c968be26
diff --git a/libs/binder/Binder.cpp b/libs/binder/Binder.cpp
index b5ea60f..1dc6233 100644
--- a/libs/binder/Binder.cpp
+++ b/libs/binder/Binder.cpp
@@ -202,6 +202,17 @@
     proxy->withLock(doWithLock);
 }
 
+sp<IBinder> IBinder::lookupOrCreateWeak(const void* objectID, object_make_func make,
+                                        const void* makeArgs) {
+    BBinder* local = localBinder();
+    if (local) {
+        return local->lookupOrCreateWeak(objectID, make, makeArgs);
+    }
+    BpBinder* proxy = this->remoteBinder();
+    LOG_ALWAYS_FATAL_IF(proxy == nullptr, "binder object must be either local or remote");
+    return proxy->lookupOrCreateWeak(objectID, make, makeArgs);
+}
+
 // ---------------------------------------------------------------------------
 
 class BBinder::RpcServerLink : public IBinder::DeathRecipient {
@@ -378,6 +389,14 @@
     doWithLock();
 }
 
+sp<IBinder> BBinder::lookupOrCreateWeak(const void* objectID, object_make_func make,
+                                        const void* makeArgs) {
+    Extras* e = getOrCreateExtras();
+    LOG_ALWAYS_FATAL_IF(!e, "no memory");
+    AutoMutex _l(e->mLock);
+    return e->mObjects.lookupOrCreateWeak(objectID, make, makeArgs);
+}
+
 BBinder* BBinder::localBinder()
 {
     return this;
diff --git a/libs/binder/BpBinder.cpp b/libs/binder/BpBinder.cpp
index b6d35ef..d9b7231 100644
--- a/libs/binder/BpBinder.cpp
+++ b/libs/binder/BpBinder.cpp
@@ -100,6 +100,36 @@
     return value;
 }
 
+namespace {
+struct Tag {
+    wp<IBinder> binder;
+};
+} // namespace
+
+static void cleanWeak(const void* /* id */, void* obj, void* /* cookie */) {
+    delete static_cast<Tag*>(obj);
+}
+
+sp<IBinder> BpBinder::ObjectManager::lookupOrCreateWeak(const void* objectID, object_make_func make,
+                                                        const void* makeArgs) {
+    entry_t& e = mObjects[objectID];
+    if (e.object != nullptr) {
+        if (auto attached = static_cast<Tag*>(e.object)->binder.promote()) {
+            return attached;
+        }
+    } else {
+        e.object = new Tag;
+        LOG_ALWAYS_FATAL_IF(!e.object, "no more memory");
+    }
+    sp<IBinder> newObj = make(makeArgs);
+
+    static_cast<Tag*>(e.object)->binder = newObj;
+    e.cleanupCookie = nullptr;
+    e.func = cleanWeak;
+
+    return newObj;
+}
+
 void BpBinder::ObjectManager::kill()
 {
     const size_t N = mObjects.size();
@@ -516,6 +546,12 @@
     doWithLock();
 }
 
+sp<IBinder> BpBinder::lookupOrCreateWeak(const void* objectID, object_make_func make,
+                                         const void* makeArgs) {
+    AutoMutex _l(mLock);
+    return mObjects.lookupOrCreateWeak(objectID, make, makeArgs);
+}
+
 BpBinder* BpBinder::remoteBinder()
 {
     return this;
diff --git a/libs/binder/include/binder/Binder.h b/libs/binder/include/binder/Binder.h
index 46223bb..88d9ca1 100644
--- a/libs/binder/include/binder/Binder.h
+++ b/libs/binder/include/binder/Binder.h
@@ -59,6 +59,8 @@
     virtual void*       findObject(const void* objectID) const final;
     virtual void* detachObject(const void* objectID) final;
     void withLock(const std::function<void()>& doWithLock);
+    sp<IBinder> lookupOrCreateWeak(const void* objectID, IBinder::object_make_func make,
+                                   const void* makeArgs);
 
     virtual BBinder*    localBinder();
 
diff --git a/libs/binder/include/binder/BpBinder.h b/libs/binder/include/binder/BpBinder.h
index 19ad5e6..4172cc5 100644
--- a/libs/binder/include/binder/BpBinder.h
+++ b/libs/binder/include/binder/BpBinder.h
@@ -72,6 +72,8 @@
     virtual void*       findObject(const void* objectID) const final;
     virtual void* detachObject(const void* objectID) final;
     void withLock(const std::function<void()>& doWithLock);
+    sp<IBinder> lookupOrCreateWeak(const void* objectID, IBinder::object_make_func make,
+                                   const void* makeArgs);
 
     virtual BpBinder*   remoteBinder();
 
@@ -96,6 +98,8 @@
                      IBinder::object_cleanup_func func);
         void* find(const void* objectID) const;
         void* detach(const void* objectID);
+        sp<IBinder> lookupOrCreateWeak(const void* objectID, IBinder::object_make_func make,
+                                       const void* makeArgs);
 
         void kill();
 
@@ -104,9 +108,9 @@
         ObjectManager& operator=(const ObjectManager&);
 
         struct entry_t {
-            void* object;
-            void* cleanupCookie;
-            IBinder::object_cleanup_func func;
+            void* object = nullptr;
+            void* cleanupCookie = nullptr;
+            IBinder::object_cleanup_func func = nullptr;
         };
 
         std::map<const void*, entry_t> mObjects;
diff --git a/libs/binder/include/binder/IBinder.h b/libs/binder/include/binder/IBinder.h
index 43fc5ff..83aaca7 100644
--- a/libs/binder/include/binder/IBinder.h
+++ b/libs/binder/include/binder/IBinder.h
@@ -284,6 +284,9 @@
 
     virtual BBinder*        localBinder();
     virtual BpBinder*       remoteBinder();
+    typedef sp<IBinder> (*object_make_func)(const void* makeArgs);
+    sp<IBinder> lookupOrCreateWeak(const void* objectID, object_make_func make,
+                                   const void* makeArgs);
 
 protected:
     virtual          ~IBinder();
diff --git a/libs/binder/tests/binderBinderUnitTest.cpp b/libs/binder/tests/binderBinderUnitTest.cpp
index ce2770f..b6aed0d 100644
--- a/libs/binder/tests/binderBinderUnitTest.cpp
+++ b/libs/binder/tests/binderBinderUnitTest.cpp
@@ -15,10 +15,11 @@
  */
 
 #include <binder/Binder.h>
-#include <binder/IBinder.h>
+#include <binder/IInterface.h>
 #include <gtest/gtest.h>
 
 using android::BBinder;
+using android::IBinder;
 using android::OK;
 using android::sp;
 
@@ -48,3 +49,49 @@
     binder->setExtension(ext);
     EXPECT_EQ(ext, binder->getExtension());
 }
+
+struct MyCookie {
+    bool* deleted;
+};
+
+class UniqueBinder : public BBinder {
+public:
+    UniqueBinder(const void* c) : cookie(reinterpret_cast<const MyCookie*>(c)) {
+        *cookie->deleted = false;
+    }
+    ~UniqueBinder() { *cookie->deleted = true; }
+    const MyCookie* cookie;
+};
+
+static sp<IBinder> make(const void* arg) {
+    return sp<UniqueBinder>::make(arg);
+}
+
+TEST(Binder, LookupOrCreateWeak) {
+    auto binder = sp<BBinder>::make();
+    bool deleted;
+    MyCookie cookie = {&deleted};
+    sp<IBinder> createdBinder = binder->lookupOrCreateWeak(kObjectId1, make, &cookie);
+    EXPECT_NE(binder, createdBinder);
+
+    sp<IBinder> lookedUpBinder = binder->lookupOrCreateWeak(kObjectId1, make, &cookie);
+    EXPECT_EQ(createdBinder, lookedUpBinder);
+    EXPECT_FALSE(deleted);
+}
+
+TEST(Binder, LookupOrCreateWeakDropSp) {
+    auto binder = sp<BBinder>::make();
+    bool deleted1 = false;
+    bool deleted2 = false;
+    MyCookie cookie1 = {&deleted1};
+    MyCookie cookie2 = {&deleted2};
+    sp<IBinder> createdBinder = binder->lookupOrCreateWeak(kObjectId1, make, &cookie1);
+    EXPECT_NE(binder, createdBinder);
+
+    createdBinder.clear();
+    EXPECT_TRUE(deleted1);
+
+    sp<IBinder> lookedUpBinder = binder->lookupOrCreateWeak(kObjectId1, make, &cookie2);
+    EXPECT_EQ(&cookie2, sp<UniqueBinder>::cast(lookedUpBinder)->cookie);
+    EXPECT_FALSE(deleted2);
+}