From 73632517cd769da73a062459d056b82ec7bc6a54 Mon Sep 17 00:00:00 2001 From: zhun Date: Tue, 13 Jan 2026 17:14:45 +0000 Subject: [PATCH] update cybergym/oss-fuzz subset --- src/cybergym/server/__main__.py | 10 +- src/cybergym/server/server_utils.py | 151 +++++++--------------------- 2 files changed, 40 insertions(+), 121 deletions(-) diff --git a/src/cybergym/server/__main__.py b/src/cybergym/server/__main__.py index d171dd6..63509c1 100644 --- a/src/cybergym/server/__main__.py +++ b/src/cybergym/server/__main__.py @@ -17,7 +17,6 @@ from cybergym.task.types import DEFAULT_SALT SALT = DEFAULT_SALT LOG_DIR = Path("./logs") DB_PATH = Path("./poc.db") -OSS_FUZZ_PATH = Path("./oss-fuzz-data") API_KEY = "cybergym-030a0cd7-5908-4862-8ab9-91f2bfc7b56d" API_KEY_NAME = "X-API-Key" @@ -66,7 +65,7 @@ def submit_vul(db: SessionDep, metadata: Annotated[str, Form()], file: Annotated except Exception: raise HTTPException(status_code=400, detail="Invalid metadata format") from None payload.data = file.file.read() - res = submit_poc(db, payload, mode="vul", log_dir=LOG_DIR, salt=SALT, oss_fuzz_path=OSS_FUZZ_PATH) + res = submit_poc(db, payload, mode="vul", log_dir=LOG_DIR, salt=SALT) res = _post_process_result(res, payload.require_flag) return res @@ -78,7 +77,7 @@ def submit_fix(db: SessionDep, metadata: Annotated[str, Form()], file: Annotated except Exception: raise HTTPException(status_code=400, detail="Invalid metadata format") from None payload.data = file.file.read() - res = submit_poc(db, payload, mode="fix", log_dir=LOG_DIR, salt=SALT, oss_fuzz_path=OSS_FUZZ_PATH) + res = submit_poc(db, payload, mode="fix", log_dir=LOG_DIR, salt=SALT) res = _post_process_result(res, payload.require_flag) return res @@ -101,7 +100,7 @@ def verify_all_pocs_for_agent_id(db: SessionDep, query: VerifyPocs): raise HTTPException(status_code=404, detail="No records found for this agent_id") for record in records: - run_poc_id(db, LOG_DIR, record.poc_id, oss_fuzz_path=OSS_FUZZ_PATH) + run_poc_id(db, LOG_DIR, record.poc_id) return { "message": f"All {len(records)} PoCs for this agent_id have been verified", @@ -120,7 +119,6 @@ if __name__ == "__main__": parser.add_argument("--salt", type=str, default=SALT, help="Salt for checksum") parser.add_argument("--log_dir", type=Path, default=LOG_DIR, help="Directory to store logs") parser.add_argument("--db_path", type=Path, default=DB_PATH, help="Path to SQLite DB") - parser.add_argument("--cybergym_oss_fuzz_path", type=Path, default=OSS_FUZZ_PATH, help="Path to OSS-Fuzz") args = parser.parse_args() SALT = args.salt @@ -129,6 +127,4 @@ if __name__ == "__main__": DB_PATH = Path(args.db_path) - OSS_FUZZ_PATH = Path(args.cybergym_oss_fuzz_path) - uvicorn.run(app, host=args.host, port=args.port) diff --git a/src/cybergym/server/server_utils.py b/src/cybergym/server/server_utils.py index a9b0e79..153772b 100644 --- a/src/cybergym/server/server_utils.py +++ b/src/cybergym/server/server_utils.py @@ -1,6 +1,5 @@ import hashlib -import json -import os +import shlex from enum import IntEnum from pathlib import Path from typing import Literal @@ -40,19 +39,45 @@ def _post_process_result(res: dict, require_flag: bool = False): return res -def run_arvo_container( +def _image_and_command_from_task_id(task_id: str, mode: str) -> tuple[str, list[str]]: + if task_id.startswith("arvo:"): + arvo_id = get_arvo_id(task_id) + image = f"n132/arvo:{arvo_id}-{mode}" + command = ["/bin/arvo"] + elif task_id.startswith("oss-fuzz:"): + oss_fuzz_id = get_oss_fuzz_id(task_id) + image = f"cybergym/oss-fuzz:{oss_fuzz_id}-{mode}" + command = ["/usr/local/bin/run_poc"] + elif task_id.startswith("oss-fuzz-latest:"): + raise HTTPException(status_code=400, detail="oss-fuzz-latest does not support this operation") + else: + raise HTTPException(status_code=400, detail="Invalid task_id") + return image, command + + +def is_integer(s): + try: + int(s) + return True + except ValueError: + return False + + + +def run_container( + task_id: str, poc_path: Path, - arvo_id: str, mode: Literal["vul", "fix"], docker_timeout: int = DEFAULT_DOCKER_TIMEOUT, cmd_timeout: int = DEFAULT_CMD_TIMEOUT, ): + image, cmd = _image_and_command_from_task_id(task_id, mode) + cmd = ["/bin/bash", "-c", f"timeout -s SIGKILL {cmd_timeout} {shlex.join(cmd)} 2>&1"] client = docker.from_env() container = None try: - cmd = ["/bin/bash", "-c", f"timeout -s SIGKILL {cmd_timeout} /bin/arvo 2>&1"] container = client.containers.run( - image=f"n132/arvo:{arvo_id}-{mode}", + image=image, command=cmd, volumes={str(poc_path.absolute()): {"bind": "/tmp/poc", "mode": "ro"}}, # noqa: S108 detach=True, @@ -74,109 +99,7 @@ def run_arvo_container( if container: container.remove(force=True) - return exit_code, docker_output - - -def is_integer(s): - try: - int(s) - return True - except ValueError: - return False - - -def run_oss_fuzz_container( - poc_path: Path, - oss_fuzz_id: str, - mode: Literal["vul", "fix"], - oss_fuzz_path: Path, - docker_timeout: int = DEFAULT_DOCKER_TIMEOUT, - cmd_timeout: int = DEFAULT_CMD_TIMEOUT, -): - client = docker.from_env() - container = None - try: - if is_integer(oss_fuzz_id): - out_dir = Path(oss_fuzz_path, f"{oss_fuzz_id}-{mode}", "out") - else: - if mode == "fix": - raise HTTPException(status_code=400, detail="Fix mode is not supported for oss-fuzz-latest") - project, index = oss_fuzz_id.rsplit("-", 1) - out_dir = Path(oss_fuzz_path, project, "out") - volumes = {str(poc_path.absolute()): {"bind": "/testcase", "mode": "ro"}} - for filename in os.listdir(out_dir): - host_path = str(Path(out_dir, filename).absolute()) - container_path = os.path.join("/out", filename) - volumes[host_path] = {"bind": container_path, "mode": "ro"} - if is_integer(oss_fuzz_id): - meta_file = os.path.join(oss_fuzz_path, f"{oss_fuzz_id}-{mode}", "metadata.json") - with open(meta_file) as f: - metadata = json.load(f) - fuzzer_name = metadata["fuzz_target"] - else: - project, index = oss_fuzz_id.rsplit("-", 1) - meta_file = os.path.join(oss_fuzz_path, project, "metadata.json") - with open(meta_file) as f: - metadata = json.load(f) - fuzzer_name = metadata["fuzz_targets"][int(index)] - - cmd = ["/bin/bash", "-c", f"timeout -s SIGKILL {cmd_timeout} reproduce {fuzzer_name} 2>&1"] - container = client.containers.run( - image="cybergym/oss-fuzz-base-runner", - command=cmd, - volumes=volumes, - detach=True, - ) - out = container.logs(stdout=True, stderr=False, stream=True, follow=True) - exit_code = container.wait(timeout=docker_timeout)["StatusCode"] - if exit_code == 137: # Process killed by timeout - exit_code = CustomExitCode.Timeout - docker_output = b"" - else: - docker_output = b"".join(out) - except requests.exceptions.ReadTimeout: - raise HTTPException(status_code=500, detail="Timeout waiting for the program") from None - except DockerException as e: - raise HTTPException(status_code=500, detail=f"Running error: {e}") from None - except Exception as e: - raise HTTPException(status_code=500, detail=f"Unexpected error: {e}") from None - finally: - if container: - container.remove(force=True) - - return exit_code, docker_output - - -def run_container( - task_id: str, - poc_path: Path, - mode: Literal["vul", "fix"], - docker_timeout: int = DEFAULT_DOCKER_TIMEOUT, - cmd_timeout: int = DEFAULT_CMD_TIMEOUT, - **kwargs, -): - if task_id.startswith("arvo:"): - arvo_id = get_arvo_id(task_id) - return run_arvo_container( - poc_path, - arvo_id, - mode, - docker_timeout=docker_timeout, - cmd_timeout=cmd_timeout, - ) - elif task_id.startswith("oss-fuzz:") or task_id.startswith("oss-fuzz-latest:"): - oss_fuzz_id = get_oss_fuzz_id(task_id) - oss_fuzz_path = kwargs.get("oss_fuzz_path") - return run_oss_fuzz_container( - poc_path, - oss_fuzz_id, - mode, - oss_fuzz_path, - docker_timeout=docker_timeout, - cmd_timeout=cmd_timeout, - ) - else: - raise HTTPException(status_code=400, detail="Invalid task_id") + return exit_code, docker_output def get_poc_storage_path(poc_id: str, log_dir: Path): @@ -184,7 +107,7 @@ def get_poc_storage_path(poc_id: str, log_dir: Path): return log_dir / poc_id[:2] / poc_id[2:4] / poc_id -def submit_poc(db: Session, payload: Payload, mode: str, log_dir: Path, salt: str, oss_fuzz_path: Path | None = None): +def submit_poc(db: Session, payload: Payload, mode: str, log_dir: Path, salt: str): # TODO: limit output size for return if not verify_task(payload.task_id, payload.agent_id, payload.checksum, salt=salt): raise HTTPException(status_code=400, detail="Invalid checksum") @@ -239,7 +162,7 @@ def submit_poc(db: Session, payload: Payload, mode: str, log_dir: Path, salt: st ) # Run the PoC - exit_code, docker_output = run_container(payload.task_id, poc_bin_file, mode, oss_fuzz_path=oss_fuzz_path) + exit_code, docker_output = run_container(payload.task_id, poc_bin_file, mode) output_file = poc_dir / f"output.{mode}" with open(output_file, "wb") as f: f.write(docker_output) @@ -255,7 +178,7 @@ def submit_poc(db: Session, payload: Payload, mode: str, log_dir: Path, salt: st return res -def run_poc_id(db: Session, log_dir: Path, poc_id: str, rerun: bool = False, oss_fuzz_path: Path | None = None): +def run_poc_id(db: Session, log_dir: Path, poc_id: str, rerun: bool = False): records = db.query(PoCRecord).filter_by(poc_id=poc_id).all() if len(records) != 1: raise HTTPException(status_code=500, detail=f"{len(records)} PoC records for same poc_id found") @@ -268,7 +191,7 @@ def run_poc_id(db: Session, log_dir: Path, poc_id: str, rerun: bool = False, oss if rerun or record.vul_exit_code is None: # Run the PoC - exit_code, docker_output = run_container(record.task_id, poc_path, "vul", oss_fuzz_path=oss_fuzz_path) + exit_code, docker_output = run_container(record.task_id, poc_path, "vul") with open(poc_dir / "output.vul", "wb") as f: f.write(docker_output) update_poc_output(db, record, "vul", exit_code) @@ -279,7 +202,7 @@ def run_poc_id(db: Session, log_dir: Path, poc_id: str, rerun: bool = False, oss if rerun or record.fix_exit_code is None: # Run the PoC - exit_code, docker_output = run_container(record.task_id, poc_path, "fix", oss_fuzz_path=oss_fuzz_path) + exit_code, docker_output = run_container(record.task_id, poc_path, "fix") with open(poc_dir / "output.fix", "wb") as f: f.write(docker_output) update_poc_output(db, record, "fix", exit_code)