[#47] Differentiate TaintMap.Source by call site address.

KeenSecurityLab · Jul 18, 2022 · 4326030 · 4326030
1 parent 952d8ba
commit 4326030
Show file tree

Hide file tree

Showing 6 changed files with 37 additions and 13 deletions.
diff --git a/src/main/java/com/bai/env/Context.java b/src/main/java/com/bai/env/Context.java
@@ -247,7 +247,7 @@ public void prepareMainAbsEnv(AbsEnv absEnv, Function mainFunction) {
         } else {
             offset = entryLocal.getBase();
         }
-        long taints = TaintMap.getTaints(this, GlobalState.eEntryFunction);
+        long taints = TaintMap.getTaints(null, this, GlobalState.eEntryFunction);
         int unit = GlobalState.arch.getDefaultPointerSize();
         for (int i = 0; i < TAINT_ARGV_COUNT; i++) {
             absEnv.set(ALoc.getALoc(entryLocal, offset + ((long) i * unit), unit), KSet.getTop(taints), true);

diff --git a/src/main/java/com/bai/env/TaintMap.java b/src/main/java/com/bai/env/TaintMap.java
@@ -1,5 +1,6 @@
 package com.bai.env;
 
+import ghidra.program.model.address.Address;
 import ghidra.program.model.listing.Function;
 import com.bai.util.Logging;
 import java.util.ArrayList;
@@ -18,14 +19,20 @@ public class TaintMap {
      */
     public static class Source {
 
+        private final Address callSite;
         private final Context context;
         private final Function function;
 
-        public Source(Context context, Function function) {
+        public Source(Address callSite, Context context, Function function) {
+            this.callSite = callSite;
             this.context = context;
             this.function = function;
         }
 
+        public Address getCallSite() {
+            return callSite;
+        }
+
         public Context getContext() {
             return context;
         }
@@ -43,12 +50,12 @@ public boolean equals(Object o) {
                 return false;
             }
             Source source = (Source) o;
-            return Objects.equals(context, source.context) && Objects.equals(function, source.function);
+            return Objects.equals(callSite, source.callSite) && Objects.equals(context, source.context) && Objects.equals(function, source.function);
         }
 
         @Override
         public int hashCode() {
-            return Objects.hash(context, function);
+            return Objects.hash(callSite, context, function);
         }
     }
 
@@ -64,13 +71,13 @@ public static void reset() {
         taintSourceToIdMap.clear();
     }
 
-    protected static int getTaintId(Context context, Function function) {
+    protected static int getTaintId(Address callSite, Context context, Function function) {
         if (taintId >= MAX_TAINT_CNT) {
             Logging.error("Taint id number reach " + MAX_TAINT_CNT
                     + "this may lead to false positive.");
             taintId = taintId % MAX_TAINT_CNT;
         }
-        Source src = new Source(context, function);
+        Source src = new Source(callSite, context, function);
         Integer id = taintSourceToIdMap.get(src);
         if (id != null) {
             return id;
@@ -99,12 +106,13 @@ public static List<Source> getTaintSourceList(long taints) {
 
     /**
      * Get a taint bitmap for a taint source consisting of a context and a function
+     * @param callSite Call site address of the Function component
      * @param context Context component for a taint source
      * @param function Function component for a taint source
      * @return A taint bitmap for the information of a taint source
      */
-    public static long getTaints(Context context, Function function) {
-        return 1L << getTaintId(context, function);
+    public static long getTaints(Address callSite, Context context, Function function) {
+        return 1L << getTaintId(callSite, context, function);
     }
 
     /**

diff --git a/src/main/java/com/bai/env/funcs/externalfuncs/GetenvFunction.java b/src/main/java/com/bai/env/funcs/externalfuncs/GetenvFunction.java
@@ -1,10 +1,13 @@
 package com.bai.env.funcs.externalfuncs;
 
+import static com.bai.util.Utils.getAddress;
+
 import com.bai.env.ALoc;
 import com.bai.env.AbsEnv;
 import com.bai.env.Context;
 import com.bai.env.KSet;
 import com.bai.env.TaintMap;
+import ghidra.program.model.address.Address;
 import ghidra.program.model.data.PointerDataType;
 import ghidra.program.model.listing.Function;
 import ghidra.program.model.pcode.PcodeOp;
@@ -33,7 +36,8 @@ public void invoke(PcodeOp pcode, AbsEnv inOutEnv, AbsEnv tmpEnv, Context contex
         if (retALoc == null) {
             return;
         }
-        long taints = TaintMap.getTaints(context, callFunc);
+        Address callAddress = getAddress(pcode);
+        long taints = TaintMap.getTaints(callAddress, context, callFunc);
         inOutEnv.set(retALoc, KSet.getTop(taints), true);
     }
 }
diff --git a/src/main/java/com/bai/env/funcs/externalfuncs/InputFunctionBase.java b/src/main/java/com/bai/env/funcs/externalfuncs/InputFunctionBase.java
@@ -1,11 +1,14 @@
 package com.bai.env.funcs.externalfuncs;
 
+import static com.bai.util.Utils.getAddress;
+
 import com.bai.env.ALoc;
 import com.bai.env.AbsEnv;
 import com.bai.env.AbsVal;
 import com.bai.env.Context;
 import com.bai.env.KSet;
 import com.bai.env.TaintMap;
+import ghidra.program.model.address.Address;
 import ghidra.program.model.data.IntegerDataType;
 import ghidra.program.model.data.PointerDataType;
 import ghidra.program.model.listing.Function;
@@ -48,17 +51,18 @@ public void invoke(PcodeOp pcode, AbsEnv inOutEnv, AbsEnv tmpEnv, Context contex
         if (retALoc == null) {
             return;
         }
+        Address callAddress = getAddress(pcode);
         if (taintedBufParamIndex == -1) {
             if (isReturnNewTaint) {
-                long newTaints = TaintMap.getTaints(context, callFunc);
+                long newTaints = TaintMap.getTaints(callAddress, context, callFunc);
                 inOutEnv.set(retALoc, KSet.getTop(newTaints), true);
             }
             return;
         }
         KSet res = new KSet(retALoc.getLen() * 8);
         for (ALoc bufALoc : getParamALocs(callFunc, taintedBufParamIndex, inOutEnv)) {
             KSet bufPtrKSet = inOutEnv.get(bufALoc);
-            long newTaints = TaintMap.getTaints(context, callFunc);
+            long newTaints = TaintMap.getTaints(callAddress, context, callFunc);
 
             if (!bufPtrKSet.isNormal()) {
                 bufPtrKSet = KSet.getTop(newTaints);

diff --git a/src/main/java/com/bai/env/funcs/externalfuncs/InputVarArgsFunctionBase.java b/src/main/java/com/bai/env/funcs/externalfuncs/InputVarArgsFunctionBase.java
@@ -1,5 +1,7 @@
 package com.bai.env.funcs.externalfuncs;
 
+import static com.bai.util.Utils.getAddress;
+
 import com.bai.env.ALoc;
 import com.bai.env.AbsEnv;
 import com.bai.env.AbsVal;
@@ -9,6 +11,7 @@
 import com.bai.env.region.Reg;
 import com.bai.util.GlobalState;
 import com.bai.util.Utils;
+import ghidra.program.model.address.Address;
 import ghidra.program.model.data.ParameterDefinition;
 import ghidra.program.model.lang.PrototypeModel;
 import ghidra.program.model.listing.Function;
@@ -81,7 +84,8 @@ private void taintVarnodeWithTop(Varnode varnode, AbsEnv absEnv, long taints) {
 
     public void invoke(PcodeOp pcodeOp, AbsEnv inOutEnv, AbsEnv tmpEnv, Context context, Function callFunc) {
         super.invoke(pcodeOp, inOutEnv, tmpEnv, context, callFunc);
-        long newTaints = TaintMap.getTaints(context, callFunc);
+        Address callAddress = getAddress(pcodeOp);
+        long newTaints = TaintMap.getTaints(callAddress, context, callFunc);
         taintVarArgs(pcodeOp, inOutEnv, callFunc, newTaints);
     }
 }
diff --git a/src/main/java/com/bai/env/funcs/externalfuncs/RandFunction.java b/src/main/java/com/bai/env/funcs/externalfuncs/RandFunction.java
@@ -1,10 +1,13 @@
 package com.bai.env.funcs.externalfuncs;
 
+import static com.bai.util.Utils.getAddress;
+
 import com.bai.env.ALoc;
 import com.bai.env.AbsEnv;
 import com.bai.env.Context;
 import com.bai.env.KSet;
 import com.bai.env.TaintMap;
+import ghidra.program.model.address.Address;
 import ghidra.program.model.data.IntegerDataType;
 import ghidra.program.model.listing.Function;
 import ghidra.program.model.pcode.PcodeOp;
@@ -32,7 +35,8 @@ public void invoke(PcodeOp pcode, AbsEnv inOutEnv, AbsEnv tmpEnv, Context contex
         if (retALoc == null) {
             return;
         }
-        long taints = TaintMap.getTaints(context, callFunc);
+        Address callAddress = getAddress(pcode);
+        long taints = TaintMap.getTaints(callAddress, context, callFunc);
         inOutEnv.set(retALoc, KSet.getTop(taints), true);
     }