From 6bba9263f9e2bea39803470d856fb7a14dac4b43 Mon Sep 17 00:00:00 2001
From: Heikki Tampio <heikki.tampio@tuni.fi>
Date: Wed, 28 Jun 2023 11:08:23 +0300
Subject: [PATCH] multiplex: Improve thread safety in reception_flow

---
 src/reception_flow.cc | 31 ++++++++++++++++++++-----------
 src/reception_flow.hh |  1 +
 2 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/src/reception_flow.cc b/src/reception_flow.cc
index 73dacde..3fd00e4 100644
--- a/src/reception_flow.cc
+++ b/src/reception_flow.cc
@@ -247,9 +247,12 @@ uvgrtp::frame::rtp_frame* uvgrtp::reception_flow::pull_frame(std::shared_ptr<std
     frames_mtx_.lock();
     if (!frames_.empty()) {
         frame = frames_.front();
-        if (frame->header.ssrc == remote_ssrc.get()->load()) {
+        if ( frame && frame->header.ssrc == remote_ssrc.get()->load()) {
             frames_.erase(frames_.begin());
         }
+        else {
+            frame = nullptr;
+        }
     }
     frames_mtx_.unlock();
     return frame;
@@ -269,23 +272,25 @@ uvgrtp::frame::rtp_frame* uvgrtp::reception_flow::pull_frame(ssize_t timeout_ms,
     if (should_stop_ || frames_.empty())
         return nullptr;
     // Check if the source ssrc in the frame matches the remote ssrc that we want to pull frames from
-    bool found_frame = false;
+    uvgrtp::frame::rtp_frame* frame = nullptr;
     frames_mtx_.lock();
-    auto frame = frames_.front();
-    if (frame->header.ssrc == remote_ssrc.get()->load()) {
-        frames_.pop_front();
-        found_frame = true;
+    if (!frames_.empty()) {
+        frame = frames_.front();
+        if (frame && frame->header.ssrc == remote_ssrc.get()->load()) {
+            frames_.pop_front();
+        }
+        else {
+            frame = nullptr;
+        }
     }
     frames_mtx_.unlock();
-    if (found_frame) {
-        return frame;
-    }
-    return nullptr;
+    return frame;
 }
 
 rtp_error_t uvgrtp::reception_flow::new_install_handler(int type, std::shared_ptr<std::atomic<std::uint32_t>> remote_ssrc,
     std::function<rtp_error_t(void*, int, uint8_t*, size_t, frame::rtp_frame** out)> handler, void* args)
 {
+    handlers_mutex_.lock();
     switch (type) {
         case 1: {
             packet_handlers_[remote_ssrc].rtp.handler = handler;
@@ -322,13 +327,16 @@ rtp_error_t uvgrtp::reception_flow::new_install_handler(int type, std::shared_pt
             break;
         }
     }
+    handlers_mutex_.unlock();
     return RTP_OK;
 }
 
 rtp_error_t uvgrtp::reception_flow::new_install_getter(std::shared_ptr<std::atomic<std::uint32_t>> remote_ssrc,
     std::function<rtp_error_t(uvgrtp::frame::rtp_frame**)> getter)
 {
+    handlers_mutex_.lock();
     packet_handlers_[remote_ssrc].getter = getter;
+    handlers_mutex_.unlock();
     return RTP_OK;
 }
 
@@ -702,10 +710,11 @@ int uvgrtp::reception_flow::clear_stream_from_flow(std::shared_ptr<std::atomic<s
     if (hooks_.find(remote_ssrc) != hooks_.end()) {
         hooks_.erase(remote_ssrc);
     }
+    handlers_mutex_.lock();
     if (packet_handlers_.find(remote_ssrc) != packet_handlers_.end()) {
         packet_handlers_.erase(remote_ssrc);
     }
-
+    handlers_mutex_.unlock();
     // If all the data structures are empty, return 1 which means that there is no streams left for this reception_flow
     // and it can be safely deleted
     if (hooks_.empty() && packet_handlers_.empty()) {
diff --git a/src/reception_flow.hh b/src/reception_flow.hh
index 5dc2be6..85e22bb 100644
--- a/src/reception_flow.hh
+++ b/src/reception_flow.hh
@@ -205,6 +205,7 @@ namespace uvgrtp {
             std::map<std::shared_ptr<std::atomic<std::uint32_t>>, handler> packet_handlers_;
 
             std::vector<Buffer> ring_buffer_;
+            std::mutex handlers_mutex_;
             std::mutex ring_mutex_;
             // these uphold the ring buffer details
             std::atomic<ssize_t> ring_read_index_;