[Python] Select existing path node (with flow state) in the isAdditionalFlowState predicate

[jackfromeast]

I am trying to track how many times the get operation is performed on the return object. However, my current CodeQL query is unable to correctly distinguish between different numbers of get operations.

The functions below demonstrate my expected flow states:

def test1(obj, key):        # Source: Key, SourceKeyFlowState
  for _ in range(2):
    obj = obj.get(key)      # if key has SourceKeyFlowState, obj.get(key) -> ObjectLayerOneFlowState 
                            # if key has SourceKeyFlowState and obj has ObjectLayerOneFlowState, obj.get(key) -> ObjectLayerTwoFlowState

  return obj                # Sink: obj, ObjectLayerTwoFlowState

def test2(obj, key):        # Source: Key, SourceKeyFlowState
  for _ in range(1):
    obj = obj.get(key)      # if key has SourceKeyFlowState, obj.get(key) -> ObjectLayerOneFlowState 
                            # if key has SourceKeyFlowState and obj has ObjectLayerOneFlowState, obj.get(key) -> ObjectLayerTwoFlowState

  return obj                # Sink: obj, ObjectLayerOneFlowState

Expected Taint Flows:

• test1: SourceKeyFlowState → ObjectFlowState-MoreThanOne
• test2: SourceKeyFlowState → ObjectFlowState-One

However, my current CodeQL query does not correctly distinguish between these cases and instead selects all possible four flows.

import python
import semmle.python.ApiGraphs
import codeql.dataflow.DataFlow
import semmle.python.dataflow.new.DataFlow
import semmle.python.dataflow.new.TaintTracking
import semmle.python.dataflow.new.internal.DataFlowPublic
import semmle.python.dataflow.new.internal.TaintTrackingPublic

module TempTest {

abstract class FlowState extends string {
  bindingset[this]
  FlowState() { any() }
}

class SourceKeyFlowState extends TempTest::FlowState {
  SourceKeyFlowState() { this = "SourceKeyFlowState" }
}

class Occurrence extends string {
  Occurrence() { 
    this = "One" or
    this = "MoreThanOne"
  }
}

class ObjectFlowState extends TempTest::FlowState {
  Occurrence occur;

  ObjectFlowState() { this = "ObjectFlowState" + "-" + occur }
}

module TrackingGetOperationConfiguration implements DataFlow::StateConfigSig {
  class FlowState = TempTest::FlowState;

  predicate isSource(DataFlow::Node source, FlowState state) {
    exists (Function func, Parameter param | 
      func.getArg(1) = param and
      source.asExpr() = param and
      func.getName().matches("test%")
    ) and
    (
      state instanceof SourceKeyFlowState
    )
  }

  predicate isSink(DataFlow::Node sink, FlowState state) {
    exists (Function func, Return ret | 
      ret.getScope() = func.getEvaluatingScope() and
      ret.contains(sink.asExpr()) and
      func.getName().matches("test%")
    ) and
    (
      state instanceof ObjectFlowState
    )
  }

  predicate isAdditionalFlowStep(DataFlow::Node fromNode, FlowState fromState, DataFlow::Node toNode, FlowState toState) {
    exists( MethodCallNode callNode, Call call |
      callNode.asExpr() = call and
      callNode.getMethodName() = "get" and
      callNode.asExpr() = toNode.asExpr() and
      call.getArg(0) = fromNode.asExpr()
    ) and 
    fromState instanceof SourceKeyFlowState and
    toState instanceof ObjectFlowState
  }
}

module TrackingGetOperationFlow = DataFlow::GlobalWithState<TrackingGetOperationConfiguration>;
module Flow = TrackingGetOperationFlow; // For shortening the name

predicate run(Flow::PathNode source, Flow::PathNode sink, FlowState state) {
  Flow::flowPath(source, sink) and
  sink.getState() = state
}

}

I think the key problem here is that in the isAdditionalFlowStep predicate, I cannot select the path node beside the fromNode and toNode, and use its flow state information to determine the flow state of toNode. In my case, the taint propagation step should take into account both the flow state of the key (fromNode) and the flow state of the base object to correctly determine the flow state of toNode.

Could anyone provide suggestions on how to fix this issue? Any insights would be greatly appreciated!

[superboy-zjc]

I have the same question!

[aschackmull]

You cannot inspect flow to both obj and key at the same time in order to determine the output state at obj.get(key) in that way - this is because the calculated flow is a path that goes from one node to another one step at a time. However, you can do something that's close and might actually get you what you want: You can use two additional steps: One from key to the output that updates the state from SourceKeyFlowState to ObjectFlowState-One and one from the qualifier obj to the output that updates the state from ObjectFlowState-One to ObjectFlowState-MoreThanOne. Your current step definition mostly achieves the former, except that you've given the resulting state as toState instanceof ObjectFlowState which includes both One and MoreThanOne. You probably want to change that to state = "ObjectFlowState-One". And then supplement with the other step. The two steps might look something like this:

  predicate isAdditionalFlowStep(DataFlow::Node fromNode, FlowState fromState, DataFlow::Node toNode, FlowState toState) {
    exists( MethodCallNode callNode, Call call |
      callNode.asExpr() = call and
      callNode.getMethodName() = "get" and
      callNode.asExpr() = toNode.asExpr()
    |
      call.getArg(0) = fromNode.asExpr() and
      fromState instanceof SourceKeyFlowState and
      toState = "ObjectFlowState-One"
      or
      call.getQualifier() = fromNode.asExpr() and
      fromState = "ObjectFlowState-One" and
      toState = "ObjectFlowState-MoreThanOne"
    )
  }

Note that you cannot expect the out-of-the-box analysis to distinguish for _ in range(1): and for _ in range(2):. Determining the number of times a loop is executed in general is halting-problem territory.