Python Class Variable data flow

[michal-sladecek-cognite]

Hi,
I'm having trouble with data flow in following python code:

def source():
    return "SOURCE"

class Test():
    VAR = source()

def sink(s):
    print(s)


VAR = source()

sink(VAR)               # Found data flow
sink(source())          # Also found
sink(Test().VAR).       # Query can't detect this

The query I'm using is:

class ClassVariableFlow extends DataFlow::Configuration {
    ClassVariableFlow() { this = "ClassVariableFlow"}

    override predicate isSource(DataFlow::Node source) {
        exists ( FunctionValue f |
            f.getName() = "source" and
            f.getACall() = source.asCfgNode()
        )
    }
    override predicate isSink(DataFlow::Node sink) {   
        exists( FunctionValue f |
            f.getName() = "sink" and
            f.getACall().getAnArg() = sink.asCfgNode()
        )
    }
}


from ClassVariableFlow config,DataFlow::Node source,DataFlow::Node sink
where config.hasFlow(source, sink)
select source,sink

The query should detect all data flows from function source() into argument of function sink(), however it works only for the first two calls of sink - it doesn't work when function source() is called inside of the class, and then passed as a class variable into sink(). Can you recommend me some changes to the query that would make dataflow through class variables also visible?

[yoff]

Hi, thank you very much for alerting us to this hole in our analysis! It looks like we do indeed not handle this flow. I think what would be needed is to add a jump step along the lines of

predicate readClassVarStep(DataFlow::EssaNode fromNode, DataFlow::AttrRead toNode) {
  exists(ClassDef classDef, SsaVariable c | classDef.defines(c.getVariable()) |
    fromNode.getVar().getScope() = classDef.getDefinedClass() and
    toNode.getObject().(DataFlow::CallCfgNode).getFunction().(DataFlow::CfgNode).getNode() =
      c.getAUse()
  )
}

and then a corresponding one for writing to class variables.

I think you cannot add jump steps, but you could add

  override predicate isAdditionalFlowStep(DataFlow::Node fromNode, DataFlow::Node toNode) {
    readClassVarStep(fromNode, toNode)
  }

to your configuration.

Alternatively, depending on your use case, you could perhaps use instance variables instead. We do support those; you can get an idea of what we handle from the field flow tests here and here.

[michal-sladecek-cognite]

Thank you, this works! Is it somehow possible for me to apply this additionalFlowStep on all the queries I'm running by default, so that it will also be applied when searching for the vulnerabilities that are in standard CodeQL library?

[yoff]

We will add this to the standard analysis shortly :-) Until then, I am afraid you have to apply it to each one...you might create a subclass of configuration with it in, or define something to import to make it more ergonomic..

[michal-sladecek-cognite]

Alright, thank you!

[michal-sladecek-cognite]

I stumbled on another problem, even with the additional steps the query finds no results if the class definition is in another file. I've been trying to change the predicates and also readClassVarStep to use the API graph, but I can't find the way to get it working.

Example:
package1/classdef.py :

def source():
    return "SOURCE"

class Test():
    VAR = source()

main.py:

from package1.classdef import Test

def sink(s):
    print(s)

sink(Test().VAR)

[yoff]

Right, for that I think we need the proper solution via field flow. I am afraid that you cannot add that from the outside as it goes via read steps and store steps rather than flow steps.