fix(core, analyzer): Fix method boundary instructions

core/opentaint-dataflow-core/opentaint-dataflow/src/main/kotlin/org/opentaint/dataflow/ap/ifds/trace/MethodForwardTraceResolver.kt

@@ -111,6 +111,7 @@ class MethodForwardTraceResolver(
 
         val initialEdge = ZeroToFact(methodEntryPoint, initialStatement, initialFact)
         if (!startAtStatement) {
+            builder.handleEdgeToMethodEnd(initialEdge, initialEdge, EdgeReason.Unchanged)
             builder.propagateEdgeToSuccessors(initialEdge, initialEdge, EdgeReason.Unchanged)
         } else {
             builder.enqueue(initialEdge)

core/opentaint-dataflow-core/opentaint-jvm-dataflow/src/main/kotlin/org/opentaint/dataflow/jvm/ap/ifds/LambdaAnonymousClassFeature.kt

@@ -53,6 +53,7 @@ import org.opentaint.ir.impl.features.classpaths.virtual.JIRVirtualParameter
 import org.opentaint.ir.impl.types.JIRClassTypeImpl
 import org.opentaint.ir.impl.types.JIRTypedFieldImpl
 import org.opentaint.ir.impl.types.substition.JIRSubstitutorImpl
+import org.opentaint.jvm.graph.JMethodBoundaryInstFeature
 import java.util.Objects
 import java.util.concurrent.ConcurrentHashMap
 
@@ -308,7 +309,9 @@ class LambdaAnonymousClassFeature : JIRClasspathExtFeature {
         parameters: List<JIRVirtualParameter>,
         private val instructions: JIRInstList<JIRInst>
     ) : JIRVirtualMethodImpl(name, returnType = returnType, parameters = parameters, description = description) {
-        override val instList: JIRInstList<JIRInst> get() = instructions
+        override val instList: JIRInstList<JIRInst> by lazy {
+            JMethodBoundaryInstFeature.transformInstList(this, instructions)
+        }
 
         override fun hashCode(): Int = originalMethod.hashCode()
 
@@ -364,7 +367,9 @@ class LambdaAnonymousClassFeature : JIRClasspathExtFeature {
         parameters: List<JIRVirtualParameter>,
         private val instructions: JIRInstList<JIRInst>
     ) : JIRVirtualMethodImpl(name, returnType = returnType, parameters = parameters, description = description) {
-        override val instList: JIRInstList<JIRInst> get() = instructions
+        override val instList: JIRInstList<JIRInst> by lazy {
+            JMethodBoundaryInstFeature.transformInstList(this, instructions)
+        }
 
         override fun hashCode(): Int = Objects.hash(name, enclosingClass)
 

core/opentaint-dataflow-core/opentaint-jvm-dataflow/src/main/kotlin/org/opentaint/dataflow/jvm/ap/ifds/analysis/JIRAnalysisManager.kt

@@ -50,6 +50,8 @@ import org.opentaint.ir.api.jvm.cfg.JIRImmediate
 import org.opentaint.ir.api.jvm.cfg.JIRInst
 import org.opentaint.ir.api.jvm.cfg.JIRReturnInst
 import org.opentaint.jvm.graph.JApplicationGraph
+import org.opentaint.jvm.graph.JMethodExitExceptionalInst
+import org.opentaint.jvm.graph.JMethodExitNormalInst
 import org.opentaint.util.analysis.ApplicationGraph
 import java.util.concurrent.ConcurrentHashMap
 import java.util.concurrent.ConcurrentLinkedQueue
@@ -112,6 +114,11 @@ class JIRAnalysisManager(
         val jirContextForEmptyMethod = contextForEmptyMethod as? JIRMethodAnalysisContext
 
         val method = entryPointStatement.location.method
+        val instructions = method.instList.instructions
+        check(instructions.isEmpty() || instructions.last() is JMethodExitExceptionalInst) {
+            "Method $method is analysed without method boundary instructions"
+        }
+
         val localVariableReachability = jirContextForEmptyMethod?.localVariableReachability
             ?: JIRLocalVariableReachability(method, graph, this)
 
@@ -303,7 +310,7 @@ class JIRAnalysisManager(
     override val methodContextSerializer = JIRMethodContextSerializer(cp)
 
     override fun isTraceRequiredInstruction(inst: CommonInst): Boolean =
-        inst is JIRReturnInst
+        inst is JIRReturnInst || inst is JMethodExitNormalInst
 
     override fun onInstructionReached(inst: CommonInst) {
 

core/opentaint-dataflow-core/opentaint-jvm-dataflow/src/main/kotlin/org/opentaint/dataflow/jvm/ap/ifds/analysis/JIRMethodSequentFlowFunction.kt

@@ -39,6 +39,8 @@ import org.opentaint.ir.api.jvm.cfg.JIRInst
 import org.opentaint.ir.api.jvm.cfg.JIRReturnInst
 import org.opentaint.ir.api.jvm.cfg.JIRThrowInst
 import org.opentaint.ir.api.jvm.cfg.JIRValue
+import org.opentaint.jvm.graph.JMethodExitExceptionalInst
+import org.opentaint.jvm.graph.JMethodExitNormalInst
 import org.opentaint.util.onSome
 
 class JIRMethodSequentFlowFunction(
@@ -157,17 +159,29 @@ class JIRMethodSequentFlowFunction(
 
             is JIRReturnInst -> {
                 val access = currentInst.returnValue?.let { accessPathBase(it) }
+                simpleAssign(AccessPathBase.Return, access, factAp, { unchanged() }) {
+                    propagateFact(it, TraceInfo.Flow)
+                }
+            }
+
+            is JIRThrowInst -> {
+                val access = accessPathBase(currentInst.throwable)
+                simpleAssign(AccessPathBase.Exception, access, factAp, { unchanged() }) {
+                    propagateFact(it, TraceInfo.Flow)
+                }
+            }
+
+            is JMethodExitNormalInst -> {
                 propagateExitFact(
                     initialFacts, AccessPathBase.Return,
-                    access, factAp, unchanged, propagateFactWithRefinement, sideEffect
+                    factAp, unchanged, propagateFactWithRefinement, sideEffect
                 )
             }
 
-            is JIRThrowInst -> {
-                val access = accessPathBase(currentInst.throwable)
+            is JMethodExitExceptionalInst -> {
                 propagateExitFact(
                     initialFacts, AccessPathBase.Exception,
-                    access, factAp, unchanged, propagateFactWithRefinement, sideEffect
+                    factAp, unchanged, propagateFactWithRefinement, sideEffect
                 )
             }
 
@@ -180,33 +194,15 @@ class JIRMethodSequentFlowFunction(
     private fun propagateExitFact(
         initialFacts: Set<InitialFactAp>?,
         exitBase: AccessPathBase,
-        access: AccessPathBase?,
         factAp: FinalFactAp,
         unchanged: () -> Unit,
         propagateFactWithRefinement: (FactRefiner, FinalFactAp, TraceInfo) -> Unit,
         sideEffect: (Sequent.SideEffect) -> Unit
     ) {
         val refiner = FactRefiner()
-
-        val currentFacts = mutableListOf<FinalFactAp>()
-
-        simpleAssign(
-            exitBase, access, factAp,
-            unchanged = {
-                currentFacts += it
-            },
-            propagateFact = {
-                propagateFactWithRefinement(refiner, it, TraceInfo.Flow)
-                currentFacts += it
-            }
-        )
-
         val resultFacts = mutableListOf<Pair<FinalFactAp, TraceInfo>>()
-
-        currentFacts.forEach { currentFact ->
-            resultFacts += currentFact to TraceInfo.Flow
-            resultFacts += applyMethodExitSourceRules(exitBase, currentFact, refiner)
-        }
+        resultFacts += factAp to TraceInfo.Flow
+        resultFacts += applyMethodExitSourceRules(exitBase, factAp, refiner)
 
         while (resultFacts.isNotEmpty()) {
             val (resultFact, factTrace) = resultFacts.removeLast()
@@ -217,7 +213,7 @@ class JIRMethodSequentFlowFunction(
             val propagatedFact = resultFact.dropFinalFacts(factsToDrop)
                 ?.dropArgumentsLocalTaintMarks(initialFacts != null && initialFacts.isEmpty())
 
-            if (propagatedFact == factAp) {
+            if (propagatedFact == factAp && !refiner.hasRefinement) {
                 unchanged()
             } else if (propagatedFact != null) {
                 propagateFactWithRefinement(refiner, propagatedFact, factTrace)
@@ -625,7 +621,7 @@ class JIRMethodSequentFlowFunction(
     }
 
     private fun applyUnconditionalSinks() = with(analysisContext.taint) {
-        if (currentInst !is JIRReturnInst) return
+        if (currentInst !is JMethodExitNormalInst) return
 
         val sinkRules = sinkRulesForMethodExit(currentInst, fact = null, initialFacts = null).toList()
         sinkRules.forEach {
@@ -662,7 +658,7 @@ class JIRMethodSequentFlowFunction(
     }
 
     private fun MutableSet<Sequent>.applyUnconditionalSources() {
-        if (currentInst is JIRReturnInst) {
+        if (currentInst is JMethodExitNormalInst) {
             applyMethodExitSourceRules(AccessPathBase.Return, fact = null, refiner = null).forEach { (fact, trace) ->
                 this += Sequent.ZeroToFact(fact, trace)
             }

core/opentaint-dataflow-core/opentaint-jvm-dataflow/src/main/kotlin/org/opentaint/dataflow/jvm/ap/ifds/trace/JIRMethodSequentPrecondition.kt

@@ -32,6 +32,8 @@ import org.opentaint.ir.api.jvm.cfg.JIRInst
 import org.opentaint.ir.api.jvm.cfg.JIRReturnInst
 import org.opentaint.ir.api.jvm.cfg.JIRThrowInst
 import org.opentaint.ir.api.jvm.cfg.JIRValue
+import org.opentaint.jvm.graph.JMethodExitExceptionalInst
+import org.opentaint.jvm.graph.JMethodExitNormalInst
 import org.opentaint.util.maybeFlatMap
 
 class JIRMethodSequentPrecondition(
@@ -42,43 +44,40 @@ class JIRMethodSequentPrecondition(
 
     override fun factPrecondition(
         fact: InitialFactAp,
-    ): Set<SequentPrecondition> {
-        if (currentInst !is JIRAssignInst && currentInst !is JIRReturnInst && currentInst !is JIRThrowInst) {
-            return setOf(SequentPrecondition.Unchanged)
+    ): Set<SequentPrecondition> = when (currentInst) {
+        is JMethodExitNormalInst, is JMethodExitExceptionalInst -> methodExitPrecondition(fact)
+
+        is JIRAssignInst, is JIRReturnInst, is JIRThrowInst -> {
+            val results = mutableSetOf<SequentPrecondition>()
+            results.computeFactPrecondition(fact)
+            results
         }
 
-        val results = mutableSetOf<SequentPrecondition>()
-        results.computeFactPrecondition(fact, applyExitSourceRules = true)
+        else -> setOf(SequentPrecondition.Unchanged)
+    }
+
+    private fun methodExitPrecondition(fact: InitialFactAp): Set<SequentPrecondition> {
+        val results = mutableSetOf<SequentPrecondition>(SequentPrecondition.Unchanged)
+        results.methodExitSourcePrecondition(fact)
         return results
     }
 
-    private fun MutableSet<SequentPrecondition>.computeFactPrecondition(
-        fact: InitialFactAp,
-        applyExitSourceRules: Boolean
-    ) {
-        val factPrecondition = computePrecondition(fact, applyExitSourceRules)
+    private fun MutableSet<SequentPrecondition>.computeFactPrecondition(fact: InitialFactAp) {
+        val factPrecondition = computePrecondition(fact)
         this += factPrecondition.ifEmpty { setOf(SequentPrecondition.Unchanged) }
 
         analysisContext.aliasAnalysis?.forEachPossibleAliasAtStatement(currentInst, fact) { aliasedFact ->
-            this += computePrecondition(aliasedFact, applyExitSourceRules)
+            this += computePrecondition(aliasedFact)
         }
     }
 
-    private fun computePrecondition(
-        aliasedFact: InitialFactAp,
-        applyExitSourceRules: Boolean
-    ): Set<SequentPrecondition> {
+    private fun computePrecondition(aliasedFact: InitialFactAp): Set<SequentPrecondition> {
         val precondition = mutableSetOf<SequentPrecondition>()
         preconditionForFact(aliasedFact)?.let {
             precondition += PreconditionFactsForInitialFact(aliasedFact, it)
         }
 
         precondition.unconditionalSourcesPrecondition(aliasedFact)
-
-        if (applyExitSourceRules) {
-            precondition.methodExitSourcePrecondition(aliasedFact)
-        }
-
         return precondition
     }
 
@@ -112,6 +111,8 @@ class JIRMethodSequentPrecondition(
                 return listOf(fact.rebase(base))
             }
 
+            is JMethodExitNormalInst, is JMethodExitExceptionalInst -> return listOf(fact)
+
             else -> return null
         }
     }
@@ -335,7 +336,7 @@ class JIRMethodSequentPrecondition(
                         }
 
                         val preFact = factCube.facts.single()
-                        computeFactPrecondition(preFact, applyExitSourceRules = false)
+                        computeFactPrecondition(preFact)
                     }
                 }
             )

core/src/main/kotlin/org/opentaint/jvm/sast/project/spring/SpringGeneratedMethod.kt

@@ -6,6 +6,7 @@ import org.opentaint.ir.api.jvm.cfg.JIRInst
 import org.opentaint.ir.api.jvm.cfg.JIRInstList
 import org.opentaint.ir.impl.features.classpaths.virtual.JIRVirtualMethodImpl
 import org.opentaint.ir.impl.features.classpaths.virtual.JIRVirtualParameter
+import org.opentaint.jvm.graph.JMethodBoundaryInstFeature
 import java.util.Objects
 
 class SpringGeneratedMethod(
@@ -21,7 +22,9 @@ class SpringGeneratedMethod(
     parameters = parameters,
     description = description
 ) {
-    override val instList: JIRInstList<JIRInst> get() = instructions
+    override val instList: JIRInstList<JIRInst> by lazy {
+        JMethodBoundaryInstFeature.transformInstList(this, instructions)
+    }
 
     override fun hashCode(): Int = Objects.hash(name, enclosingClass)
 

core/src/main/kotlin/org/opentaint/jvm/sast/sarif/JirSarifGenerator.kt

@@ -218,7 +218,7 @@ class JirSarifGenerator(
         relevantLocations: List<List<IntermediateLocation>>?
     ): IntermediateLocation? {
         if (TraceMessageBuilder.isGeneratedLocation(statement)) {
-            val normalLocation = TraceMessageBuilder.tryResolveNormalGeneratedLocation(statement)
+            val normalLocation = TraceMessageBuilder.tryResolveNormalGeneratedLocation(statement, relevantLocations)
                 ?: return null
             return statementLocation(normalLocation, type, relevantLocations)
         }

core/src/main/kotlin/org/opentaint/jvm/sast/sarif/TraceMessageBuilder.kt

@@ -37,6 +37,8 @@ import org.opentaint.ir.api.jvm.cfg.JIRValue
 import org.opentaint.ir.approximation.JIREnrichedVirtualMethod
 import org.opentaint.jvm.graph.JMethodBoundaryInst
 import org.opentaint.jvm.graph.JMethodEnterInst
+import org.opentaint.jvm.graph.JMethodExitExceptionalInst
+import org.opentaint.jvm.graph.JMethodExitNormalInst
 import org.opentaint.jvm.sast.project.spring.GeneratedSpringRegistry
 import org.opentaint.jvm.sast.project.spring.SpringGeneratedMethod
 import org.opentaint.semgrep.pattern.Mark
@@ -1134,8 +1136,17 @@ class TraceMessageBuilder(
             return false
         }
 
-        fun tryResolveNormalGeneratedLocation(stmt: CommonInst): CommonInst? {
+        fun tryResolveNormalGeneratedLocation(
+            stmt: CommonInst,
+            relevantLocations: List<List<IntermediateLocation>>?
+        ): CommonInst? {
             if (stmt is JMethodEnterInst) return stmt.location.method.instList.first()
+            if (stmt is JMethodExitNormalInst) {
+                return tryResolveExitLocation(stmt, relevantLocations) { it is JIRReturnInst }
+            }
+            if (stmt is JMethodExitExceptionalInst) {
+                return tryResolveExitLocation(stmt, relevantLocations) { it is JIRThrowInst }
+            }
             if (stmt is JMethodBoundaryInst) return null
             val locationMethod = stmt.location.method
             if (locationMethod is JIRLambdaMethod) {
@@ -1146,6 +1157,14 @@ class TraceMessageBuilder(
             return null
         }
 
+        private inline fun tryResolveExitLocation(
+            stmt: JMethodBoundaryInst,
+            relevantLocations: List<List<IntermediateLocation>>?,
+            isExit: (CommonInst) -> Boolean
+        ): CommonInst? =
+            relevantLocations?.firstOrNull()?.lastOrNull { isExit(it.inst) }?.inst
+                ?: stmt.location.method.instList.instructions.lastOrNull { isExit(it) }
+
         fun isAbnormalLocation(stmt: CommonInst): Boolean =
             stmt is JIRInst && stmt.lineNumber == 0