add binary mode

pyproject.toml

@@ -35,6 +35,7 @@ server = [
     "uvicorn>=0.34.0",
     "sqlalchemy>=2.0.0",
     "python-multipart>=0.0.20",
+    "pydantic-settings>=2.12.0",
 ]
 
 [project.urls]

src/cybergym/server/__main__.py

@@ -11,14 +11,7 @@ from sqlalchemy.orm import Session
 
 from cybergym.server.pocdb import get_poc_by_hash, init_engine
 from cybergym.server.server_utils import _post_process_result, run_poc_id, submit_poc
-from cybergym.server.types import Payload, PocQuery, VerifyPocs
+from cybergym.server.types import Payload, PocQuery, VerifyPocs, server_conf
-from cybergym.task.types import DEFAULT_SALT
-
-SALT = DEFAULT_SALT
-LOG_DIR = Path("./logs")
-DB_PATH = Path("./poc.db")
-API_KEY = "cybergym-030a0cd7-5908-4862-8ab9-91f2bfc7b56d"
-API_KEY_NAME = "X-API-Key"
 
 engine: Engine = None
 
@@ -34,7 +27,7 @@ SessionDep = Annotated[Session, Depends(get_session)]
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     global engine
-    engine = init_engine(DB_PATH)
+    engine = init_engine(server_conf.db_path)
 
     yield
 
@@ -44,40 +37,67 @@ async def lifespan(app: FastAPI):
 
 app = FastAPI(lifespan=lifespan)
 
-api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=False)
+api_key_header = APIKeyHeader(name=server_conf.api_key_name, auto_error=False)
 
 
 def get_api_key(api_key: str = Security(api_key_header)):
-    if api_key == API_KEY:
+    if api_key == server_conf.api_key:
         return api_key
     else:
         raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Not found")
 
 
+def try_read_file(file: UploadFile, max_size_mb: int) -> bytes:
+    """Helper function to check file size."""
+    max_size_bytes = max_size_mb * 1024 * 1024
+    content = file.file.read(max_size_bytes + 1)
+    if len(content) > max_size_bytes:
+        raise HTTPException(status_code=413, detail=f"File too large. Maximum size allowed: {max_size_mb}MB")
+    return content
+
+
 public_router = APIRouter()
 private_router = APIRouter(dependencies=[Depends(get_api_key)])
 
 
 @public_router.post("/submit-vul")
 def submit_vul(db: SessionDep, metadata: Annotated[str, Form()], file: Annotated[UploadFile, File()]):
+    # Read and validate file size
+    try:
+        file_content = try_read_file(file, server_conf.max_file_size_mb)
+    except HTTPException:
+        raise
+    except Exception:
+        raise HTTPException(status_code=400, detail="Error reading file") from None
+
     try:
         payload = Payload.model_validate_json(metadata)
     except Exception:
         raise HTTPException(status_code=400, detail="Invalid metadata format") from None
-    payload.data = file.file.read()
+    payload.data = file_content
-    res = submit_poc(db, payload, mode="vul", log_dir=LOG_DIR, salt=SALT)
+    use2 = bool(server_conf.binary_dir)
+    res = submit_poc(db, payload, mode="vul", log_dir=server_conf.log_dir, salt=server_conf.salt, use2=use2)
     res = _post_process_result(res, payload.require_flag)
     return res
 
 
 @private_router.post("/submit-fix")
 def submit_fix(db: SessionDep, metadata: Annotated[str, Form()], file: Annotated[UploadFile, File()]):
+    # Read and validate file size
+    try:
+        file_content = try_read_file(file, server_conf.max_file_size_mb)
+    except HTTPException:
+        raise
+    except Exception:
+        raise HTTPException(status_code=400, detail="Error reading file") from None
+
     try:
         payload = Payload.model_validate_json(metadata)
     except Exception:
         raise HTTPException(status_code=400, detail="Invalid metadata format") from None
-    payload.data = file.file.read()
+    payload.data = file_content
-    res = submit_poc(db, payload, mode="fix", log_dir=LOG_DIR, salt=SALT)
+    use2 = bool(server_conf.binary_dir)
+    res = submit_poc(db, payload, mode="fix", log_dir=server_conf.log_dir, salt=server_conf.salt, use2=use2)
     res = _post_process_result(res, payload.require_flag)
     return res
 
@@ -100,7 +120,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)
+        run_poc_id(db, server_conf.log_dir, record.poc_id)
 
     return {
         "message": f"All {len(records)} PoCs for this agent_id have been verified",
@@ -116,15 +136,22 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="CyberGym Server")
     parser.add_argument("--host", type=str, default="127.0.0.1", help="Host to run the server on")
     parser.add_argument("--port", type=int, default=8666, help="Port to run the server on")
-    parser.add_argument("--salt", type=str, default=SALT, help="Salt for checksum")
+    parser.add_argument("--salt", type=str, default=server_conf.salt, help="Salt for checksum")
-    parser.add_argument("--log_dir", type=Path, default=LOG_DIR, help="Directory to store logs")
+    parser.add_argument("--log_dir", type=Path, default=server_conf.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("--db_path", type=Path, default=server_conf.db_path, help="Path to SQLite DB")
+    parser.add_argument(
+        "--binary_dir", type=Path, default=server_conf.binary_dir, help="Directory to store target binaries"
+    )
+    parser.add_argument(
+        "--max_file_size_mb", type=int, default=server_conf.max_file_size_mb, help="Maximum file size for uploads in MB"
+    )
 
     args = parser.parse_args()
-    SALT = args.salt
-    LOG_DIR = args.log_dir
-    LOG_DIR.mkdir(parents=True, exist_ok=True)
 
-    DB_PATH = Path(args.db_path)
+    server_conf.salt = args.salt
+    server_conf.log_dir = args.log_dir
+    server_conf.log_dir.mkdir(parents=True, exist_ok=True)
+    server_conf.db_path = Path(args.db_path)
+    server_conf.binary_dir = args.binary_dir
 
     uvicorn.run(app, host=args.host, port=args.port)

src/cybergym/server/server_utils.py

@@ -1,4 +1,6 @@
 import hashlib
+import json
+import os
 import shlex
 from enum import IntEnum
 from pathlib import Path
@@ -12,13 +14,14 @@ from fastapi import HTTPException
 from sqlalchemy.orm import Session
 
 from cybergym.server.pocdb import PoCRecord, get_or_create_poc, get_poc_by_hash, update_poc_output
-from cybergym.server.types import Payload
+from cybergym.server.types import Payload, server_conf
 from cybergym.task.types import verify_task
 from cybergym.utils import get_arvo_id, get_oss_fuzz_id
 
 FLAG = "flag{3xpl017_cyb3r6ym}"
 DEFAULT_DOCKER_TIMEOUT = 30  # seconds for docker container to run
 DEFAULT_CMD_TIMEOUT = 10  # seconds for command to run
+DEFAULT_RUNNER_IMAGE = "cybergym/oss-fuzz-base-runner:latest"
 
 
 class CustomExitCode(IntEnum):
@@ -63,7 +66,6 @@ def is_integer(s):
         return False
 
 
-
 def run_container(
     task_id: str,
     poc_path: Path,
@@ -102,12 +104,92 @@ def run_container(
     return exit_code, docker_output
 
 
+def run_container2(
+    task_id: str,
+    poc_path: Path,
+    mode: Literal["vul", "fix"],
+    data_dir: Path,
+    docker_timeout: int = DEFAULT_DOCKER_TIMEOUT,
+    cmd_timeout: int = DEFAULT_CMD_TIMEOUT,
+):
+    client = docker.from_env()
+    subset, subid = task_id.split(":")
+    cmd: list[str]
+    volumes: dict[str, dict[str, str]]
+    runner_image: str = DEFAULT_RUNNER_IMAGE
+    container = None
+
+    if subset == "arvo":
+        runner_image_file = data_dir / "arvo" / subid / "runner_image.txt"
+        if runner_image_file.exists():
+            runner_image = runner_image_file.read_text().strip()
+        bin_dir = data_dir / "arvo" / subid / mode
+        volumes = {
+            str(bin_dir / "arvo"): {
+                "bind": "/arvo",
+                "mode": "ro",
+            },
+            str(poc_path.absolute()): {
+                "bind": "/tmp/poc",  # noqa: S108
+                "mode": "ro",
+            },
+            str(bin_dir / "libs"): {
+                "bind": "/out-libs",
+                "mode": "ro",
+            },
+        }
+        cmd = ["env", "LD_LIBRARY_PATH=/out-libs", "/bin/bash", "/arvo"]
+    elif subset == "oss-fuzz":
+        if not is_integer(subid):
+            raise HTTPException(status_code=400, detail="Invalid task_id format for oss-fuzz")
+        oss_fuzz_path = data_dir / "oss-fuzz"
+        out_dir = oss_fuzz_path / f"{subid}-{mode}/out"
+        meta_file = oss_fuzz_path / f"{subid}-{mode}/metadata.json"
+        with open(meta_file) as f:
+            metadata = json.load(f)
+        fuzzer_name = metadata["fuzz_target"]
+        volumes = {str(poc_path.absolute()): {"bind": "/testcase", "mode": "ro"}}
+        for subfile in out_dir.iterdir():
+            host_path = str(subfile.absolute())
+            container_path = os.path.join("/out", subfile.name)
+            volumes[host_path] = {"bind": container_path, "mode": "ro"}
+        cmd = ["reproduce", fuzzer_name]
+    else:
+        raise HTTPException(status_code=400, detail="Invalid task_id format")
+
+    try:
+        container = client.containers.run(
+            image=runner_image,
+            command=["/bin/bash", "-c", f"timeout -s SIGKILL {cmd_timeout} {shlex.join(cmd)} 2>&1"],
+            detach=True,
+            volumes=volumes,
+        )
+        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 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):
+def submit_poc(db: Session, payload: Payload, mode: str, log_dir: Path, salt: str, use2: bool = False):
     # 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")
@@ -162,7 +244,10 @@ 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)
+    if use2:
+        exit_code, docker_output = run_container2(payload.task_id, poc_bin_file, mode, data_dir=server_conf.binary_dir)
+    else:
+        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)

src/cybergym/server/types.py

@@ -1,4 +1,26 @@
-from pydantic import BaseModel
+from pathlib import Path
+
+from pydantic import BaseModel, Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+from cybergym.task.types import DEFAULT_SALT
+
+
+class ServerConfig(BaseSettings):
+    salt: str = Field(default=DEFAULT_SALT, description="Salt for checksum")
+    log_dir: Path = Field(default=Path("./logs"), description="Directory to store logs")
+    db_path: Path = Field(default=Path("./poc.db"), description="Path to SQLite DB")
+    binary_dir: Path | None = Field(default=None, description="Directory to store target binaries")
+    api_key: str = Field(
+        default="cybergym-030a0cd7-5908-4862-8ab9-91f2bfc7b56d", description="API key for authentication"
+    )
+    api_key_name: str = Field(default="X-API-Key", description="Name of the API key header")
+    max_file_size_mb: int = Field(default=10, description="Maximum file size for uploads in MB")
+
+    model_config = SettingsConfigDict(env_prefix="CYBERGYM_")
+
+
+server_conf = ServerConfig()
 
 
 class Payload(BaseModel):