2 files changed, 27 insertions, 23 deletions

diff --git a/core/java/android/os/Binder.java b/core/java/android/os/Binder.java
index ed75491b8e21..86dc20c497e6 100644
--- a/core/java/android/os/Binder.java
+++ b/core/java/android/os/Binder.java
@@ -149,6 +149,11 @@ public class Binder implements IBinder {
     private static volatile boolean sStackTrackingEnabled = false;
 
     /**
+     * The extension binder object
+     */
+    private IBinder mExtension = null;
+
+    /**
      * Enable Binder IPC stack tracking. If enabled, every binder transaction will be logged to
      * {@link TransactionTracker}.
      *
@@ -1234,7 +1239,9 @@ public class Binder implements IBinder {
 
     /** @hide */
     @Override
-    public final native @Nullable IBinder getExtension();
+    public final @Nullable IBinder getExtension() {
+        return mExtension;
+    }
 
     /**
      * Set the binder extension.
@@ -1242,7 +1249,12 @@ public class Binder implements IBinder {
      *
      * @hide
      */
-    public final native void setExtension(@Nullable IBinder extension);
+    public final void setExtension(@Nullable IBinder extension) {
+        mExtension = extension;
+        setExtensionNative(extension);
+    }
+
+    private final native void setExtensionNative(@Nullable IBinder extension);
 
     /**
      * Default implementation rewinds the parcels and calls onTransact. On
diff --git a/core/jni/android_util_Binder.cpp b/core/jni/android_util_Binder.cpp
index 7f545164dbc3..5a0b21ac43b1 100644
--- a/core/jni/android_util_Binder.cpp
+++ b/core/jni/android_util_Binder.cpp
@@ -73,6 +73,7 @@ static struct bindernative_offsets_t
     jmethodID mExecTransact;
     jmethodID mGetInterfaceDescriptor;
     jmethodID mTransactionCallback;
+    jmethodID mGetExtension;
 
     // Object state.
     jfieldID mObject;
@@ -488,8 +489,12 @@ public:
             if (mVintf) {
                 ::android::internal::Stability::markVintf(b.get());
             }
-            if (mExtension != nullptr) {
-                b.get()->setExtension(mExtension);
+            if (mSetExtensionCalled) {
+                jobject javaIBinderObject = env->CallObjectMethod(obj, gBinderOffsets.mGetExtension);
+                sp<IBinder> extensionFromJava = ibinderForJavaObject(env, javaIBinderObject);
+                if (extensionFromJava != nullptr) {
+                    b.get()->setExtension(extensionFromJava);
+                }
             }
             mBinder = b;
             ALOGV("Creating JavaBinder %p (refs %p) for Object %p, weakCount=%" PRId32 "\n",
@@ -515,21 +520,12 @@ public:
         mVintf = false;
     }
 
-    sp<IBinder> getExtension() {
-        AutoMutex _l(mLock);
-        sp<JavaBBinder> b = mBinder.promote();
-        if (b != nullptr) {
-            return b.get()->getExtension();
-        }
-        return mExtension;
-    }
-
     void setExtension(const sp<IBinder>& extension) {
         AutoMutex _l(mLock);
-        mExtension = extension;
+        mSetExtensionCalled = true;
         sp<JavaBBinder> b = mBinder.promote();
         if (b != nullptr) {
-            b.get()->setExtension(mExtension);
+            b.get()->setExtension(extension);
         }
     }
 
@@ -541,8 +537,7 @@ private:
     // is too much binder state here, we can think about making JavaBBinder an
     // sp here (avoid recreating it)
     bool            mVintf = false;
-
-    sp<IBinder>     mExtension;
+    bool            mSetExtensionCalled = false;
 };
 
 // ----------------------------------------------------------------------------
@@ -1254,10 +1249,6 @@ static void android_os_Binder_blockUntilThreadAvailable(JNIEnv* env, jobject cla
     return IPCThreadState::self()->blockUntilThreadAvailable();
 }
 
-static jobject android_os_Binder_getExtension(JNIEnv* env, jobject obj) {
-    JavaBBinderHolder* jbh = (JavaBBinderHolder*) env->GetLongField(obj, gBinderOffsets.mObject);
-    return javaObjectForIBinder(env, jbh->getExtension());
-}
 
 static void android_os_Binder_setExtension(JNIEnv* env, jobject obj, jobject extensionObject) {
     JavaBBinderHolder* jbh = (JavaBBinderHolder*) env->GetLongField(obj, gBinderOffsets.mObject);
@@ -1300,8 +1291,7 @@ static const JNINativeMethod gBinderMethods[] = {
     { "getNativeBBinderHolder", "()J", (void*)android_os_Binder_getNativeBBinderHolder },
     { "getNativeFinalizer", "()J", (void*)android_os_Binder_getNativeFinalizer },
     { "blockUntilThreadAvailable", "()V", (void*)android_os_Binder_blockUntilThreadAvailable },
-    { "getExtension", "()Landroid/os/IBinder;", (void*)android_os_Binder_getExtension },
-    { "setExtension", "(Landroid/os/IBinder;)V", (void*)android_os_Binder_setExtension },
+    { "setExtensionNative", "(Landroid/os/IBinder;)V", (void*)android_os_Binder_setExtension },
 };
 // clang-format on
 
@@ -1318,6 +1308,8 @@ static int int_register_android_os_Binder(JNIEnv* env)
     gBinderOffsets.mTransactionCallback =
             GetStaticMethodIDOrDie(env, clazz, "transactionCallback", "(IIII)V");
     gBinderOffsets.mObject = GetFieldIDOrDie(env, clazz, "mObject", "J");
+    gBinderOffsets.mGetExtension = GetMethodIDOrDie(env, clazz, "getExtension",
+                                                        "()Landroid/os/IBinder;");
 
     return RegisterMethodsOrDie(
         env, kBinderPathName,