resume the pattern of Alloc function's annotation

SVF-tools · Nov 27, 2023 · 434f138 · 434f138
1 parent 9dc0447
commit 434f138
Show file tree

Hide file tree

Showing 2 changed files with 136 additions and 115 deletions.
diff --git a/svf-llvm/lib/SymbolTableBuilder.cpp b/svf-llvm/lib/SymbolTableBuilder.cpp
@@ -671,10 +671,10 @@ void SymbolTableBuilder::analyzeObjType(ObjTypeInfo* typeinfo, const Value* val)
 
 /*!
  * Analyze byte size of heap alloc function (e.g. malloc/calloc/...)
- * 1) __attribute__((annotate("ALLOC_RET"), annotate("Arg0")))
+ * 1) __attribute__((annotate("ALLOC_RET"), annotate("AllocSize:Arg0")))
      void* safe_malloc(unsigned long size).
      Byte Size is the size(Arg0)
-   2)__attribute__((annotate("ALLOC_RET"), annotate("Arg0"), annotate("Arg1")))
+   2)__attribute__((annotate("ALLOC_RET"), annotate("AllocSize:Arg0*Arg1")))
     char* safecalloc(int a, int b)
     Byte Size is a(Arg0) * b(Arg1)
    3)__attribute__((annotate("ALLOC_RET"), annotate("UNKNOWN")))
@@ -684,44 +684,65 @@ void SymbolTableBuilder::analyzeObjType(ObjTypeInfo* typeinfo, const Value* val)
     otherwise return ByteSize 0
  */
 u32_t SymbolTableBuilder::analyzeHeapAllocByteSize(const Value* val) {
-    if (const llvm::CallInst* callInst = llvm::dyn_cast<llvm::CallInst>(val))
+    if(const llvm::CallInst* callInst = llvm::dyn_cast<llvm::CallInst>(val))
     {
-        if (const llvm::Function* calledFunction = callInst->getCalledFunction())
+        if (const llvm::Function* calledFunction =
+                callInst->getCalledFunction())
         {
             const SVFFunction* svfFunction =
                 LLVMModuleSet::getLLVMModuleSet()->getSVFFunction(
                     calledFunction);
-            std::vector<const llvm::Value*> args;
-            for (const std::string& annot : svfFunction->getAnnotations())
+            std::vector<const Value*> args;
+            // Heap alloc functions have annoation like "AllocSize:Arg1"
+            for (std::string annotation : svfFunction->getAnnotations())
             {
-                if (annot == "UNKNOWN")
+                if (annotation.find("AllocSize:") != std::string::npos)
                 {
-                    return 0;
-                }
-                if (annot.rfind("Arg", 0) == 0)
-                {
-                    u32_t argIndex;
-                    std::istringstream(annot.substr(3)) >> argIndex;
-                    if (argIndex < callInst->getNumOperands() - 1)
+                    std::string allocSize = annotation.substr(10);
+                    std::stringstream ss(allocSize);
+                    std::string token;
+                    // Analyaze annotation string and attract Arg list
+                    while (std::getline(ss, token, '*'))
                     {
-                        args.push_back(callInst->getArgOperand(argIndex));
+                        if (token.rfind("Arg", 0) == 0)
+                        {
+                            u32_t argIndex;
+                            std::istringstream(token.substr(3)) >> argIndex;
+                            if (argIndex < callInst->getNumOperands() - 1)
+                            {
+                                args.push_back(
+                                    callInst->getArgOperand(argIndex));
+                            }
+                        }
                     }
                 }
             }
             u64_t product = 1;
-            for (const llvm::Value* arg : args)
+            if (args.size() > 0)
             {
-                if (const llvm::ConstantInt* constIntArg =
-                        llvm::dyn_cast<llvm::ConstantInt>(arg))
-                {
-                    product *= constIntArg->getZExtValue();
-                }
-                else
+                // for annotations like "AllocSize:Arg0*Arg1"
+                for (const llvm::Value* arg : args)
                 {
-                    return 0;
+                    if (const llvm::ConstantInt* constIntArg =
+                            llvm::dyn_cast<llvm::ConstantInt>(arg))
+                    {
+                        // Multiply the constant Value if all Args are const
+                        product *= constIntArg->getZExtValue();
+                    }
+                    else
+                    {
+                        // if Arg list has non-const value, return 0 to indicate it is non const byte size
+                        return 0;
+                    }
                 }
+                // If all the Args are const, return product
+                return product;
+            }
+            else
+            {
+                // for annotations like "AllocSize:UNKNOWN"
+                return 0;
             }
-            return product;
         }
     }
     // if it is not CallInst or CallInst has no CalledFunction, return 0 to indicate it is non const byte size