3 files changed, 43 insertions, 16 deletions

diff --git a/services/core/jni/tvinput/JTvInputHal.cpp b/services/core/jni/tvinput/JTvInputHal.cpp
index dc05462c1a1c..80427b346f1a 100644
--- a/services/core/jni/tvinput/JTvInputHal.cpp
+++ b/services/core/jni/tvinput/JTvInputHal.cpp
@@ -25,7 +25,7 @@ JTvInputHal::JTvInputHal(JNIEnv* env, jobject thiz, std::shared_ptr<ITvInputWrap
     mThiz = env->NewWeakGlobalRef(thiz);
     mTvInput = tvInput;
     mLooper = looper;
-    mTvInputCallback = ::ndk::SharedRefBase::make<TvInputCallback>(this);
+    mTvInputCallback = std::shared_ptr<TvInputCallbackWrapper>(new TvInputCallbackWrapper(this));
     mTvInput->setCallback(mTvInputCallback);
 }
 
@@ -443,18 +443,23 @@ void JTvInputHal::NotifyTvMessageHandler::handleMessage(const Message& message)
     }
 }
 
-JTvInputHal::TvInputCallback::TvInputCallback(JTvInputHal* hal) {
+JTvInputHal::TvInputCallbackWrapper::TvInputCallbackWrapper(JTvInputHal* hal) {
+    aidlTvInputCallback = ::ndk::SharedRefBase::make<AidlTvInputCallback>(hal);
+    hidlTvInputCallback = sp<HidlTvInputCallback>::make(hal);
+}
+
+JTvInputHal::AidlTvInputCallback::AidlTvInputCallback(JTvInputHal* hal) {
     mHal = hal;
 }
 
-::ndk::ScopedAStatus JTvInputHal::TvInputCallback::notify(const AidlTvInputEvent& event) {
+::ndk::ScopedAStatus JTvInputHal::AidlTvInputCallback::notify(const AidlTvInputEvent& event) {
     mHal->mLooper->sendMessage(new NotifyHandler(mHal,
                                                  TvInputEventWrapper::createEventWrapper(event)),
                                static_cast<int>(event.type));
     return ::ndk::ScopedAStatus::ok();
 }
 
-::ndk::ScopedAStatus JTvInputHal::TvInputCallback::notifyTvMessageEvent(
+::ndk::ScopedAStatus JTvInputHal::AidlTvInputCallback::notifyTvMessageEvent(
         const AidlTvMessageEvent& event) {
     const std::string DEVICE_ID_SUBTYPE = "device_id";
     ::ndk::ScopedAStatus status = ::ndk::ScopedAStatus::ok();
@@ -487,11 +492,14 @@ JTvInputHal::ITvInputWrapper::ITvInputWrapper(std::shared_ptr<AidlITvInput>& aid
       : mIsHidl(false), mAidlTvInput(aidlTvInput) {}
 
 ::ndk::ScopedAStatus JTvInputHal::ITvInputWrapper::setCallback(
-        const std::shared_ptr<TvInputCallback>& in_callback) {
+        const std::shared_ptr<TvInputCallbackWrapper>& in_callback) {
     if (mIsHidl) {
-        return hidlSetCallback(in_callback);
+        in_callback->aidlTvInputCallback = nullptr;
+        return hidlSetCallback(in_callback == nullptr ? nullptr : in_callback->hidlTvInputCallback);
     } else {
-        return mAidlTvInput->setCallback(in_callback);
+        in_callback->hidlTvInputCallback = nullptr;
+        return mAidlTvInput->setCallback(in_callback == nullptr ? nullptr
+                                                                : in_callback->aidlTvInputCallback);
     }
 }
 
diff --git a/services/core/jni/tvinput/JTvInputHal.h b/services/core/jni/tvinput/JTvInputHal.h
index 6026a107c67f..2ef94ac4a3b0 100644
--- a/services/core/jni/tvinput/JTvInputHal.h
+++ b/services/core/jni/tvinput/JTvInputHal.h
@@ -168,23 +168,39 @@ private:
         JTvInputHal* mHal;
     };
 
-    class TvInputCallback : public HidlITvInputCallback, public BnTvInputCallback {
+    class AidlTvInputCallback : public BnTvInputCallback {
     public:
-        explicit TvInputCallback(JTvInputHal* hal);
+        explicit AidlTvInputCallback(JTvInputHal* hal);
         ::ndk::ScopedAStatus notify(const AidlTvInputEvent& event) override;
         ::ndk::ScopedAStatus notifyTvMessageEvent(const AidlTvMessageEvent& event) override;
+
+    private:
+        JTvInputHal* mHal;
+    };
+
+    class HidlTvInputCallback : public HidlITvInputCallback {
+    public:
+        explicit HidlTvInputCallback(JTvInputHal* hal);
         Return<void> notify(const HidlTvInputEvent& event) override;
 
     private:
         JTvInputHal* mHal;
     };
 
+    class TvInputCallbackWrapper {
+    public:
+        explicit TvInputCallbackWrapper(JTvInputHal* hal);
+        std::shared_ptr<AidlTvInputCallback> aidlTvInputCallback;
+        sp<HidlTvInputCallback> hidlTvInputCallback;
+    };
+
     class ITvInputWrapper {
     public:
         ITvInputWrapper(std::shared_ptr<AidlITvInput>& aidlTvInput);
         ITvInputWrapper(sp<HidlITvInput>& hidlTvInput);
 
-        ::ndk::ScopedAStatus setCallback(const std::shared_ptr<TvInputCallback>& in_callback);
+        ::ndk::ScopedAStatus setCallback(
+                const std::shared_ptr<TvInputCallbackWrapper>& in_callback);
         ::ndk::ScopedAStatus getStreamConfigurations(int32_t in_deviceId,
                                                      std::vector<AidlTvStreamConfig>* _aidl_return);
         ::ndk::ScopedAStatus openStream(int32_t in_deviceId, int32_t in_streamId,
@@ -198,7 +214,7 @@ private:
         ::ndk::ScopedAStatus getAidlInterfaceVersion(int32_t* _aidl_return);
 
     private:
-        ::ndk::ScopedAStatus hidlSetCallback(const std::shared_ptr<TvInputCallback>& in_callback);
+        ::ndk::ScopedAStatus hidlSetCallback(const sp<HidlTvInputCallback>& in_callback);
         ::ndk::ScopedAStatus hidlGetStreamConfigurations(
                 int32_t in_deviceId, std::vector<AidlTvStreamConfig>* _aidl_return);
         ::ndk::ScopedAStatus hidlOpenStream(int32_t in_deviceId, int32_t in_streamId,
@@ -229,7 +245,7 @@ private:
     KeyedVector<int, KeyedVector<int, Connection> > mConnections;
 
     std::shared_ptr<ITvInputWrapper> mTvInput;
-    std::shared_ptr<TvInputCallback> mTvInputCallback;
+    std::shared_ptr<TvInputCallbackWrapper> mTvInputCallback;
 };
 
 } // namespace android
diff --git a/services/core/jni/tvinput/TvInputHal_hidl.cpp b/services/core/jni/tvinput/TvInputHal_hidl.cpp
index 37cf8445920f..cdd926622dcc 100644
--- a/services/core/jni/tvinput/TvInputHal_hidl.cpp
+++ b/services/core/jni/tvinput/TvInputHal_hidl.cpp
@@ -59,7 +59,11 @@ JTvInputHal::TvInputEventWrapper JTvInputHal::TvInputEventWrapper::createEventWr
     return event;
 }
 
-Return<void> JTvInputHal::TvInputCallback::notify(const HidlTvInputEvent& event) {
+JTvInputHal::HidlTvInputCallback::HidlTvInputCallback(JTvInputHal* hal) {
+    mHal = hal;
+}
+
+Return<void> JTvInputHal::HidlTvInputCallback::notify(const HidlTvInputEvent& event) {
     mHal->mLooper->sendMessage(new NotifyHandler(mHal,
                                                  TvInputEventWrapper::createEventWrapper(event)),
                                static_cast<int>(event.type));
@@ -70,9 +74,8 @@ JTvInputHal::ITvInputWrapper::ITvInputWrapper(sp<HidlITvInput>& hidlTvInput)
       : mIsHidl(true), mHidlTvInput(hidlTvInput) {}
 
 ::ndk::ScopedAStatus JTvInputHal::ITvInputWrapper::hidlSetCallback(
-        const std::shared_ptr<TvInputCallback>& in_callback) {
-    mHidlTvInput->setCallback(in_callback == nullptr ? nullptr
-                                                     : sp<TvInputCallback>(in_callback.get()));
+        const sp<HidlTvInputCallback>& in_callback) {
+    mHidlTvInput->setCallback(in_callback);
     return ::ndk::ScopedAStatus::ok();
 }