From 55667a90586546bfaf0ee7a1f2a1443032c6397d Mon Sep 17 00:00:00 2001
From: SuRuoyan <suruoyan1@huawei.com>
Date: Wed, 19 Mar 2025 17:33:59 +0800
Subject: [PATCH] fix media codec UAF

Signed-off-by: SuRuoyan <suruoyan1@huawei.com>
---
 .../modules/media_codec/media_codec.cpp             | 13 +++++++++----
 .../media_engine/modules/media_codec/media_codec.h  |  7 +++----
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/services/media_engine/modules/media_codec/media_codec.cpp b/services/media_engine/modules/media_codec/media_codec.cpp
index 2b340c95e..06ae9e2ef 100644
--- a/services/media_engine/modules/media_codec/media_codec.cpp
+++ b/services/media_engine/modules/media_codec/media_codec.cpp
@@ -46,18 +46,23 @@ namespace OHOS {
 namespace Media {
 class InputBufferAvailableListener : public IConsumerListener {
 public:
-    explicit InputBufferAvailableListener(MediaCodec *mediaCodec)
+    explicit InputBufferAvailableListener(std::shared_ptr<MediaCodec> mediaCodec)
     {
         mediaCodec_ = mediaCodec;
     }
 
     void OnBufferAvailable() override
     {
-        mediaCodec_->ProcessInputBuffer();
+        auto realPtr = mediaCodec_.lock();
+        if (realPtr != nullptr) {
+            realPtr->ProcessInputBuffer();
+        } else {
+            MEDIA_LOG_W("mediacodec was released, can not process input buffer");
+        }
     }
 
 private:
-    MediaCodec *mediaCodec_;
+    std::weak_ptr<MediaCodec> mediaCodec_;
 };
 
 MediaCodec::MediaCodec()
@@ -621,7 +626,7 @@ int32_t MediaCodec::PrepareInputBufferQueue()
     }
     FALSE_RETURN_V_MSG_E(inputBufferQueue_ != nullptr, (int32_t)Status::ERROR_UNKNOWN, "inputBufferQueue_ is nullptr");
     inputBufferQueueConsumer_ = inputBufferQueue_->GetConsumer();
-    sptr<IConsumerListener> listener = new InputBufferAvailableListener(this);
+    sptr<IConsumerListener> listener = new InputBufferAvailableListener(shared_from_this());
     FALSE_RETURN_V_MSG_E(inputBufferQueueConsumer_ != nullptr, (int32_t)Status::ERROR_UNKNOWN,
                          "inputBufferQueueConsumer_ is nullptr");
     inputBufferQueueConsumer_->SetBufferAvailableListener(listener);
diff --git a/services/media_engine/modules/media_codec/media_codec.h b/services/media_engine/modules/media_codec/media_codec.h
index 85caa9708..f61237885 100644
--- a/services/media_engine/modules/media_codec/media_codec.h
+++ b/services/media_engine/modules/media_codec/media_codec.h
@@ -77,7 +77,7 @@ public:
     virtual void OnOutputFormatChanged(const std::shared_ptr<Meta> &format) = 0;
 };
 
-class MediaCodec : public Plugins::DataCallback {
+class MediaCodec : public std::enable_shared_from_this<MediaCodec>, public Plugins::DataCallback {
 public:
     MediaCodec();
 
@@ -89,8 +89,7 @@ public:
 
     int32_t Configure(const std::shared_ptr<Meta> &meta);
 
-    __attribute__((no_sanitize("cfi"))) int32_t SetOutputBufferQueue(
-        const sptr<AVBufferQueueProducer> &bufferQueueProducer);
+    int32_t SetOutputBufferQueue(const sptr<AVBufferQueueProducer> &bufferQueueProducer);
 
     int32_t SetCodecCallback(const std::shared_ptr<CodecCallback> &codecCallback);
 
@@ -144,7 +143,7 @@ private:
     Status AttachDrmBufffer(std::shared_ptr<AVBuffer> &drmInbuf, std::shared_ptr<AVBuffer> &drmOutbuf,
         uint32_t size);
     Status DrmAudioCencDecrypt(std::shared_ptr<AVBuffer> &filledInputBuffer);
-    __attribute__((no_sanitize("cfi"))) Status HandleOutputBuffer(uint32_t eosStatus);
+    Status HandleOutputBuffer(uint32_t eosStatus);
 
     int32_t PrepareInputBufferQueue();
 
-- 
Gitee