From 8154bddd53e8db762a073ac6c4f264bb863e7fcc Mon Sep 17 00:00:00 2001
From: fwx1246009 <fuzijian5@huawei.com>
Date: Fri, 30 Jun 2023 10:31:29 +0800
Subject: [PATCH] Solve the UAF problem

Issue: I7GF0Q
Test: UT

Signed-off-by: fwx1246009 <fuzijian5@huawei.com>
Change-Id: Ifa122b3846d2e64295a90f81df46bf6c63158c6a
---
 base/src/event_demultiplexer.cpp | 45 ++++++++++++++++++++------------
 base/src/event_demultiplexer.h   |  3 ++-
 base/src/event_handler.h         |  4 ++-
 base/src/timer.cpp               |  8 ------
 base/src/timer_event_handler.h   |  2 +-
 5 files changed, 35 insertions(+), 27 deletions(-)

diff --git a/base/src/event_demultiplexer.cpp b/base/src/event_demultiplexer.cpp
index 87df8a7..8a578b1 100644
--- a/base/src/event_demultiplexer.cpp
+++ b/base/src/event_demultiplexer.cpp
@@ -70,7 +70,7 @@ uint32_t EventDemultiplexer::UpdateEventHandler(EventHandler* handler)
     std::lock_guard<std::recursive_mutex> lock(mutex_);
     auto itor = eventHandlers_.find(handler->GetHandle());
     if (itor == eventHandlers_.end()) {
-        eventHandlers_.insert(std::make_pair(handler->GetHandle(), handler));
+        eventHandlers_.insert(std::make_pair(handler->GetHandle(), handler->shared_from_this()));
         return Update(EPOLL_CTL_ADD, handler);
     }
 
@@ -79,7 +79,7 @@ uint32_t EventDemultiplexer::UpdateEventHandler(EventHandler* handler)
         return Update(EPOLL_CTL_DEL, handler);
     }
 
-    if (handler != itor->second) {
+    if (handler != itor->second.get()) {
         return TIMER_ERR_DEAL_FAILED;
     }
     return Update(EPOLL_CTL_MOD, handler);
@@ -102,23 +102,36 @@ uint32_t EventDemultiplexer::Update(int operation, EventHandler* handler)
 
 void EventDemultiplexer::Polling(int timeout /* ms */)
 {
+    std::vector<std::shared_ptr<EventHandler>> taskque;
     std::vector<struct epoll_event> epollEvents(maxEvents_);
-    int nfds = epoll_wait(epollFd_, &epollEvents[0], static_cast<int>(epollEvents.size()), timeout);
-    if (nfds == 0) {
-        return;
-    }
-    if (nfds == -1) {
-        UTILS_LOGE("epoll_wait failed.");
-        return;
+    int nfds = 0;
+    {
+        std::lock_guard<std::recursive_mutex> lock(mutex_);
+        if (eventHandlers_.size() == 0) {
+            return;
+        }
+
+        nfds = epoll_wait(epollFd_, &epollEvents[0], static_cast<int>(epollEvents.size()), timeout);
+        if (nfds == 0) {
+            return;
+        }
+        if (nfds == -1) {
+            UTILS_LOGE("epoll_wait failed.");
+            return;
+        }
+
+        for (int idx = 0; idx < nfds; ++idx) {
+            void* ptr = epollEvents[idx].data.ptr;
+            auto handler = reinterpret_cast<EventHandler*>(ptr);
+            if (handler != nullptr) {
+                taskque.emplace_back(handler->shared_from_this());
+            }
+        }
     }
 
-    for (int idx = 0; idx < nfds; ++idx) {
+    for (int idx = 0; idx < nfds && idx < taskque.size(); ++idx) {
         uint32_t events = epollEvents[idx].events;
-        void* ptr = epollEvents[idx].data.ptr;
-        auto handler = reinterpret_cast<EventHandler*>(ptr);
-        if (handler != nullptr) {
-            handler->HandleEvents(Epoll2Reactor(events));
-        }
+        taskque[idx]->HandleEvents(Epoll2Reactor(events));
     }
 
     if (nfds == maxEvents_) {
@@ -149,4 +162,4 @@ uint32_t EventDemultiplexer::Reactor2Epoll(uint32_t reactorEvent)
 }
 
 }
-}
+}
\ No newline at end of file
diff --git a/base/src/event_demultiplexer.h b/base/src/event_demultiplexer.h
index 5c25e56..13fe966 100644
--- a/base/src/event_demultiplexer.h
+++ b/base/src/event_demultiplexer.h
@@ -17,6 +17,7 @@
 #define UTILS_EVENT_DEMULTIPLEXER_H
 
 #include <mutex>
+#include <memory>
 #include <cstdint>
 #include <map>
 
@@ -47,7 +48,7 @@ private:
     int epollFd_;
     int maxEvents_;
     std::recursive_mutex mutex_;
-    std::map<int, EventHandler*> eventHandlers_; // guard by mutex_
+    std::map<int, std::shared_ptr<EventHandler>> eventHandlers_; // guard by mutex_
 };
 
 }
diff --git a/base/src/event_handler.h b/base/src/event_handler.h
index af75cbe..511bb6f 100644
--- a/base/src/event_handler.h
+++ b/base/src/event_handler.h
@@ -19,13 +19,15 @@
 #include <cstdint>
 #include <map>
 #include <functional>
+#include <unistd.h>
+#include <memory>
 
 namespace OHOS {
 namespace Utils {
 
 class EventReactor;
 
-class EventHandler {
+class EventHandler : public std::enable_shared_from_this<EventHandler> {
 public:
     using Callback = std::function<void()>;
 
diff --git a/base/src/timer.cpp b/base/src/timer.cpp
index 4fc9a1d..bb4db90 100644
--- a/base/src/timer.cpp
+++ b/base/src/timer.cpp
@@ -49,14 +49,6 @@ void Timer::Shutdown(bool useJoin)
     }
 
     reactor_->SwitchOff();
-    if (timeoutMs_ == -1) {
-        std::lock_guard<std::mutex> lock(mutex_);
-        if (intervalToTimers_.empty()) {
-            UTILS_LOGI("no event for epoll wait, use detach to shutdown");
-            thread_.detach();
-            return;
-        }
-    }
     if (!useJoin) {
         thread_.detach();
         return;
diff --git a/base/src/timer_event_handler.h b/base/src/timer_event_handler.h
index 3201480..1fd0bd6 100644
--- a/base/src/timer_event_handler.h
+++ b/base/src/timer_event_handler.h
@@ -55,7 +55,7 @@ private:
     uint32_t       interval_;
     EventReactor*  reactor_;
 
-    std::unique_ptr<EventHandler> handler_;
+    std::shared_ptr<EventHandler> handler_;
     TimerCallback                 callback_;
 };
 
-- 
Gitee