From aef19c6f9794da3c7be83553812ee88b3790109e Mon Sep 17 00:00:00 2001
From: "Anthony J. Thibault" <tony@highfidelity.io>
Date: Tue, 15 Aug 2017 18:28:51 -0700
Subject: [PATCH] Rig: Use a registry to prevent use after free
 crashes/corruption

Create a global registry to hold all the currently active Rig instances.
Use this registry and it's mutex to prevent accessing the rig after it has already been destroyed,
or is in the process of being destroyed on the main thread.
---
 libraries/animation/src/Rig.cpp | 39 +++++++++++++++++++++++++++++++--
 libraries/animation/src/Rig.h   |  6 +++--
 2 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/libraries/animation/src/Rig.cpp b/libraries/animation/src/Rig.cpp
index 3a31ccd25f..6222850d52 100644
--- a/libraries/animation/src/Rig.cpp
+++ b/libraries/animation/src/Rig.cpp
@@ -32,6 +32,10 @@
 #include "AnimUtil.h"
 #include "IKTarget.h"
 
+static int nextRigId = 1;
+static std::map<int, Rig*> rigRegistry;
+static std::mutex rigRegistryMutex;
+
 static bool isEqual(const glm::vec3& u, const glm::vec3& v) {
     const float EPSILON = 0.0001f;
     return glm::length(u - v) / glm::length(u) <= EPSILON;
@@ -49,6 +53,26 @@ const glm::vec3 DEFAULT_RIGHT_EYE_POS(-0.3f, 0.9f, 0.0f);
 const glm::vec3 DEFAULT_LEFT_EYE_POS(0.3f, 0.9f, 0.0f);
 const glm::vec3 DEFAULT_HEAD_POS(0.0f, 0.75f, 0.0f);
 
+Rig::Rig() {
+    // Ensure thread-safe access to the rigRegistry.
+    std::lock_guard<std::mutex> guard(rigRegistryMutex);
+
+    // Insert this newly allocated rig into the rig registry
+    _rigId = nextRigId;
+    rigRegistry[_rigId] = this;
+    nextRigId++;
+}
+
+Rig::~Rig() {
+    // Ensure thread-safe access to the rigRegstry, but also prevent the rig from being deleted
+    // while Rig::animationStateHandlerResult is being invoked on a script thread.
+    std::lock_guard<std::mutex> guard(rigRegistryMutex);
+    auto iter = rigRegistry.find(_rigId);
+    if (iter != rigRegistry.end()) {
+        rigRegistry.erase(iter);
+    }
+}
+
 void Rig::overrideAnimation(const QString& url, float fps, bool loop, float firstFrame, float lastFrame) {
 
     UserAnimState::ClipNodeEnum clipNodeEnum;
@@ -933,8 +957,19 @@ void Rig::updateAnimationStateHandlers() { // called on avatar update thread (wh
         int identifier = data.key();
         StateHandler& value = data.value();
         QScriptValue& function = value.function;
-        auto handleResult = [this, identifier](QScriptValue result) { // called in script thread to get the result back to us.
-            animationStateHandlerResult(identifier, result);
+        int rigId = _rigId;
+        auto handleResult = [rigId, identifier](QScriptValue result) { // called in script thread to get the result back to us.
+            // Hold the rigRegistryMutex to ensure thread-safe access to the rigRegistry, but
+            // also to prevent the rig from being deleted while this lambda is being executed.
+            std::lock_guard<std::mutex> guard(rigRegistryMutex);
+
+            // if the rig pointer is in the registry, it has not been deleted yet.
+            auto iter = rigRegistry.find(rigId);
+            if (iter != rigRegistry.end()) {
+                Rig* rig = iter->second;
+                assert(rig);
+                rig->animationStateHandlerResult(identifier, result);
+            }
         };
         // invokeMethod makes a copy of the args, and copies of AnimVariantMap do copy the underlying map, so this will correctly capture
         // the state of _animVars and allow continued changes to _animVars in this thread without conflict.
diff --git a/libraries/animation/src/Rig.h b/libraries/animation/src/Rig.h
index ca55635250..eabc62ab75 100644
--- a/libraries/animation/src/Rig.h
+++ b/libraries/animation/src/Rig.h
@@ -97,8 +97,8 @@ public:
         Hover
     };
 
-    Rig() {}
-    virtual ~Rig() {}
+    Rig();
+    virtual ~Rig();
 
     void destroyAnimGraph();
 
@@ -372,6 +372,8 @@ protected:
 
     glm::vec3 _prevLeftHandPoleVector { -Vectors::UNIT_Z };
     bool _prevLeftHandPoleVectorValid { false };
+
+    int _rigId;
 };
 
 #endif /* defined(__hifi__Rig__) */