diff --git a/src/main/java/com/bai/env/Context.java b/src/main/java/com/bai/env/Context.java index 16bb1c6..d34dda5 100644 --- 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 index 8d6965a..2ba1d1d 100644 --- 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 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 index cb900e7..80b398a 100644 --- 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 index 197caa8..5189560 100644 --- 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,9 +51,10 @@ 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; @@ -58,7 +62,7 @@ public void invoke(PcodeOp pcode, AbsEnv inOutEnv, AbsEnv tmpEnv, Context contex 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 index 969d200..30a5ad1 100644 --- 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 index 21adaed..5055f0a 100644 --- 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); }