"""CodeQuery primitives testing""" import pytest import subprocess from pathlib import Path from buttercup.common.challenge_task import ChallengeTask from buttercup.program_model.codequery import CodeQuery, CodeQueryPersistent from buttercup.common.task_meta import TaskMeta def setup_dirs(tmp_path: Path) -> Path: """Create a mock challenge task directory structure.""" # Create the main directories oss_fuzz = tmp_path / "fuzz-tooling" / "my-oss-fuzz" source = tmp_path / "src" / "my-source" diffs = tmp_path / "diff" / "my-diff" oss_fuzz.mkdir(parents=True, exist_ok=True) source.mkdir(parents=True, exist_ok=True) diffs.mkdir(parents=True, exist_ok=True) # Create a mock project.yaml file project_yaml_path = oss_fuzz / "projects" / "example_project" / "project.yaml" project_yaml_path.parent.mkdir(parents=True, exist_ok=True) project_yaml_path.write_text("language: c\n") # Create some mock patch files (diffs / "patch1.diff").write_text("mock patch 1") (diffs / "patch2.diff").write_text("mock patch 2") # Create a mock helper.py file helper_path = oss_fuzz / "infra/helper.py" helper_path.parent.mkdir(parents=True, exist_ok=True) helper_path.write_text("import sys;\nsys.exit(0)\n") # Create a mock test.txt file (source / "test.c").write_text("int main() { return 0; }") (source / "test2.c").write_text("""int function2(int a, int b) { int c = a + b; return c; } """) (source / "test3.c").write_text("""int function3(int a, int b) { int c = a + b; return c; } int function4(char *s) { return strlen(s); } """) # Create task metadata TaskMeta( project_name="example_project", focus="my-source", task_id="task-id-challenge-task", ).save(tmp_path) return tmp_path @pytest.fixture def task_dir(tmp_path: Path) -> Path: return setup_dirs(tmp_path / "task_rw") @pytest.fixture def task_dir_ro(tmp_path: Path) -> Path: return setup_dirs(tmp_path / "task_ro") @pytest.fixture def mock_challenge_task(task_dir: Path) -> ChallengeTask: """Create a mock challenge task""" return ChallengeTask(task_dir, local_task_dir=task_dir) @pytest.fixture def mock_challenge_task_ro(task_dir_ro: Path) -> ChallengeTask: """Create a mock challenge task""" return ChallengeTask(task_dir_ro, local_task_dir=task_dir_ro) def test_get_functions_simple(mock_challenge_task: ChallengeTask): """Test that we can get the main function""" codequery = CodeQuery(mock_challenge_task) main_functions = codequery.get_functions("main") assert len(main_functions) == 1 assert main_functions[0].name == "main" assert len(main_functions[0].bodies) == 1 assert main_functions[0].bodies[0].body == "int main() { return 0; }" def test_get_functions_file(mock_challenge_task: ChallengeTask): """Test that we can get the main function from a specific file""" codequery = CodeQuery(mock_challenge_task) main_functions = codequery.get_functions("main", Path("test.c")) assert len(main_functions) == 1 assert main_functions[0].name == "main" assert len(main_functions[0].bodies) == 1 assert main_functions[0].bodies[0].body == "int main() { return 0; }" def test_get_functions_multiple(mock_challenge_task: ChallengeTask): """Test that we can get multiple functions from a file""" codequery = CodeQuery(mock_challenge_task) function3 = codequery.get_functions("function3", Path("test3.c")) assert len(function3) == 1 assert function3[0].name == "function3" assert ( function3[0].bodies[0].body == "int function3(int a, int b) {\n int c = a + b;\n return c;\n}" ) function4 = codequery.get_functions("function4", Path("test3.c")) assert len(function4) == 1 assert function4[0].name == "function4" assert ( function4[0].bodies[0].body == "int function4(char *s) {\n return strlen(s);\n}" ) def test_keep_status( mock_challenge_task: ChallengeTask, mock_challenge_task_ro: ChallengeTask, tmp_path: Path, ): """Test that we can access the same db from different instances""" wdir = tmp_path wdir.mkdir(parents=True, exist_ok=True) codequery = CodeQueryPersistent(mock_challenge_task, work_dir=wdir) assert codequery.get_functions("main") assert mock_challenge_task.task_dir.exists() codequery2 = CodeQueryPersistent(mock_challenge_task_ro, work_dir=wdir) assert codequery2.get_functions("main") assert codequery2.challenge.task_dir == codequery.challenge.task_dir assert mock_challenge_task.task_dir.exists() assert mock_challenge_task_ro.task_dir.exists() with mock_challenge_task_ro.get_rw_copy( mock_challenge_task_ro.task_dir.parent ) as nd_challenge: codequery3 = CodeQueryPersistent(nd_challenge, work_dir=wdir) assert codequery3.get_functions("main") assert codequery3.challenge.task_dir == codequery.challenge.task_dir assert mock_challenge_task.task_dir.exists() assert mock_challenge_task_ro.task_dir.exists() with mock_challenge_task.get_rw_copy( mock_challenge_task.task_dir.parent ) as nd_challenge: codequery4 = CodeQueryPersistent(nd_challenge, work_dir=wdir) assert codequery4.get_functions("main") assert codequery4.challenge.task_dir == codequery.challenge.task_dir assert mock_challenge_task.task_dir.exists() assert mock_challenge_task_ro.task_dir.exists() @pytest.fixture def libjpeg_oss_fuzz_task(tmp_path: Path) -> ChallengeTask: """Create a challenge task using a real OSS-Fuzz repository.""" # Clone real oss-fuzz repo into temp dir oss_fuzz_dir = tmp_path / "fuzz-tooling" oss_fuzz_dir.mkdir(parents=True) source_dir = tmp_path / "src" source_dir.mkdir(parents=True) subprocess.run( [ "git", "-C", str(oss_fuzz_dir), "clone", "https://github.com/google/oss-fuzz.git", ], check=True, ) # Restore libjpeg-turbo project directory to specific commit subprocess.run( [ "git", "-C", str(oss_fuzz_dir / "oss-fuzz"), "checkout", "7e664533834b558a859b0f8eb1f2c2caf676c12a", "--", "projects/libjpeg-turbo", ], check=True, ) # Download libpng source code libjpeg_url = "https://github.com/libjpeg-turbo/libjpeg-turbo" # Checkout specific libjpeg commit for reproducibility subprocess.run(["git", "-C", str(source_dir), "clone", libjpeg_url], check=True) subprocess.run( [ "git", "-C", str(source_dir / "libjpeg-turbo"), "checkout", "6d91e950c871103a11bac2f10c63bf998796c719", ], check=True, ) # Create task metadata TaskMeta( project_name="libjpeg-turbo", focus="libjpeg-turbo", task_id="task-id-libjpeg-turbo", ).save(tmp_path) return ChallengeTask( read_only_task_dir=tmp_path, local_task_dir=tmp_path, ) @pytest.mark.integration def test_libjpeg_indexing(libjpeg_oss_fuzz_task: ChallengeTask): """Test that we can index libjpeg""" codequery = CodeQuery(libjpeg_oss_fuzz_task) functions = codequery.get_functions("jpeg_read_header") assert len(functions) == 1 assert functions[0].name == "jpeg_read_header" parse_switches = codequery.get_functions("parse_switches") assert len(parse_switches) == 3 parse_switches.sort(key=lambda x: x.file_path) assert parse_switches[0].file_path == Path("src/libjpeg-turbo/cjpeg.c") assert parse_switches[0].file_path.name == "cjpeg.c" assert len(parse_switches[0].bodies) == 1 assert ( "parse_switches(j_compress_ptr cinfo, int argc, char **argv," in parse_switches[0].bodies[0].body ) assert parse_switches[1].file_path.name == "djpeg.c" assert parse_switches[1].name == "parse_switches" assert len(parse_switches[1].bodies) == 1 assert "/* Parse optional switches." in parse_switches[1].bodies[0].body assert ( """LOCAL(int) parse_switches(j_decompress_ptr cinfo, int argc, char **argv, int last_file_arg_seen, boolean for_real) """ in parse_switches[1].bodies[0].body ) assert ( ' } else if (keymatch(arg, "crop", 2)) {' in parse_switches[1].bodies[0].body ) assert ( "return argn; /* return index of next arg (file name) */" in parse_switches[1].bodies[0].body ) assert parse_switches[2].file_path.name == "jpegtran.c" assert len(parse_switches[2].bodies) == 1 assert ( "parse_switches(j_compress_ptr cinfo, int argc, char **argv," in parse_switches[2].bodies[0].body ) @pytest.fixture def selinux_oss_fuzz_task(tmp_path: Path) -> ChallengeTask: """Create a challenge task using a real OSS-Fuzz repository.""" oss_fuzz_dir = tmp_path / "fuzz-tooling" oss_fuzz_dir.mkdir(parents=True) source_dir = tmp_path / "src" source_dir.mkdir(parents=True) subprocess.run( [ "git", "-C", str(oss_fuzz_dir), "clone", "https://github.com/google/oss-fuzz.git", ], check=True, ) # Restore libjpeg-turbo project directory to specific commit subprocess.run( [ "git", "-C", str(oss_fuzz_dir / "oss-fuzz"), "checkout", "ef2f42b3b10af381d3d55cc901fde0729e54573b", "--", "projects/selinux", ], check=True, ) # Download selinux source code url = "https://github.com/SELinuxProject/selinux" subprocess.run(["git", "-C", str(source_dir), "clone", url], check=True) subprocess.run( [ "git", "-C", str(source_dir / "selinux"), "checkout", "c35919a703302bd571476f245d856174a1fe1926", ], check=True, ) # Create task metadata TaskMeta(project_name="selinux", focus="selinux", task_id="task-id-selinux").save( tmp_path ) return ChallengeTask( read_only_task_dir=tmp_path, local_task_dir=tmp_path, ) @pytest.mark.integration def test_selinux_indexing(selinux_oss_fuzz_task: ChallengeTask): """Test that we can index selinux and files inside oss-fuzz repo""" codequery = CodeQuery(selinux_oss_fuzz_task) functions = codequery.get_functions("LLVMFuzzerTestOneInput") assert len(functions) == 1 assert functions[0].name == "LLVMFuzzerTestOneInput" assert functions[0].file_path == Path( "fuzz-tooling/oss-fuzz/projects/selinux/secilc-fuzzer.c" ) assert len(functions[0].bodies) == 1 assert ( "if (sepol_policydb_optimize(pdb) != SEPOL_OK)" in functions[0].bodies[0].body )