Use some WeakSymbolicVariable to avoid memleak

JonathanSalwan · Sep 24, 2019 · 6441891 · 6441891
1 parent 16bc17d
commit 6441891
Show file tree

Hide file tree

Showing 5 changed files with 61 additions and 37 deletions.
diff --git a/src/libtriton/api/api.cpp b/src/libtriton/api/api.cpp
@@ -635,19 +635,19 @@ namespace triton {
   }
 
 
-  const triton::engines::symbolic::SharedSymbolicVariable& API::convertExpressionToSymbolicVariable(triton::usize exprId, triton::uint32 symVarSize, const std::string& symVarComment) {
+  triton::engines::symbolic::SharedSymbolicVariable API::convertExpressionToSymbolicVariable(triton::usize exprId, triton::uint32 symVarSize, const std::string& symVarComment) {
     this->checkSymbolic();
     return this->symbolic->convertExpressionToSymbolicVariable(exprId, symVarSize, symVarComment);
   }
 
 
-  const triton::engines::symbolic::SharedSymbolicVariable& API::convertMemoryToSymbolicVariable(const triton::arch::MemoryAccess& mem, const std::string& symVarComment) {
+  triton::engines::symbolic::SharedSymbolicVariable API::convertMemoryToSymbolicVariable(const triton::arch::MemoryAccess& mem, const std::string& symVarComment) {
     this->checkSymbolic();
     return this->symbolic->convertMemoryToSymbolicVariable(mem, symVarComment);
   }
 
 
-  const triton::engines::symbolic::SharedSymbolicVariable& API::convertRegisterToSymbolicVariable(const triton::arch::Register& reg, const std::string& symVarComment) {
+  triton::engines::symbolic::SharedSymbolicVariable API::convertRegisterToSymbolicVariable(const triton::arch::Register& reg, const std::string& symVarComment) {
     this->checkSymbolic();
     return this->symbolic->convertRegisterToSymbolicVariable(reg, symVarComment);
   }
@@ -707,7 +707,7 @@ namespace triton {
   }
 
 
-  const triton::engines::symbolic::SharedSymbolicVariable& API::newSymbolicVariable(triton::uint32 varSize, const std::string& comment) {
+  triton::engines::symbolic::SharedSymbolicVariable API::newSymbolicVariable(triton::uint32 varSize, const std::string& comment) {
     this->checkSymbolic();
     return this->symbolic->newSymbolicVariable(triton::engines::symbolic::UNDEFINED_VARIABLE, 0, varSize, comment);
   }
@@ -837,13 +837,13 @@ namespace triton {
   }
 
 
-  const triton::engines::symbolic::SharedSymbolicVariable& API::getSymbolicVariableFromId(triton::usize symVarId) const {
+  triton::engines::symbolic::SharedSymbolicVariable API::getSymbolicVariableFromId(triton::usize symVarId) const {
     this->checkSymbolic();
     return this->symbolic->getSymbolicVariableFromId(symVarId);
   }
 
 
-  const triton::engines::symbolic::SharedSymbolicVariable& API::getSymbolicVariableFromName(const std::string& symVarName) const {
+  triton::engines::symbolic::SharedSymbolicVariable API::getSymbolicVariableFromName(const std::string& symVarName) const {
     this->checkSymbolic();
     return this->symbolic->getSymbolicVariableFromName(symVarName);
   }
@@ -957,7 +957,7 @@ namespace triton {
   }
 
 
-  const std::unordered_map<triton::usize, triton::engines::symbolic::SharedSymbolicVariable>& API::getSymbolicVariables(void) const {
+  std::unordered_map<triton::usize, triton::engines::symbolic::SharedSymbolicVariable> API::getSymbolicVariables(void) const {
     this->checkSymbolic();
     return this->symbolic->getSymbolicVariables();
   }

diff --git a/src/libtriton/engines/symbolic/symbolicEngine.cpp b/src/libtriton/engines/symbolic/symbolicEngine.cpp
@@ -204,32 +204,53 @@ namespace triton {
 
 
       /* Returns the symbolic variable otherwise raises an exception */
-      const SharedSymbolicVariable& SymbolicEngine::getSymbolicVariableFromId(triton::usize symVarId) const {
+      SharedSymbolicVariable SymbolicEngine::getSymbolicVariableFromId(triton::usize symVarId) const {
         auto it = this->symbolicVariables.find(symVarId);
         if (it == this->symbolicVariables.end())
-          throw triton::exceptions::SymbolicEngine("SymbolicEngine::getSymbolicVariableFromId(): Unregistred variable.");
-        return it->second;
+          throw triton::exceptions::SymbolicEngine("SymbolicEngine::getSymbolicVariableFromId(): Unregistred symbolic variable.");
+
+        if (auto node = it->second.lock())
+          return node;
+        else
+          throw triton::exceptions::SymbolicEngine("SymbolicEngine::getSymbolicVariableFromId(): This symbolic variable is dead.");
       }
 
 
       /* Returns the symbolic variable otherwise returns nullptr */
-      const SharedSymbolicVariable& SymbolicEngine::getSymbolicVariableFromName(const std::string& symVarName) const {
+      SharedSymbolicVariable SymbolicEngine::getSymbolicVariableFromName(const std::string& symVarName) const {
         /*
          * FIXME: When there is a ton of symvar, this loop takes a while to go through.
          *        What about adding two maps {id:symvar} and {string:symvar}? See #648.
          */
         for (auto& sv: this->symbolicVariables) {
-          if (sv.second->getName() == symVarName)
-            return sv.second;
+          if (auto symVar = sv.second.lock()) {
+            if (symVar->getName() == symVarName) {
+              return symVar;
+            }
+          }
         }
-
-        throw triton::exceptions::SymbolicEngine("SymbolicEngine::getSymbolicVariableFromName(): Unregistred variable.");
+        throw triton::exceptions::SymbolicEngine("SymbolicEngine::getSymbolicVariableFromName(): Unregistred or dead symbolic variable.");
       }
 
 
       /* Returns all symbolic variables */
-      const std::unordered_map<triton::usize, SharedSymbolicVariable>& SymbolicEngine::getSymbolicVariables(void) const {
-        return this->symbolicVariables;
+      std::unordered_map<triton::usize, SharedSymbolicVariable> SymbolicEngine::getSymbolicVariables(void) const {
+        // Copy and clean up dead weak ref
+        std::unordered_map<triton::usize, SharedSymbolicVariable> ret;
+        std::vector<triton::usize> toRemove;
+
+        for (auto& kv : this->symbolicVariables) {
+          if (auto sp = kv.second.lock()) {
+            ret[kv.first] = sp;
+          } else {
+            toRemove.push_back(kv.first);
+          }
+        }
+
+        for (auto id : toRemove)
+          this->symbolicVariables.erase(id);
+
+        return ret;
       }
 
 
@@ -444,7 +465,7 @@ namespace triton {
        * convertExpressionToSymbolicVariable(43, 8)
        * #43 = SymVar_4
        */
-      const SharedSymbolicVariable& SymbolicEngine::convertExpressionToSymbolicVariable(triton::usize exprId, triton::uint32 symVarSize, const std::string& symVarComment) {
+      SharedSymbolicVariable SymbolicEngine::convertExpressionToSymbolicVariable(triton::usize exprId, triton::uint32 symVarSize, const std::string& symVarComment) {
         const SharedSymbolicExpression& expression = this->getSymbolicExpressionFromId(exprId);
         const SharedSymbolicVariable& symVar       = this->newSymbolicVariable(UNDEFINED_VARIABLE, 0, symVarSize, symVarComment);
         const triton::ast::SharedAbstractNode& tmp = this->astCtxt->variable(symVar);
@@ -459,7 +480,7 @@ namespace triton {
 
 
       /* The memory size is used to define the symbolic variable's size. */
-      const SharedSymbolicVariable& SymbolicEngine::convertMemoryToSymbolicVariable(const triton::arch::MemoryAccess& mem, const std::string& symVarComment) {
+      SharedSymbolicVariable SymbolicEngine::convertMemoryToSymbolicVariable(const triton::arch::MemoryAccess& mem, const std::string& symVarComment) {
         triton::uint64 memAddr          = mem.getAddress();
         triton::uint32 symVarSize       = mem.getSize();
         triton::uint512 cv              = this->architecture->getConcreteMemoryValue(mem);
@@ -505,7 +526,7 @@ namespace triton {
       }
 
 
-      const SharedSymbolicVariable& SymbolicEngine::convertRegisterToSymbolicVariable(const triton::arch::Register& reg, const std::string& symVarComment) {
+      SharedSymbolicVariable SymbolicEngine::convertRegisterToSymbolicVariable(const triton::arch::Register& reg, const std::string& symVarComment) {
         const triton::arch::Register& parent  = this->architecture->getRegister(reg.getParent());
         triton::uint32 symVarSize             = reg.getBitSize();
         triton::uint512 cv                    = this->architecture->getConcreteRegisterValue(reg);
@@ -543,15 +564,15 @@ namespace triton {
 
 
       /* Adds a new symbolic variable */
-      const SharedSymbolicVariable& SymbolicEngine::newSymbolicVariable(triton::engines::symbolic::variable_e type, triton::uint64 origin, triton::uint32 size, const std::string& comment) {
+      SharedSymbolicVariable SymbolicEngine::newSymbolicVariable(triton::engines::symbolic::variable_e type, triton::uint64 origin, triton::uint32 size, const std::string& comment) {
         triton::usize uniqueId = this->getUniqueSymVarId();
 
         SharedSymbolicVariable symVar = std::make_shared<SymbolicVariable>(type, origin, uniqueId, size, comment);
         if (symVar == nullptr)
           throw triton::exceptions::SymbolicEngine("SymbolicEngine::newSymbolicVariable(): Cannot allocate a new symbolic variable");
 
         this->symbolicVariables[uniqueId] = symVar;
-        return this->symbolicVariables[uniqueId];
+        return symVar;
       }
 
 

diff --git a/src/libtriton/includes/triton/api.hpp b/src/libtriton/includes/triton/api.hpp
@@ -337,13 +337,13 @@ namespace triton {
         TRITON_EXPORT triton::uint512 getSymbolicRegisterValue(const triton::arch::Register& reg);
 
         //! [**symbolic api**] - Converts a symbolic expression to a symbolic variable. `symVarSize` must be in bits.
-        TRITON_EXPORT const triton::engines::symbolic::SharedSymbolicVariable& convertExpressionToSymbolicVariable(triton::usize exprId, triton::uint32 symVarSize, const std::string& symVarComment="");
+        TRITON_EXPORT triton::engines::symbolic::SharedSymbolicVariable convertExpressionToSymbolicVariable(triton::usize exprId, triton::uint32 symVarSize, const std::string& symVarComment="");
 
         //! [**symbolic api**] - Converts a symbolic memory expression to a symbolic variable.
-        TRITON_EXPORT const triton::engines::symbolic::SharedSymbolicVariable& convertMemoryToSymbolicVariable(const triton::arch::MemoryAccess& mem, const std::string& symVarComment="");
+        TRITON_EXPORT triton::engines::symbolic::SharedSymbolicVariable convertMemoryToSymbolicVariable(const triton::arch::MemoryAccess& mem, const std::string& symVarComment="");
 
         //! [**symbolic api**] - Converts a symbolic register expression to a symbolic variable.
-        TRITON_EXPORT const triton::engines::symbolic::SharedSymbolicVariable& convertRegisterToSymbolicVariable(const triton::arch::Register& reg, const std::string& symVarComment="");
+        TRITON_EXPORT triton::engines::symbolic::SharedSymbolicVariable convertRegisterToSymbolicVariable(const triton::arch::Register& reg, const std::string& symVarComment="");
 
         //! [**symbolic api**] - Returns the AST corresponding to the operand.
         TRITON_EXPORT triton::ast::SharedAbstractNode getOperandAst(const triton::arch::OperandWrapper& op);
@@ -373,7 +373,7 @@ namespace triton {
         TRITON_EXPORT triton::engines::symbolic::SharedSymbolicExpression newSymbolicExpression(const triton::ast::SharedAbstractNode& node, const std::string& comment="");
 
         //! [**symbolic api**] - Returns a new symbolic variable.
-        TRITON_EXPORT const triton::engines::symbolic::SharedSymbolicVariable& newSymbolicVariable(triton::uint32 varSize, const std::string& comment="");
+        TRITON_EXPORT triton::engines::symbolic::SharedSymbolicVariable newSymbolicVariable(triton::uint32 varSize, const std::string& comment="");
 
         //! [**symbolic api**] - Removes the symbolic expression corresponding to the id.
         TRITON_EXPORT void removeSymbolicExpression(triton::usize symExprId);
@@ -406,10 +406,10 @@ namespace triton {
         TRITON_EXPORT triton::engines::symbolic::SharedSymbolicExpression getSymbolicExpressionFromId(triton::usize symExprId) const;
 
         //! [**symbolic api**] - Returns the symbolic variable corresponding to the symbolic variable id.
-        TRITON_EXPORT const triton::engines::symbolic::SharedSymbolicVariable& getSymbolicVariableFromId(triton::usize symVarId) const;
+        TRITON_EXPORT triton::engines::symbolic::SharedSymbolicVariable getSymbolicVariableFromId(triton::usize symVarId) const;
 
         //! [**symbolic api**] - Returns the symbolic variable corresponding to the symbolic variable name.
-        TRITON_EXPORT const triton::engines::symbolic::SharedSymbolicVariable& getSymbolicVariableFromName(const std::string& symVarName) const;
+        TRITON_EXPORT triton::engines::symbolic::SharedSymbolicVariable getSymbolicVariableFromName(const std::string& symVarName) const;
 
         //! [**symbolic api**] - Returns the logical conjunction vector of path constraints.
         TRITON_EXPORT const std::vector<triton::engines::symbolic::PathConstraint>& getPathConstraints(void) const;
@@ -466,7 +466,7 @@ namespace triton {
         TRITON_EXPORT std::unordered_map<triton::usize, triton::engines::symbolic::SharedSymbolicExpression> getSymbolicExpressions(void) const;
 
         //! [**symbolic api**] - Returns all symbolic variables as a map of <SymVarId : SymVar>
-        TRITON_EXPORT const std::unordered_map<triton::usize, triton::engines::symbolic::SharedSymbolicVariable>& getSymbolicVariables(void) const;
+        TRITON_EXPORT std::unordered_map<triton::usize, triton::engines::symbolic::SharedSymbolicVariable> getSymbolicVariables(void) const;
 
         //! [**symbolic api**] - Gets the concrete value of a symbolic variable.
         TRITON_EXPORT const triton::uint512& getConcreteVariableValue(const triton::engines::symbolic::SharedSymbolicVariable& symVar) const;

diff --git a/src/libtriton/includes/triton/symbolicEngine.hpp b/src/libtriton/includes/triton/symbolicEngine.hpp
@@ -78,7 +78,7 @@ namespace triton {
            * **item1**: variable id<br>
            * **item2**: symbolic variable
            */
-          std::unordered_map<triton::usize, SharedSymbolicVariable> symbolicVariables;
+          mutable std::unordered_map<triton::usize, WeakSymbolicVariable> symbolicVariables;
 
           /*! \brief The map of symbolic expressions
            *
@@ -167,22 +167,22 @@ namespace triton {
           TRITON_EXPORT void removeSymbolicExpression(triton::usize symExprId);
 
           //! Adds a symbolic variable.
-          TRITON_EXPORT const SharedSymbolicVariable& newSymbolicVariable(triton::engines::symbolic::variable_e type, triton::uint64 source, triton::uint32 size, const std::string& comment="");
+          TRITON_EXPORT SharedSymbolicVariable newSymbolicVariable(triton::engines::symbolic::variable_e type, triton::uint64 source, triton::uint32 size, const std::string& comment="");
 
           //! Converts a symbolic expression to a symbolic variable. `symVarSize` must be in bits.
-          TRITON_EXPORT const SharedSymbolicVariable& convertExpressionToSymbolicVariable(triton::usize exprId, triton::uint32 symVarSize, const std::string& symVarComment="");
+          TRITON_EXPORT SharedSymbolicVariable convertExpressionToSymbolicVariable(triton::usize exprId, triton::uint32 symVarSize, const std::string& symVarComment="");
 
           //! Converts a symbolic memory expression to a symbolic variable.
-          TRITON_EXPORT const SharedSymbolicVariable& convertMemoryToSymbolicVariable(const triton::arch::MemoryAccess& mem, const std::string& symVarComment="");
+          TRITON_EXPORT SharedSymbolicVariable convertMemoryToSymbolicVariable(const triton::arch::MemoryAccess& mem, const std::string& symVarComment="");
 
           //! Converts a symbolic register expression to a symbolic variable.
-          TRITON_EXPORT const SharedSymbolicVariable& convertRegisterToSymbolicVariable(const triton::arch::Register& reg, const std::string& symVarComment="");
+          TRITON_EXPORT SharedSymbolicVariable convertRegisterToSymbolicVariable(const triton::arch::Register& reg, const std::string& symVarComment="");
 
           //! Returns the symbolic variable corresponding to the symbolic variable id.
-          TRITON_EXPORT const SharedSymbolicVariable& getSymbolicVariableFromId(triton::usize symVarId) const;
+          TRITON_EXPORT SharedSymbolicVariable getSymbolicVariableFromId(triton::usize symVarId) const;
 
           //! Returns the symbolic variable corresponding to the symbolic variable name.
-          TRITON_EXPORT const SharedSymbolicVariable& getSymbolicVariableFromName(const std::string& symVarName) const;
+          TRITON_EXPORT SharedSymbolicVariable getSymbolicVariableFromName(const std::string& symVarName) const;
 
           //! Returns the symbolic expression corresponding to an id.
           TRITON_EXPORT SharedSymbolicExpression getSymbolicExpressionFromId(triton::usize symExprId) const;
@@ -266,7 +266,7 @@ namespace triton {
           TRITON_EXPORT std::unordered_map<triton::usize, SharedSymbolicExpression> getSymbolicExpressions(void) const;
 
           //! Returns all symbolic variables.
-          TRITON_EXPORT const std::unordered_map<triton::usize, SharedSymbolicVariable>& getSymbolicVariables(void) const;
+          TRITON_EXPORT std::unordered_map<triton::usize, SharedSymbolicVariable> getSymbolicVariables(void) const;
 
           //! Concretizes all symbolic memory references.
           TRITON_EXPORT void concretizeAllMemory(void);

diff --git a/src/libtriton/includes/triton/symbolicVariable.hpp b/src/libtriton/includes/triton/symbolicVariable.hpp
@@ -45,6 +45,9 @@ namespace triton {
       //! Shared Symbolic variable
       using SharedSymbolicVariable = std::shared_ptr<triton::engines::symbolic::SymbolicVariable>;
 
+      //! Weak Symbolic variable
+      using WeakSymbolicVariable = std::weak_ptr<triton::engines::symbolic::SymbolicVariable>;
+
       /*! \class SymbolicVariable
           \brief The symbolic variable class. */
       class SymbolicVariable {