Merge pull request #4 from sunblaze-ucb/feat/ratelimiter

Feat/ratelimiter

.gitignore

@@ -3,6 +3,8 @@
 .vscode/
 /logs/
 /cybergym_data/
+/cybergym-server-data/
+/server/
 /oss-fuzz-data/
 *.7z
 cybergym-oss-fuzz-data.7z.part-*

pyproject.toml

@@ -54,7 +54,7 @@ include = [
 
 [tool.ruff]
 # Enable flake8-bugbear (`B`) rules.
-select = ["E", "F", "B", "I", "N", "UP", "S", "A", "C4", "DTZ", "T10", "ISC"]
-ignore = ["E501"]
+lint.select = ["E", "F", "B", "I", "N", "UP", "S", "A", "C4", "DTZ", "T10", "ISC"]
+lint.ignore = ["E501"]
 
 line-length = 120

src/cybergym/server/__main__.py

@@ -1,19 +1,69 @@
 import argparse
+import logging
 from contextlib import asynccontextmanager
 from pathlib import Path
+import time
 from typing import Annotated
 
 import uvicorn
-from fastapi import APIRouter, Depends, FastAPI, File, Form, HTTPException, Security, UploadFile, status
+from fastapi import APIRouter, Depends, FastAPI, File, Form, HTTPException, Request, Security, UploadFile, status
 from fastapi.security import APIKeyHeader
 from sqlalchemy import Engine
 from sqlalchemy.orm import Session
 
 from cybergym.server.pocdb import get_poc_by_hash, init_engine
+from cybergym.server.rate_limiter import RateLimiter
 from cybergym.server.server_utils import _post_process_result, run_poc_id, submit_poc
 from cybergym.server.types import Payload, PocQuery, VerifyPocs, server_conf
+from cybergym.task.mask import load_mask_map
+from cybergym.task.types import verify_task
+
+LOG_FORMAT = "%(asctime)s [%(levelname)s] %(name)s: %(message)s"
+LOG_DATE_FORMAT = "%Y-%m-%d %H:%M:%S"
+
+
+def make_log_config(log_file: str) -> dict:
+    return {
+        "version": 1,
+        "disable_existing_loggers": False,
+        "formatters": {
+            "default": {
+                "format": LOG_FORMAT,
+                "datefmt": LOG_DATE_FORMAT,
+            },
+        },
+        "handlers": {
+            "console": {
+                "class": "logging.StreamHandler",
+                "formatter": "default",
+                "stream": "ext://sys.stderr",
+            },
+            "file": {
+                "class": "logging.handlers.RotatingFileHandler",
+                "formatter": "default",
+                "filename": log_file,
+                "maxBytes": 10 * 1024 * 1024,  # 10 MB
+                "backupCount": 5,
+            },
+        },
+        "loggers": {
+            "uvicorn": {"handlers": ["console", "file"], "level": "INFO", "propagate": False},
+            "uvicorn.error": {"handlers": ["console", "file"], "level": "INFO", "propagate": False},
+            "uvicorn.access": {"handlers": ["console", "file"], "level": "INFO", "propagate": False},
+            "cybergym.server": {"handlers": ["console", "file"], "level": "INFO", "propagate": False},
+        },
+    }
+
+
+logger = logging.getLogger("cybergym.server")
+if not logger.handlers:
+    _handler = logging.StreamHandler()
+    _handler.setFormatter(logging.Formatter(LOG_FORMAT, datefmt=LOG_DATE_FORMAT))
+    logger.addHandler(_handler)
+    logger.setLevel(logging.INFO)
 
 engine: Engine = None
+rate_limiter: RateLimiter = None
 
 
 def get_session():
@@ -26,17 +76,39 @@ SessionDep = Annotated[Session, Depends(get_session)]
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    global engine
+    global engine, rate_limiter
+    logger.info("Starting server: db_path=%s, log_dir=%s", server_conf.db_path, server_conf.log_dir)
     engine = init_engine(server_conf.db_path)
+    rate_limiter = RateLimiter(
+        max_requests=server_conf.rate_limit_max_requests, window_seconds=server_conf.rate_limit_window_seconds
+    )
+    logger.info("Server ready")
 
     yield
 
+    logger.info("Shutting down server")
     if engine:
         engine.dispose()
 
 
 app = FastAPI(lifespan=lifespan)
 
+
+@app.middleware("http")
+async def log_requests(request: Request, call_next):
+    start = time.perf_counter()
+    response = await call_next(request)
+    elapsed_ms = (time.perf_counter() - start) * 1000
+    logger.info(
+        "%s %s -> %d (%.1fms)",
+        request.method,
+        request.url.path,
+        response.status_code,
+        elapsed_ms,
+    )
+    return response
+
+
 api_key_header = APIKeyHeader(name=server_conf.api_key_name, auto_error=False)
 
 
@@ -68,18 +140,29 @@ def submit_vul(db: SessionDep, metadata: Annotated[str, Form()], file: Annotated
     except HTTPException:
         raise
     except Exception:
+        logger.warning("Failed to read uploaded file")
         raise HTTPException(status_code=400, detail="Error reading file") from None
 
     try:
         payload = Payload.model_validate_json(metadata)
     except Exception:
+        logger.warning("Invalid metadata in submit-vul request")
         raise HTTPException(status_code=400, detail="Invalid metadata format") from None
+
+    logger.info("submit-vul: agent=%s task=%s file_size=%d", payload.agent_id, payload.task_id, len(file_content))
+
+    if not verify_task(payload.task_id, payload.agent_id, payload.checksum, salt=server_conf.salt):
+        raise HTTPException(status_code=400, detail="Invalid checksum")
+
+    rate_limiter.check(payload.agent_id)
+
     payload.data = file_content
     binary_only_mode = bool(server_conf.binary_dir)
     res = submit_poc(
         db, payload, mode="vul", log_dir=server_conf.log_dir, salt=server_conf.salt, binary_only_mode=binary_only_mode
     )
     res = _post_process_result(res, payload.require_flag)
+    logger.info("submit-vul done: agent=%s task=%s exit_code=%s", payload.agent_id, payload.task_id, res["exit_code"])
     return res
 
 
@@ -91,23 +174,30 @@ def submit_fix(db: SessionDep, metadata: Annotated[str, Form()], file: Annotated
     except HTTPException:
         raise
     except Exception:
+        logger.warning("Failed to read uploaded file")
         raise HTTPException(status_code=400, detail="Error reading file") from None
 
     try:
         payload = Payload.model_validate_json(metadata)
     except Exception:
+        logger.warning("Invalid metadata in submit-fix request")
         raise HTTPException(status_code=400, detail="Invalid metadata format") from None
+
+    logger.info("submit-fix: agent=%s task=%s file_size=%d", payload.agent_id, payload.task_id, len(file_content))
+
     payload.data = file_content
     binary_only_mode = bool(server_conf.binary_dir)
     res = submit_poc(
         db, payload, mode="fix", log_dir=server_conf.log_dir, salt=server_conf.salt, binary_only_mode=binary_only_mode
     )
     res = _post_process_result(res, payload.require_flag)
+    logger.info("submit-fix done: agent=%s task=%s exit_code=%s", payload.agent_id, payload.task_id, res["exit_code"])
     return res
 
 
 @private_router.post("/query-poc")
 def query_db(db: SessionDep, query: PocQuery):
+    logger.info("query-poc: agent=%s task=%s", query.agent_id, query.task_id)
     records = get_poc_by_hash(db, query.agent_id, query.task_id)
     if not records:
         raise HTTPException(status_code=404, detail="Record not found")
@@ -119,13 +209,19 @@ def verify_all_pocs_for_agent_id(db: SessionDep, query: VerifyPocs):
     """
     Verify all PoCs for a given agent_id.
     """
+    logger.info("verify-agent-pocs: agent=%s", query.agent_id)
     records = get_poc_by_hash(db, query.agent_id)
     if not records:
         raise HTTPException(status_code=404, detail="No records found for this agent_id")
 
     for record in records:
+        if record.vul_exit_code in [0, 300]:
+            continue  # Skip PoCs that did not trigger a crash
+        logger.info("Re-verifying poc_id=%s task=%s", record.poc_id, record.task_id)
         run_poc_id(db, server_conf.log_dir, record.poc_id, binary_only_mode=bool(server_conf.binary_dir))
+        time.sleep(0.5)  # Small delay to avoid overwhelming the docker
 
+    logger.info("verify-agent-pocs done: agent=%s count=%d", query.agent_id, len(records))
     return {
         "message": f"All {len(records)} PoCs for this agent_id have been verified",
         "poc_ids": [record.poc_id for record in records],
@@ -141,6 +237,9 @@ if __name__ == "__main__":
     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=server_conf.salt, help="Salt for checksum")
+    parser.add_argument(
+        "--mask_map_path", type=Path, default=server_conf.mask_map_path, help="Path to task ID mask mapping JSON file"
+    )
     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=server_conf.db_path, help="Path to SQLite DB")
     parser.add_argument(
@@ -149,14 +248,34 @@ if __name__ == "__main__":
     parser.add_argument(
         "--max_file_size_mb", type=int, default=server_conf.max_file_size_mb, help="Maximum file size for uploads in MB"
     )
+    parser.add_argument(
+        "--rate_limit_max_requests",
+        type=int,
+        default=server_conf.rate_limit_max_requests,
+        help="Max requests per agent per window",
+    )
+    parser.add_argument(
+        "--rate_limit_window_seconds",
+        type=int,
+        default=server_conf.rate_limit_window_seconds,
+        help="Rate limit window in seconds",
+    )
 
     args = parser.parse_args()
 
     server_conf.salt = args.salt
+    server_conf.mask_map_path = args.mask_map_path
     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
     server_conf.max_file_size_mb = args.max_file_size_mb
+    server_conf.rate_limit_max_requests = args.rate_limit_max_requests
+    server_conf.rate_limit_window_seconds = args.rate_limit_window_seconds
 
-    uvicorn.run(app, host=args.host, port=args.port)
+    if server_conf.mask_map_path:
+        load_mask_map(server_conf.mask_map_path)
+
+    uvicorn.run(
+        app, host=args.host, port=args.port, log_config=make_log_config(str(server_conf.log_dir / "server.log"))
+    )

src/cybergym/server/pocdb.py

@@ -1,7 +1,7 @@
 import datetime
 from pathlib import Path
 
-from sqlalchemy import Column, DateTime, Engine, Integer, String, UniqueConstraint, create_engine
+from sqlalchemy import Column, DateTime, Engine, Integer, String, UniqueConstraint, create_engine, event
 from sqlalchemy.orm import DeclarativeBase, Session
 
 
@@ -90,7 +90,19 @@ def get_poc_by_hash(
 
 def init_engine(db_path: Path) -> Engine:
     engine = create_engine(
-        f"sqlite:///{db_path}", echo=False, connect_args={"check_same_thread": False}, pool_size=64, max_overflow=64
+        f"sqlite:///{db_path}",
+        echo=False,
+        connect_args={"check_same_thread": False, "timeout": 30},
+        pool_size=8,
+        max_overflow=16,
     )
+
+    @event.listens_for(engine, "connect")
+    def set_sqlite_pragma(dbapi_connection, connection_record):
+        cursor = dbapi_connection.cursor()
+        cursor.execute("PRAGMA journal_mode=WAL")
+        cursor.execute("PRAGMA synchronous=NORMAL")
+        cursor.close()
+
     Base.metadata.create_all(engine)
     return engine

src/cybergym/server/rate_limiter.py

@@ -0,0 +1,44 @@
+import threading
+import time
+from collections import defaultdict, deque
+
+from fastapi import HTTPException
+
+
+class RateLimiter:
+    """Sliding-window rate limiter keyed by agent_id."""
+
+    def __init__(self, max_requests: int, window_seconds: int):
+        self.max_requests = max_requests
+        self.window_seconds = window_seconds
+        self._requests: dict[str, deque[float]] = defaultdict(deque)
+        self._lock = threading.Lock()
+        self._last_purge: float = 0.0
+
+    def check(self, agent_id: str):
+        """Raise HTTP 429 if the agent has exceeded the rate limit."""
+        now = time.monotonic()
+        cutoff = now - self.window_seconds
+
+        with self._lock:
+            if now - self._last_purge >= self.window_seconds:
+                self._purge_stale(cutoff)
+                self._last_purge = now
+
+            q = self._requests[agent_id]
+            # Evict expired entries
+            while q and q[0] <= cutoff:
+                q.popleft()
+            if len(q) >= self.max_requests:
+                raise HTTPException(
+                    status_code=429,
+                    detail=f"Rate limit exceeded for agent {agent_id}. "
+                    f"Max {self.max_requests} requests per {self.window_seconds}s.",
+                )
+            q.append(now)
+
+    def _purge_stale(self, cutoff: float):
+        """Remove agent buckets whose entries have all expired."""
+        stale = [aid for aid, q in self._requests.items() if not q or q[-1] <= cutoff]
+        for aid in stale:
+            del self._requests[aid]

src/cybergym/server/server_utils.py

@@ -15,11 +15,12 @@ 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, server_conf
+from cybergym.task.mask import _reverse_map, unmask_task_id
 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_DOCKER_TIMEOUT = 60  # seconds for docker container to run
 DEFAULT_CMD_TIMEOUT = 10  # seconds for command to run
 DEFAULT_RUNNER_IMAGE = "cybergym/oss-fuzz-base-runner:latest"
 
@@ -29,7 +30,7 @@ class CustomExitCode(IntEnum):
 
 
 CUSTOM_ERROR_MESSAGES = {
-    CustomExitCode.Timeout: "Timeout waiting for the program",
+    CustomExitCode.Timeout: "Timeout waiting for the target binary, not crashed",
 }
 
 
@@ -78,12 +79,16 @@ def run_container(
     client = docker.from_env()
     container = None
     try:
-        container = client.containers.run(
+        # Split create + start so we always have a container ref for cleanup.
+        # containers.run(detach=True) internally does create+start, but if start
+        # fails the container ref is lost and the container leaks in "created" state.
+        container = client.containers.create(
             image=image,
             command=cmd,
+            network_mode="none",
             volumes={str(poc_path.absolute()): {"bind": "/tmp/poc", "mode": "ro"}},  # noqa: S108
-            detach=True,
         )
+        container.start()
         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
@@ -163,12 +168,13 @@ def run_container_binary(
         raise HTTPException(status_code=400, detail="Invalid task_id format")
 
     try:
-        container = client.containers.run(
+        container = client.containers.create(
             image=runner_image,
             command=["/bin/bash", "-c", f"timeout -s SIGKILL {cmd_timeout} {shlex.join(cmd)} 2>&1"],
-            detach=True,
+            network_mode="none",
             volumes=volumes,
         )
+        container.start()
         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
@@ -196,16 +202,26 @@ def get_poc_storage_path(poc_id: str, log_dir: Path):
 
 def submit_poc(db: Session, payload: Payload, mode: str, log_dir: Path, salt: str, binary_only_mode: bool = False):
     # TODO: limit output size for return
+    # Verify checksum with masked task_id (agent computed checksum with what it sees)
     if not verify_task(payload.task_id, payload.agent_id, payload.checksum, salt=salt):
         raise HTTPException(status_code=400, detail="Invalid checksum")
 
+    # Unmask to get real task_id for internal use (container, DB)
+    if _reverse_map:
+        try:
+            real_task_id = unmask_task_id(payload.task_id)
+        except ValueError:
+            raise HTTPException(status_code=400, detail="Invalid task_id") from None
+    else:
+        real_task_id = payload.task_id
+
     decoded = payload.data
 
     # Compute hash of PoC
     poc_hash = hashlib.sha256(decoded).hexdigest()
 
-    # Check if PoC already exists for this agent/task/hash
-    existings = get_poc_by_hash(db, payload.agent_id, payload.task_id, poc_hash)
+    # Check if PoC already exists for this agent/task/hash (DB stores real task_id)
+    existings = get_poc_by_hash(db, payload.agent_id, real_task_id, poc_hash)
     poc_id = uuid4().hex
     if existings:
         if len(existings) > 1:
@@ -224,7 +240,7 @@ def submit_poc(db: Session, payload: Payload, mode: str, log_dir: Path, salt: st
             except Exception:
                 output = ""
             res = {
-                "task_id": payload.task_id,
+                "task_id": payload.task_id,  # return masked to agent
                 "exit_code": exit_code,
                 "output": output,
                 "poc_id": poc_id,
@@ -238,23 +254,23 @@ def submit_poc(db: Session, payload: Payload, mode: str, log_dir: Path, salt: st
     with open(poc_bin_file, "wb") as f:
         f.write(decoded)
 
-    # Insert or update DB record
+    # Insert or update DB record (store real task_id for admin analysis)
     record = get_or_create_poc(
         db,
         agent_id=payload.agent_id,
-        task_id=payload.task_id,
+        task_id=real_task_id,
         poc_id=poc_id,
         poc_hash=poc_hash,
         poc_length=len(decoded),
     )
 
-    # Run the PoC
+    # Run the PoC with real task_id (resolves docker image)
     if binary_only_mode:
         exit_code, docker_output = run_container_binary(
-            payload.task_id, poc_bin_file, mode, data_dir=server_conf.binary_dir
+            real_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)
+        exit_code, docker_output = run_container(real_task_id, poc_bin_file, mode)
     output_file = poc_dir / f"output.{mode}"
     with open(output_file, "wb") as f:
         f.write(docker_output)
@@ -262,7 +278,7 @@ def submit_poc(db: Session, payload: Payload, mode: str, log_dir: Path, salt: st
     update_poc_output(db, record, mode, exit_code)
 
     res = {
-        "task_id": payload.task_id,
+        "task_id": payload.task_id,  # return masked to agent
         "exit_code": exit_code,
         "output": docker_output.decode("utf-8"),
         "poc_id": poc_id,

src/cybergym/server/types.py

@@ -8,6 +8,7 @@ from cybergym.task.types import DEFAULT_SALT
 
 class ServerConfig(BaseSettings):
     salt: str = Field(default=DEFAULT_SALT, description="Salt for checksum")
+    mask_map_path: Path | None = Field(default=None, description="Path to task ID mask mapping JSON file")
     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")
@@ -16,6 +17,8 @@ class ServerConfig(BaseSettings):
     )
     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")
+    rate_limit_max_requests: int = Field(default=20, description="Maximum number of requests for rate limiting")
+    rate_limit_window_seconds: int = Field(default=60, description="Time window in seconds for rate limiting")
 
     model_config = SettingsConfigDict(env_prefix="CYBERGYM_")
 

src/cybergym/task/arvo_task.py

@@ -2,6 +2,7 @@ import logging
 import shutil
 from pathlib import Path
 
+from cybergym.task.mask import mask_task_id
 from cybergym.utils import get_arvo_id
 
 from .types import Task, TaskConfig, TaskDifficulty, generate_agent_id_and_checksum
@@ -104,14 +105,17 @@ def generate_arvo_task(config: TaskConfig) -> Task:
     arvo_id = get_arvo_id(config.task_id)
     arvo_dir = config.data_dir / "arvo" / arvo_id
 
-    # Create a unique agent ID and checksum
-    agent_id, checksum = generate_agent_id_and_checksum(config.task_id, config.salt, config.agent_id)
+    # Mask task_id so the agent never sees the real identifier
+    agent_facing_id = mask_task_id(config.task_id) if config.mask_map_path else config.task_id
 
-    # Prepare the output directory
+    # Checksum is computed with the agent-facing ID
+    agent_id, checksum = generate_agent_id_and_checksum(agent_facing_id, config.salt, config.agent_id)
+
+    # Embed agent-facing ID (not real task_id) into submit.sh
     prepare_arvo_files(
         config.out_dir,
         arvo_dir,
-        config.task_id,
+        agent_facing_id,
         config.server,
         agent_id,
         checksum,

src/cybergym/task/gen_task.py

@@ -3,6 +3,7 @@ import logging
 from pathlib import Path
 
 from cybergym.task.arvo_task import generate_arvo_task
+from cybergym.task.mask import load_mask_map
 from cybergym.task.oss_fuzz_task import generate_oss_fuzz_latest_task, generate_oss_fuzz_task
 from cybergym.task.types import DEFAULT_SALT, Task, TaskConfig, TaskDifficulty, TaskType
 
@@ -13,6 +14,9 @@ def generate_task(config: TaskConfig) -> Task:
     """
     Generate a task based on the task type.
     """
+    if config.mask_map_path:
+        load_mask_map(config.mask_map_path)
+
     task_type = config.task_id.split(":")[0]
     if task_type not in TaskType:
         raise ValueError(f"Unsupported task type: {task_type}")
@@ -64,6 +68,12 @@ def init_parser(parser):
         choices=[str(lvl) for lvl in TaskDifficulty],
         help="Difficulty level of the task.",
     )
+    parser.add_argument(
+        "--mask-map",
+        type=Path,
+        default=None,
+        help="Path to task ID mask mapping JSON file.",
+    )
     parser.add_argument(
         "--with-flag",
         action="store_true",
@@ -93,6 +103,7 @@ def main(raw_args=None):
         server=args.server,
         difficulty=args.difficulty,
         salt=DEFAULT_SALT,
+        mask_map_path=args.mask_map,
         with_flag=args.with_flag,
     )
 

src/cybergym/task/mask.py

@@ -0,0 +1,39 @@
+import json
+from pathlib import Path
+from uuid import uuid4
+
+_forward_map: dict[str, str] = {}  # real_task_id -> masked_id
+_reverse_map: dict[str, str] = {}  # masked_id -> real_task_id
+
+
+def load_mask_map(path: Path):
+    """Load a task ID mapping from a JSON file."""
+    with open(path) as f:
+        data = json.load(f)
+    # Mutate in-place so existing references (from `from mask import _reverse_map`) stay valid
+    _forward_map.clear()
+    _forward_map.update(data)
+    _reverse_map.clear()
+    _reverse_map.update({v: k for k, v in data.items()})
+
+
+def mask_task_id(task_id: str) -> str:
+    """Look up the masked ID for a real task_id."""
+    if task_id not in _forward_map:
+        raise ValueError(f"Task ID not in mask map: {task_id}")
+    return _forward_map[task_id]
+
+
+def unmask_task_id(masked_id: str) -> str:
+    """Look up the real task_id for a masked ID."""
+    if masked_id not in _reverse_map:
+        raise ValueError(f"Masked ID not in mask map: {masked_id}")
+    return _reverse_map[masked_id]
+
+
+def generate_mask_map(task_ids: list[str], output_path: Path):
+    """Generate a mapping file for a list of task IDs using 12-char UUIDs."""
+    mapping = {task_id: uuid4().hex[:12] for task_id in task_ids}
+    with open(output_path, "w") as f:
+        json.dump(mapping, f, indent=2)
+    return mapping

src/cybergym/task/oss_fuzz_task.py

@@ -1,23 +1,23 @@
 from cybergym.task.arvo_task import prepare_arvo_files
+from cybergym.task.mask import mask_task_id
 from cybergym.task.types import Task, TaskConfig, generate_agent_id_and_checksum
 from cybergym.utils import get_oss_fuzz_id
 
 
 def generate_oss_fuzz_task(config: TaskConfig) -> Task:
     """
-    Generate an ARVO task.
+    Generate an OSS-Fuzz task.
     """
     ossfuzz_id = get_oss_fuzz_id(config.task_id)
     ossfuzz_dir = config.data_dir / "oss-fuzz" / ossfuzz_id
 
-    # Create a unique agent ID and checksum
-    agent_id, checksum = generate_agent_id_and_checksum(config.task_id, config.salt, config.agent_id)
+    agent_facing_id = mask_task_id(config.task_id) if config.mask_map_path else config.task_id
+    agent_id, checksum = generate_agent_id_and_checksum(agent_facing_id, config.salt, config.agent_id)
 
-    # Prepare the output directory
     prepare_arvo_files(
         config.out_dir,
         ossfuzz_dir,
-        config.task_id,
+        agent_facing_id,
         config.server,
         agent_id,
         checksum,
@@ -37,19 +37,18 @@ def generate_oss_fuzz_task(config: TaskConfig) -> Task:
 
 def generate_oss_fuzz_latest_task(config: TaskConfig) -> Task:
     """
-    Generate an ARVO task.
+    Generate an OSS-Fuzz-Latest task.
     """
     ossfuzz_id = get_oss_fuzz_id(config.task_id)
     ossfuzz_dir = config.data_dir / "oss-fuzz-latest" / ossfuzz_id
 
-    # Create a unique agent ID and checksum
-    agent_id, checksum = generate_agent_id_and_checksum(config.task_id, config.salt, config.agent_id)
+    agent_facing_id = mask_task_id(config.task_id) if config.mask_map_path else config.task_id
+    agent_id, checksum = generate_agent_id_and_checksum(agent_facing_id, config.salt, config.agent_id)
 
-    # Prepare the output directory
     prepare_arvo_files(
         config.out_dir,
         ossfuzz_dir,
-        config.task_id,
+        agent_facing_id,
         config.server,
         agent_id,
         checksum,

src/cybergym/task/types.py

@@ -39,6 +39,7 @@ class TaskConfig(BaseModel):
     server: str
     difficulty: TaskDifficulty
     salt: str = DEFAULT_SALT
+    mask_map_path: Path | None = None
     agent_id: str | None = None
     with_flag: bool = False