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()
    )
  }