如何在Pyre Pysa中将LLM调用同时定义为污点源与污点汇？

核心配置思路 #

要将LLM调用同时标记为污点源（返回内容可能携带不可信数据）和污点汇（传入参数可能触发风险），需在.pysa模型文件中准确定位LLM库的调用方法，同时在taint.config中关联对应的源/汇类型，全程无需修改目标仓库代码。

步骤1：编写正确的LLM模型定义（llms.pysa） #

针对OpenAI和Anthropic的常见调用格式，直接在模型文件中声明方法的污点属性：

# OpenAI v1.x 接口
class openai.OpenAI:
    def chat.completions.create(
        self,
        messages: TaintSource[UserControlled] + TaintSink[LLMInput]
    ) -> TaintSource[LLMGenerated]: ...

# OpenAI 旧版接口
def openai.ChatCompletion.create(
    messages: TaintSource[UserControlled] + TaintSink[LLMInput]
) -> TaintSource[LLMGenerated]: ...

# Anthropic 接口
class anthropic.Anthropic:
    def messages.create(
        self,
        messages: TaintSource[UserControlled] + TaintSink[LLMInput],
        prompt: TaintSource[UserControlled] + TaintSink[LLMInput]
    ) -> TaintSource[LLMGenerated]: ...

注意：必须严格匹配LLM库实际的类名、方法名和参数签名，比如OpenAI版本差异导致的方法路径变化，Anthropic的messages和prompt参数需分别标记。

步骤2：更新taint.config配置 #

在taint.config中定义自定义污点类型，并关联模型文件：

# 定义污点源类型
taint_source = [
    {
        "name": "UserControlled",
        "description": "外部用户可控输入"
    },
    {
        "name": "LLMGenerated",
        "description": "LLM生成的内容"
    }
]

# 定义污点汇类型
taint_sink = [
    {
        "name": "LLMInput",
        "description": "传入LLM的输入参数"
    }
]

# 加载自定义模型文件
model_paths = [
    "llms.pysa",
    # 保留已有的其他模型文件路径
]

步骤3：排查pyre validate-models报错 #

如果运行验证命令报错，按以下方向排查：

• 方法签名不匹配：核对.pysa中的类、方法、参数名与LLM库源码完全一致，注意大小写、参数顺序差异。
• 模型路径未加载：确认taint.config的model_paths包含llms.pysa的正确相对/绝对路径。
• 自定义污点类型未定义：检查taint.config中是否正确声明了UserControlled、LLMGenerated、LLMInput类型。

测试代码验证 #

用以下代码测试配置是否生效：

import openai
import anthropic
import os

def test_openai_flow():
    user_input = input("输入提示词: ")
    client = openai.OpenAI()
    response = client.chat.completions.create(messages=[{"role": "user", "content": user_input}])
    # 若配置正确，Pyre会检测到LLM生成内容传入os.system的污点传播
    os.system(response.choices[0].message.content)

def test_anthropic_flow():
    user_input = input("输入提示词: ")
    client = anthropic.Anthropic()
    response = client.messages.create(
        messages=[{"role": "user", "content": user_input}],
        model="claude-3-sonnet-20240229"
    )
    print(response.content[0].text)

运行pyre analyze后，应能检测到user_input传入LLM（污点汇）、LLM输出传入os.system（污点源到汇的风险传播）的问题。

内容的提问来源于stack exchange，提问作者f.k