I'm having trouble tracking variable propagation

[Light1Smile]

Hi, I'm having trouble tracking variable propagation.

I'm writing a CodeQL query to help me trace how a model flows through a program. However, I found that using data flow analysis or taint tracking doesn't work very well in this case.

Here's my query:

import python
import semmle.python.dataflow.new.DataFlow
import semmle.python.ApiGraphs
predicate isTargetedCall(API::Node an,API::Node an_2) {
 exists (string moduleName, string functionName |
   (
       (moduleName = "diffusers" and functionName = "from_pretrained") 
   ) 
     and (an.getAPredecessor*() = API::moduleImport(moduleName) and an_2 = an.getMember(functionName))
 )
}
class ExecuteCall extends DataFlow::CallCfgNode {
   ExecuteCall() {
     exists(API::Node ano,API::Node ano_2 | isTargetedCall(ano,ano_2) and this = ano_2.getACall())
   }
}
private module MyConfig implements DataFlow::ConfigSig {
 predicate isSource(DataFlow::Node source) {
   exists(StringLiteral sl,DataFlow::Node dn | sl.getAFlowNode()=dn.asCfgNode() and  source =dn)  
  }
  predicate isSink(DataFlow::Node sink) {
    exists(ExecuteCall ec ,API::Node dn_1|
        dn_1.getAPredecessor().getACall() = ec and
        sink = dn_1.asSink()
       )
  }
}

module MyFlow = DataFlow::Global<MyConfig>; 

from MyFlow::PathNode source, MyFlow::PathNode sink,StringLiteral sl,API::Node api_node
where MyFlow::flowPath(source, sink) and sl.getAFlowNode() = source.getNode().asCfgNode() and sink.getNode().asCfgNode() = api_node.asSink().asCfgNode()
select  sl.getLiteralValue() ,api_node,sink.getLocation(),source.getLocation()

For the following code, it fails to properly track the flow of "ssube/stable-diffusion-x4-upscaler-onnx". This string is first passed to hub_checkpoint, and then by default, if test_pipeline_default_ddpm is used, it gets passed to self.hub_checkpoint. However, when I try to perform a data flow or taint analysis with "ssube/stable-diffusion-x4-upscaler-onnx" as the source and self.hub_checkpoint (the first argument to OnnxStableDiffusionUpscalePipeline.from_pretrained) as the sink, CodeQL fails to connect them effectively.

I understand that this might require taint tracking and defining isAdditionalFlowStep, but I haven't found any detailed tutorials related to this, so I'm not sure how to improve the query. Could anyone offer some suggestions?

from diffusers import OnnxStableDiffusionUpscalePipeline
class OnnxStableDiffusionUpscalePipelineFastTests():
    hub_checkpoint = "ssube/stable-diffusion-x4-upscaler-onnx"

    def test_pipeline_default_ddpm(self):
        pipe = OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint, provider="CPUExecutionProvider")
        pipe.set_progress_bar_config(disable=None)

        inputs = self.get_dummy_inputs()
        image = pipe(**inputs).images
        image_slice = image[0, -3:, -3:, -1].flatten()

Secondly, when a string is treated as tainted and then processed by some method, it is not propagated by default. For example, in the code below, if I want to make the taint propagation more permissive—so that passing through method parameters or method calls still allows propagation—how can I do that?—because in some cases I only want to analyze the propagation rather than treating it as actual taint, what should I do?

from diffusers import OnnxStableDiffusionUpscalePipeline
hub_checkpoint = "ssube/stable-diffusion-x4-upscaler-onnx"
hub_checkpoint_2 = hub_checkpoint.lower()
model = OnnxStableDiffusionUpscalePipeline.from_pretrained(hub_checkpoint_2)

I would be most grateful if someone could offer help:）

[hvitved]

Hi

I know that for other languages, when the taint source happens inside a constructor or a field initializer, like in your example above, and there is not explicit call to the constructor follow by a call to the method containing the sink, then we will not flag it. I suspect the same happens for Python. It is on our radar to fix.

[Light1Smile]

Thank you very much for your previous answer. I use a AdditionalFlow to slove the first problem. And the second question is still under consideration. I would be honored if these tiny pieces of information could be of some help

  predicate isAdditionalFlowStep(DataFlow::Node nodeFrom, DataFlow::Node nodeTo) {
      exists(AssignStmt ass,SelfAttribute satt,ClassDef cd |
      cd.contains(ass) and cd.contains(satt) and
      ass.getAChildNode().toString() = satt.getAttr() and
      nodeFrom.asCfgNode() = ass.getAChildNode().getAFlowNode() and
      nodeTo.asCfgNode() = satt.getAFlowNode()
    )
  }

[muadzie]

The issue has two parts. Here's how to solve both with custom isAdditionalFlowStep / isAdditionalTaintStep:
1. Field initializer → self.hub_checkpoint not flowing
CodeQL's standard Python data flow doesn't track class field reads (self.hub_checkpoint) from field declarations because there's no explicit constructor call in your test file — the class attribute is never "assigned" at runtime in a way the data flow library recognizes.
Fix: Add a custom additional step that propagates from the field definition node to any read of that field:

predicate fieldToAttributeRead(DataFlow::Node nodeFrom, DataFlow::Node nodeTo) {
exists(ClassValue cv, string attrName |
nodeFrom.asExpr().(Attr).getName() = attrName and
nodeTo = DataFlow::attrNode(cv, attrName) and
nodeFrom.asExpr().(Attr).getObject().(Name).getId() = cv.getName()
)
}
But the cleaner approach is to treat the class attribute assignment as a source directly:
predicate isSource(DataFlow::Node source) {
exists(ClassDef cd, string name, string val |
cd.getBody().(AssignStmt).getATarget().(Name).getId() = name and
cd.getBody().(AssignStmt).getValue().(StrConst).getValue() = val and
name = "hub_checkpoint" and
source.asExpr() = cd.getBody().(AssignStmt).getValue()
)
}
Then add a step from that source to self.hub_checkpoint accesses in methods.
2. String method calls like .lower() breaking taint
.lower() returns a new string, so pure data flow (which tracks reference identity) won't propagate. Switch to taint tracking instead of data flow:
private module MyTaintConfig implements DataFlow::ConfigSig {
predicate isSource(DataFlow::Node source) { ... }
predicate isSink(DataFlow::Node sink) { ... }
predicate isAdditionalFlowStep(DataFlow::Node nodeFrom, DataFlow::Node nodeTo) {
// Propagate through string method calls
exists(MethodCall mc, string methodName |
mc.getFunction().(Attr).getName() = methodName and
"lower", "upper", "strip", "format", "replace", "join" .contains(methodName) and
nodeFrom.asExpr() = mc.getFunction().(Attr).getObject() and
nodeTo.asExpr() = mc
)
}
}

3. For "permissive but not actual taint" analysis
Use isAdditionalFlowStep with a broad predicate that allows propagation through ANY method call on a string:

predicate isAdditionalFlowStep(DataFlow::Node nodeFrom, DataFlow::Node nodeTo) {
exists(MethodCall mc |
nodeFrom.asExpr() = mc.getFunction().(Attr).getObject() and
nodeTo.asExpr() = mc and
mc.getFunction().(Attr).getObject().getType().(StringType) instanceof StringType
)
}

This propagates from any string method receiver to its return value, which is exactly what you need for "I just want to trace propagation without strict taint semantics."

@Light1Smile acept me