iohex/cybergym
1
0
Fork 0
mirror of https://github.com/sunblaze-ucb/cybergym.git synced 2026-10-12 08:42:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

update cybergym/oss-fuzz subset

Browse Source
This commit is contained in:
zhun
2026-01-13 17:14:45 +00:00
parent df48a05709
commit 73632517cd
2 changed files with 40 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
src/cybergym/server/__main__.py
View File

@@ -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)

151
src/cybergym/server/server_utils.py
View File

@@ -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)