From fc1bc3a606651bf41bcd09025086088748e71db1 Mon Sep 17 00:00:00 2001
From: Tom Hvitved <hvitved@github.com>
Date: Wed, 20 Nov 2024 12:58:30 +0100
Subject: [PATCH] Address review comments

---
 .../dataflow/internal/DataFlowImplCommon.qll  | 45 +++++++++----------
 1 file changed, 21 insertions(+), 24 deletions(-)

diff --git a/shared/dataflow/codeql/dataflow/internal/DataFlowImplCommon.qll b/shared/dataflow/codeql/dataflow/internal/DataFlowImplCommon.qll
index cd2af16340513..e88e5defa6625 100644
--- a/shared/dataflow/codeql/dataflow/internal/DataFlowImplCommon.qll
+++ b/shared/dataflow/codeql/dataflow/internal/DataFlowImplCommon.qll
@@ -1435,27 +1435,19 @@ module MakeImplCommon<LocationSig Location, InputSig<Location> Lang> {
           string model, CachedCallContextSensitivity::CcNoCall ctx
         ) {
           exists(
-            DataFlowCall call, DataFlowCallable callable, ArgNode arg, string model1, string model2,
-            CachedCallContextSensitivity::CcNoCall ctx1, CachedCallContextSensitivity::CcNoCall ctx2
+            ArgNode arg, string model1, string model2, CachedCallContextSensitivity::CcNoCall ctx1,
+            CachedCallContextSensitivity::CcNoCall ctx2
           |
             model = mergeModels(model1, model2) and
-            (
-              // call may restrict the set of call sites that can be returned to
-              ctx2.(CachedCallContextSensitivity::CcReturn).isReturn(callable, call)
-              or
-              // call does not restrict the set of call sites that can be returned to
-              not exists(CachedCallContextSensitivity::CcReturn ret | ret.isReturn(callable, call)) and
-              CachedCallContextSensitivity::viableImplNotCallContextReducedReverse(ctx2)
-            ) and
             ctx = mergeContexts(ctx1, ctx2)
           |
             // flow through: no prior read
             parameterValueFlowArg(p, arg, mustBeNone, model1, ctx1) and
-            argumentValueFlowsThrough(call, callable, arg, read, node, model2)
+            argumentValueFlowsThrough(arg, read, node, model2, ctx2)
             or
             // flow through: no read inside method
             parameterValueFlowArg(p, arg, read, model1, ctx1) and
-            argumentValueFlowsThrough(call, callable, arg, mustBeNone, node, model2)
+            argumentValueFlowsThrough(arg, mustBeNone, node, model2, ctx2)
           )
         }
 
@@ -1470,27 +1462,32 @@ module MakeImplCommon<LocationSig Location, InputSig<Location> Lang> {
 
         pragma[nomagic]
         private predicate argumentValueFlowsThrough0(
-          DataFlowCall call, DataFlowCallable callable, ArgNode arg, ReturnKind kind,
-          ReadStepTypesOption read, string model
+          DataFlowCall call, ArgNode arg, ReturnKind kind, ReadStepTypesOption read, string model,
+          CachedCallContextSensitivity::CcNoCall outerCtx
         ) {
-          exists(ParamNode param, CachedCallContextSensitivity::CcNoCall ctx |
+          exists(
+            ParamNode param, DataFlowCallable callable,
+            CachedCallContextSensitivity::CcNoCall innerCtx
+          |
             viableParamArg(call, param, arg) and
-            parameterValueFlowReturn(param, kind, read, model, ctx) and
-            callable = nodeGetEnclosingCallable(param)
+            parameterValueFlowReturn(param, kind, read, model, innerCtx) and
+            callable = nodeGetEnclosingCallable(param) and
+            outerCtx = CachedCallContextSensitivity::getCallContextReturn(callable, call)
           |
-            CachedCallContextSensitivity::viableImplNotCallContextReducedReverse(ctx)
+            CachedCallContextSensitivity::viableImplNotCallContextReducedReverse(innerCtx)
             or
-            call = CachedCallContextSensitivity::viableImplCallContextReducedReverse(callable, ctx)
+            call =
+              CachedCallContextSensitivity::viableImplCallContextReducedReverse(callable, innerCtx)
           )
         }
 
         pragma[nomagic]
         private predicate argumentValueFlowsThrough(
-          DataFlowCall call, DataFlowCallable callable, ArgNode arg, ReadStepTypesOption read,
-          Node out, string model
+          ArgNode arg, ReadStepTypesOption read, Node out, string model,
+          CachedCallContextSensitivity::CcNoCall ctx
         ) {
-          exists(ReturnKind kind |
-            argumentValueFlowsThrough0(call, callable, arg, kind, read, model) and
+          exists(DataFlowCall call, ReturnKind kind |
+            argumentValueFlowsThrough0(call, arg, kind, read, model, ctx) and
             out = getAnOutNode(call, kind)
           |
             // normal flow through
@@ -1515,7 +1512,7 @@ module MakeImplCommon<LocationSig Location, InputSig<Location> Lang> {
         predicate argumentValueFlowsThrough(
           ArgNode arg, ReadStepTypesOption read, Node out, string model
         ) {
-          argumentValueFlowsThrough(_, _, arg, read, out, model)
+          argumentValueFlowsThrough(arg, read, out, model, _)
         }
 
         /**