From 318f2219309750e38e7521d924a3c6458322e6af Mon Sep 17 00:00:00 2001
From: Wataru Ishida <ishida@nel-america.com>
Date: Tue, 11 Jun 2019 15:21:36 -0700
Subject: [PATCH] mux: fix race condition of notification oid conversion

Signed-off-by: Wataru Ishida <ishida@nel-america.com>
---
 tai_mux/mux.cpp | 26 +++++++++++++++++++-------
 tai_mux/mux.hpp |  9 ++++++++-
 2 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/tai_mux/mux.cpp b/tai_mux/mux.cpp
index 37487cb..03b7b88 100644
--- a/tai_mux/mux.cpp
+++ b/tai_mux/mux.cpp
@@ -247,7 +247,6 @@ void Multiplexier::notify(notification_context* ctx, tai_object_id_t real_oid, t
     if ( oid == TAI_NULL_OBJECT_ID ) {
         return;
     }
-    auto key = std::pair<tai_object_id_t, tai_attr_id_t>(oid, ctx->notify_id);
     tai_attribute_t dst;
     tai_alloc_info_t info;
     info.reference = src;
@@ -272,7 +271,10 @@ void Multiplexier::notify(notification_context* ctx, tai_object_id_t real_oid, t
         goto err;
     }
 
-    ctx->handler.notify(ctx->handler.context, oid, &dst);
+    {
+        std::unique_lock<std::mutex> lk(ctx->mutex);
+        ctx->handler.notify(ctx->handler.context, oid, &dst);
+    }
 err:
     tai_metadata_free_attr_value(meta, &dst, nullptr);
 }
@@ -292,6 +294,7 @@ tai_status_t Multiplexier::set_attributes( std::function<tai_status_t(ModuleAdap
     std::vector<tai_attribute_t> attrs;
     tai_alloc_info_t info;
     auto t = object_type_query(oid);
+    std::vector<notification_key> n_to_clear;
     for ( auto i = 0; i < attr_count; i++ ) {
         auto meta = tai_metadata_get_attr_metadata(t, attr_list[i].id);
         if ( meta == nullptr ) {
@@ -312,8 +315,19 @@ tai_status_t Multiplexier::set_attributes( std::function<tai_status_t(ModuleAdap
             goto err;
         }
         attrs.emplace_back(attr);
+        // we can't delete m_notification_map[key] before calling f() since
+        // the adapter can still call the notification callback.
+        // store the keys to be removed into n_to_clear and delete them after calling f()
+        if ( meta->attrvaluetype == TAI_ATTR_VALUE_TYPE_NOTIFICATION && attr.value.notification.notify == nullptr ) {
+            n_to_clear.emplace_back(notification_key(oid, attr.id));
+        }
     }
     ret = f(m_adapter, id, attrs.size(), attrs.data());
+
+    for ( auto& key : n_to_clear ) {
+        delete m_notification_map[key];
+        m_notification_map.erase(key);
+    }
 err:
     if ( free_attributes(t, attrs) < 0 ) {
         return TAI_STATUS_FAILURE;
@@ -365,7 +379,7 @@ tai_status_t Multiplexier::convert_oid(tai_object_id_t oid, const tai_attribute_
         }
         return oid;
     };
-    auto key = std::pair<tai_object_id_t, tai_attr_id_t>(oid, src->id);
+    auto key = notification_key(oid, src->id);
 
     switch (meta->attrvaluetype) {
     case TAI_ATTR_VALUE_TYPE_OID:
@@ -405,14 +419,12 @@ tai_status_t Multiplexier::convert_oid(tai_object_id_t oid, const tai_attribute_
                 dst->value.notification.notify = n->handler.notify;
             }
         } else {
-            if ( src->value.notification.notify == nullptr ) {
-                delete m_notification_map[key];
-                m_notification_map.erase(key);
-            } else {
+            if ( src->value.notification.notify != nullptr ) {
                 if ( m_notification_map.find(key) == m_notification_map.end() ) {
                     m_notification_map[key] = new notification_context();
                 }
                 auto n = m_notification_map[key];
+                std::unique_lock<std::mutex> lk(n->mutex);
                 n->mux = this;
                 n->adapter = adapter;
                 n->handler = src->value.notification;
diff --git a/tai_mux/mux.hpp b/tai_mux/mux.hpp
index 11b1c72..9324c4b 100644
--- a/tai_mux/mux.hpp
+++ b/tai_mux/mux.hpp
@@ -9,6 +9,7 @@
 #include "tai.h"
 #include <sstream>
 #include <bitset>
+#include <mutex>
 
 class OIDAllocator {
     public:
@@ -36,8 +37,11 @@ struct notification_context {
     ModuleAdapter* adapter;
     tai_notification_handler_t handler;
     tai_attr_id_t notify_id;
+    std::mutex mutex;
 };
 
+using notification_key = std::pair<tai_object_id_t, tai_attr_id_t>;
+
 // singleton
 class Multiplexier {
     public:
@@ -57,6 +61,9 @@ class Multiplexier {
             if ( m_pa != nullptr ) {
                 delete m_pa;
             }
+            for ( auto& n : m_notification_map ) {
+                delete n.second;
+            }
         }
 
         ModuleAdapter * get_module_adapter(const std::string& location) {
@@ -133,7 +140,7 @@ class Multiplexier {
         void operator = (const Multiplexier&){}
         PlatformAdapter *m_pa;
         std::map<tai_object_id_t, std::pair<tai_object_id_t, ModuleAdapter* >> m_map;
-        std::map<std::pair<tai_object_id_t, tai_attr_id_t>, notification_context*> m_notification_map;
+        std::map<notification_key, notification_context*> m_notification_map;
 
         OIDAllocator m_oid_allocator;