update cybergym/oss-fuzz subset

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)

src/cybergym/server/server_utils.py

@@ -1,6 +1,5 @@
 import hashlib
-import json
+import shlex
-import os
 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,
@@ -77,114 +102,12 @@ def run_arvo_container(
     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")
-
-
 def get_poc_storage_path(poc_id: str, log_dir: Path):
     # logs/ab/cd/1234/...
     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)