testgen need to ^ help: add ; here

Author: joyguoguo

fuzz/collect_fuzz_python.py

@@ -15,9 +15,16 @@
 from difflib import SequenceMatcher
 from itertools import islice
 from datetime import datetime
+import re
 
 def build_image(repos: list[str], jobs: int):
-    """构建每个仓库对应的OSS-Fuzz项目的Docker镜像"""
+    """
+    构建每个仓库对应的OSS-Fuzz项目的Docker镜像
+    
+    Args:
+        repos (list[str]): 仓库路径列表
+        jobs (int): 并行任务数
+    """
     logging.info(f"Building Docker images for {len(repos)} OSS-Fuzz projects")
     log_dir = os.path.abspath("fuzz_pipeline_log")
     os.makedirs(log_dir, exist_ok=True)
@@ -26,6 +33,8 @@ def _build_cmd(path: str):
         project_name = os.path.basename(path.rstrip("/"))
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
         log_file = os.path.join(log_dir, f"{project_name}_{timestamp}.log")
+
+        logging.info(f"Start building {project_name}, logging to {log_file}")
         return subprocess.Popen(
             f"yes | python3 infra/helper.py build_image {project_name}",
             cwd=os.path.abspath(os.path.join(path, "../../")),
@@ -37,7 +46,13 @@ def _build_cmd(path: str):
     _ = parallel_subprocess(repos, jobs, _build_cmd, on_exit=None)
 
 def build_fuzzer(repos: list[str], jobs: int):
-    """对构建成功的项目并行构建模糊测试器"""
+    """
+    对构建成功的项目并行构建模糊测试器
+    
+    Args:
+        repos (list[str]): 仓库路径列表
+        jobs (int): 并行任务数
+    """
     logging.info(f"Building fuzzers for {len(repos)} OSS-Fuzz projects")
     log_dir = os.path.abspath("fuzz_pipeline_log")
     os.makedirs(log_dir, exist_ok=True)
@@ -46,6 +61,8 @@ def _build_cmd(path: str):
         project_name = os.path.basename(path.rstrip("/"))
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
         log_file = os.path.join(log_dir, f"{project_name}_fuzzer_{timestamp}.log")
+
+        logging.info(f"Start building fuzzers for {project_name}, logging to {log_file}")
         return subprocess.Popen(
             f"python3 infra/helper.py build_fuzzers --sanitizer address {project_name}",
             cwd=os.path.abspath(os.path.join(path, "../../")),
@@ -57,7 +74,16 @@ def _build_cmd(path: str):
     _ = parallel_subprocess(repos, jobs, _build_cmd, on_exit=None)
 
 def discover_targets(project_name: str, oss_fuzz_dir: Path) -> list[str]:
-    """发现模糊测试目标"""
+    """
+    发现模糊测试目标
+    
+    Args:
+        project_name (str): 项目名称
+        oss_fuzz_dir (Path): OSS-Fuzz根目录
+        
+    Returns:
+        list[str]: 目标名称列表
+    """
     out_dir = oss_fuzz_dir / "build" / "out" / project_name
     targets = []
 
@@ -77,7 +103,16 @@ def discover_targets(project_name: str, oss_fuzz_dir: Path) -> list[str]:
     return targets
 
 def fuzz_one_target(target: tuple[str, str], timeout: int):
-    """对单个模糊测试目标执行模糊测试"""
+    """
+    对单个模糊测试目标执行模糊测试
+    
+    Args:
+        target (tuple[str, str]): (仓库路径, 目标名称)
+        timeout (int): 超时时间(秒)
+        
+    Returns:
+        subprocess.Popen: 子进程对象
+    """
     repo_path, target_name = target
     project_name = os.path.basename(repo_path)
     oss_fuzz_root = os.path.dirname(os.path.dirname(repo_path))
@@ -103,7 +138,14 @@ def fuzz_one_target(target: tuple[str, str], timeout: int):
         return None
 
 def fuzz_repos(repos: list[str], jobs: int, timeout: int = 60):
-    """对一组仓库执行模糊测试"""
+    """
+    对一组仓库执行模糊测试
+    
+    Args:
+        repos (list[str]): 仓库路径列表
+        jobs (int): 并行任务数
+        timeout (int): 超时时间(秒)
+    """
     logging.info("Discovering fuzz targets")
 
     # 获取所有目标
@@ -130,17 +172,26 @@ def fuzz_repos(repos: list[str], jobs: int, timeout: int = 60):
     parallel_subprocess(targets, jobs, lambda p: fuzz_one_target(p, timeout), on_exit=None)
 
 def generate_test_template(target_name: str, repo_path: str):
-    """为单个目标生成测试模板"""
+    """
+    为单个目标生成测试模板
+    
+    Args:
+        target_name (str): 目标名称
+        repo_path (str): 仓库路径
+        
+    Returns:
+        str: 模板文件路径
+    """
     template_dir = pjoin(repo_path, "tests-gen")
     os.makedirs(template_dir, exist_ok=True)
     template_path = pjoin(template_dir, f"{target_name}.rs")
 
-    # 基本测试模板
+    # 基本测试模板 - 使用字节数组而不是字节字符
     template = f"""
     #[test]
     fn test_{target_name}() {{
         // 测试逻辑将在这里生成
-        let input = []; // 模糊测试输入将替换这里
+        let input = b""; // 模糊测试输入将替换这里
         let result = process_input(&input);
         assert!(result.is_ok());
     }}
@@ -152,7 +203,13 @@ def generate_test_template(target_name: str, repo_path: str):
     return template_path
 
 def transform_repos(repos: list[str], jobs: int):
-    """为所有目标生成测试模板"""
+    """
+    为所有目标生成测试模板
+    
+    Args:
+        repos (list[str]): 仓库路径列表
+        jobs (int): 并行任务数
+    """
     logging.info("Generating test templates")
 
     def _transform_repo(repo: str):
@@ -164,18 +221,67 @@ def _transform_repo(repo: str):
     with ProcessingPool(jobs) as p:
         return list(p.map(_transform_repo, repos))
 
-def substitute_input(template: str, input_data: str, idx: int) -> str:
-    """将模糊测试输入替换到测试模板中"""
-    return template.replace(
-        'let input = []; // 模糊测试输入将替换这里',
-        f"let input = {input_data};"
-    ).replace(
-        f"fn test_",
-        f"fn test_{idx}_"
+def escape_special_chars(input_data: str) -> str:
+    """
+    转义输入数据中的特殊字符
+    
+    Args:
+        input_data (str): 原始输入数据
+        
+    Returns:
+        str: 转义后的输入数据
+    """
+    # 转义反斜杠和双引号
+    escaped = input_data.replace('\\', '\\\\').replace('"', '\\"')
+    
+    # 处理非ASCII字符
+    if any(ord(c) > 127 for c in escaped):
+        # 如果包含非ASCII字符，使用字节数组表示
+        byte_array = [str(b) for b in input_data.encode()]
+        return f"b\"\" // Original: {input_data}\n        let input = vec![{', '.join(byte_array)}];"
+    
+    return f"b\"{escaped}\""
+
+def substitute_input(template: str, input_data: str, idx: int, target_name: str) -> str:
+    """
+    将模糊测试输入替换到测试模板中
+    
+    Args:
+        template (str): 模板内容
+        input_data (str): 输入数据
+        idx (int): 测试索引
+        target_name (str): 目标名称
+        
+    Returns:
+        str: 替换后的测试代码
+    """
+    # 转义特殊字符并处理非ASCII字符
+    escaped_input = escape_special_chars(input_data)
+    
+    # 替换输入占位符
+    new_template = template.replace(
+        'let input = b""; // 模糊测试输入将替换这里',
+        escaped_input
+    )
+    
+    # 替换函数名避免重复
+    return new_template.replace(
+        f"fn test_{target_name}()",
+        f"fn test_{target_name}_{idx}()"
     )
 
 def has_similar(selected: list[str], x: str, thresh: float = 0.8) -> bool:
-    """检查字符串是否与已选列表中的任何字符串足够相似"""
+    """
+    检查字符串是否与已选列表中的任何字符串足够相似
+    
+    Args:
+        selected (list[str]): 已选字符串列表
+        x (str): 待检查字符串
+        thresh (float): 相似度阈值
+        
+    Returns:
+        bool: 是否相似
+    """
     def similar(a, b):
         return SequenceMatcher(None, a, b).ratio()
     return any(similar(x, y) > thresh for y in selected)
@@ -188,21 +294,45 @@ def substitute_one_repo(
     max_len: int,
     sim_thresh: float,
 ):
-    """处理单个仓库，将模糊测试输入替换到测试模板中"""
+    """
+    处理单个仓库，将模糊测试输入替换到测试模板中
+    
+    Args:
+        repo (str): 仓库路径
+        targets (list[str]): 目标列表
+        n_fuzz (int): 使用的输入数量
+        strategy (str): 选择策略
+        max_len (int): 最大长度
+        sim_thresh (float): 相似度阈值
+    """
     template_dir = pjoin(repo, "tests-gen")
     input_dir = pjoin(repo, "fuzz_inputs")
 
-    for t in targets:
-        template_path = pjoin(template_dir, f"{t}.rs")
-        input_path = pjoin(input_dir, t)
+    for target_name in targets:  # 使用target_name作为循环变量
+        template_path = pjoin(template_dir, f"{target_name}.rs")
+        input_path = pjoin(input_dir, target_name)
 
         try:
+            if not os.path.exists(template_path):
+                logging.warning(f"Template file not found: {template_path}")
+                continue
+                
+            if not os.path.exists(input_path):
+                logging.warning(f"Input file not found: {input_path}")
+                continue
+                
             with open(template_path) as f_template:
                 template = f_template.read()
 
             with open(input_path, "r") as f_input:
-                all_inputs = [i for i in f_input.read().splitlines() if i]
+                all_inputs = [line.strip() for line in f_input if line.strip()]
 
+            if not all_inputs:
+                logging.warning(f"No valid inputs found for {target_name}")
+                continue
+                
+            logging.info(f"Loaded {len(all_inputs)} inputs for {target_name}")
+
             # 选择输入策略
             if strategy == "shuffle":
                 random.shuffle(all_inputs)
@@ -221,20 +351,20 @@ def substitute_one_repo(
 
             # 生成测试用例
             tests = [
-                substitute_input(template, input_data, i)
+                substitute_input(template, input_data, i, target_name)  # 传递target_name
                 for i, input_data in enumerate(inputs)
             ]
 
             # 写入生成的测试文件
-            generated_path = pjoin(template_dir, f"{t}.inputs.rs")
+            generated_path = pjoin(template_dir, f"{target_name}.inputs.rs")
             with open(generated_path, "w") as f:
                 f.write("\n".join(tests))
 
             # 格式化代码
             subprocess.run(["rustfmt", generated_path], check=False)
 
         except Exception as e:
-            logging.error(f"Error processing {t}: {e}")
+            logging.error(f"Error processing {target_name}: {e}")
 
 def testgen_repos(
     repos: list[str],
@@ -244,7 +374,17 @@ def testgen_repos(
     max_len: int = 100,
     sim_thresh: float = 0.8,
 ):
-    """从模糊测试输入生成测试用例"""
+    """
+    从模糊测试输入生成测试用例
+    
+    Args:
+        repos (list[str]): 仓库路径列表
+        jobs (int): 并行任务数
+        n_fuzz (int): 使用的输入数量
+        strategy (str): 选择策略
+        max_len (int): 最大长度
+        sim_thresh (float): 相似度阈值
+    """
     # 首先获取所有目标
     targets_list = []
     for repo in repos:
@@ -275,7 +415,20 @@ def main(
     max_len: int = 100,
     sim_thresh: float = 0.8,
 ):
-    """主函数，控制整个模糊测试流程"""
+    """
+    主函数，控制整个模糊测试流程
+    
+    Args:
+        repo_id (str): 项目ID文件路径
+        repo_root (str): 项目根目录
+        timeout (int): 超时时间
+        jobs (int): 并行任务数
+        pipeline (str): 流程类型
+        n_fuzz (int): 使用的输入数量
+        strategy (str): 选择策略
+        max_len (int): 最大长度
+        sim_thresh (float): 相似度阈值
+    """
     try:
         with open(repo_id, "r") as f:
             repo_id_list = [line.strip() for line in f if line.strip()]