admin/trailofbits-buttercup
1
0
Fork 0
mirror of https://github.com/trailofbits/buttercup synced 2026-06-21 14:11:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

style: apply ruff auto-fixes and formatting across entire codebase (#309)

Browse Source 
* style: apply ruff auto-fixes and formatting across entire codebase

Applied safe auto-fixes from ruff v0.12.9 with --select ALL to improve code quality:
- Reorder imports (stdlib → third-party → local)
- Use modern type hints (collections.abc.Generator instead of typing.Generator)
- Add trailing commas for better diffs
- Format multi-line function parameters for readability
- Add strict=False to zip() calls for explicit behavior
- Simplify redundant elif to if after return statements
- Consistent code formatting with ruff format

These are all mechanical, non-controversial changes that improve code consistency
without altering functionality. Changes affect 180 files across all modules:
common, fuzzer, orchestrator, patcher, program-model, and seed-gen.

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>

* re-applt ruff after merge

---------

Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Michael D Brown <michael.brown@trailofbits.com>
This commit is contained in:
Dan Guido
2025-08-22 10:30:12 -04:00
committed by GitHub
parent e4bb6390a7
commit 42691e50b4
180 changed files with 3266 additions and 2662 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/src/buttercup/common/challenge_task.py
+63 -49
View File
@@ -1,36 +1,43 @@
from __future__ import annotations
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, Any, Callable, TypeVar, cast
from os import PathLike
from functools import wraps, cached_property
import contextlib
import logging
import shlex
import os
from contextlib import contextmanager
import tempfile
import shutil
import uuid
import subprocess
import re
import shlex
import shutil
import subprocess
import tempfile
import uuid
from collections.abc import Callable, Iterator
from contextlib import contextmanager
from dataclasses import dataclass, field
from functools import cached_property, wraps
from os import PathLike
from pathlib import Path
from typing import Any, TypeVar, cast
from packaging.version import Version
from buttercup.common import node_local
from buttercup.common.constants import ARCHITECTURE
from buttercup.common.stack_parsing import get_crash_token
from buttercup.common.task_meta import TaskMeta
from buttercup.common.utils import copyanything, get_diffs
from buttercup.common.stack_parsing import get_crash_token
from typing import Iterator
import buttercup.common.node_local as node_local
from buttercup.common.constants import ARCHITECTURE
from packaging.version import Version
logger = logging.getLogger(__name__)
@contextmanager
def create_tmp_dir(
challenge: ChallengeTask, work_dir: Path | None, delete: bool = True, prefix: str | None = None
challenge: ChallengeTask,
work_dir: Path | None,
delete: bool = True,
prefix: str | None = None,
) -> Iterator[Path]:
"""Create a temporary directory inside a working dir and either keep or
delete it after use."""
delete it after use.
"""
if work_dir:
work_dir.mkdir(parents=True, exist_ok=True)
@@ -60,8 +67,6 @@ def create_tmp_dir(
class ChallengeTaskError(Exception):
"""Base class for Challenge Task errors."""
pass
FAILURE_ERR_RESULT = 201
TIMEOUT_ERR_RESULT = 124
@@ -104,7 +109,7 @@ class ReproduceResult:
"""Determine if the fuzzer at least ran"""
return bool(
(self.command_result.output and b"INFO: Seed: " in self.command_result.output)
or (self.command_result.error and b"INFO: Seed: " in self.command_result.error)
or (self.command_result.error and b"INFO: Seed: " in self.command_result.error),
)
# This is intended to encapsulate heuristics for determining if a run caused a crash
@@ -189,7 +194,8 @@ class ChallengeTask:
def _local_ro_dir(self, path: PathLike[str] | str) -> Path:
"""Return the local path to the read-only task directory.
If the path doesn't exist, it will be downloaded from the remote storage"""
If the path doesn't exist, it will be downloaded from the remote storage
"""
lp = Path(path)
if not lp.exists():
try:
@@ -226,7 +232,9 @@ class ChallengeTask:
return self._find_first_dir(self.OSS_FUZZ_DIR)
def _task_dir_compose_path(
self, subpath_method: Callable[[], Path | None], raise_on_none: bool = False
self,
subpath_method: Callable[[], Path | None],
raise_on_none: bool = False,
) -> Path | None:
subpath = subpath_method()
if subpath is None:
@@ -322,7 +330,7 @@ class ChallengeTask:
raise ChallengeTaskError("Challenge Task is read-only, cannot perform this operation")
return func(self, *args, **kwargs)
return cast(F, wrapper)
return cast("F", wrapper)
def _add_optional_arg(self, cmd: list[str], flag: str, arg: Any | None) -> None:
if arg is not None:
@@ -364,7 +372,11 @@ class ChallengeTask:
return current_line
def _run_cmd(
self, cmd: list[str], cwd: Path | None = None, log: bool = True, env_helper: Dict[str, str] | None = None
self,
cmd: list[str],
cwd: Path | None = None,
log: bool = True,
env_helper: dict[str, str] | None = None,
) -> CommandResult:
try:
if env_helper:
@@ -419,7 +431,7 @@ class ChallengeTask:
logger.exception(f"Command failed (cwd={cwd}): {' '.join(cmd)}")
return CommandResult(success=False, returncode=None, error=str(e).encode(), output=None)
def _run_helper_cmd(self, cmd: list[str], env_helper: Dict[str, str] | None = None) -> CommandResult:
def _run_helper_cmd(self, cmd: list[str], env_helper: dict[str, str] | None = None) -> CommandResult:
oss_fuzz_subpath = self.get_oss_fuzz_subpath()
if oss_fuzz_subpath is None:
raise ChallengeTaskError("OSS-Fuzz path not found")
@@ -432,7 +444,7 @@ class ChallengeTask:
try:
result = self._run_helper_cmd(grep_cmd)
except Exception as e:
logger.exception(f"[task {self.task_dir}] Error grep'ing for base-runner version: {str(e)}")
logger.exception(f"[task {self.task_dir}] Error grep'ing for base-runner version: {e!s}")
return None
if not result.success:
return None
@@ -448,7 +460,7 @@ class ChallengeTask:
base_runner_str = m.group(1).strip(":v")
return Version(base_runner_str)
except Exception as e:
logger.exception(f"[task {self.task_dir}] Error parsing base-runner version: {str(e)}")
logger.exception(f"[task {self.task_dir}] Error parsing base-runner version: {e!s}")
return None
@cached_property
@@ -469,7 +481,7 @@ class ChallengeTask:
if image.startswith("gcr.io/oss-fuzz"):
logger.info(f"Using oss-fuzz container org: {result}")
break
elif image.startswith("ghcr.io/aixcc-finals"):
if image.startswith("ghcr.io/aixcc-finals"):
result = "aixcc-afc"
logger.info(f"Using aixcc-afc container org: {result}")
break
@@ -482,8 +494,7 @@ class ChallengeTask:
return f"{self.oss_fuzz_container_org}/{self.project_name}"
def container_src_dir(self) -> str:
"""
Name of the src directory in the container (e.g. /src/FreeRDP -> FreeRDP).
"""Name of the src directory in the container (e.g. /src/FreeRDP -> FreeRDP).
This assumes that the src directory is the same as the workdir.
"""
return self.workdir_from_dockerfile().parts[-1]
@@ -497,7 +508,8 @@ class ChallengeTask:
always_build_image: bool = False,
) -> CommandResult:
"""Execute a command inside a docker container. If not specified, the
docker container is the oss-fuzz one."""
docker container is the oss-fuzz one.
"""
return self.exec_docker_cmd_rw(cmd, mount_dirs, container_image, always_build_image=always_build_image)
def exec_docker_cmd_rw(
@@ -568,8 +580,8 @@ class ChallengeTask:
architecture: str | None = ARCHITECTURE,
engine: str | None = None,
sanitizer: str | None = None,
env: Dict[str, str] | None = None,
env_helper: Dict[str, str] | None = None,
env: dict[str, str] | None = None,
env_helper: dict[str, str] | None = None,
) -> CommandResult:
logger.info(
"Building fuzzers for project %s | architecture=%s | engine=%s | sanitizer=%s | env=%s | use_source_dir=%s",
@@ -614,8 +626,8 @@ class ChallengeTask:
engine: str | None = None,
sanitizer: str | None = None,
pull_latest_base_image: bool = True,
env: Dict[str, str] | None = None,
env_helper: Dict[str, str] | None = None,
env: dict[str, str] | None = None,
env_helper: dict[str, str] | None = None,
) -> CommandResult:
check_build_res = self.check_build(architecture=architecture, engine=engine, sanitizer=sanitizer, env=env)
if check_build_res.success:
@@ -643,7 +655,7 @@ class ChallengeTask:
engine: str | None = None,
sanitizer: str | None = None,
pull_latest_base_image: bool = True,
env: Dict[str, str] | None = None,
env: dict[str, str] | None = None,
) -> CommandResult:
env_helper = {
"OSS_FUZZ_SAVE_CONTAINERS_NAME": container_name,
@@ -666,7 +678,7 @@ class ChallengeTask:
architecture: str | None = ARCHITECTURE,
engine: str | None = None,
sanitizer: str | None = None,
env: Dict[str, str] | None = None,
env: dict[str, str] | None = None,
) -> CommandResult:
logger.info(
"Checking build for project %s | architecture=%s | engine=%s | sanitizer=%s | env=%s",
@@ -695,7 +707,7 @@ class ChallengeTask:
fuzzer_args: list[str] | None = None,
*,
architecture: str | None = ARCHITECTURE,
env: Dict[str, str] | None = None,
env: dict[str, str] | None = None,
) -> ReproduceResult:
logger.info(
"Reproducing POV for project %s | fuzzer_name=%s | crash_path=%s | fuzzer_args=%s | architecture=%s | env=%s",
@@ -743,7 +755,7 @@ class ChallengeTask:
architecture: str | None = ARCHITECTURE,
engine: str | None = None,
sanitizer: str | None = None,
env: Dict[str, str] | None = None,
env: dict[str, str] | None = None,
) -> CommandResult:
logger.info(
"Running fuzzer for project %s | harness_name=%s | fuzzer_args=%s | corpus_dir=%s | architecture=%s | engine=%s | sanitizer=%s | env=%s",
@@ -778,7 +790,7 @@ class ChallengeTask:
harness_name: str,
corpus_dir: str,
architecture: str | None = ARCHITECTURE,
env: Dict[str, str] | None = None,
env: dict[str, str] | None = None,
) -> CommandResult:
logger.info(
"Running coverage for project %s | harness_name=%s | corpus_dir=%s | architecture=%s | env=%s",
@@ -839,17 +851,17 @@ class ChallengeTask:
return True
except FileNotFoundError as e:
logger.error(f"[task {self.task_dir}] File not found: {str(e)}")
raise ChallengeTaskError(f"[task {self.task_dir}] File not found: {str(e)}") from e
logger.error(f"[task {self.task_dir}] File not found: {e!s}")
raise ChallengeTaskError(f"[task {self.task_dir}] File not found: {e!s}") from e
except subprocess.CalledProcessError as e:
logger.error(f"[task {self.task_dir}] Error applying diff: {str(e)}")
logger.error(f"[task {self.task_dir}] Error applying diff: {e!s}")
logger.debug(f"[task {self.task_dir}] Error returncode: {e.returncode}")
logger.debug(f"[task {self.task_dir}] Error stdout: {e.stdout}")
logger.debug(f"[task {self.task_dir}] Error stderr: {e.stderr}")
raise ChallengeTaskError(f"[task {self.task_dir}] Error applying diff: {str(e)}") from e
raise ChallengeTaskError(f"[task {self.task_dir}] Error applying diff: {e!s}") from e
except Exception as e:
logger.exception(f"[task {self.task_dir}] Error applying diff: {str(e)}")
raise ChallengeTaskError(f"[task {self.task_dir}] Error applying diff: {str(e)}") from e
logger.exception(f"[task {self.task_dir}] Error applying diff: {e!s}")
raise ChallengeTaskError(f"[task {self.task_dir}] Error applying diff: {e!s}") from e
@contextmanager
def get_rw_copy(self, work_dir: PathLike | None, delete: bool = True) -> Iterator[ChallengeTask]:
@@ -859,6 +871,7 @@ class ChallengeTask:
Example:
with task.get_rw_copy(work_dir) as local_task:
local_task.build_fuzzers()
"""
work_dir = Path(work_dir) if work_dir else Path(node_local.scratch_path())
work_dir = work_dir / self.task_meta.task_id
@@ -906,7 +919,7 @@ class ChallengeTask:
raise ChallengeTaskError("Failed to commit task") from e
logger.error(
f"Failed to commit task {self.local_task_dir} to {new_name}. Retrying with a random suffix..."
f"Failed to commit task {self.local_task_dir} to {new_name}. Retrying with a random suffix...",
)
suffix = None
@@ -915,7 +928,8 @@ class ChallengeTask:
@read_write_decorator
def restore(self) -> None:
"""Restore the task from the original read-only task directory (if
different from the local task directory)."""
different from the local task directory).
"""
if self.read_only_task_dir == self.local_task_dir:
raise ChallengeTaskError("Task cannot be restored, it doesn't have a local task directory")
common/src/buttercup/common/challenge_task_cli.py
+20 -21
View File
@@ -1,11 +1,12 @@
from pydantic_settings import BaseSettings, CliSubCommand, get_subcommand, CliImplicitFlag
from collections.abc import Iterator
from contextlib import contextmanager
from pathlib import Path
from pydantic import BaseModel
from pydantic_settings import BaseSettings, CliImplicitFlag, CliSubCommand, get_subcommand
from buttercup.common.challenge_task import ChallengeTask, CommandResult, ReproduceResult
from buttercup.common.logger import setup_package_logger
from pathlib import Path
from typing import Dict
from contextlib import contextmanager
from typing import Iterator
class BuildImageCommand(BaseModel):
@@ -22,7 +23,7 @@ class BuildFuzzersCommand(BaseModel):
architecture: str | None = None
engine: str | None = None
sanitizer: str | None = None
env: Dict[str, str] | None = None
env: dict[str, str] | None = None
use_cache: bool = True
@@ -30,7 +31,7 @@ class CheckBuildCommand(BaseModel):
architecture: str | None = None
engine: str | None = None
sanitizer: str | None = None
env: Dict[str, str] | None = None
env: dict[str, str] | None = None
class ReproducePovCommand(BaseModel):
@@ -38,7 +39,7 @@ class ReproducePovCommand(BaseModel):
crash_path: Path
fuzzer_args: list[str] | None = None
architecture: str | None = None
env: Dict[str, str] | None = None
env: dict[str, str] | None = None
class Settings(BaseSettings):
@@ -68,7 +69,7 @@ def handle_subcommand(task: ChallengeTask, subcommand: BaseModel | None) -> Comm
cache=subcommand.cache,
architecture=subcommand.architecture,
)
elif isinstance(subcommand, BuildFuzzersCommand):
if isinstance(subcommand, BuildFuzzersCommand):
if subcommand.use_cache:
return task.build_fuzzers_with_cache(
architecture=subcommand.architecture,
@@ -76,21 +77,20 @@ def handle_subcommand(task: ChallengeTask, subcommand: BaseModel | None) -> Comm
sanitizer=subcommand.sanitizer,
env=subcommand.env,
)
else:
return task.build_fuzzers(
architecture=subcommand.architecture,
engine=subcommand.engine,
sanitizer=subcommand.sanitizer,
env=subcommand.env,
)
elif isinstance(subcommand, CheckBuildCommand):
return task.build_fuzzers(
architecture=subcommand.architecture,
engine=subcommand.engine,
sanitizer=subcommand.sanitizer,
env=subcommand.env,
)
if isinstance(subcommand, CheckBuildCommand):
return task.check_build(
architecture=subcommand.architecture,
engine=subcommand.engine,
sanitizer=subcommand.sanitizer,
env=subcommand.env,
)
elif isinstance(subcommand, ReproducePovCommand):
if isinstance(subcommand, ReproducePovCommand):
return task.reproduce_pov(
fuzzer_name=subcommand.fuzzer_name,
crash_path=subcommand.crash_path,
@@ -98,10 +98,9 @@ def handle_subcommand(task: ChallengeTask, subcommand: BaseModel | None) -> Comm
architecture=subcommand.architecture,
env=subcommand.env,
).command_result
elif isinstance(subcommand, ApplyPatchCommand):
if isinstance(subcommand, ApplyPatchCommand):
return CommandResult(success=task.apply_patch_diff(diff_file=subcommand.diff_file))
else:
raise ValueError(f"Unknown subcommand: {subcommand}")
raise ValueError(f"Unknown subcommand: {subcommand}")
@contextmanager
common/src/buttercup/common/constants.py
+2 -3
View File
@@ -5,10 +5,9 @@ def _detect_architecture() -> str:
machine = platform.machine().lower()
if machine in ("x86_64", "amd64"):
return "x86_64"
elif machine in ("aarch64", "arm64"):
if machine in ("aarch64", "arm64"):
return "aarch64"
else:
raise RuntimeError(f"Unsupported architecture: {machine}")
raise RuntimeError(f"Unsupported architecture: {machine}")
CORPUS_DIR_NAME = "buttercup_corpus"
common/src/buttercup/common/corpus.py
+15 -13
View File
@@ -1,15 +1,17 @@
import logging
from typing import BinaryIO, List
import buttercup.common.node_local as node_local
from buttercup.common.constants import CORPUS_DIR_NAME, CRASH_DIR_NAME
import os
import hashlib
import logging
import os
import shutil
import subprocess
from pathlib import Path
from redis import Redis
from buttercup.common.sets import MergedCorpusSet
import tempfile
from pathlib import Path
from typing import BinaryIO
from redis import Redis
from buttercup.common import node_local
from buttercup.common.constants import CORPUS_DIR_NAME, CRASH_DIR_NAME
from buttercup.common.sets import MergedCorpusSet
logger = logging.getLogger(__name__)
@@ -97,7 +99,7 @@ class InputDir:
return len(name) == 64 and all(c in "0123456789abcdef" for c in name)
@classmethod
def hash_corpus(cls, path: str) -> List[str]:
def hash_corpus(cls, path: str) -> list[str]:
hashed_files = []
for file in os.listdir(path):
# If the file is already a hash, skip it
@@ -130,7 +132,7 @@ class InputDir:
"--exclude=*",
str(src_path) + "/",
str(dst_path) + "/",
]
],
)
def sync_to_remote(self) -> None:
@@ -139,11 +141,11 @@ class InputDir:
self._do_sync(self.path, self.remote_path)
def sync_specific_files_to_remote(self, files: list[str]) -> None:
"""
Sync only specific files to remote storage.
"""Sync only specific files to remote storage.
Args:
files: List of filenames (basename only, not full path) to sync to remote
"""
self.hash_new_corpus()
os.makedirs(self.remote_path, exist_ok=True)
@@ -166,7 +168,7 @@ class InputDir:
f"--files-from={file_list.name}",
str(self.path) + "/",
str(self.remote_path) + "/",
]
],
)
def sync_from_remote(self) -> None:
common/src/buttercup/common/coverage_monitor.py
+32 -20
View File
@@ -8,12 +8,13 @@ from pathlib import Path
from typing import Any
from redis import Redis
from buttercup.common.maps import CoverageMap, FunctionCoverage, HarnessWeights
# Add matplotlib import for visualization
try:
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import matplotlib.pyplot as plt
MATPLOTLIB_AVAILABLE = True
except ImportError:
@@ -22,12 +23,12 @@ except ImportError:
# Move _print_coverage_metrics to be a free function
def print_coverage_metrics(func_coverage_list: list[FunctionCoverage], snapshot_count: int) -> None:
"""
Print coverage metrics for the given list of function coverage objects.
"""Print coverage metrics for the given list of function coverage objects.
Args:
func_coverage_list: List of FunctionCoverage objects
snapshot_count: The current snapshot count
"""
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
total_functions = len(func_coverage_list)
@@ -50,8 +51,7 @@ def print_coverage_metrics(func_coverage_list: list[FunctionCoverage], snapshot_
def coverage_data_equal(old_data: list[dict], new_data: list[dict]) -> bool:
"""
Compare two coverage data lists to check if they are equal.
"""Compare two coverage data lists to check if they are equal.
Args:
old_data: Previous coverage data
@@ -59,6 +59,7 @@ def coverage_data_equal(old_data: list[dict], new_data: list[dict]) -> bool:
Returns:
True if coverage data is the same, False otherwise
"""
if len(old_data) != len(new_data):
return False
@@ -67,7 +68,7 @@ def coverage_data_equal(old_data: list[dict], new_data: list[dict]) -> bool:
old_sorted = sorted(old_data, key=lambda x: x["function_name"])
new_sorted = sorted(new_data, key=lambda x: x["function_name"])
for old_item, new_item in zip(old_sorted, new_sorted):
for old_item, new_item in zip(old_sorted, new_sorted, strict=False):
if (
old_item["function_name"] != new_item["function_name"]
or old_item["total_lines"] != new_item["total_lines"]
@@ -106,14 +107,14 @@ class CoverageMonitor:
}
def monitor_coverage(self, duration_seconds: int | None = None) -> str:
"""
Monitor function coverage over time and save results to a file for all packages and harnesses.
"""Monitor function coverage over time and save results to a file for all packages and harnesses.
Args:
duration_seconds: How long to run the monitor. If None, run indefinitely.
Returns:
Path to the output file.
"""
coverage_snapshots = []
start_time = time.time()
@@ -243,8 +244,7 @@ class CoverageMonitor:
def _extract_metrics(
coverage_snapshots: list[dict],
) -> tuple[list[datetime], list[int], list[int], list[int], list[float]]:
"""
Extract metrics from coverage snapshots for analysis and visualization.
"""Extract metrics from coverage snapshots for analysis and visualization.
Returns:
Tuple containing:
@@ -253,6 +253,7 @@ class CoverageMonitor:
- total_lines: List of total line counts for each snapshot
- covered_lines: List of covered line counts for each snapshot
- coverage_percentages: List of coverage percentages for each snapshot
"""
timestamps = []
function_counts = []
@@ -286,11 +287,11 @@ class CoverageMonitor:
covered_lines: list[int],
coverage_percentages: list[float],
) -> str | None:
"""
Create a visualization of coverage metrics over time.
"""Create a visualization of coverage metrics over time.
Returns:
Path to the generated image file, or None if visualization failed.
"""
if not MATPLOTLIB_AVAILABLE:
print("Matplotlib is not available. Install with 'pip install matplotlib' to enable visualization.")
@@ -355,10 +356,12 @@ class CoverageMonitor:
@staticmethod
def analyze_coverage_file(
file_path: str, harness_key: str | None = None, visualize: bool = False, list_only: bool = False
file_path: str,
harness_key: str | None = None,
visualize: bool = False,
list_only: bool = False,
) -> None:
"""
Analyze a previously recorded coverage file and print summary statistics.
"""Analyze a previously recorded coverage file and print summary statistics.
Args:
file_path: Path to the coverage data file.
@@ -366,9 +369,10 @@ class CoverageMonitor:
If None, analyze all harnesses in the file.
visualize: Whether to generate a visualization of the data.
list_only: If True, only list the available harnesses without analyzing.
"""
try:
with open(file_path, "r") as f:
with open(file_path) as f:
coverage_snapshots = json.load(f)
if not coverage_snapshots:
@@ -384,7 +388,7 @@ class CoverageMonitor:
print(f"Coverage file: {file_path}")
print(
f"Time range: {datetime.fromtimestamp(coverage_snapshots[0]['timestamp']).strftime('%Y-%m-%d %H:%M:%S')} - "
f"{datetime.fromtimestamp(coverage_snapshots[-1]['timestamp']).strftime('%Y-%m-%d %H:%M:%S')}"
f"{datetime.fromtimestamp(coverage_snapshots[-1]['timestamp']).strftime('%Y-%m-%d %H:%M:%S')}",
)
print(f"Total snapshots: {len(coverage_snapshots)}")
print("\nAvailable harnesses:")
@@ -527,7 +531,9 @@ def main() -> None:
monitor_parser.add_argument("--task-id", help="Task ID to filter harnesses (optional)")
monitor_parser.add_argument("--interval", type=int, default=10, help="Interval between snapshots in seconds")
monitor_parser.add_argument(
"--duration", type=int, help="Duration to run the monitor in seconds (default: run indefinitely)"
"--duration",
type=int,
help="Duration to run the monitor in seconds (default: run indefinitely)",
)
# Analyze command
@@ -535,7 +541,10 @@ def main() -> None:
analyze_parser.add_argument("file_path", help="Path to the coverage data file")
analyze_parser.add_argument("--harness-key", help="Harness key to analyze in format 'package_name-harness_name'")
analyze_parser.add_argument(
"--visualize", "-v", action="store_true", help="Generate visualization of coverage metrics"
"--visualize",
"-v",
action="store_true",
help="Generate visualization of coverage metrics",
)
analyze_parser.add_argument("--list", "-l", action="store_true", help="List available harnesses in the file")
@@ -553,7 +562,10 @@ def main() -> None:
elif args.command == "analyze":
CoverageMonitor.analyze_coverage_file(
args.file_path, harness_key=args.harness_key, visualize=args.visualize, list_only=args.list
args.file_path,
harness_key=args.harness_key,
visualize=args.visualize,
list_only=args.list,
)
else:
common/src/buttercup/common/default_task_loop.py
+8 -7
View File
@@ -1,11 +1,12 @@
from abc import ABC, abstractmethod
from redis import Redis
from buttercup.common.utils import serve_loop
from buttercup.common.datastructures.msg_pb2 import WeightedHarness, BuildOutput, BuildType
from buttercup.common.maps import HarnessWeights, BuildMap
import random
import logging
import random
from abc import ABC, abstractmethod
from redis import Redis
from buttercup.common.datastructures.msg_pb2 import BuildOutput, BuildType, WeightedHarness
from buttercup.common.maps import BuildMap, HarnessWeights
from buttercup.common.utils import serve_loop
logger = logging.getLogger(__name__)
common/src/buttercup/common/llm.py
+8 -6
View File
@@ -1,15 +1,15 @@
import functools
import logging
import os
from enum import Enum
from typing import Any
import requests
import functools
from langchain.callbacks.base import BaseCallbackHandler
from langchain_core.language_models import BaseChatModel
from langchain_core.runnables import ConfigurableField
from langchain_openai.chat_models import ChatOpenAI
from langfuse.callback import CallbackHandler
from langchain.callbacks.base import BaseCallbackHandler
from langchain_core.runnables import ConfigurableField
import logging
logger = logging.getLogger(__name__)
@@ -62,7 +62,9 @@ def langfuse_auth_check() -> bool:
try:
response = requests.post(
f"{langfuse_host}/api/public/ingestion", timeout=2, auth=(langfuse_public_key, langfuse_secret_key)
f"{langfuse_host}/api/public/ingestion",
timeout=2,
auth=(langfuse_public_key, langfuse_secret_key),
)
return response.status_code == 400 # expect that we authenticate, but the request is invalid
except requests.RequestException:
common/src/buttercup/common/logger.py
+11 -5
View File
@@ -1,17 +1,20 @@
from opentelemetry._logs import set_logger_provider
import os
from opentelemetry._logs import set_logger_provider
if os.environ.get("OTEL_EXPORTER_OTLP_PROTOCOL") == "grpc":
from opentelemetry.exporter.otlp.proto.grpc._log_exporter import OTLPLogExporter
else:
from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter # type: ignore
import logging
import tempfile
from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler
from opentelemetry.sdk._logs.export import BatchLogRecordProcessor
from opentelemetry.sdk.resources import Resource
from buttercup.common.telemetry import crs_instance_id, service_instance_id
import logging
import tempfile
_is_initialized = False
PACKAGE_LOGGER_NAME = "buttercup"
@@ -30,7 +33,10 @@ class MaxLengthFormatter(logging.Formatter):
def setup_package_logger(
application_name: str, logger_name: str, log_level: str = "info", max_line_length: int | None = None
application_name: str,
logger_name: str,
log_level: str = "info",
max_line_length: int | None = None,
) -> logging.Logger:
global _is_initialized
@@ -47,7 +53,7 @@ def setup_package_logger(
"service.name": application_name,
"service.instance.id": service_instance_id,
"crs.instance.id": crs_instance_id,
}
},
)
# Initialize the LoggerProvider with the created resource.
common/src/buttercup/common/maps.py
+15 -7
View File
@@ -1,8 +1,11 @@
from typing import Generic, TypeVar, Type, Iterator
from buttercup.common.datastructures.msg_pb2 import WeightedHarness, BuildOutput, FunctionCoverage, BuildType
from redis import Redis
from bson.json_util import dumps, CANONICAL_JSON_OPTIONS
from collections.abc import Iterator
from typing import Generic, TypeVar
from bson.json_util import CANONICAL_JSON_OPTIONS, dumps
from google.protobuf.message import Message
from redis import Redis
from buttercup.common.datastructures.msg_pb2 import BuildOutput, BuildType, FunctionCoverage, WeightedHarness
from buttercup.common.sets import RedisSet
MsgType = TypeVar("MsgType", bound=Message)
@@ -10,7 +13,7 @@ MSG_FIELD_NAME = b"msg"
class RedisMap(Generic[MsgType]):
def __init__(self, redis: Redis, hash_name: str, msg_builder: Type[MsgType]):
def __init__(self, redis: Redis, hash_name: str, msg_builder: type[MsgType]):
self.redis = redis
self.msg_builder = msg_builder
self.hash_name = hash_name
@@ -52,7 +55,8 @@ class BuildMap:
def _build_output_key(self, task_id: str, build_type: BuildType, san: str, internal_patch_id: str) -> str:
return dumps(
[task_id, BUILD_SAN_MAP_NAME, build_type, san, internal_patch_id], json_options=CANONICAL_JSON_OPTIONS
[task_id, BUILD_SAN_MAP_NAME, build_type, san, internal_patch_id],
json_options=CANONICAL_JSON_OPTIONS,
)
def add_build(self, build: BuildOutput) -> None:
@@ -78,7 +82,11 @@ class BuildMap:
return builds
def get_build_from_san(
self, task_id: str, build_type: BuildType, san: str, internal_patch_id: str = ""
self,
task_id: str,
build_type: BuildType,
san: str,
internal_patch_id: str = "",
) -> BuildOutput | None:
if internal_patch_id != "":
assert build_type == BuildType.PATCH, "internal_patch_id is only valid for PATCH builds"
common/src/buttercup/common/node_local.py
+7 -6
View File
@@ -1,14 +1,14 @@
# Store the node local path for subsequent use
from contextlib import contextmanager
import logging
import os
from pathlib import Path
import shutil
import tarfile
from tempfile import NamedTemporaryFile
import tempfile
from typing import Any, Iterator, TypeAlias
from contextlib import AbstractContextManager
from collections.abc import Iterator
from contextlib import AbstractContextManager, contextmanager
from pathlib import Path
from tempfile import NamedTemporaryFile
from typing import Any, TypeAlias
logger = logging.getLogger(__name__)
@@ -188,6 +188,7 @@ def dir_to_remote_archive(local_path: NodeLocalPath) -> RemotePath:
def lopen(local_path: NodeLocalPath, mode: str) -> Any:
"""Open a file in the node local storage
If it doesn't exist, it will be downloaded from the remote storage"""
If it doesn't exist, it will be downloaded from the remote storage
"""
make_locally_available(local_path)
return open(local_path, mode)
common/src/buttercup/common/project_yaml.py
+5 -3
View File
@@ -1,8 +1,10 @@
from dataclasses import dataclass, field
import yaml
from buttercup.common.challenge_task import ChallengeTask
from enum import Enum
import yaml
from buttercup.common.challenge_task import ChallengeTask
class Language(str, Enum):
C = "c"
@@ -45,7 +47,7 @@ class ProjectYaml:
"""Language field but with a more consistent naming convention."""
if self.language.lower() in ["c", "c++", "cpp"]:
return Language.C
elif self.language.lower() in ["java", "jvm"]:
if self.language.lower() in ["java", "jvm"]:
return Language.JAVA
raise ValueError(f"Unsupported language: {self.language}")
common/src/buttercup/common/queues.py
+89 -46
View File
@@ -1,31 +1,32 @@
from __future__ import annotations
import logging
import os
import uuid
from collections.abc import Callable
from dataclasses import dataclass, field
from redis import Redis, RedisError
from google.protobuf.message import Message
from enum import Enum
from functools import wraps
from typing import Any, Generic, Literal, TypeVar, cast, overload
from google.protobuf.message import Message
from redis import Redis, RedisError
from buttercup.common.datastructures.msg_pb2 import (
BuildRequest,
BuildOutput,
Crash,
TaskDownload,
TaskReady,
TaskDelete,
Patch,
BuildRequest,
ConfirmedVulnerability,
IndexRequest,
Crash,
IndexOutput,
TracedCrash,
IndexRequest,
Patch,
POVReproduceRequest,
POVReproduceResponse,
TaskDelete,
TaskDownload,
TaskReady,
TracedCrash,
)
import logging
from typing import Type, Generic, TypeVar, Literal, overload, Callable, cast
import uuid
import os
from enum import Enum
from typing import Any
TIMES_DELIVERED_FIELD = "times_delivered"
@@ -86,9 +87,7 @@ MsgType = TypeVar("MsgType", bound=Message)
@dataclass
class RQItem(Generic[MsgType]):
"""
A single item in a reliable queue.
"""
"""A single item in a reliable queue."""
item_id: str
deserialized: MsgType
@@ -96,13 +95,11 @@ class RQItem(Generic[MsgType]):
@dataclass
class ReliableQueue(Generic[MsgType]):
"""
A queue that is reliable and can be used to process tasks in a distributed environment.
"""
"""A queue that is reliable and can be used to process tasks in a distributed environment."""
redis: Redis
queue_name: str
msg_builder: Type[MsgType]
msg_builder: type[MsgType]
group_name: str | None = None
task_timeout_ms: int = 180000
reader_name: str | None = None
@@ -113,7 +110,7 @@ class ReliableQueue(Generic[MsgType]):
def __post_init__(self) -> None:
if self.reader_name is None:
self.reader_name = f"rqueue_{str(uuid.uuid4())}"
self.reader_name = f"rqueue_{uuid.uuid4()!s}"
if self.group_name is not None:
# Create consumer group if it doesn't exist
@@ -125,9 +122,10 @@ class ReliableQueue(Generic[MsgType]):
pass
else:
logger.exception(
"Failed to create consumer group %s for queue %s", self.group_name, self.queue_name
"Failed to create consumer group %s for queue %s",
self.group_name,
self.queue_name,
)
pass
def size(self) -> int:
return self.redis.xlen(self.queue_name)
@@ -145,7 +143,7 @@ class ReliableQueue(Generic[MsgType]):
return func(self, *args, **kwargs)
return cast(F, wrapper)
return cast("F", wrapper)
@_ensure_group_name
def pop(self) -> RQItem[MsgType] | None:
@@ -205,7 +203,7 @@ class ReliableQueue(Generic[MsgType]):
if pending is None or len(pending) == 0:
return 0
return cast(int, pending[0][TIMES_DELIVERED_FIELD])
return cast("int", pending[0][TIMES_DELIVERED_FIELD])
@_ensure_group_name
def claim_item(self, item_id: str, min_idle_time: int = 0) -> None:
@@ -215,7 +213,7 @@ class ReliableQueue(Generic[MsgType]):
@dataclass
class QueueConfig:
queue_name: QueueNames
msg_builder: Type
msg_builder: type
task_timeout_ms: int
group_names: list[GroupNames] = field(default_factory=list)
@@ -305,81 +303,126 @@ class QueueFactory:
POV_REPRODUCER_RESPONSES_TASK_TIMEOUT_MS,
[GroupNames.ORCHESTRATOR],
),
}
},
)
@overload
def create(
self, queue_name: Literal[QueueNames.BUILD], group_name: GroupNames, **kwargs: Any
self,
queue_name: Literal[QueueNames.BUILD],
group_name: GroupNames,
**kwargs: Any,
) -> ReliableQueue[BuildRequest]: ...
@overload
def create(
self, queue_name: Literal[QueueNames.BUILD_OUTPUT], group_name: GroupNames, **kwargs: Any
self,
queue_name: Literal[QueueNames.BUILD_OUTPUT],
group_name: GroupNames,
**kwargs: Any,
) -> ReliableQueue[BuildOutput]: ...
@overload
def create(
self, queue_name: Literal[QueueNames.DOWNLOAD_TASKS], group_name: GroupNames, **kwargs: Any
self,
queue_name: Literal[QueueNames.DOWNLOAD_TASKS],
group_name: GroupNames,
**kwargs: Any,
) -> ReliableQueue[TaskDownload]: ...
@overload
def create(
self, queue_name: Literal[QueueNames.READY_TASKS], group_name: GroupNames, **kwargs: Any
self,
queue_name: Literal[QueueNames.READY_TASKS],
group_name: GroupNames,
**kwargs: Any,
) -> ReliableQueue[TaskReady]: ...
@overload
def create(
self, queue_name: Literal[QueueNames.CRASH], group_name: GroupNames, **kwargs: Any
self,
queue_name: Literal[QueueNames.CRASH],
group_name: GroupNames,
**kwargs: Any,
) -> ReliableQueue[Crash]: ...
@overload
def create(
self, queue_name: Literal[QueueNames.TRACED_VULNERABILITIES], group_name: GroupNames, **kwargs: Any
self,
queue_name: Literal[QueueNames.TRACED_VULNERABILITIES],
group_name: GroupNames,
**kwargs: Any,
) -> ReliableQueue[TracedCrash]: ...
@overload
def create(
self, queue_name: Literal[QueueNames.CONFIRMED_VULNERABILITIES], group_name: GroupNames, **kwargs: Any
self,
queue_name: Literal[QueueNames.CONFIRMED_VULNERABILITIES],
group_name: GroupNames,
**kwargs: Any,
) -> ReliableQueue[ConfirmedVulnerability]: ...
@overload
def create(
self, queue_name: Literal[QueueNames.DELETE_TASK], group_name: GroupNames, **kwargs: Any
self,
queue_name: Literal[QueueNames.DELETE_TASK],
group_name: GroupNames,
**kwargs: Any,
) -> ReliableQueue[TaskDelete]: ...
@overload
def create(
self, queue_name: Literal[QueueNames.PATCHES], group_name: GroupNames, **kwargs: Any
self,
queue_name: Literal[QueueNames.PATCHES],
group_name: GroupNames,
**kwargs: Any,
) -> ReliableQueue[Patch]: ...
@overload
def create(
self, queue_name: Literal[QueueNames.INDEX], group_name: GroupNames, **kwargs: Any
self,
queue_name: Literal[QueueNames.INDEX],
group_name: GroupNames,
**kwargs: Any,
) -> ReliableQueue[IndexRequest]: ...
@overload
def create(
self, queue_name: Literal[QueueNames.INDEX_OUTPUT], group_name: GroupNames, **kwargs: Any
self,
queue_name: Literal[QueueNames.INDEX_OUTPUT],
group_name: GroupNames,
**kwargs: Any,
) -> ReliableQueue[IndexOutput]: ...
@overload
def create(
self, queue_name: Literal[QueueNames.POV_REPRODUCER_REQUESTS], group_name: GroupNames, **kwargs: Any
self,
queue_name: Literal[QueueNames.POV_REPRODUCER_REQUESTS],
group_name: GroupNames,
**kwargs: Any,
) -> ReliableQueue[POVReproduceRequest]: ...
@overload
def create(
self, queue_name: Literal[QueueNames.POV_REPRODUCER_RESPONSES], group_name: GroupNames, **kwargs: Any
self,
queue_name: Literal[QueueNames.POV_REPRODUCER_RESPONSES],
group_name: GroupNames,
**kwargs: Any,
) -> ReliableQueue[POVReproduceResponse]: ...
@overload
def create(
self, queue_name: QueueNames, group_name: GroupNames | None = None, **kwargs: Any
self,
queue_name: QueueNames,
group_name: GroupNames | None = None,
**kwargs: Any,
) -> ReliableQueue[MsgType]: ...
def create(
self, queue_name: QueueNames, group_name: GroupNames | None = None, **kwargs: Any
self,
queue_name: QueueNames,
group_name: GroupNames | None = None,
**kwargs: Any,
) -> ReliableQueue[MsgType]:
if queue_name not in self._config:
raise ValueError(f"Invalid queue name: {queue_name}")
common/src/buttercup/common/reproduce_multiple.py
+14 -8
View File
@@ -1,17 +1,21 @@
from buttercup.common.datastructures.msg_pb2 import BuildOutput
from buttercup.common.challenge_task import ReproduceResult, ChallengeTask
from pathlib import Path
from typing import Generator
from contextlib import contextmanager
import contextlib
import logging
from collections.abc import Generator
from contextlib import contextmanager
from pathlib import Path
from buttercup.common.challenge_task import ChallengeTask, ReproduceResult
from buttercup.common.datastructures.msg_pb2 import BuildOutput
logger = logging.getLogger(__name__)
class ReproduceMultiple:
def __init__(
self, wdir: Path, build_outputs: list[BuildOutput], build_cache: list[ChallengeTask] | None = None
self,
wdir: Path,
build_outputs: list[BuildOutput],
build_cache: list[ChallengeTask] | None = None,
) -> None:
self.build_outputs = build_outputs
self.wdir = wdir
@@ -32,11 +36,13 @@ class ReproduceMultiple:
pass
def attempt_reproduce(
self, pov: Path, harness_name: str
self,
pov: Path,
harness_name: str,
) -> Generator[tuple[BuildOutput, ReproduceResult], None, None]:
if self.builds_cache is None:
raise RuntimeError("Build cache is not populated")
for build, task in zip(self.build_outputs, self.builds_cache):
for build, task in zip(self.build_outputs, self.builds_cache, strict=False):
yield (build, task.reproduce_pov(harness_name, pov))
def get_first_crash(self, pov: Path, harness_name: str) -> tuple[BuildOutput, ReproduceResult] | None:
common/src/buttercup/common/sarif_store.py
+20 -19
View File
@@ -1,6 +1,7 @@
from __future__ import annotations
from typing import Dict, Any, List
from typing import Any
from pydantic import BaseModel, Field
from redis import Redis
@@ -8,11 +9,11 @@ from redis import Redis
class SARIFBroadcastDetail(BaseModel):
"""Model for SARIF broadcast details, matches the model in types.py"""
metadata: Dict[str, Any] = Field(
metadata: dict[str, Any] = Field(
...,
description="String to string map containing data that should be attached to outputs like log messages and OpenTelemetry trace attributes for traceability",
)
sarif: Dict[str, Any] = Field(..., description="SARIF Report compliant with provided schema")
sarif: dict[str, Any] = Field(..., description="SARIF Report compliant with provided schema")
sarif_id: str
task_id: str
@@ -21,47 +22,47 @@ class SARIFStore:
"""Store and retrieve SARIF objects in Redis"""
def __init__(self, redis: Redis):
"""
Initialize the SARIF store with a Redis connection.
"""Initialize the SARIF store with a Redis connection.
Args:
redis: Redis connection
"""
self.redis = redis
self.key_prefix = "sarif:"
def _get_key(self, task_id: str) -> str:
"""
Get the Redis key for a task_id.
"""Get the Redis key for a task_id.
Args:
task_id: Task ID
Returns:
Redis key
"""
return f"{self.key_prefix}{task_id.lower()}"
def _decode_key(self, key: str | bytes) -> str:
"""
Decode a Redis key if it's bytes, otherwise return as is.
"""Decode a Redis key if it's bytes, otherwise return as is.
Args:
key: Redis key, either bytes or string
Returns:
Decoded key as string
"""
if isinstance(key, bytes):
return key.decode("utf-8")
return key
def store(self, sarif_detail: SARIFBroadcastDetail) -> None:
"""
Store a SARIF broadcast detail in Redis.
"""Store a SARIF broadcast detail in Redis.
Args:
sarif_detail: The SARIF broadcast detail to store
"""
task_id = sarif_detail.task_id
key = self._get_key(task_id)
@@ -73,12 +74,12 @@ class SARIFStore:
# Add to the list for this task
self.redis.rpush(key, sarif_json)
def get_all(self) -> List[SARIFBroadcastDetail]:
"""
Retrieve all SARIF objects from Redis.
def get_all(self) -> list[SARIFBroadcastDetail]:
"""Retrieve all SARIF objects from Redis.
Returns:
List of SARIF broadcast details
"""
# Get all SARIF keys in Redis
all_keys = self.redis.keys(f"{self.key_prefix}*")
@@ -97,15 +98,15 @@ class SARIFStore:
return result
def get_by_task_id(self, task_id: str) -> List[SARIFBroadcastDetail]:
"""
Retrieve all SARIF objects for a specific task.
def get_by_task_id(self, task_id: str) -> list[SARIFBroadcastDetail]:
"""Retrieve all SARIF objects for a specific task.
Args:
task_id: Task ID
Returns:
List of SARIF broadcast details for this task
"""
key = self._get_key(task_id)
sarif_list = self.redis.lrange(key, 0, -1)
@@ -118,14 +119,14 @@ class SARIFStore:
return result
def delete_by_task_id(self, task_id: str) -> int:
"""
Remove all SARIF objects for a specific task.
"""Remove all SARIF objects for a specific task.
Args:
task_id: Task ID
Returns:
Number of removed keys (0 or 1)
"""
key = self._get_key(task_id)
return self.redis.delete(key)
common/src/buttercup/common/sarif_utility.py
+8 -9
View File
@@ -1,21 +1,20 @@
"""
Utility script for SARIF storage and retrieval operations.
"""
"""Utility script for SARIF storage and retrieval operations."""
import argparse
import json
from redis import Redis
from buttercup.common.sarif_store import SARIFStore
def list_all_sarifs(redis_url: str, verbose: bool = False) -> None:
"""
List all SARIF objects in the database.
"""List all SARIF objects in the database.
Args:
redis_url: Redis URL
verbose: Whether to print the full SARIF object
"""
redis_client = Redis.from_url(redis_url)
sarif_store = SARIFStore(redis_client)
@@ -32,13 +31,13 @@ def list_all_sarifs(redis_url: str, verbose: bool = False) -> None:
def list_task_sarifs(redis_url: str, task_id: str, verbose: bool = False) -> None:
"""
List all SARIF objects for a specific task.
"""List all SARIF objects for a specific task.
Args:
redis_url: Redis URL
task_id: Task ID
verbose: Whether to print the full SARIF object
"""
redis_client = Redis.from_url(redis_url)
sarif_store = SARIFStore(redis_client)
@@ -55,12 +54,12 @@ def list_task_sarifs(redis_url: str, task_id: str, verbose: bool = False) -> Non
def delete_task_sarifs(redis_url: str, task_id: str) -> None:
"""
Delete all SARIF objects for a specific task.
"""Delete all SARIF objects for a specific task.
Args:
redis_url: Redis URL
task_id: Task ID
"""
redis_client = Redis.from_url(redis_url)
sarif_store = SARIFStore(redis_client)
common/src/buttercup/common/sets.py
+21 -16
View File
@@ -1,12 +1,12 @@
from typing import Iterator
import json
import random
from collections.abc import Generator, Iterator
from contextlib import contextmanager
from functools import lru_cache
from bson.json_util import CANONICAL_JSON_OPTIONS, dumps
from redis import Redis
from redis.exceptions import ResponseError
from bson.json_util import dumps, CANONICAL_JSON_OPTIONS
from contextlib import contextmanager
import random
import json
from functools import lru_cache
from typing import Generator
# Import POVReproduceRequest for the refactored PoVReproduceStatus
from buttercup.common.datastructures.msg_pb2 import POVReproduceRequest, POVReproduceResponse
@@ -64,7 +64,7 @@ class RedisLock:
@contextmanager
def acquire(self) -> Generator[None, None, None]:
if not self.redis.set(self.key, "1", ex=self.lock_timeout_seconds, nx=True):
raise FailedToAcquireLock()
raise FailedToAcquireLock
try:
yield
finally:
@@ -116,6 +116,7 @@ class PoVReproduceStatus:
Returns:
False if mitigated (didn't crash), True if non-mitigated (did crash), None if not in final states
"""
pipeline = self.redis.pipeline()
pipeline.sismember(POV_REPRODUCE_MITIGATED_SET_NAME, key)
@@ -124,10 +125,9 @@ class PoVReproduceStatus:
if result[0]:
return False # Mitigated - didn't crash
elif result[1]:
if result[1]:
return True # Non-mitigated - did crash
else:
return None # Not in final states
return None # Not in final states
def _make_key(self, request: POVReproduceRequest) -> str:
"""Create a unique key from a POVReproduceRequest by serializing it to string."""
@@ -144,6 +144,7 @@ class PoVReproduceStatus:
Returns:
None if pending, POVReproduceResponse if completed
"""
key = self._make_key(request)
@@ -161,13 +162,13 @@ class PoVReproduceStatus:
if result[0]:
return None # Pending
elif result[1]:
if result[1]:
return POVReproduceResponse(request=request, did_crash=False) # Completed and mitigated
elif result[2]:
if result[2]:
return POVReproduceResponse(request=request, did_crash=True) # Completed and not mitigated
else: # First time, schedule it for testing
self.redis.sadd(POV_REPRODUCE_PENDING_SET_NAME, key)
return None
# First time, schedule it for testing
self.redis.sadd(POV_REPRODUCE_PENDING_SET_NAME, key)
return None
def mark_mitigated(self, request: POVReproduceRequest) -> bool:
"""Mark a POV reproduction as mitigated (patch successfully prevented the crash).
@@ -177,6 +178,7 @@ class PoVReproduceStatus:
Returns:
True if the item was moved from pending to mitigated, False if item wasn't pending.
"""
key = self._make_key(request)
moved_count = self.redis.smove(POV_REPRODUCE_PENDING_SET_NAME, POV_REPRODUCE_MITIGATED_SET_NAME, key)
@@ -190,6 +192,7 @@ class PoVReproduceStatus:
Returns:
True if the item was moved from pending to non-mitigated, False if item wasn't pending.
"""
key = self._make_key(request)
moved_count = self.redis.smove(POV_REPRODUCE_PENDING_SET_NAME, POV_REPRODUCE_NON_MITIGATED_SET_NAME, key)
@@ -203,6 +206,7 @@ class PoVReproduceStatus:
Returns:
True if the item was moved from pending to expired, False if item wasn't pending.
"""
# NOTE: This function isn't strictly needed. We could just have keys expire after a certain time.
# However, this allows us to track which items have expired.
@@ -215,6 +219,7 @@ class PoVReproduceStatus:
Returns:
POVReproduceRequest if one is available, None otherwise
"""
pending_set = self.redis.smembers(POV_REPRODUCE_PENDING_SET_NAME)
if len(pending_set) == 0:
common/src/buttercup/common/stack_parsing.py
+6 -4
View File
@@ -1,9 +1,11 @@
import re
from buttercup.common.clusterfuzz_parser import StackParser, CrashInfo
import logging
from buttercup.common.sets import RedisSet
import re
from bson.json_util import CANONICAL_JSON_OPTIONS, dumps
from redis import Redis
from bson.json_util import dumps, CANONICAL_JSON_OPTIONS
from buttercup.common.clusterfuzz_parser import CrashInfo, StackParser
from buttercup.common.sets import RedisSet
logger = logging.getLogger(__name__)
common/src/buttercup/common/task_meta.py
+4 -2
View File
@@ -1,7 +1,7 @@
from dataclasses import dataclass, asdict
import json
from pathlib import Path
from dataclasses import asdict, dataclass
from os import PathLike
from pathlib import Path
@dataclass
@@ -27,6 +27,7 @@ class TaskMeta:
Raises:
FileNotFoundError: If the file doesn't exist
JSONDecodeError: If the file contains invalid JSON
"""
path = Path(directory) / cls.METADATA_FILENAME
with path.open() as f:
@@ -38,6 +39,7 @@ class TaskMeta:
Args:
directory: Directory where to save the metadata file
"""
path = Path(directory) / self.METADATA_FILENAME
with path.open("w") as f:
common/src/buttercup/common/task_registry.py
+25 -10
View File
@@ -1,11 +1,14 @@
from functools import lru_cache
from buttercup.common.datastructures.msg_pb2 import Task
from redis import Redis
from buttercup.common.queues import HashNames
from dataclasses import dataclass
from google.protobuf import text_format
import builtins
import time
from typing import Set, Iterator
from collections.abc import Iterator
from dataclasses import dataclass
from functools import lru_cache
from google.protobuf import text_format
from redis import Redis
from buttercup.common.datastructures.msg_pb2 import Task
from buttercup.common.queues import HashNames
# Redis set keys for tracking task states
CANCELLED_TASKS_SET = "cancelled_tasks"
@@ -68,6 +71,7 @@ class TaskRegistry:
Returns:
Task object if found, None otherwise
"""
prepared_key = self._prepare_key(task_id)
task_bytes = self.redis.hget(self.hash_name, prepared_key)
@@ -87,6 +91,7 @@ class TaskRegistry:
Args:
task_id: The ID of the task to delete
"""
prepared_key = self._prepare_key(task_id)
@@ -104,6 +109,7 @@ class TaskRegistry:
Args:
task_or_id: Either a Task object or a task ID string to be added to the cancelled set
"""
# Extract task_id based on the type of input
task_id = task_or_id.task_id if isinstance(task_or_id, Task) else task_or_id
@@ -119,6 +125,7 @@ class TaskRegistry:
Returns:
True if the task is in the cancelled tasks set, False otherwise
"""
# Get task_id
task_id = task_or_id.task_id if isinstance(task_or_id, Task) else task_or_id
@@ -140,6 +147,7 @@ class TaskRegistry:
Returns:
True if the task is expired (deadline has passed), False otherwise.
Returns False if the task doesn't exist.
"""
@lru_cache(maxsize=1000)
@@ -166,6 +174,7 @@ class TaskRegistry:
Returns:
list[Task]: List of active tasks
"""
# Iterate through all tasks, filtering out cancelled and expired ones
# The cancelled flag is already set correctly by the __iter__ method
@@ -176,13 +185,14 @@ class TaskRegistry:
Returns:
list[str]: List of task IDs that are in the cancelled tasks set
"""
# Get all cancelled task IDs from the Redis set
cancelled_ids = self.redis.smembers(CANCELLED_TASKS_SET)
# Decode bytes to strings if needed
return [task_id.decode("utf-8") if isinstance(task_id, bytes) else task_id for task_id in cancelled_ids]
def should_stop_processing(self, task_or_id: str | Task, cancelled_ids: Set[str] | None = None) -> bool:
def should_stop_processing(self, task_or_id: str | Task, cancelled_ids: builtins.set[str] | None = None) -> bool:
"""Check if a task should no longer be processed due to cancellation or expiration.
Args:
@@ -193,8 +203,8 @@ class TaskRegistry:
Returns:
bool: True if the task should not be processed (is cancelled or expired),
False otherwise
"""
"""
# Extract task_id for cancelled IDs check
task_id = task_or_id.task_id if isinstance(task_or_id, Task) else task_or_id
@@ -223,6 +233,7 @@ class TaskRegistry:
Args:
task_or_id: Either a Task object or a task ID string to be added to the successful set
"""
# Extract task_id based on the type of input
task_id = task_or_id.task_id if isinstance(task_or_id, Task) else task_or_id
@@ -238,6 +249,7 @@ class TaskRegistry:
Returns:
True if the task is in the successful tasks set, False otherwise
"""
# Get task_id
task_id = task_or_id.task_id if isinstance(task_or_id, Task) else task_or_id
@@ -255,6 +267,7 @@ class TaskRegistry:
Args:
task_or_id: Either a Task object or a task ID string to be added to the errored set
"""
# Extract task_id based on the type of input
task_id = task_or_id.task_id if isinstance(task_or_id, Task) else task_or_id
@@ -270,6 +283,7 @@ class TaskRegistry:
Returns:
True if the task is in the errored tasks set, False otherwise
"""
# Get task_id
task_id = task_or_id.task_id if isinstance(task_or_id, Task) else task_or_id
@@ -282,9 +296,10 @@ class TaskRegistry:
def task_registry_cli() -> None:
"""CLI for the task registry"""
from pydantic_settings import BaseSettings
from typing import Annotated
from pydantic import Field
from pydantic_settings import BaseSettings
class TaskRegistrySettings(BaseSettings):
redis_url: Annotated[str, Field(default="redis://localhost:6379", description="Redis URL")]
common/src/buttercup/common/telemetry.py
+4 -4
View File
@@ -1,14 +1,14 @@
import os
import logging
from enum import Enum
import os
import uuid
from enum import Enum
from typing import Any
import openlit
import opentelemetry.attributes
from opentelemetry import trace
from opentelemetry.trace import Span, Tracer, Status, StatusCode
from langchain_core.prompt_values import ChatPromptValue
from opentelemetry import trace
from opentelemetry.trace import Span, Status, StatusCode, Tracer
# Monkey patch the _clean_attribute function to handle ChatPromptValue
_clean_attribute_orig = opentelemetry.attributes._clean_attribute
common/src/buttercup/common/util_cli.py
+22 -20
View File
@@ -1,26 +1,27 @@
import logging
from pathlib import Path
from typing import Annotated
from uuid import uuid4
from google.protobuf.text_format import Parse
from pydantic import BaseModel, Field
from pydantic_settings import BaseSettings, CliPositionalArg, CliSubCommand, get_subcommand
from redis import Redis
from buttercup.common.datastructures.msg_pb2 import (
BuildOutput,
BuildType,
SubmissionEntry,
SubmissionResult,
WeightedHarness,
)
from buttercup.common.logger import setup_package_logger
from buttercup.common.maps import (
BuildMap,
HarnessWeights,
)
from buttercup.common.datastructures.msg_pb2 import (
BuildOutput,
BuildType,
WeightedHarness,
SubmissionEntry,
SubmissionResult,
)
from buttercup.common.queues import QueueFactory, QueueNames, ReliableQueue
from buttercup.common.task_registry import TaskRegistry
from uuid import uuid4
from redis import Redis
from pydantic_settings import BaseSettings, CliSubCommand, CliPositionalArg, get_subcommand
from pydantic import BaseModel
from typing import Annotated
from pydantic import Field
from pathlib import Path
from google.protobuf.text_format import Parse
import logging
logger = logging.getLogger(__name__)
@@ -229,7 +230,9 @@ def handle_subcommand(redis: Redis, command: BaseModel | None) -> None:
if task_id not in result:
result[task_id] = TaskResult(
task_id=task_id, project_name=task.project_name, mode=str(task.task_type)
task_id=task_id,
project_name=task.project_name,
mode=str(task.task_type),
)
c = next((c for c in submission.crashes if c.result == SubmissionResult.PASSED), None)
if c:
@@ -240,9 +243,8 @@ def handle_subcommand(redis: Redis, command: BaseModel | None) -> None:
result[task_id].n_patches += 1
assert c is not None
result[task_id].patched_vulnerabilities.append(c.competition_pov_id)
else:
if c:
result[task_id].non_patched_vulnerabilities.append(c.competition_pov_id)
elif c:
result[task_id].non_patched_vulnerabilities.append(c.competition_pov_id)
b = next((b for b in submission.bundles), None)
if b:
common/src/buttercup/common/utils.py
+6 -5
View File
@@ -1,12 +1,13 @@
import shutil
import errno
import logging
import os
from typing import Any, Callable
from pathlib import Path
from os import PathLike
import time
import shutil
import threading
import time
from collections.abc import Callable
from os import PathLike
from pathlib import Path
from typing import Any
logger = logging.getLogger(__name__)
common/tests/test_challenge_task.py
+50 -33
View File
@@ -1,18 +1,20 @@
from pathlib import Path
import pytest
from unittest.mock import MagicMock, patch
from dirty_equals import IsStr
import subprocess
import os
import base64
import os
import subprocess
import tempfile
from pathlib import Path
from unittest.mock import MagicMock, patch
import pytest
from dirty_equals import IsStr
from buttercup.common.challenge_task import (
ChallengeTask,
ChallengeTaskError,
ReproduceResult,
CommandResult,
ReproduceResult,
)
from buttercup.common.task_meta import TaskMeta
import tempfile
@pytest.fixture
@@ -373,7 +375,7 @@ def libjpeg_oss_fuzz_task_dir(tmp_path: Path) -> Path:
metadata={"task_id": "task-id-libjpeg-turbo", "round_id": "testing", "team_id": "tob"},
).save(tmp_path)
yield tmp_path
return tmp_path
@pytest.fixture
@@ -717,7 +719,10 @@ def mock_node_local(monkeypatch, tmp_path: Path):
@patch("buttercup.common.node_local._get_root_path")
@patch("buttercup.common.node_local.remote_archive_to_dir")
def test_challenge_task_with_node_local_storage_existing(
mock_remote_archive_to_dir, mock_get_root_path, mock_node_local_storage, task_dir
mock_remote_archive_to_dir,
mock_get_root_path,
mock_node_local_storage,
task_dir,
):
"""Test ChallengeTask behavior when using node_local and path exists."""
mock_get_root_path.return_value = Path(mock_node_local_storage)
@@ -770,7 +775,10 @@ def test_challenge_task_with_node_local_storage_existing(
@patch("buttercup.common.node_local._get_root_path")
@patch("buttercup.common.node_local.remote_archive_to_dir")
def test_challenge_task_with_node_local_storage_download(
mock_remote_archive_to_dir, mock_get_root_path, mock_node_local_storage, task_dir
mock_remote_archive_to_dir,
mock_get_root_path,
mock_node_local_storage,
task_dir,
):
"""Test ChallengeTask behavior when using node_local and path doesn't exist."""
mock_get_root_path.return_value = Path(mock_node_local_storage)
@@ -866,7 +874,7 @@ def mock_node_local_storage(tmp_path: Path):
metadata={"task_id": "task-id-challenge-task", "round_id": "testing", "team_id": "tob"},
).save(local_task_path)
yield node_data_dir
return node_data_dir
@pytest.mark.integration
@@ -891,7 +899,7 @@ def test_reproduce_result_stacktrace():
# Test case 1: Short string (under 1MB limit)
short_output = b"INFO: Seed: 12345\nRunning normally\n==ERROR: AddressSanitizer: heap-buffer-overflow"
result1 = ReproduceResult(
command_result=CommandResult(success=False, returncode=1, output=short_output, error=None)
command_result=CommandResult(success=False, returncode=1, output=short_output, error=None),
)
stacktrace1 = result1.stacktrace()
assert stacktrace1 is not None
@@ -906,7 +914,7 @@ def test_reproduce_result_stacktrace():
)
result2 = ReproduceResult(
command_result=CommandResult(success=False, returncode=1, output=large_content, error=None)
command_result=CommandResult(success=False, returncode=1, output=large_content, error=None),
)
stacktrace2 = result2.stacktrace()
assert stacktrace2 is not None
@@ -938,7 +946,7 @@ def test_reproduce_result_stacktrace():
exact_size_content = prefix + b"B" * (MAX_OUTPUT_LEN - len(prefix) - len(suffix)) + suffix
result4 = ReproduceResult(
command_result=CommandResult(success=False, returncode=1, output=exact_size_content, error=None)
command_result=CommandResult(success=False, returncode=1, output=exact_size_content, error=None),
)
stacktrace4 = result4.stacktrace()
assert stacktrace4 is not None
@@ -954,7 +962,7 @@ def test_reproduce_result_stacktrace():
)
result5 = ReproduceResult(
command_result=CommandResult(success=False, returncode=1, output=over_limit_content, error=None)
command_result=CommandResult(success=False, returncode=1, output=over_limit_content, error=None),
)
stacktrace5 = result5.stacktrace()
assert stacktrace5 is not None
@@ -969,7 +977,7 @@ def test_reproduce_result_stacktrace():
)
result6 = ReproduceResult(
command_result=CommandResult(success=False, returncode=1, output=very_large_content, error=None)
command_result=CommandResult(success=False, returncode=1, output=very_large_content, error=None),
)
stacktrace6 = result6.stacktrace()
assert stacktrace6 is not None
@@ -989,15 +997,18 @@ def test_reproduce_result_methods():
# Test case 1: Successful run, no crash
result1 = ReproduceResult(
command_result=CommandResult(
success=True, returncode=0, output=b"INFO: Seed: 12345\nRunning normally", error=None
)
success=True,
returncode=0,
output=b"INFO: Seed: 12345\nRunning normally",
error=None,
),
)
assert result1.did_run() is True
assert result1.did_crash() is False
# Test case 2: Failed run, no crash (fuzzer didn't start)
result2 = ReproduceResult(
command_result=CommandResult(success=False, returncode=1, output=b"Error: Could not start fuzzer", error=None)
command_result=CommandResult(success=False, returncode=1, output=b"Error: Could not start fuzzer", error=None),
)
assert result2.did_run() is False
assert result2.did_crash() is False
@@ -1009,7 +1020,7 @@ def test_reproduce_result_methods():
returncode=1,
output=b"INFO: Seed: 12345\nRunning normally\n==ERROR: AddressSanitizer: heap-buffer-overflow",
error=None,
)
),
)
assert result3.did_run() is True
assert result3.did_crash() is True
@@ -1021,7 +1032,7 @@ def test_reproduce_result_methods():
returncode=1,
output=None,
error=b"INFO: Seed: 12345\nRunning normally\n==ERROR: AddressSanitizer: heap-buffer-overflow",
)
),
)
assert result4.did_run() is True
assert result4.did_crash() is True
@@ -1029,8 +1040,11 @@ def test_reproduce_result_methods():
# Test case 5: Run with None returncode
result5 = ReproduceResult(
command_result=CommandResult(
success=False, returncode=None, output=b"INFO: Seed: 12345\nRunning normally", error=None
)
success=False,
returncode=None,
output=b"INFO: Seed: 12345\nRunning normally",
error=None,
),
)
assert result5.did_run() is True
assert result5.did_crash() is False
@@ -1042,7 +1056,7 @@ def test_reproduce_result_methods():
returncode=124, # TIMEOUT_ERR_RESULT
output=b"INFO: Seed: 12345\nRunning normally\n==ERROR: AddressSanitizer: heap-buffer-overflow\nTimeout occurred",
error=None,
)
),
)
assert result6.did_run() is True
assert result6.did_crash() is True # Should detect crash due to crash token in stacktrace
@@ -1054,7 +1068,7 @@ def test_reproduce_result_methods():
returncode=124, # TIMEOUT_ERR_RESULT
output=b"INFO: Seed: 12345\nRunning normally\nTimeout occurred\nNo crash detected",
error=None,
)
),
)
assert result7.did_run() is True
assert result7.did_crash() is False # Should not detect crash due to no crash token
@@ -1102,7 +1116,7 @@ subprocess command returned a non-zero exit status: 1"""
returncode=124, # TIMEOUT_ERR_RESULT
output=b"INFO: Seed: 12345\nRunning normally\n" + output + b"\nTimeout occurred",
error=b"",
)
),
)
assert result8.did_run() is True
assert result8.did_crash() is True # Should detect crash due to crash token in error output
@@ -1114,7 +1128,7 @@ subprocess command returned a non-zero exit status: 1"""
returncode=124, # TIMEOUT_ERR_RESULT
output=None,
error=None,
)
),
)
assert result9.did_run() is False
assert result9.did_crash() is False # Should not detect crash due to no output and no crash token
@@ -1126,7 +1140,7 @@ subprocess command returned a non-zero exit status: 1"""
returncode=124, # TIMEOUT_ERR_RESULT
output=b"INFO: Seed: 12345\nTimeout occurred",
error=None,
)
),
)
assert result10.did_run() is True
assert result10.did_crash() is False # Should not detect crash due to no crash token
@@ -1138,7 +1152,7 @@ subprocess command returned a non-zero exit status: 1"""
returncode=201, # FAILURE_ERR_RESULT
output=b"INFO: Seed: 12345\nRunning normally\n==ERROR: AddressSanitizer: heap-buffer-overflow",
error=None,
)
),
)
assert result11.did_run() is True
assert result11.did_crash() is False # Should not detect crash due to FAILURE_ERR_RESULT
@@ -1150,7 +1164,7 @@ subprocess command returned a non-zero exit status: 1"""
returncode=124, # TIMEOUT_ERR_RESULT
output=b"INFO: Seed: 12345\nRunning normally\n" + output + b"\nTimeout occurred",
error=None,
)
),
)
assert result12.did_run() is True
assert result12.did_crash() is True # Should detect crash due to UBSan error in stacktrace
@@ -1162,7 +1176,7 @@ subprocess command returned a non-zero exit status: 1"""
returncode=124, # TIMEOUT_ERR_RESULT
output=b"INFO: Seed: 12345\nRunning normally\n" + output + b"\n" + output + b"\nTimeout occurred",
error=None,
)
),
)
assert result13.did_run() is True
assert result13.did_crash() is True # Should detect crash due to multiple crash patterns in stacktrace
@@ -1265,7 +1279,10 @@ def test_apply_patch_diff_git_apply_failure(challenge_task: ChallengeTask):
with patch("subprocess.run") as mock_run:
# Simulate git apply failure
mock_run.side_effect = subprocess.CalledProcessError(
returncode=1, cmd=["patch", "-p1"], output="", stderr="patch does not apply"
returncode=1,
cmd=["patch", "-p1"],
output="",
stderr="patch does not apply",
)
with pytest.raises(ChallengeTaskError, match="Error applying diff"):
common/tests/test_corpus.py
+5 -3
View File
@@ -1,10 +1,12 @@
import pytest
import os
import shutil
import tempfile
from pathlib import Path
from unittest.mock import patch
from buttercup.common.corpus import InputDir, Corpus
import pytest
from buttercup.common.corpus import Corpus, InputDir
@pytest.fixture
@@ -81,7 +83,7 @@ def test_local_corpus_size_with_mixed_exceptions(temp_dir, mock_node_local):
if self.name == "test_file_2" and call_count[self.name] == 1:
raise FileNotFoundError(f"Simulated transient error for {self.name}")
# test_file_4 will always fail
elif self.name == "test_file_4":
if self.name == "test_file_4":
raise PermissionError(f"Simulated permission error for {self.name}")
# Other files work normally
return original_lstat(self)
common/tests/test_coverage_map.py
+2 -1
View File
@@ -1,7 +1,8 @@
import pytest
from redis import Redis
from buttercup.common.maps import CoverageMap
from buttercup.common.datastructures.msg_pb2 import FunctionCoverage
from buttercup.common.maps import CoverageMap
@pytest.fixture
common/tests/test_get_fuzz_targets.py
+1 -1
View File
@@ -3,7 +3,7 @@ import tempfile
import unittest
from unittest.mock import patch
from buttercup.common.clusterfuzz_utils import get_fuzz_targets, EXTRA_BUILD_DIR
from buttercup.common.clusterfuzz_utils import EXTRA_BUILD_DIR, get_fuzz_targets
class TestGetFuzzTargets(unittest.TestCase):
common/tests/test_map.py
+2 -1
View File
@@ -1,7 +1,8 @@
import pytest
from redis import Redis
from buttercup.common.maps import RedisMap
from buttercup.common.datastructures.msg_pb2 import WeightedHarness
from buttercup.common.maps import RedisMap
@pytest.fixture
common/tests/test_multiple_reproduce.py
+5 -3
View File
@@ -1,9 +1,11 @@
from unittest.mock import MagicMock, patch
from pathlib import Path
from unittest.mock import MagicMock, patch
import pytest
from buttercup.common.reproduce_multiple import ReproduceMultiple
from buttercup.common.challenge_task import CommandResult, ReproduceResult
from buttercup.common.datastructures.msg_pb2 import BuildOutput
from buttercup.common.challenge_task import ReproduceResult, CommandResult
from buttercup.common.reproduce_multiple import ReproduceMultiple
@pytest.fixture
common/tests/test_node_local.py
+29 -19
View File
@@ -1,24 +1,28 @@
from pathlib import Path
import pytest
import os
from unittest.mock import patch, mock_open, MagicMock
from contextlib import contextmanager
from pathlib import Path
from unittest.mock import MagicMock, mock_open, patch
import pytest
from buttercup.common.node_local import (
_get_root_path,
TmpDir,
temp_dir,
rename_atomically,
remote_path as remote_path_func,
remote_archive_path as remote_archive_path_func,
scratch_path,
scratch_dir,
local_scratch_file,
remote_scratch_file,
_get_root_path,
dir_to_remote_archive,
local_scratch_file,
lopen,
remote_scratch_file,
rename_atomically,
scratch_dir,
scratch_path,
temp_dir,
)
from buttercup.common.node_local import (
remote_archive_path as remote_archive_path_func,
)
from buttercup.common.node_local import (
remote_path as remote_path_func,
)
# Use this root path for all tests
TEST_ROOT_PATH = Path("/test/node/data/dir")
@@ -217,13 +221,15 @@ class TestNodeLocal:
mock_scratch_context.__enter__.return_value = mock_scratch_file
with patch(
"buttercup.common.node_local.local_scratch_file", return_value=mock_scratch_context
"buttercup.common.node_local.local_scratch_file",
return_value=mock_scratch_context,
) as mock_local_scratch:
# Mock copyfileobj to avoid actual copying
with patch("shutil.copyfileobj") as mock_copy:
# Mock rename_atomically
with patch(
"buttercup.common.node_local.rename_atomically", return_value=local_path
"buttercup.common.node_local.rename_atomically",
return_value=local_path,
) as mock_rename:
# Mock path operations
with patch.object(Path, "mkdir"):
@@ -265,7 +271,8 @@ class TestNodeLocal:
# Mock the archive path
with patch(
"buttercup.common.node_local.remote_archive_path", return_value=remote_archive_path_val
"buttercup.common.node_local.remote_archive_path",
return_value=remote_archive_path_val,
) as mock_rpath:
# Mock file operations
with patch("builtins.open", mock_open(read_data=b"test data")) as mocked_open:
@@ -276,7 +283,8 @@ class TestNodeLocal:
mock_scratch_context.__enter__.return_value = mock_scratch_file
with patch(
"buttercup.common.node_local.local_scratch_file", return_value=mock_scratch_context
"buttercup.common.node_local.local_scratch_file",
return_value=mock_scratch_context,
) as mock_local_scratch:
# Mock copyfileobj to avoid actual copying
with patch("shutil.copyfileobj") as mock_copy:
@@ -294,11 +302,13 @@ class TestNodeLocal:
mock_tmp_dir_context.__enter__.return_value = mock_tmp_dir
with patch(
"buttercup.common.node_local.scratch_dir", return_value=mock_tmp_dir_context
"buttercup.common.node_local.scratch_dir",
return_value=mock_tmp_dir_context,
) as mock_scratch_dir:
# Mock rename_atomically
with patch(
"buttercup.common.node_local.rename_atomically", return_value=local_path
"buttercup.common.node_local.rename_atomically",
return_value=local_path,
) as mock_rename:
# Mock path.exists
with patch.object(Path, "exists", return_value=False):
common/tests/test_reliable_queue.py
+11 -9
View File
@@ -1,17 +1,19 @@
import pytest
import time
from redis import Redis
import pytest
from google.protobuf.struct_pb2 import Struct
from redis import Redis
from buttercup.common.datastructures.msg_pb2 import BuildOutput, BuildRequest
from buttercup.common.queues import (
ReliableQueue,
RQItem,
BUILD_OUTPUT_TASK_TIMEOUT_MS,
BUILD_TASK_TIMEOUT_MS,
GroupNames,
QueueFactory,
QueueNames,
GroupNames,
BUILD_TASK_TIMEOUT_MS,
BUILD_OUTPUT_TASK_TIMEOUT_MS,
ReliableQueue,
RQItem,
)
from buttercup.common.datastructures.msg_pb2 import BuildRequest, BuildOutput
GROUP_NAME = "test_group"
QUEUE_NAME = "test_queue"
@@ -210,7 +212,7 @@ def test_queue_factory(redis_client):
sanitizer="test_sanitizer",
task_dir="test_task_dir",
task_id="test_task_id",
)
),
)
item = queue.pop()
common/tests/test_sarif_store.py
+4 -2
View File
@@ -1,7 +1,9 @@
import pytest
import json
import pytest
from redis import Redis
from buttercup.common.sarif_store import SARIFStore, SARIFBroadcastDetail
from buttercup.common.sarif_store import SARIFBroadcastDetail, SARIFStore
@pytest.fixture
common/tests/test_set.py
+4 -2
View File
@@ -1,8 +1,10 @@
from unittest.mock import MagicMock
import pytest
from redis import Redis
from buttercup.common.sets import RedisSet, PoVReproduceStatus
from buttercup.common.datastructures.msg_pb2 import POVReproduceRequest, POVReproduceResponse
from unittest.mock import MagicMock
from buttercup.common.sets import PoVReproduceStatus, RedisSet
@pytest.fixture
common/tests/test_stack_parser.py
+12 -10
View File
File diff suppressed because one or more lines are too long
common/tests/test_task_registry.py
+19 -16
View File
@@ -1,11 +1,11 @@
import pytest
from unittest.mock import Mock
from buttercup.common.task_registry import TaskRegistry, CANCELLED_TASKS_SET, SUCCEEDED_TASKS_SET, ERRORED_TASKS_SET
from buttercup.common.datastructures.msg_pb2 import Task, SourceDetail
import time
from typing import Set
from unittest.mock import Mock, patch
import pytest
from redis import Redis
from unittest.mock import patch
from buttercup.common.datastructures.msg_pb2 import SourceDetail, Task
from buttercup.common.task_registry import CANCELLED_TASKS_SET, ERRORED_TASKS_SET, SUCCEEDED_TASKS_SET, TaskRegistry
@pytest.fixture
@@ -166,7 +166,10 @@ def test_iter_tasks_with_different_types(task_registry, redis_client):
# Create and add two different tasks
full_task = Task(task_id="full123", task_type=Task.TaskType.TASK_TYPE_FULL, message_id="msg_full", cancelled=False)
delta_task = Task(
task_id="delta456", task_type=Task.TaskType.TASK_TYPE_DELTA, message_id="msg_delta", cancelled=True
task_id="delta456",
task_type=Task.TaskType.TASK_TYPE_DELTA,
message_id="msg_delta",
cancelled=True,
)
# Setup Redis mock
@@ -319,7 +322,7 @@ def test_is_expired(task_registry, redis_client):
def mock_hget(hash_name, key):
if key == "expired-task":
return expired_task.SerializeToString()
elif key == "live-task":
if key == "live-task":
return live_task.SerializeToString()
return None
@@ -432,7 +435,7 @@ def test_get_cancelled_task_ids(task_registry, redis_client):
def test_should_stop_processing_with_cancelled_ids(task_registry, redis_client):
"""Test that should_stop_processing handles cancelled_ids correctly."""
# Create a set of cancelled IDs
cancelled_ids: Set[str] = {"cancelled-task-1", "cancelled-task-2"}
cancelled_ids: set[str] = {"cancelled-task-1", "cancelled-task-2"}
# Create tasks to test with
current_time = int(time.time())
@@ -444,9 +447,9 @@ def test_should_stop_processing_with_cancelled_ids(task_registry, redis_client):
def mock_hget(hash_name, key):
if key == "cancelled-task-1":
return cancelled_task1.SerializeToString()
elif key == "cancelled-task-2":
if key == "cancelled-task-2":
return cancelled_task2.SerializeToString()
elif key == "active-task":
if key == "active-task":
return active_task.SerializeToString()
return None
@@ -479,7 +482,7 @@ def test_should_stop_processing_no_cancelled_ids(task_registry, redis_client):
def mock_hget(hash_name, key):
if key == "active-task":
return active_task.SerializeToString()
elif key == "cancelled-task":
if key == "cancelled-task":
return cancelled_task.SerializeToString()
return None
@@ -517,7 +520,7 @@ def test_should_stop_processing_expired_task(task_registry, redis_client):
def mock_hget(hash_name, key):
if key == "expired-task":
return expired_task.SerializeToString()
elif key == "active-task":
if key == "active-task":
return active_task.SerializeToString()
return None
@@ -659,9 +662,9 @@ def test_is_expired_with_delta(task_registry, redis_client):
def mock_hget(hash_name, key):
if key == "just-expired":
return just_expired_task.SerializeToString()
elif key == "soon-expired":
if key == "soon-expired":
return soon_expired_task.SerializeToString()
elif key == "future-task":
if key == "future-task":
return future_task.SerializeToString()
return None
@@ -732,7 +735,7 @@ def test_is_expired_with_delta_edge_cases(task_registry, redis_client):
def mock_hget(hash_name, key):
if key == "one-second-expired":
return one_second_expired.SerializeToString()
elif key == "one-second-future":
if key == "one-second-future":
return one_second_future.SerializeToString()
return None
fuzzer/src/buttercup/fuzzing_infra/builder_bot.py
+26 -25
View File
@@ -1,22 +1,21 @@
from buttercup.common.queues import QueueNames, GroupNames
from redis import Redis
from buttercup.common.queues import QueueFactory
from buttercup.common.datastructures.msg_pb2 import BuildType, BuildOutput, BuildRequest
from buttercup.common.logger import setup_package_logger
from dataclasses import dataclass, field
from buttercup.common.queues import ReliableQueue
import logging
import tempfile
from buttercup.common.utils import serve_loop
from buttercup.common.challenge_task import ChallengeTask, ChallengeTaskError
from dataclasses import dataclass, field
from pathlib import Path
from buttercup.fuzzing_infra.settings import BuilderBotSettings
import buttercup.common.node_local as node_local
from buttercup.common.telemetry import init_telemetry
from opentelemetry import trace
from opentelemetry.trace import Status, StatusCode
from buttercup.common.telemetry import set_crs_attributes, CRSActionCategory
from redis import Redis
from buttercup.common import node_local
from buttercup.common.challenge_task import ChallengeTask, ChallengeTaskError
from buttercup.common.datastructures.msg_pb2 import BuildOutput, BuildRequest, BuildType
from buttercup.common.logger import setup_package_logger
from buttercup.common.queues import GroupNames, QueueFactory, QueueNames, ReliableQueue
from buttercup.common.task_registry import TaskRegistry
from buttercup.common.telemetry import CRSActionCategory, init_telemetry, set_crs_attributes
from buttercup.common.utils import serve_loop
from buttercup.fuzzing_infra.settings import BuilderBotSettings
logger = logging.getLogger(__name__)
@@ -44,13 +43,13 @@ class BuilderBot:
def _apply_challenge_diff(self, task: ChallengeTask, msg: BuildRequest) -> bool:
if msg.apply_diff and task.is_delta_mode():
logger.info(
f"Applying diff for {msg.task_id} | {msg.engine} | {msg.sanitizer} | {BuildType.Name(msg.build_type)} | diff {msg.apply_diff}"
f"Applying diff for {msg.task_id} | {msg.engine} | {msg.sanitizer} | {BuildType.Name(msg.build_type)} | diff {msg.apply_diff}",
)
try:
res = task.apply_patch_diff()
if not res:
logger.warning(
f"No diffs for {msg.task_id} | {msg.engine} | {msg.sanitizer} | {BuildType.Name(msg.build_type)} | diff {msg.apply_diff}"
f"No diffs for {msg.task_id} | {msg.engine} | {msg.sanitizer} | {BuildType.Name(msg.build_type)} | diff {msg.apply_diff}",
)
return False
except ChallengeTaskError:
@@ -74,13 +73,13 @@ class BuilderBot:
logger.debug("Patch written to %s", patch_file.name)
logger.info(
f"Applying patch for {msg.task_id} | {msg.engine} | {msg.sanitizer} | {BuildType.Name(msg.build_type)} | diff {msg.apply_diff}"
f"Applying patch for {msg.task_id} | {msg.engine} | {msg.sanitizer} | {BuildType.Name(msg.build_type)} | diff {msg.apply_diff}",
)
try:
res = task.apply_patch_diff(Path(patch_file.name))
if not res:
logger.info(
f"Failed to apply patch for {msg.task_id} | {msg.engine} | {msg.sanitizer} | {BuildType.Name(msg.build_type)} | diff {msg.apply_diff}"
f"Failed to apply patch for {msg.task_id} | {msg.engine} | {msg.sanitizer} | {BuildType.Name(msg.build_type)} | diff {msg.apply_diff}",
)
return False
except ChallengeTaskError:
@@ -103,7 +102,7 @@ class BuilderBot:
msg = rqit.deserialized
logger.info(
f"Received build request for {msg.task_id} | {msg.engine} | {msg.sanitizer} | {BuildType.Name(msg.build_type)} | diff {msg.apply_diff}"
f"Received build request for {msg.task_id} | {msg.engine} | {msg.sanitizer} | {BuildType.Name(msg.build_type)} | diff {msg.apply_diff}",
)
# Check if task should not be processed (expired or cancelled)
@@ -129,7 +128,7 @@ class BuilderBot:
if not self._apply_challenge_diff(task, msg):
if self._build_requests_queue.times_delivered(rqit.item_id) > self.max_tries:
logger.error(
f"Max tries reached for {msg.task_id} | {msg.engine} | {msg.sanitizer} | {BuildType.Name(msg.build_type)} | diff {msg.apply_diff}"
f"Max tries reached for {msg.task_id} | {msg.engine} | {msg.sanitizer} | {BuildType.Name(msg.build_type)} | diff {msg.apply_diff}",
)
self._build_requests_queue.ack_item(rqit.item_id)
@@ -138,7 +137,7 @@ class BuilderBot:
if not self._apply_patch(task, msg):
if self._build_requests_queue.times_delivered(rqit.item_id) > self.max_tries:
logger.error(
f"Max tries reached for {msg.task_id} | {msg.engine} | {msg.sanitizer} | {BuildType.Name(msg.build_type)} | diff {msg.apply_diff} | patch {msg.internal_patch_id}"
f"Max tries reached for {msg.task_id} | {msg.engine} | {msg.sanitizer} | {BuildType.Name(msg.build_type)} | diff {msg.apply_diff} | patch {msg.internal_patch_id}",
)
self._build_requests_queue.ack_item(rqit.item_id)
@@ -154,12 +153,14 @@ class BuilderBot:
task_metadata=dict(origin_task.task_meta.metadata),
)
res = task.build_fuzzers_with_cache(
engine=msg.engine, sanitizer=msg.sanitizer, pull_latest_base_image=self.allow_pull
engine=msg.engine,
sanitizer=msg.sanitizer,
pull_latest_base_image=self.allow_pull,
)
if not res.success:
logger.error(
f"Could not build fuzzer {msg.task_id} | {msg.engine} | {msg.sanitizer} | {BuildType.Name(msg.build_type)} | diff {msg.apply_diff}"
f"Could not build fuzzer {msg.task_id} | {msg.engine} | {msg.sanitizer} | {BuildType.Name(msg.build_type)} | diff {msg.apply_diff}",
)
span.set_status(Status(StatusCode.ERROR))
return True
@@ -168,7 +169,7 @@ class BuilderBot:
task.commit()
logger.info(
f"Pushing build output for {msg.task_id} | {msg.engine} | {msg.sanitizer} | {BuildType.Name(msg.build_type)} | diff {msg.apply_diff}"
f"Pushing build output for {msg.task_id} | {msg.engine} | {msg.sanitizer} | {BuildType.Name(msg.build_type)} | diff {msg.apply_diff}",
)
node_local.dir_to_remote_archive(task.task_dir)
self._build_outputs_queue.push(
@@ -180,10 +181,10 @@ class BuilderBot:
build_type=msg.build_type,
apply_diff=msg.apply_diff,
internal_patch_id=msg.internal_patch_id,
)
),
)
logger.info(
f"Acked build request for {msg.task_id} | {msg.engine} | {msg.sanitizer} | {BuildType.Name(msg.build_type)} | diff {msg.apply_diff}"
f"Acked build request for {msg.task_id} | {msg.engine} | {msg.sanitizer} | {BuildType.Name(msg.build_type)} | diff {msg.apply_diff}",
)
self._build_requests_queue.ack_item(rqit.item_id)
return True
fuzzer/src/buttercup/fuzzing_infra/corpus_merger.py
+52 -56
View File
@@ -1,31 +1,28 @@
from buttercup.common import node_local
from buttercup.fuzzing_infra.runner import Runner, Conf, FuzzConfiguration
from dataclasses import dataclass
import os
from buttercup.common.datastructures.msg_pb2 import BuildType, WeightedHarness
from buttercup.common.datastructures.aliases import BuildType as BuildTypeHint
from buttercup.common.corpus import Corpus
from buttercup.common.maps import HarnessWeights, BuildMap
from buttercup.common.utils import serve_loop, setup_periodic_zombie_reaper
from buttercup.common.logger import setup_package_logger
from redis import Redis
from typing import List
from os import PathLike
import random
from buttercup.common.datastructures.msg_pb2 import BuildOutput
import datetime
import logging
from buttercup.common.challenge_task import ChallengeTask
from buttercup.fuzzing_infra.settings import FuzzerBotSettings
from buttercup.common.sets import MergedCorpusSetLock
from buttercup.common.constants import ADDRESS_SANITIZER
from buttercup.common.sets import FailedToAcquireLock
from buttercup.common.sets import MERGING_LOCK_TIMEOUT_SECONDS
from buttercup.common.telemetry import init_telemetry
import os
import random
import shutil
from dataclasses import dataclass
from os import PathLike
from opentelemetry import trace
from opentelemetry.trace import Status, StatusCode
from buttercup.common.telemetry import set_crs_attributes, CRSActionCategory
import datetime
import shutil
from redis import Redis
from buttercup.common import node_local
from buttercup.common.challenge_task import ChallengeTask
from buttercup.common.constants import ADDRESS_SANITIZER
from buttercup.common.corpus import Corpus
from buttercup.common.datastructures.aliases import BuildType as BuildTypeHint
from buttercup.common.datastructures.msg_pb2 import BuildOutput, BuildType, WeightedHarness
from buttercup.common.logger import setup_package_logger
from buttercup.common.maps import BuildMap, HarnessWeights
from buttercup.common.sets import MERGING_LOCK_TIMEOUT_SECONDS, FailedToAcquireLock, MergedCorpusSetLock
from buttercup.common.telemetry import CRSActionCategory, init_telemetry, set_crs_attributes
from buttercup.common.utils import serve_loop, setup_periodic_zombie_reaper
from buttercup.fuzzing_infra.runner import Conf, FuzzConfiguration, Runner
from buttercup.fuzzing_infra.settings import FuzzerBotSettings
logger = logging.getLogger(__name__)
@@ -38,9 +35,7 @@ logger = logging.getLogger(__name__)
@dataclass
class FinalCorpus:
"""
Represents the corpus after the merge operation has been performed.
"""
"""Represents the corpus after the merge operation has been performed."""
def __init__(self, corpus: Corpus, push_remotely: set[str], delete_locally: set[str]):
self._corpus = corpus
@@ -48,9 +43,7 @@ class FinalCorpus:
self._delete_locally = delete_locally
def push_remotely(self) -> int:
"""
Push the files to remote storage.
"""
"""Push the files to remote storage."""
n = 0
if self._push_remotely:
n = len(self._push_remotely)
@@ -59,9 +52,7 @@ class FinalCorpus:
return n
def delete_locally(self) -> int:
"""
Delete the files from local storage.
"""
"""Delete the files from local storage."""
n = 0
for file in self._delete_locally:
try:
@@ -76,9 +67,7 @@ class FinalCorpus:
@dataclass
class PartitionedCorpus:
"""
Represents the corpus split into local and remote parts.
"""
"""Represents the corpus split into local and remote parts."""
corpus: Corpus
local_dir: PathLike[str]
@@ -118,8 +107,7 @@ class PartitionedCorpus:
logger.debug(f"Error copying file {file} to remote directory: {e}. Copied from remote storage instead.")
def to_final(self) -> FinalCorpus:
"""
Returns a FinalCorpus object that represents the corpus after the merge operation has been performed.
"""Returns a FinalCorpus object that represents the corpus after the merge operation has been performed.
NOTE: This should be called after the merge operation has been performed.
Will rehash any files in the remote_directory as the merge operation may have changed the file names.
@@ -148,8 +136,7 @@ class PartitionedCorpus:
@dataclass
class BaseCorpus:
"""
Represents the initial corpus state, before any merge operations have been performed.
"""Represents the initial corpus state, before any merge operations have been performed.
- local_dir: PathLike directory for the local corpus
- remote_dir: PathLike directory for the remote corpus
@@ -163,8 +150,7 @@ class BaseCorpus:
max_local_files: int = 500
def partition_corpus(self) -> PartitionedCorpus:
"""
1. Collect the remote corpus files
"""1. Collect the remote corpus files
2. Collect the list of files only available remotely
3. Partition the corpus into two sets,
- files that are in the remote corpus,
@@ -192,7 +178,12 @@ class BaseCorpus:
class MergerBot:
def __init__(
self, redis: Redis, timeout_seconds: int, python: str, crs_scratch_dir: str, max_local_files: int = 500
self,
redis: Redis,
timeout_seconds: int,
python: str,
crs_scratch_dir: str,
max_local_files: int = 500,
):
self.redis = redis
self.runner = Runner(Conf(timeout_seconds))
@@ -202,7 +193,7 @@ class MergerBot:
self.builds = BuildMap(redis)
self.max_local_files = max_local_files
def required_builds(self) -> List[BuildTypeHint]:
def required_builds(self) -> list[BuildTypeHint]:
return [BuildType.FUZZER]
def _run_merge_operation(
@@ -215,8 +206,7 @@ class MergerBot:
remote_files: set[str],
corp: Corpus,
) -> None:
"""
Run the merge operation to find which local files add coverage.
"""Run the merge operation to find which local files add coverage.
Args:
task: The WeightedHarness object
@@ -229,6 +219,7 @@ class MergerBot:
Returns:
No return value - files that add coverage will be moved to remote_dir
"""
with node_local.scratch_dir() as td:
tsk = ChallengeTask(read_only_task_dir=build.task_dir, python_path=self.python)
@@ -268,8 +259,7 @@ class MergerBot:
span.set_status(Status(StatusCode.OK))
def run_task(self, task: WeightedHarness, builds: list[BuildOutput]) -> bool:
"""
Strategy:
"""Strategy:
Given a task/WeightedHarness, we want to merge the local corpus into the remote corpus if it adds coverage
- acquire a lock on the merged corpus set, if not possible, return and move on to next task
- ensure all of the remotely stored corpus files are available locally
@@ -282,7 +272,6 @@ class MergerBot:
- remove any files only in L from the local corpus (as we know those don't add any coverage)
- release the lock on the merged corpus set
"""
logger.debug(f"Running merge pass for {task.harness_name} | {task.package_name} | {task.task_id}")
build = next(iter([b for b in builds if b.sanitizer == ADDRESS_SANITIZER]), None)
@@ -298,7 +287,10 @@ class MergerBot:
# We need to acquire a lock to ensure that we dont double remove a conflict
try:
with MergedCorpusSetLock(
self.redis, task.task_id, task.harness_name, MERGING_LOCK_TIMEOUT_SECONDS
self.redis,
task.task_id,
task.harness_name,
MERGING_LOCK_TIMEOUT_SECONDS,
).acquire():
# Create scratch directories for remote (R) and local-only (L) corpus parts, and copy files
with node_local.scratch_dir() as remote_dir, node_local.scratch_dir() as local_dir:
@@ -309,12 +301,12 @@ class MergerBot:
# If L is empty, the node is up to date
if not partitioned_corpus.local_only_files:
logger.debug(
f"Skipping merge for {task.harness_name} | {task.package_name} | {task.task_id} because local corpus is up to date"
f"Skipping merge for {task.harness_name} | {task.package_name} | {task.task_id} because local corpus is up to date",
)
return False # We did not do any work
logger.info(
f"Found {len(partitioned_corpus.local_only_files)} files only in local corpus for {task.harness_name}. Will run merge operation."
f"Found {len(partitioned_corpus.local_only_files)} files only in local corpus for {task.harness_name}. Will run merge operation.",
)
try:
@@ -344,14 +336,14 @@ class MergerBot:
remove_count = final_corpus.delete_locally()
if remove_count > 0:
logger.info(
f"Removed {remove_count} files from local corpus {corp.path} that don't add coverage"
f"Removed {remove_count} files from local corpus {corp.path} that don't add coverage",
)
return True # We did work
except FailedToAcquireLock:
logger.debug(
f"Skipping merge for {task.harness_name} | {task.package_name} | {task.task_id} because another worker is already merging"
f"Skipping merge for {task.harness_name} | {task.package_name} | {task.task_id} because another worker is already merging",
)
except Exception as e:
logger.error(f"Error merging corpus: {e}")
@@ -402,7 +394,11 @@ def main() -> None:
logger.info(f"Starting merger (crs_scratch_dir: {args.crs_scratch_dir})")
merger = MergerBot(
Redis.from_url(args.redis_url), args.timeout, args.python, args.crs_scratch_dir, args.max_local_files
Redis.from_url(args.redis_url),
args.timeout,
args.python,
args.crs_scratch_dir,
args.max_local_files,
)
merger.run()
fuzzer/src/buttercup/fuzzing_infra/coverage_bot.py
+31 -28
View File
@@ -1,36 +1,34 @@
from functools import lru_cache
import logging
import os
import random
from buttercup.common.logger import setup_package_logger
from buttercup.common.default_task_loop import TaskLoop
from buttercup.common.datastructures.msg_pb2 import BuildType, WeightedHarness, FunctionCoverage
from buttercup.common.datastructures.aliases import BuildType as BuildTypeHint
from buttercup.common.maps import CoverageMap
from typing import List, Generator
from redis import Redis
from buttercup.common.datastructures.msg_pb2 import BuildOutput
from buttercup.common.corpus import Corpus
from buttercup.fuzzing_infra.coverage_runner import CoverageRunner, CoveredFunction
from buttercup.fuzzing_infra.settings import CoverageBotSettings
from buttercup.common.challenge_task import ChallengeTask
from buttercup.common.utils import setup_periodic_zombie_reaper
import shutil
import buttercup.common.node_local as node_local
from collections.abc import Generator
from contextlib import contextmanager
from buttercup.common.telemetry import init_telemetry
from functools import lru_cache
from opentelemetry import trace
from opentelemetry.trace import Status, StatusCode
from buttercup.common.telemetry import set_crs_attributes, CRSActionCategory
from redis import Redis
from buttercup.common import node_local
from buttercup.common.challenge_task import ChallengeTask
from buttercup.common.corpus import Corpus
from buttercup.common.datastructures.aliases import BuildType as BuildTypeHint
from buttercup.common.datastructures.msg_pb2 import BuildOutput, BuildType, FunctionCoverage, WeightedHarness
from buttercup.common.default_task_loop import TaskLoop
from buttercup.common.logger import setup_package_logger
from buttercup.common.maps import CoverageMap
from buttercup.common.telemetry import CRSActionCategory, init_telemetry, set_crs_attributes
from buttercup.common.utils import setup_periodic_zombie_reaper
from buttercup.fuzzing_infra.coverage_runner import CoverageRunner, CoveredFunction
from buttercup.fuzzing_infra.settings import CoverageBotSettings
logger = logging.getLogger(__name__)
@lru_cache(maxsize=10)
def get_processed_coverage(corpus_path: str) -> set[str]:
"""
Get the set of processed coverage files in the corpus.
"""
"""Get the set of processed coverage files in the corpus."""
return set()
@@ -53,7 +51,7 @@ class CoverageBot(TaskLoop):
logger.info(f"Coverage bot initialized with sample_size: {sample_size}")
super().__init__(redis, timer_seconds)
def required_builds(self) -> List[BuildTypeHint]:
def required_builds(self) -> list[BuildTypeHint]:
return [BuildType.COVERAGE]
@contextmanager
@@ -67,6 +65,7 @@ class CoverageBot(TaskLoop):
Returns:
A context manager yielding a temporary directory containing symlinks
to the sampled corpus files, or the original corpus path if sample_size is 0.
"""
# Get list of input files from corpus
input_files = os.listdir(corpus.path)
@@ -78,7 +77,7 @@ class CoverageBot(TaskLoop):
# If sample_size is 0, use the entire corpus directly without sampling
if self.sample_size == 0:
logger.info(
f"Using entire non-processed corpus ({len(input_files)} files) in {corpus.path} (sample_size=0)"
f"Using entire non-processed corpus ({len(input_files)} files) in {corpus.path} (sample_size=0)",
)
yield (corpus.path, input_files)
return
@@ -128,7 +127,7 @@ class CoverageBot(TaskLoop):
with self._sample_corpus(corpus) as (sampled_corpus_path, remaining_files):
if len(remaining_files) == 0:
logger.info(
f"No files to process for {task.harness_name} | {corpus.path} | {local_tsk.project_name}"
f"No files to process for {task.harness_name} | {corpus.path} | {local_tsk.project_name}",
)
return
@@ -154,7 +153,7 @@ class CoverageBot(TaskLoop):
if func_coverage is None:
logger.error(
f"No function coverage found for {task.harness_name} | {corpus.path} | {local_tsk.project_name}"
f"No function coverage found for {task.harness_name} | {corpus.path} | {local_tsk.project_name}",
)
span.set_status(Status(StatusCode.ERROR))
return
@@ -162,7 +161,7 @@ class CoverageBot(TaskLoop):
get_processed_coverage(corpus.path).update(remaining_files)
logger.info(
f"Coverage for {task.harness_name} | {corpus.path} | {local_tsk.project_name} | processed {len(func_coverage)} functions"
f"Coverage for {task.harness_name} | {corpus.path} | {local_tsk.project_name} | processed {len(func_coverage)} functions",
)
self._submit_function_coverage(func_coverage, task.harness_name, task.package_name, task.task_id)
@@ -179,16 +178,20 @@ class CoverageBot(TaskLoop):
return function_coverage.covered_lines > old_function_coverage.covered_lines # type: ignore[no-any-return]
def _submit_function_coverage(
self, func_coverage: list[CoveredFunction], harness_name: str, package_name: str, task_id: str
self,
func_coverage: list[CoveredFunction],
harness_name: str,
package_name: str,
task_id: str,
) -> None:
"""
Store function coverage in Redis.
"""Store function coverage in Redis.
Args:
func_coverage: List of dictionaries containing function coverage metrics
harness_name: Name of the harness
package_name: Name of the package
task_id: Task ID
"""
coverage_map = CoverageMap(self.redis, harness_name, package_name, task_id)
fuzzer/src/buttercup/fuzzing_infra/coverage_runner.py
+20 -16
View File
@@ -1,13 +1,15 @@
from buttercup.common.challenge_task import ChallengeTask
import argparse
import subprocess
import json
import logging
import subprocess
from dataclasses import dataclass
from buttercup.common.project_yaml import ProjectYaml, Language
from bs4 import BeautifulSoup, Tag
from typing import Any, cast
from bs4 import BeautifulSoup, Tag
from buttercup.common.challenge_task import ChallengeTask
from buttercup.common.project_yaml import Language, ProjectYaml
logger = logging.getLogger(__name__)
@@ -26,8 +28,7 @@ class CoverageRunner:
@staticmethod
def _process_function_coverage(coverage_data: dict[str, Any]) -> list[CoveredFunction]:
"""
Process the LLVM coverage data to extract function-level line coverage.
"""Process the LLVM coverage data to extract function-level line coverage.
Returns a dictionary mapping function names to their line coverage metrics.
@@ -78,7 +79,7 @@ class CoverageRunner:
total_line_count,
covered_line_count,
function.get("filenames", []),
)
),
)
return function_coverage
@@ -105,29 +106,29 @@ class CoverageRunner:
jacoco_path = build_dir / "dumps" / f"{harness_name}.xml"
if not jacoco_path.exists():
logger.error(
f"Failed to find jacoco file for {harness_name} | {corpus_dir} | {self.tool.project_name} | in {jacoco_path}"
f"Failed to find jacoco file for {harness_name} | {corpus_dir} | {self.tool.project_name} | in {jacoco_path}",
)
return None
# parse the jacoco file
with open(jacoco_path, "r") as f:
with open(jacoco_path) as f:
soup = BeautifulSoup(f, "xml")
covered_functions = []
for target_class in soup.find_all("class"):
target_class = cast(Tag, target_class)
target_class = cast("Tag", target_class)
file_paths = []
source_file_name = target_class.get("sourcefilename")
if source_file_name is not None:
file_paths.append(source_file_name)
for method in target_class.find_all("method"):
method = cast(Tag, method)
method = cast("Tag", method)
method_name = method.get("name")
if method_name is None:
continue
method_name = str(method_name)
for ctr in method.find_all("counter"):
ctr = cast(Tag, ctr)
ctr = cast("Tag", ctr)
if ctr.get("type") == "LINE":
covered_attr = ctr.get("covered")
missed_attr = ctr.get("missed")
@@ -138,8 +139,11 @@ class CoverageRunner:
if covered_lines > 0:
covered_functions.append(
CoveredFunction(
method_name, total_lines, covered_lines, [str(f) for f in file_paths]
)
method_name,
total_lines,
covered_lines,
[str(f) for f in file_paths],
),
)
return covered_functions
@@ -155,7 +159,7 @@ class CoverageRunner:
profdata_path = package_path / "dumps" / "merged.profdata"
if not profdata_path.exists():
logger.error(
f"Failed to find profdata for {harness_name} | {corpus_dir} | {self.tool.project_name} | in {profdata_path}"
f"Failed to find profdata for {harness_name} | {corpus_dir} | {self.tool.project_name} | in {profdata_path}",
)
return None
@@ -163,7 +167,7 @@ class CoverageRunner:
coverage_file = package_path / "dumps" / "coverage.json"
harness_path = package_path / harness_name
args = [self.llvm_cov_path, "export", "-format=text", f"--instr-profile={profdata_path}", harness_path]
ret = subprocess.run(args, stdout=subprocess.PIPE)
ret = subprocess.run(args, check=False, stdout=subprocess.PIPE)
if ret.returncode != 0:
logger.error(
"Failed to convert profdata to json for %s | %s | %s | in %s (return code: %s)",
fuzzer/src/buttercup/fuzzing_infra/fuzzer_bot.py
+25 -22
View File
@@ -1,26 +1,26 @@
from buttercup.fuzzing_infra.runner import Runner, Conf, FuzzConfiguration
from buttercup.common.datastructures.msg_pb2 import BuildType, WeightedHarness, Crash
from buttercup.common.datastructures.aliases import BuildType as BuildTypeHint
from buttercup.common.queues import QueueFactory, QueueNames
from buttercup.common.corpus import Corpus, CrashDir
from buttercup.common import stack_parsing
from buttercup.common.stack_parsing import CrashSet
from buttercup.common.logger import setup_package_logger
from buttercup.common.utils import setup_periodic_zombie_reaper
from redis import Redis
from clusterfuzz.fuzz import engine
from buttercup.common.default_task_loop import TaskLoop
from typing import List
import random
from buttercup.common.datastructures.msg_pb2 import BuildOutput
import logging
from buttercup.common.challenge_task import ChallengeTask
from buttercup.fuzzing_infra.settings import FuzzerBotSettings
from buttercup.common.telemetry import init_telemetry, CRSActionCategory, set_crs_attributes
import random
from pathlib import Path
from clusterfuzz.fuzz import engine
from opentelemetry import trace
from opentelemetry.trace.status import Status, StatusCode
from redis import Redis
from buttercup.common import stack_parsing
from buttercup.common.challenge_task import ChallengeTask
from buttercup.common.corpus import Corpus, CrashDir
from buttercup.common.datastructures.aliases import BuildType as BuildTypeHint
from buttercup.common.datastructures.msg_pb2 import BuildOutput, BuildType, Crash, WeightedHarness
from buttercup.common.default_task_loop import TaskLoop
from buttercup.common.logger import setup_package_logger
from buttercup.common.node_local import scratch_dir
from pathlib import Path
from buttercup.common.queues import QueueFactory, QueueNames
from buttercup.common.stack_parsing import CrashSet
from buttercup.common.telemetry import CRSActionCategory, init_telemetry, set_crs_attributes
from buttercup.common.utils import setup_periodic_zombie_reaper
from buttercup.fuzzing_infra.runner import Conf, FuzzConfiguration, Runner
from buttercup.fuzzing_infra.settings import FuzzerBotSettings
logger = logging.getLogger(__name__)
@@ -44,7 +44,7 @@ class FuzzerBot(TaskLoop):
self.max_pov_size = max_pov_size
super().__init__(redis, timer_seconds)
def required_builds(self) -> List[BuildTypeHint]:
def required_builds(self) -> list[BuildTypeHint]:
return [BuildType.FUZZER]
def run_task(self, task: WeightedHarness, builds: dict[BuildTypeHint, BuildOutput]) -> None:
@@ -86,7 +86,10 @@ class FuzzerBot(TaskLoop):
crash_set = CrashSet(self.redis)
crash_dir = CrashDir(
self.crs_scratch_dir, task.task_id, task.harness_name, count_limit=self.crash_dir_count_limit
self.crs_scratch_dir,
task.task_id,
task.harness_name,
count_limit=self.crash_dir_count_limit,
)
for crash_ in result.crashes:
crash: engine.Crash = crash_
@@ -111,7 +114,7 @@ class FuzzerBot(TaskLoop):
crash.stacktrace,
):
logger.info(
f"Crash {crash.input_path}|{crash.reproduce_args}|{crash.crash_time} already in set"
f"Crash {crash.input_path}|{crash.reproduce_args}|{crash.crash_time} already in set",
)
logger.debug(f"Crash stacktrace: {crash.stacktrace}")
continue
fuzzer/src/buttercup/fuzzing_infra/orchestrator.py
+18 -14
View File
@@ -1,26 +1,30 @@
import argparse
import os
import time
from redis import Redis
from buttercup.common.clusterfuzz_utils import get_fuzz_targets
from buttercup.common.datastructures.msg_pb2 import BuildOutput, BuildType, WeightedHarness
from buttercup.common.logger import setup_package_logger
from buttercup.common.maps import BuildMap, HarnessWeights
from buttercup.common.queues import (
ReliableQueue,
RQItem,
GroupNames,
QueueFactory,
QueueNames,
GroupNames,
ReliableQueue,
RQItem,
)
from buttercup.common.maps import HarnessWeights, BuildMap
import argparse
from redis import Redis
from buttercup.common.datastructures.msg_pb2 import BuildType, BuildOutput, WeightedHarness
import time
from buttercup.common.clusterfuzz_utils import get_fuzz_targets
import os
from buttercup.common.logger import setup_package_logger
logger = setup_package_logger("fuzzer-orchestrator", __name__)
DEFAULT_WEIGHT = 1.0
def loop(
output_queue: ReliableQueue, target_list: HarnessWeights, build_map: BuildMap, sleep_time_seconds: int
output_queue: ReliableQueue,
target_list: HarnessWeights,
build_map: BuildMap,
sleep_time_seconds: int,
) -> None:
while True:
time.sleep(sleep_time_seconds)
@@ -47,7 +51,7 @@ def loop(
harness_name=os.path.basename(tgt),
package_name=deser_output.package_name,
task_id=deser_output.task_id,
)
),
)
output_queue.ack_item(output.item_id)
fuzzer/src/buttercup/fuzzing_infra/runner.py
+19 -12
View File
@@ -1,15 +1,17 @@
import argparse
import logging
import os
import typing
import uuid
from dataclasses import dataclass
from clusterfuzz.fuzz import get_engine
from clusterfuzz.fuzz.engine import Engine, FuzzResult, FuzzOptions
from buttercup.common.queues import FuzzConfiguration
from clusterfuzz.fuzz.engine import Engine, FuzzOptions, FuzzResult
from buttercup.common.logger import setup_package_logger
from buttercup.common.node_local import scratch_dir
from buttercup.common.queues import FuzzConfiguration
from buttercup.fuzzing_infra.temp_dir import patched_temp_dir, scratch_cwd
import typing
import os
from dataclasses import dataclass
import argparse
import uuid
import logging
logger = logging.getLogger(__name__)
@@ -29,11 +31,11 @@ class Runner:
job_name = f"{conf.engine}_{conf.sanitizer}"
with patched_temp_dir() as _td, scratch_cwd() as _cwd_temp:
engine = typing.cast(Engine, get_engine(conf.engine))
engine = typing.cast("Engine", get_engine(conf.engine))
target = conf.target_path
build_dir = os.path.dirname(target)
distinguisher = uuid.uuid4()
repro_dir = os.path.join(build_dir, f"repro{str(distinguisher)}")
repro_dir = os.path.join(build_dir, f"repro{distinguisher!s}")
os.makedirs(repro_dir, exist_ok=True)
os.environ["JOB_NAME"] = job_name
logger.debug(f"Calling engine.prepare with {conf.corpus_dir} | {target} | {build_dir}")
@@ -52,11 +54,16 @@ class Runner:
job_name = f"{conf.engine}_{conf.sanitizer}"
os.environ["JOB_NAME"] = job_name
with patched_temp_dir() as _td, scratch_cwd() as _cwd_temp:
engine = typing.cast(Engine, get_engine(conf.engine))
engine = typing.cast("Engine", get_engine(conf.engine))
# Temporary directory ignores crashes
with scratch_dir() as td:
engine.minimize_corpus(
conf.target_path, [], [conf.corpus_dir], output_dir, str(td.path), self.conf.timeout
conf.target_path,
[],
[conf.corpus_dir],
output_dir,
str(td.path),
self.conf.timeout,
)
fuzzer/src/buttercup/fuzzing_infra/settings.py
+3 -2
View File
@@ -1,7 +1,8 @@
from pydantic_settings import BaseSettings
from pydantic import Field
from typing import Annotated
from pydantic import Field
from pydantic_settings import BaseSettings
class Config:
cli_parse_args = True
fuzzer/src/buttercup/fuzzing_infra/stimulate_build_bot.py
+3 -1
View File
@@ -1,7 +1,9 @@
import argparse
from redis import Redis
from buttercup.common.queues import QueueFactory, QueueNames
from buttercup.common.datastructures.msg_pb2 import BuildRequest, BuildType
from buttercup.common.queues import QueueFactory, QueueNames
def main() -> None:
fuzzer/src/buttercup/fuzzing_infra/temp_dir.py
+6 -4
View File
@@ -1,12 +1,14 @@
import os
import tempfile
import contextvars
import logging
import os
import tempfile
from collections.abc import Iterator
from contextlib import contextmanager
from typing import Iterator
from unittest.mock import patch
from clusterfuzz._internal.system import environment, shell
from buttercup.common.node_local import scratch_dir, TmpDir
from buttercup.common.node_local import TmpDir, scratch_dir
logger = logging.getLogger(__name__)
fuzzer/src/buttercup/fuzzing_infra/tracer_bot.py
+14 -13
View File
@@ -1,17 +1,18 @@
import logging
import os
from pathlib import Path
from redis import Redis
from buttercup.common import node_local, stack_parsing
from buttercup.common.datastructures.msg_pb2 import TracedCrash
from buttercup.common.logger import setup_package_logger
from buttercup.common.queues import GroupNames, QueueFactory, QueueNames
from buttercup.common.task_registry import TaskRegistry
from buttercup.common.telemetry import init_telemetry
from buttercup.common.utils import serve_loop, setup_periodic_zombie_reaper
from buttercup.fuzzing_infra.settings import TracerSettings
from buttercup.fuzzing_infra.tracer_runner import TracerRunner
from buttercup.common.logger import setup_package_logger
import os
import logging
from redis import Redis
from buttercup.common.queues import QueueFactory, QueueNames, GroupNames
from buttercup.common.datastructures.msg_pb2 import TracedCrash
from buttercup.common.task_registry import TaskRegistry
from pathlib import Path
from buttercup.common import stack_parsing
from buttercup.common.utils import serve_loop, setup_periodic_zombie_reaper
import buttercup.common.node_local as node_local
from buttercup.common.telemetry import init_telemetry
logger = logging.getLogger(__name__)
@@ -68,7 +69,7 @@ class TracerBot:
TracedCrash(
crash=item.deserialized,
tracer_stacktrace=ntrace,
)
),
)
logger.info(f"Acknowledging tracer request for {item.deserialized.target.task_id}")
fuzzer/src/buttercup/fuzzing_infra/tracer_runner.py
+9 -7
View File
@@ -1,13 +1,15 @@
from buttercup.common.maps import BuildMap
from buttercup.common.challenge_task import ChallengeTask, ReproduceResult
from buttercup.common.datastructures.msg_pb2 import BuildType
from pathlib import Path
from dataclasses import dataclass
import logging
from redis import Redis
from dataclasses import dataclass
from pathlib import Path
from opentelemetry import trace
from opentelemetry.trace import Status, StatusCode
from buttercup.common.telemetry import set_crs_attributes, CRSActionCategory
from redis import Redis
from buttercup.common.challenge_task import ChallengeTask, ReproduceResult
from buttercup.common.datastructures.msg_pb2 import BuildType
from buttercup.common.maps import BuildMap
from buttercup.common.telemetry import CRSActionCategory, set_crs_attributes
logger = logging.getLogger(__name__)
fuzzer/src/buttercup/fuzzing_infra/utils_cli.py
+3 -3
View File
@@ -3,7 +3,6 @@
import argparse
import sys
from pathlib import Path
from typing import Optional
from buttercup.common.stack_parsing import get_crash_data
@@ -15,7 +14,8 @@ def parse_args() -> argparse.Namespace:
# get_crash_data subcommand
get_crash_data_parser = subparsers.add_parser(
"get_crash_data", help="Extract crash state from a stacktrace file or stdin"
"get_crash_data",
help="Extract crash state from a stacktrace file or stdin",
)
get_crash_data_parser.add_argument(
"stacktrace_file",
@@ -32,7 +32,7 @@ def parse_args() -> argparse.Namespace:
return parser.parse_args()
def read_stacktrace(file_path: Optional[Path] = None) -> str:
def read_stacktrace(file_path: Path | None = None) -> str:
"""Read stacktrace from file or stdin."""
if file_path:
return file_path.read_text()
fuzzer/tests/test_builder_bot.py
+14 -7
View File
@@ -1,11 +1,12 @@
import unittest
from unittest.mock import Mock, patch, MagicMock
from redis import Redis
from pathlib import Path
from unittest.mock import MagicMock, Mock, patch
from buttercup.fuzzing_infra.builder_bot import BuilderBot
from buttercup.common.datastructures.msg_pb2 import BuildRequest, BuildOutput, BuildType
from redis import Redis
from buttercup.common.datastructures.msg_pb2 import BuildOutput, BuildRequest, BuildType
from buttercup.common.task_registry import TaskRegistry
from buttercup.fuzzing_infra.builder_bot import BuilderBot
class TestBuilderBot(unittest.TestCase):
@@ -197,12 +198,16 @@ class TestBuilderBot(unittest.TestCase):
# Verify ChallengeTask was created with caching enabled
challenge_task_mock.assert_called_once_with(
Path("/path/to/task"), python_path=self.python, local_task_dir=Path("/path/to/task")
Path("/path/to/task"),
python_path=self.python,
local_task_dir=Path("/path/to/task"),
)
# Verify build was called
task_mock.build_fuzzers_with_cache.assert_called_once_with(
engine="libfuzzer", sanitizer="address", pull_latest_base_image=self.allow_pull
engine="libfuzzer",
sanitizer="address",
pull_latest_base_image=self.allow_pull,
)
# Verify task was committed and archived
@@ -281,7 +286,9 @@ class TestBuilderBot(unittest.TestCase):
# Verify build was called
task_mock.build_fuzzers_with_cache.assert_called_once_with(
engine="libfuzzer", sanitizer="address", pull_latest_base_image=self.allow_pull
engine="libfuzzer",
sanitizer="address",
pull_latest_base_image=self.allow_pull,
)
# Verify task was committed and archived
fuzzer/tests/test_coverage_bot.py
+7 -5
View File
@@ -1,12 +1,14 @@
import pytest
import os
import tempfile
from unittest.mock import patch, MagicMock
from unittest.mock import MagicMock, patch
import pytest
from redis import Redis
from buttercup.common.datastructures.msg_pb2 import FunctionCoverage
from buttercup.common.maps import CoverageMap
from buttercup.fuzzing_infra.coverage_bot import CoverageBot
from buttercup.fuzzing_infra.coverage_runner import CoveredFunction
from buttercup.common.maps import CoverageMap
from buttercup.common.datastructures.msg_pb2 import FunctionCoverage
@pytest.fixture
@@ -229,7 +231,7 @@ def test_should_update_function_coverage_worse_coverage(redis_client):
def test_submit_function_coverage(coverage_bot, redis_client):
# Create test data
func_coverage = [
CoveredFunction(names="test_function", total_lines=100, covered_lines=75, function_paths=["path1", "path2"])
CoveredFunction(names="test_function", total_lines=100, covered_lines=75, function_paths=["path1", "path2"]),
]
harness_name = "test_harness"
package_name = "test_package"
fuzzer/tests/test_merger_bot.py
+69 -42
View File
@@ -1,12 +1,13 @@
import unittest
from unittest.mock import Mock, patch, MagicMock
from redis import Redis
from pathlib import Path
from unittest.mock import MagicMock, Mock, patch
from redis import Redis
from buttercup.fuzzing_infra.corpus_merger import MergerBot, BaseCorpus, PartitionedCorpus
from buttercup.common.datastructures.msg_pb2 import WeightedHarness, BuildOutput
from buttercup.common.sets import FailedToAcquireLock
from buttercup.common.constants import ADDRESS_SANITIZER
from buttercup.common.datastructures.msg_pb2 import BuildOutput, WeightedHarness
from buttercup.common.sets import FailedToAcquireLock
from buttercup.fuzzing_infra.corpus_merger import BaseCorpus, MergerBot, PartitionedCorpus
class TestMergerBot(unittest.TestCase):
@@ -27,7 +28,11 @@ class TestMergerBot(unittest.TestCase):
with patch("buttercup.fuzzing_infra.corpus_merger.Runner") as runner_class_mock:
runner_class_mock.return_value = self.runner_mock
self.merger_bot = MergerBot(
self.redis_mock, self.timeout_seconds, self.python, self.crs_scratch_dir, self.max_local_files
self.redis_mock,
self.timeout_seconds,
self.python,
self.crs_scratch_dir,
self.max_local_files,
)
@patch("buttercup.fuzzing_infra.corpus_merger.Corpus")
@@ -60,7 +65,10 @@ class TestMergerBot(unittest.TestCase):
# Verify behavior
corpus_instance.hash_new_corpus.assert_called_once()
base_corpus_mock.assert_called_once_with(
corpus_instance, scratch_dir_mock().__enter__(), scratch_dir_mock().__enter__(), self.max_local_files
corpus_instance,
scratch_dir_mock().__enter__(),
scratch_dir_mock().__enter__(),
self.max_local_files,
)
base_corpus_instance.partition_corpus.assert_called_once()
@@ -72,7 +80,11 @@ class TestMergerBot(unittest.TestCase):
@patch("buttercup.fuzzing_infra.corpus_merger.node_local.scratch_dir")
@patch("buttercup.fuzzing_infra.corpus_merger.BaseCorpus")
def test_run_task_failed_to_acquire_lock(
self, base_corpus_mock, scratch_dir_mock, lock_class_mock, corpus_class_mock
self,
base_corpus_mock,
scratch_dir_mock,
lock_class_mock,
corpus_class_mock,
):
# Setup mocks
corpus_instance = corpus_class_mock.return_value
@@ -164,7 +176,10 @@ class TestMergerBot(unittest.TestCase):
# Verify behavior
corpus_instance.hash_new_corpus.assert_called_once()
base_corpus_mock.assert_called_once_with(
corpus_instance, scratch_dir_mock().__enter__(), scratch_dir_mock().__enter__(), self.max_local_files
corpus_instance,
scratch_dir_mock().__enter__(),
scratch_dir_mock().__enter__(),
self.max_local_files,
)
base_corpus_instance.partition_corpus.assert_called_once()
@@ -184,7 +199,7 @@ class TestMergerBot(unittest.TestCase):
This test is no longer applicable as the _rehash_files method no longer exists.
The file hashing functionality is now handled by the Corpus.hash_corpus method.
"""
pass # Skipping test as functionality has been moved to Corpus class
# Skipping test as functionality has been moved to Corpus class
class TestBaseCorpus(unittest.TestCase):
@@ -208,14 +223,13 @@ class TestBaseCorpus(unittest.TestCase):
if args[0] == "/corpus/path":
return f"/corpus/path/{args[1]}"
# For local_dir joins
elif str(args[0]) == "/tmp/local_dir":
if str(args[0]) == "/tmp/local_dir":
return f"/tmp/local_dir/{args[1]}"
# For remote_dir joins
elif str(args[0]) == "/tmp/remote_dir":
if str(args[0]) == "/tmp/remote_dir":
return f"/tmp/remote_dir/{args[1]}"
# Default
else:
return "/".join(str(arg) for arg in args)
return "/".join(str(arg) for arg in args)
path_join_mock.side_effect = mock_path_join
@@ -256,7 +270,8 @@ class TestBaseCorpus(unittest.TestCase):
# Check that local_only_files and remote_files sets contain the right filenames
self.assertEqual(
kwargs.get("local_only_files"), {"c123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"}
kwargs.get("local_only_files"),
{"c123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"},
)
self.assertEqual(
kwargs.get("remote_files"),
@@ -279,7 +294,13 @@ class TestPartitionedCorpus(unittest.TestCase):
@patch("os.path.exists")
@patch("random.shuffle")
def test_initialization_with_file_limit(
self, shuffle_mock, path_exists_mock, path_join_mock, shutil_copy_mock, scratch_dir_mock, corpus_mock
self,
shuffle_mock,
path_exists_mock,
path_join_mock,
shutil_copy_mock,
scratch_dir_mock,
corpus_mock,
):
"""Test that files are shuffled and limited to max_local_files."""
# Setup mocks
@@ -295,12 +316,11 @@ class TestPartitionedCorpus(unittest.TestCase):
def mock_path_join(*args):
if args[0] == "/corpus/path":
return f"/corpus/path/{args[1]}"
elif str(args[0]) == "/tmp/local_dir":
if str(args[0]) == "/tmp/local_dir":
return f"/tmp/local_dir/{args[1]}"
elif str(args[0]) == "/tmp/remote_dir":
if str(args[0]) == "/tmp/remote_dir":
return f"/tmp/remote_dir/{args[1]}"
else:
return "/".join(str(arg) for arg in args)
return "/".join(str(arg) for arg in args)
path_join_mock.side_effect = mock_path_join
@@ -357,14 +377,13 @@ class TestPartitionedCorpus(unittest.TestCase):
if args[0] == "/corpus/path":
return f"/corpus/path/{args[1]}"
# For local_dir joins
elif str(args[0]) == "/tmp/local_dir":
if str(args[0]) == "/tmp/local_dir":
return f"/tmp/local_dir/{args[1]}"
# For remote_dir joins
elif str(args[0]) == "/tmp/remote_dir":
if str(args[0]) == "/tmp/remote_dir":
return f"/tmp/remote_dir/{args[1]}"
# Default
else:
return "/".join(str(arg) for arg in args)
return "/".join(str(arg) for arg in args)
path_join_mock.side_effect = mock_path_join
@@ -418,7 +437,12 @@ class TestPartitionedCorpus(unittest.TestCase):
@patch("os.path.join")
@patch("os.path.exists")
def test_initialization_local_file_missing(
self, path_exists_mock, path_join_mock, shutil_copy_mock, scratch_dir_mock, corpus_mock
self,
path_exists_mock,
path_join_mock,
shutil_copy_mock,
scratch_dir_mock,
corpus_mock,
):
"""Test that local files that cannot be copied are removed from local_only_files."""
# Setup mocks
@@ -434,12 +458,11 @@ class TestPartitionedCorpus(unittest.TestCase):
def mock_path_join(*args):
if args[0] == "/corpus/path":
return f"/corpus/path/{args[1]}"
elif str(args[0]) == "/tmp/local_dir":
if str(args[0]) == "/tmp/local_dir":
return f"/tmp/local_dir/{args[1]}"
elif str(args[0]) == "/tmp/remote_dir":
if str(args[0]) == "/tmp/remote_dir":
return f"/tmp/remote_dir/{args[1]}"
else:
return "/".join(str(arg) for arg in args)
return "/".join(str(arg) for arg in args)
path_join_mock.side_effect = mock_path_join
@@ -472,7 +495,8 @@ class TestPartitionedCorpus(unittest.TestCase):
# Verify the missing file was removed from local_only_files
self.assertEqual(
partitioned_corpus.local_only_files, {"c123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"}
partitioned_corpus.local_only_files,
{"c123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"},
)
# The remote files should remain unchanged
@@ -484,7 +508,12 @@ class TestPartitionedCorpus(unittest.TestCase):
@patch("os.path.join")
@patch("os.path.exists")
def test_initialization_remote_file_missing(
self, path_exists_mock, path_join_mock, shutil_copy_mock, scratch_dir_mock, corpus_mock
self,
path_exists_mock,
path_join_mock,
shutil_copy_mock,
scratch_dir_mock,
corpus_mock,
):
"""Test that when a file is missing from corpus.path but available in corpus.remote_path, it's copied from remote."""
# Setup mocks
@@ -501,14 +530,13 @@ class TestPartitionedCorpus(unittest.TestCase):
def mock_path_join(*args):
if args[0] == "/corpus/path":
return f"/corpus/path/{args[1]}"
elif args[0] == "/remote/path":
if args[0] == "/remote/path":
return f"/remote/path/{args[1]}"
elif str(args[0]) == "/tmp/local_dir":
if str(args[0]) == "/tmp/local_dir":
return f"/tmp/local_dir/{args[1]}"
elif str(args[0]) == "/tmp/remote_dir":
if str(args[0]) == "/tmp/remote_dir":
return f"/tmp/remote_dir/{args[1]}"
else:
return "/".join(str(arg) for arg in args)
return "/".join(str(arg) for arg in args)
path_join_mock.side_effect = mock_path_join
@@ -520,8 +548,8 @@ class TestPartitionedCorpus(unittest.TestCase):
if "/corpus/path/" in src:
raise FileNotFoundError(f"File not found in local corpus: {src}")
# Second call with corpus.remote_path should succeed
return None
return None
return
return
shutil_copy_mock.side_effect = mock_copy
@@ -604,14 +632,13 @@ class TestPartitionedCorpus(unittest.TestCase):
if args[0] == "/corpus/path":
return f"/corpus/path/{args[1]}"
# For local_dir joins
elif str(args[0]) == "/tmp/local_dir":
if str(args[0]) == "/tmp/local_dir":
return f"/tmp/local_dir/{args[1]}"
# For remote_dir joins
elif str(args[0]) == "/tmp/remote_dir":
if str(args[0]) == "/tmp/remote_dir":
return f"/tmp/remote_dir/{args[1]}"
# Default
else:
return "/".join(str(arg) for arg in args)
return "/".join(str(arg) for arg in args)
path_join_mock.side_effect = mock_path_join
fuzzer/tests/test_temp_dir.py
+3 -3
View File
@@ -3,9 +3,9 @@ import tempfile
import threading
import time
import unittest
from unittest.mock import patch, MagicMock
from unittest.mock import MagicMock, patch
from buttercup.fuzzing_infra.temp_dir import get_temp_dir, patched_temp_dir, _scratch_path_var
from buttercup.fuzzing_infra.temp_dir import _scratch_path_var, get_temp_dir, patched_temp_dir
class TestGetTempDir(unittest.TestCase):
@@ -108,7 +108,7 @@ class TestPatchedTempDir(unittest.TestCase):
with patched_temp_dir():
with patch("buttercup.fuzzing_infra.temp_dir.shell.create_directory") as mock_create_dir:
# Import and call the clusterfuzz function - it should be patched to use our implementation
import clusterfuzz._internal.bot.fuzzers.utils as utils
from clusterfuzz._internal.bot.fuzzers import utils
result = utils.get_temp_dir()
fuzzer/tests/test_tracer_bot.py
+8 -6
View File
@@ -1,10 +1,12 @@
import pytest
from unittest.mock import patch, MagicMock
from redis import Redis
from buttercup.fuzzing_infra.tracer_bot import TracerBot
from buttercup.common.datastructures.msg_pb2 import Crash, Task, BuildOutput
from buttercup.common.task_registry import TaskRegistry
import time
from unittest.mock import MagicMock, patch
import pytest
from redis import Redis
from buttercup.common.datastructures.msg_pb2 import BuildOutput, Crash, Task
from buttercup.common.task_registry import TaskRegistry
from buttercup.fuzzing_infra.tracer_bot import TracerBot
@pytest.fixture
orchestrator/scripts/challenge.py
+2 -3
View File
@@ -2,9 +2,9 @@
import argparse
import json
import sys
import time
import urllib.request
import sys
from typing import Any
SECONDS = 1
@@ -317,11 +317,10 @@ CHALLENGE_MAP = {
def print_request(req, file=sys.stderr):
"""Print a request object."""
print(f"{req.get_method()} {req.full_url}", file=file)
headers = dict(req.headers)
max_key_len = max(len(k) for k in headers.keys()) if headers else 0
max_key_len = max(len(k) for k in headers) if headers else 0
for key, value in headers.items():
print(f"{key.ljust(max_key_len)}: {value}", file=file)
orchestrator/src/buttercup/orchestrator/api_client_factory.py
+6 -2
View File
@@ -1,12 +1,15 @@
import logging
from buttercup.orchestrator.competition_api_client.configuration import Configuration
from buttercup.orchestrator.competition_api_client.api_client import ApiClient
from buttercup.orchestrator.competition_api_client.configuration import Configuration
logger = logging.getLogger(__name__)
def create_api_client(
competition_api_url: str, competition_api_username: str, competition_api_password: str
competition_api_url: str,
competition_api_username: str,
competition_api_password: str,
) -> ApiClient:
"""Initialize the competition API client with common configuration.
@@ -15,6 +18,7 @@ def create_api_client(
Returns:
ApiClient: Configured API client instance
"""
configuration = Configuration(
host=competition_api_url,
orchestrator/src/buttercup/orchestrator/downloader/__cli__.py
+11 -10
View File
@@ -1,18 +1,19 @@
from buttercup.orchestrator.downloader.downloader import Downloader
import requests
import requests.adapters
from pydantic_settings import get_subcommand
from redis import Redis
from requests_file import FileAdapter
from buttercup.common.datastructures.msg_pb2 import SourceDetail, Task
from buttercup.common.logger import setup_package_logger
from buttercup.orchestrator.downloader.config import (
DownloaderSettings,
DownloaderServeCommand,
DownloaderProcessCommand,
DownloaderServeCommand,
DownloaderSettings,
TaskType,
)
from buttercup.common.logger import setup_package_logger
from pydantic_settings import get_subcommand
from buttercup.common.datastructures.msg_pb2 import Task, SourceDetail
from buttercup.orchestrator.downloader.downloader import Downloader
from buttercup.orchestrator.utils import response_stream_to_file
from redis import Redis
import requests.adapters
from requests_file import FileAdapter
import requests
def prepare_task(command: DownloaderProcessCommand, session: requests.Session) -> Task:
orchestrator/src/buttercup/orchestrator/downloader/config.py
+6 -6
View File
@@ -1,11 +1,11 @@
from pydantic_settings import BaseSettings, CliPositionalArg, CliSubCommand
from pydantic import BaseModel
from typing import Annotated
from pydantic import Field
from pathlib import Path
from enum import Enum
import time
import uuid
from enum import Enum
from pathlib import Path
from typing import Annotated
from pydantic import BaseModel, Field
from pydantic_settings import BaseSettings, CliPositionalArg, CliSubCommand
class DownloaderServeCommand(BaseModel):
orchestrator/src/buttercup/orchestrator/downloader/downloader.py
+23 -22
View File
@@ -1,22 +1,22 @@
import logging
import requests
import tarfile
from dataclasses import dataclass, field
import uuid
import tempfile
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional
import requests
from redis import Redis
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from buttercup.common.queues import QueueFactory, ReliableQueue, QueueNames, GroupNames
from buttercup.common.datastructures.msg_pb2 import Task, SourceDetail, TaskDownload, TaskReady
from buttercup.orchestrator.utils import response_stream_to_file
from buttercup.common import node_local
from buttercup.common.datastructures.msg_pb2 import SourceDetail, Task, TaskDownload, TaskReady
from buttercup.common.queues import GroupNames, QueueFactory, QueueNames, ReliableQueue
from buttercup.common.task_meta import TaskMeta
from redis import Redis
from buttercup.common.task_registry import TaskRegistry
from buttercup.common.utils import serve_loop
import buttercup.common.node_local as node_local
from buttercup.orchestrator.utils import response_stream_to_file
logger = logging.getLogger(__name__)
@@ -61,7 +61,7 @@ class Downloader:
"""Creates and returns the directory path for a task"""
return self.download_dir / task_id
def download_source(self, task_id: str, tmp_task_dir: Path, source: SourceDetail) -> Optional[Path]:
def download_source(self, task_id: str, tmp_task_dir: Path, source: SourceDetail) -> Path | None:
"""Downloads a source file and verifies its SHA256"""
try:
filepath = tmp_task_dir / str(uuid.uuid4())
@@ -78,18 +78,17 @@ class Downloader:
logger.info(f"[task {task_id}] Successfully downloaded source type {source.source_type} to {filepath}")
return filepath
except Exception as e:
logger.error(f"Failed to download {source.url}: {str(e)}")
logger.error(f"Failed to download {source.url}: {e!s}")
return None
def _get_source_type_dir(self, source_type: SourceDetail.SourceType) -> str:
if source_type == SourceDetail.SourceType.SOURCE_TYPE_REPO:
return "src"
elif source_type == SourceDetail.SourceType.SOURCE_TYPE_FUZZ_TOOLING:
if source_type == SourceDetail.SourceType.SOURCE_TYPE_FUZZ_TOOLING:
return "fuzz-tooling"
elif source_type == SourceDetail.SourceType.SOURCE_TYPE_DIFF:
if source_type == SourceDetail.SourceType.SOURCE_TYPE_DIFF:
return "diff"
else:
raise ValueError(f"Unknown source type: {source_type}")
raise ValueError(f"Unknown source type: {source_type}")
def extract_source(self, task_id: str, tmp_task_dir: Path, source: SourceDetail, source_file: Path) -> bool:
"""Uncompress a source file and returns the name of the main directory it contains"""
@@ -122,7 +121,7 @@ class Downloader:
logger.info(f"[task {task_id}] Successfully extracted {source_file}")
return True
except Exception as e:
logger.error(f"[task {task_id}] Failed to extract {source_file}: {str(e)}")
logger.error(f"[task {task_id}] Failed to extract {source_file}: {e!s}")
return False
def _download_and_extract_sources(self, task_id: str, tmp_task_dir: Path, sources: list) -> bool:
@@ -153,15 +152,14 @@ class Downloader:
# need to change the Task while downloading/extracting it.
if "Directory not empty" in str(e):
logger.warning(
f"Directory {final_task_dir} already exists, another process downloaded the task first, ignore it..."
f"Directory {final_task_dir} already exists, another process downloaded the task first, ignore it...",
)
return True
else:
logger.exception(f"Failed to move task directory: {str(e)}")
return False
logger.exception(f"Failed to move task directory: {e!s}")
return False
except Exception as e:
# Re-raise any other errors
logger.exception(f"Failed to move task directory: {str(e)}")
logger.exception(f"Failed to move task directory: {e!s}")
return False
def process_task(self, task: Task) -> bool:
@@ -248,6 +246,9 @@ class Downloader:
return self
def __exit__(
self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: object | None
self,
exc_type: type[BaseException] | None,
exc_val: BaseException | None,
exc_tb: object | None,
) -> None:
self.cleanup()
orchestrator/src/buttercup/orchestrator/pov_reproducer/__cli__.py
+5 -3
View File
@@ -1,8 +1,10 @@
import logging
from redis import Redis
from buttercup.common.logger import setup_package_logger
from buttercup.orchestrator.pov_reproducer.config import Settings
from buttercup.orchestrator.pov_reproducer.pov_reproducer import POVReproducer
from buttercup.common.logger import setup_package_logger
from redis import Redis
import logging
logger = logging.getLogger(__name__)
orchestrator/src/buttercup/orchestrator/pov_reproducer/config.py
+2 -1
View File
@@ -1,6 +1,7 @@
from pydantic_settings import BaseSettings
from typing import Annotated
from pydantic import Field
from pydantic_settings import BaseSettings
class Settings(BaseSettings):
orchestrator/src/buttercup/orchestrator/pov_reproducer/pov_reproducer.py
+10 -10
View File
@@ -1,17 +1,17 @@
from __future__ import annotations
from dataclasses import dataclass, field
import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional
from buttercup.common.maps import BuildMap
from buttercup.common.challenge_task import ChallengeTask
from buttercup.common.datastructures.msg_pb2 import BuildType, BuildOutput, POVReproduceRequest
from buttercup.common.sets import PoVReproduceStatus
from buttercup.common.task_registry import TaskRegistry
import buttercup.common.node_local as node_local
from redis import Redis
from buttercup.common import node_local
from buttercup.common.challenge_task import ChallengeTask
from buttercup.common.datastructures.msg_pb2 import BuildOutput, BuildType, POVReproduceRequest
from buttercup.common.maps import BuildMap
from buttercup.common.sets import PoVReproduceStatus
from buttercup.common.task_registry import TaskRegistry
from buttercup.common.utils import serve_loop
logger = logging.getLogger(__name__)
@@ -31,7 +31,7 @@ class POVReproducer:
self.registry = TaskRegistry(self.redis)
def serve_item(self) -> bool:
entry: Optional[POVReproduceRequest] = self.pov_status.get_one_pending()
entry: POVReproduceRequest | None = self.pov_status.get_one_pending()
if entry is None:
return False
@@ -54,7 +54,7 @@ class POVReproducer:
logger.info(f"Reproducing POV for {task_id} | {harness_name} | {pov_path}")
builds = BuildMap(self.redis)
build_output_with_patch: Optional[BuildOutput] = builds.get_build_from_san(
build_output_with_patch: BuildOutput | None = builds.get_build_from_san(
task_id,
BuildType.PATCH,
sanitizer,
orchestrator/src/buttercup/orchestrator/scheduler/__cli__.py
+11 -10
View File
@@ -1,16 +1,17 @@
from buttercup.orchestrator.scheduler.config import (
Settings,
ServeCommand,
ProcessBuildOutputCommand,
ProcessReadyTaskCommand,
)
from buttercup.orchestrator.scheduler.scheduler import Scheduler, Task, BuildOutput
from buttercup.common.logger import setup_package_logger
from buttercup.common.telemetry import init_telemetry
import logging
from pydantic_settings import get_subcommand
from redis import Redis
import logging
from buttercup.common.logger import setup_package_logger
from buttercup.common.telemetry import init_telemetry
from buttercup.orchestrator.scheduler.config import (
ProcessBuildOutputCommand,
ProcessReadyTaskCommand,
ServeCommand,
Settings,
)
from buttercup.orchestrator.scheduler.scheduler import BuildOutput, Scheduler, Task
logger = logging.getLogger(__name__)
orchestrator/src/buttercup/orchestrator/scheduler/cancellation.py
+12 -6
View File
@@ -1,8 +1,10 @@
import logging
from dataclasses import dataclass, field
from redis import Redis
from buttercup.common.queues import ReliableQueue, QueueFactory, RQItem, QueueNames, GroupNames
from buttercup.common.datastructures.msg_pb2 import TaskDelete
from buttercup.common.queues import GroupNames, QueueFactory, QueueNames, ReliableQueue, RQItem
from buttercup.common.task_registry import TaskRegistry
logger = logging.getLogger(__name__)
@@ -23,6 +25,7 @@ class Cancellation:
redis (Redis): Redis connection
delete_queue (ReliableQueue): Queue for processing deletion requests
registry (TaskRegistry): Registry for tracking task state
"""
redis: Redis
@@ -32,7 +35,9 @@ class Cancellation:
def __post_init__(self) -> None:
"""Initialize Redis connection, deletion queue and task registry."""
self.delete_queue = QueueFactory(self.redis).create(
QueueNames.DELETE_TASK, GroupNames.ORCHESTRATOR, block_time=None
QueueNames.DELETE_TASK,
GroupNames.ORCHESTRATOR,
block_time=None,
)
self.registry = TaskRegistry(self.redis)
@@ -45,6 +50,7 @@ class Cancellation:
Returns:
bool: True if any task was successfully marked as cancelled, False otherwise
"""
# Handle the case where 'all' is set to True
if delete_request.HasField("all") and delete_request.all:
@@ -61,7 +67,7 @@ class Cancellation:
return any_cancelled
# Handle the case where a specific task_id is provided
elif delete_request.HasField("task_id"):
if delete_request.HasField("task_id"):
task_id = delete_request.task_id
logger.info(f"Processing delete request for task {task_id}")
@@ -71,9 +77,8 @@ class Cancellation:
return True
# Neither 'all' nor 'task_id' is set
else:
logger.warning("Delete request missing both 'task_id' and 'all' fields, ignoring")
return False
logger.warning("Delete request missing both 'task_id' and 'all' fields, ignoring")
return False
def process_cancellations(self) -> bool:
"""Process one iteration of the cancellation loop.
@@ -83,6 +88,7 @@ class Cancellation:
Returns:
bool: True if any task was cancelled via delete request, False otherwise
"""
any_cancellation = False
orchestrator/src/buttercup/orchestrator/scheduler/config.py
+12 -8
View File
@@ -1,8 +1,8 @@
from pydantic_settings import BaseSettings, CliSubCommand
from pydantic import BaseModel
from typing import Annotated
from pydantic import Field
from pathlib import Path
from typing import Annotated
from pydantic import BaseModel, Field
from pydantic_settings import BaseSettings, CliSubCommand
class ServeCommand(BaseSettings):
@@ -12,16 +12,20 @@ class ServeCommand(BaseSettings):
competition_api_key_id: Annotated[str, Field(default="api_key_id", description="Competition API username")]
competition_api_key_token: Annotated[str, Field(default="api_key_token", description="Competition API password")]
competition_api_cycle_time: Annotated[
float, Field(default=10.0, description="Min seconds between competition api interactions")
float,
Field(default=10.0, description="Min seconds between competition api interactions"),
]
patch_submission_retry_limit: Annotated[
int, Field(default=60, description="Number of retries for errored patch submissions.")
int,
Field(default=60, description="Number of retries for errored patch submissions."),
]
patch_requests_per_vulnerability: Annotated[
int, Field(default=1, description="Number of patch requests per vulnerability")
int,
Field(default=1, description="Number of patch requests per vulnerability"),
]
concurrent_patch_requests_per_task: Annotated[
int, Field(default=12, description="Number of concurrent patch requests per task")
int,
Field(default=12, description="Number of concurrent patch requests per task"),
]
class Config:
orchestrator/src/buttercup/orchestrator/scheduler/sarif_matcher.py
+26 -27
View File
@@ -1,11 +1,11 @@
from buttercup.common.datastructures.msg_pb2 import TracedCrash
from buttercup.common.clusterfuzz_parser.slice import StackFrame
from buttercup.common import stack_parsing
from buttercup.orchestrator.task_server.models.types import SARIFBroadcastDetail
from typing import List, Tuple
from pathlib import Path
from dataclasses import dataclass
import logging
from dataclasses import dataclass
from pathlib import Path
from buttercup.common import stack_parsing
from buttercup.common.clusterfuzz_parser.slice import StackFrame
from buttercup.common.datastructures.msg_pb2 import TracedCrash
from buttercup.orchestrator.task_server.models.types import SARIFBroadcastDetail
logger = logging.getLogger(__name__)
@@ -13,7 +13,7 @@ logger = logging.getLogger(__name__)
@dataclass
class SarifInfo:
file: Path
lines: Tuple[int, int]
lines: tuple[int, int]
function: str | None
cwe: str | None
@@ -37,8 +37,7 @@ class SarifMatch:
def match(sarif_broadcast: SARIFBroadcastDetail, traced_crash: TracedCrash) -> SarifMatch | None:
"""
Match a SARIF broadcast with a traced crash.
"""Match a SARIF broadcast with a traced crash.
Args:
sarif_broadcast: The SARIF broadcast to match
@@ -46,8 +45,8 @@ def match(sarif_broadcast: SARIFBroadcastDetail, traced_crash: TracedCrash) -> S
Returns:
A SarifMatch object if a match is found between the SARIF and traced crash, None otherwise
"""
"""
# Extract SARIF details
sarif_infos = _sarif_detail(sarif_broadcast)
@@ -66,15 +65,15 @@ def match(sarif_broadcast: SARIFBroadcastDetail, traced_crash: TracedCrash) -> S
return None
def _sarif_detail(sarif_broadcast: SARIFBroadcastDetail) -> List[SarifInfo]:
"""
Extract detailed information from a SARIF broadcast.
def _sarif_detail(sarif_broadcast: SARIFBroadcastDetail) -> list[SarifInfo]:
"""Extract detailed information from a SARIF broadcast.
Args:
sarif_broadcast: The SARIF broadcast to extract details from
Returns:
List of SarifInfo objects containing file, line, function, and CWE information
"""
sarif_infos = []
@@ -131,16 +130,18 @@ def _sarif_detail(sarif_broadcast: SARIFBroadcastDetail) -> List[SarifInfo]:
# Create SarifInfo object and add to the list
sarif_info = SarifInfo(
file=Path(file_uri), lines=(start_line, end_line), function=function_name, cwe=cwe
file=Path(file_uri),
lines=(start_line, end_line),
function=function_name,
cwe=cwe,
)
sarif_infos.append(sarif_info)
return sarif_infos
def _match_thread_callstack(frames: List[StackFrame], sarif_infos: List[SarifInfo]) -> SarifMatch | None:
"""
Match a thread frame with SARIF information.
def _match_thread_callstack(frames: list[StackFrame], sarif_infos: list[SarifInfo]) -> SarifMatch | None:
"""Match a thread frame with SARIF information.
Args:
frames: List of stack frames from a thread
@@ -148,6 +149,7 @@ def _match_thread_callstack(frames: List[StackFrame], sarif_infos: List[SarifInf
Returns:
True if any frame matches any SARIF info, False otherwise
"""
if not frames:
return None
@@ -171,8 +173,7 @@ def _match_thread_callstack(frames: List[StackFrame], sarif_infos: List[SarifInf
def _get_frame(frame: StackFrame) -> Frame | None:
"""
Get a Frame object from a StackFrame object.
"""Get a Frame object from a StackFrame object.
NOTE: We require the filename to be present in the StackFrame object.
"""
if frame.filename is None:
@@ -181,9 +182,8 @@ def _get_frame(frame: StackFrame) -> Frame | None:
return Frame(file=Path(frame.filename), line=int(frame.fileline), function=frame.function_name)
def _match_frame(frame: Frame, sarif_infos: List[SarifInfo]) -> SarifMatch | None:
"""
Match a frame with SARIF information.
def _match_frame(frame: Frame, sarif_infos: list[SarifInfo]) -> SarifMatch | None:
"""Match a frame with SARIF information.
Tries to match based on:
1. File path and line number
@@ -196,17 +196,16 @@ def _match_frame(frame: Frame, sarif_infos: List[SarifInfo]) -> SarifMatch | Non
Returns:
True if a match is found, False otherwise
"""
def line_match(frame_line: int | str, info_lines: Tuple[int, int]) -> bool:
def line_match(frame_line: int | str, info_lines: tuple[int, int]) -> bool:
if isinstance(frame_line, str):
frame_line = int(frame_line)
return frame_line >= info_lines[0] and frame_line <= info_lines[1]
def stripped_function_match(frame_function: str, info_function: str) -> bool:
"""
Match a function name by stripping the OSS_FUZZ_ prefix.
"""
"""Match a function name by stripping the OSS_FUZZ_ prefix."""
if frame_function.startswith("OSS_FUZZ_"):
frame_function = frame_function.split("OSS_FUZZ_")[1]
return frame_function == info_function
orchestrator/src/buttercup/orchestrator/scheduler/scheduler.py
+48 -34
View File
@@ -1,31 +1,32 @@
import logging
import random
from dataclasses import dataclass, field
from pathlib import Path
from typing import Set, Union
from redis import Redis
from buttercup.common.queues import ReliableQueue, QueueFactory, RQItem, QueueNames, GroupNames
from buttercup.common.maps import HarnessWeights, BuildMap
from buttercup.common.challenge_task import ChallengeTask
from buttercup.common.datastructures.msg_pb2 import (
TaskReady,
Task,
BuildRequest,
BuildOutput,
WeightedHarness,
IndexRequest,
BuildType,
TracedCrash,
Patch,
)
from buttercup.common.project_yaml import ProjectYaml
from buttercup.orchestrator.scheduler.cancellation import Cancellation
from buttercup.orchestrator.scheduler.submissions import Submissions, CompetitionAPI
from buttercup.common.clusterfuzz_utils import get_fuzz_targets
from buttercup.orchestrator.api_client_factory import create_api_client
from buttercup.common.utils import serve_loop
from buttercup.common.datastructures.msg_pb2 import (
BuildOutput,
BuildRequest,
BuildType,
IndexRequest,
Patch,
Task,
TaskReady,
TracedCrash,
WeightedHarness,
)
from buttercup.common.maps import BuildMap, HarnessWeights
from buttercup.common.project_yaml import ProjectYaml
from buttercup.common.queues import GroupNames, QueueFactory, QueueNames, ReliableQueue, RQItem
from buttercup.common.task_registry import TaskRegistry
from buttercup.common.utils import serve_loop
from buttercup.orchestrator.api_client_factory import create_api_client
from buttercup.orchestrator.scheduler.cancellation import Cancellation
from buttercup.orchestrator.scheduler.status_checker import StatusChecker
import random
from buttercup.orchestrator.scheduler.submissions import CompetitionAPI, Submissions
logger = logging.getLogger(__name__)
@@ -53,7 +54,7 @@ class Scheduler:
build_map: BuildMap | None = field(init=False, default=None)
cancellation: Cancellation | None = field(init=False, default=None)
task_registry: TaskRegistry | None = field(init=False, default=None)
cached_cancelled_ids: Set[str] = field(init=False, default_factory=set)
cached_cancelled_ids: set[str] = field(init=False, default_factory=set)
status_checker: StatusChecker | None = field(init=False, default=None)
patches_queue: ReliableQueue | None = field(init=False, default=None)
traced_vulnerabilities_queue: ReliableQueue | None = field(init=False, default=None)
@@ -66,6 +67,7 @@ class Scheduler:
Returns:
bool: True if there were any cancelled task IDs, False otherwise
"""
if self.task_registry is None:
return False
@@ -78,7 +80,7 @@ class Scheduler:
return len(self.cached_cancelled_ids) > 0
def should_stop_processing(self, task_or_id: Union[str, Task]) -> bool:
def should_stop_processing(self, task_or_id: str | Task) -> bool:
"""Check if a task should no longer be processed due to cancellation or expiration.
Wrapper around the registry.should_stop_processing method that uses the cached
@@ -90,6 +92,7 @@ class Scheduler:
Returns:
bool: True if the task should not be processed (is cancelled or expired),
False otherwise
"""
if self.task_registry is None:
return False
@@ -99,18 +102,24 @@ class Scheduler:
if self.redis is not None:
queue_factory = QueueFactory(self.redis)
api_client = create_api_client(
self.competition_api_url, self.competition_api_key_id, self.competition_api_key_token
self.competition_api_url,
self.competition_api_key_id,
self.competition_api_key_token,
)
# Input queues are non-blocking as we're already sleeping between iterations
self.cancellation = Cancellation(redis=self.redis)
self.ready_queue = queue_factory.create(QueueNames.READY_TASKS, GroupNames.ORCHESTRATOR, block_time=None)
self.build_requests_queue = queue_factory.create(QueueNames.BUILD, block_time=None)
self.build_output_queue = queue_factory.create(
QueueNames.BUILD_OUTPUT, GroupNames.ORCHESTRATOR, block_time=None
QueueNames.BUILD_OUTPUT,
GroupNames.ORCHESTRATOR,
block_time=None,
)
self.index_queue = queue_factory.create(QueueNames.INDEX, block_time=None)
self.index_output_queue = queue_factory.create(
QueueNames.INDEX_OUTPUT, GroupNames.ORCHESTRATOR, block_time=None
QueueNames.INDEX_OUTPUT,
GroupNames.ORCHESTRATOR,
block_time=None,
)
self.harness_map = HarnessWeights(self.redis)
self.build_map = BuildMap(self.redis)
@@ -127,7 +136,9 @@ class Scheduler:
)
self.patches_queue = queue_factory.create(QueueNames.PATCHES, GroupNames.ORCHESTRATOR, block_time=None)
self.traced_vulnerabilities_queue = queue_factory.create(
QueueNames.TRACED_VULNERABILITIES, GroupNames.ORCHESTRATOR, block_time=None
QueueNames.TRACED_VULNERABILITIES,
GroupNames.ORCHESTRATOR,
block_time=None,
)
def select_preferred(self, available_options: list[str], preferred_order: list[str]) -> str:
@@ -139,6 +150,7 @@ class Scheduler:
Returns:
Selected option string
"""
for preferred in preferred_order:
if preferred in available_options:
@@ -184,7 +196,7 @@ class Scheduler:
def process_build_output(self, build_output: BuildOutput) -> list[WeightedHarness]:
"""Process a build output"""
logger.info(
f"[{build_output.task_id}] Processing build output for type {BuildType.Name(build_output.build_type)} | {build_output.engine} | {build_output.sanitizer} | {build_output.task_dir} | {build_output.apply_diff}"
f"[{build_output.task_id}] Processing build output for type {BuildType.Name(build_output.build_type)} | {build_output.engine} | {build_output.sanitizer} | {build_output.task_dir} | {build_output.apply_diff}",
)
if build_output.build_type != BuildType.FUZZER:
@@ -220,7 +232,7 @@ class Scheduler:
# Check if the task should be stopped (cancelled or expired)
if self.should_stop_processing(task_ready.task):
logger.info(
f"Skipping ready task processing for task {task_ready.task.task_id} as it is cancelled or expired"
f"Skipping ready task processing for task {task_ready.task.task_id} as it is cancelled or expired",
)
self.ready_queue.ack_item(task_ready_item.item_id)
return True
@@ -240,7 +252,7 @@ class Scheduler:
for build_req in self.process_ready_task(task_ready.task):
self.build_requests_queue.push(build_req)
logger.info(
f"[{task_ready.task.task_id}] Pushed build request of type {BuildType.Name(build_req.build_type)} | {build_req.sanitizer} | {build_req.engine} | {build_req.apply_diff}"
f"[{task_ready.task.task_id}] Pushed build request of type {BuildType.Name(build_req.build_type)} | {build_req.sanitizer} | {build_req.engine} | {build_req.apply_diff}",
)
self.ready_queue.ack_item(task_ready_item.item_id)
return True
@@ -267,12 +279,12 @@ class Scheduler:
for target in targets:
self.harness_map.push_harness(target)
logger.info(
f"Pushed {len(targets)} targets to fuzzer map for {build_output.task_id} | {build_output.engine} | {build_output.sanitizer} | {build_output.task_dir}"
f"Pushed {len(targets)} targets to fuzzer map for {build_output.task_id} | {build_output.engine} | {build_output.sanitizer} | {build_output.task_dir}",
)
return True
except Exception as e:
logger.error(
f"Failed to process build output for {build_output.task_id} | {build_output.engine} | {build_output.sanitizer} | {build_output.task_dir}: {e}"
f"Failed to process build output for {build_output.task_id} | {build_output.engine} | {build_output.sanitizer} | {build_output.task_dir}: {e}",
)
return False
@@ -289,7 +301,7 @@ class Scheduler:
# Check if the task should be stopped (cancelled or expired)
if self.should_stop_processing(build_output.task_id):
logger.info(
f"Skipping build output processing for task {build_output.task_id} as it is cancelled or expired"
f"Skipping build output processing for task {build_output.task_id} as it is cancelled or expired",
)
self.build_output_queue.ack_item(build_output_item.item_id)
return True
@@ -302,7 +314,7 @@ class Scheduler:
if res:
logger.info(
f"Acked build output {build_output.task_id} | {build_output.engine} | {build_output.sanitizer} | {build_output.task_dir} | {build_output.internal_patch_id}"
f"Acked build output {build_output.task_id} | {build_output.engine} | {build_output.sanitizer} | {build_output.task_dir} | {build_output.internal_patch_id}",
)
self.build_output_queue.ack_item(build_output_item.item_id)
return True
@@ -335,6 +347,7 @@ class Scheduler:
Returns:
bool: True if any weights were updated, False otherwise
"""
if not self.task_registry or not self.harness_map:
return False
@@ -362,7 +375,7 @@ class Scheduler:
self.harness_map.push_harness(zero_weight_harness)
logger.info(
f"Updated weight to -1.0 for cancelled/expired task {harness.task_id}, harness {harness.harness_name}"
f"Updated weight to -1.0 for cancelled/expired task {harness.task_id}, harness {harness.harness_name}",
)
any_updated = True
@@ -382,6 +395,7 @@ class Scheduler:
Returns:
bool: True if any items were processed from the queues, False otherwise
"""
assert self.traced_vulnerabilities_queue is not None
assert self.patches_queue is not None
orchestrator/src/buttercup/orchestrator/scheduler/status_checker.py
+4 -2
View File
@@ -1,6 +1,6 @@
import logging
import time
from typing import Callable
from collections.abc import Callable
logger = logging.getLogger(__name__)
@@ -13,6 +13,7 @@ class StatusChecker:
Args:
check_interval: How often to check statuses in seconds
"""
self.last_check = 0.0
self.check_interval = check_interval
@@ -33,11 +34,12 @@ class StatusChecker:
Returns:
bool: True if the check was executed, False otherwise
"""
if not self.should_check():
return False
try:
return check_fn()
except Exception as e:
logger.error(f"Failed to check statuses: {str(e)}")
logger.error(f"Failed to check statuses: {e!s}")
return False
orchestrator/src/buttercup/orchestrator/scheduler/submissions.py
+217 -207
View File
@@ -1,52 +1,53 @@
from dataclasses import field, dataclass
from functools import lru_cache
import logging
import base64
import logging
import uuid
from redis import Redis
from typing import Callable, Iterator, List, Set, Tuple
from collections.abc import Callable, Iterator
from dataclasses import dataclass, field
from functools import lru_cache
from pathlib import Path
from opentelemetry import trace
from opentelemetry.trace import Status, StatusCode
import buttercup.common.node_local as node_local
from pathlib import Path
from redis import Redis
from buttercup.common import node_local
from buttercup.common.challenge_task import ChallengeTask
from buttercup.common.clusterfuzz_parser.crash_comparer import CrashComparer
from buttercup.common.constants import ARCHITECTURE
from buttercup.common.queues import ReliableQueue, QueueFactory, QueueNames
from buttercup.common.sets import PoVReproduceStatus
from buttercup.common.datastructures.msg_pb2 import (
TracedCrash,
ConfirmedVulnerability,
SubmissionEntry,
SubmissionEntryPatch,
BuildOutput,
BuildRequest,
BuildType,
BuildOutput,
Bundle,
ConfirmedVulnerability,
CrashWithId,
Patch,
POVReproduceRequest,
POVReproduceResponse,
CrashWithId,
Bundle,
SubmissionEntry,
SubmissionEntryPatch,
SubmissionResult,
Patch,
TracedCrash,
)
from buttercup.common.sarif_store import SARIFStore, SARIFBroadcastDetail
from buttercup.common.project_yaml import ProjectYaml
from buttercup.common.queues import QueueFactory, QueueNames, ReliableQueue
from buttercup.common.sarif_store import SARIFBroadcastDetail, SARIFStore
from buttercup.common.sets import PoVReproduceStatus
from buttercup.common.stack_parsing import get_crash_data, get_inst_key
from buttercup.common.task_registry import TaskRegistry
from buttercup.common.telemetry import set_crs_attributes, CRSActionCategory
from buttercup.orchestrator.scheduler.sarif_matcher import match
from buttercup.common.telemetry import CRSActionCategory, set_crs_attributes
from buttercup.orchestrator.competition_api_client.api import BroadcastSarifAssessmentApi, BundleApi, PatchApi, PovApi
from buttercup.orchestrator.competition_api_client.api_client import ApiClient
from buttercup.orchestrator.competition_api_client.models.types_architecture import TypesArchitecture
from buttercup.orchestrator.competition_api_client.models.types_pov_submission import TypesPOVSubmission
from buttercup.orchestrator.competition_api_client.models.types_patch_submission import TypesPatchSubmission
from buttercup.orchestrator.competition_api_client.models.types_assessment import TypesAssessment
from buttercup.orchestrator.competition_api_client.models.types_bundle_submission import TypesBundleSubmission
from buttercup.orchestrator.competition_api_client.models.types_submission_status import TypesSubmissionStatus
from buttercup.orchestrator.competition_api_client.models.types_patch_submission import TypesPatchSubmission
from buttercup.orchestrator.competition_api_client.models.types_pov_submission import TypesPOVSubmission
from buttercup.orchestrator.competition_api_client.models.types_sarif_assessment_submission import (
TypesSarifAssessmentSubmission,
)
from buttercup.orchestrator.competition_api_client.models.types_assessment import TypesAssessment
from buttercup.orchestrator.competition_api_client.api_client import ApiClient
from buttercup.orchestrator.competition_api_client.api import PovApi, PatchApi, BundleApi, BroadcastSarifAssessmentApi
from buttercup.common.challenge_task import ChallengeTask
from buttercup.common.project_yaml import ProjectYaml
from buttercup.common.stack_parsing import get_crash_data, get_inst_key
from buttercup.common.clusterfuzz_parser.crash_comparer import CrashComparer
from buttercup.orchestrator.competition_api_client.models.types_submission_status import TypesSubmissionStatus
from buttercup.orchestrator.scheduler.sarif_matcher import match
logger = logging.getLogger(__name__)
@@ -68,10 +69,9 @@ def _task_id(e: SubmissionEntry | TracedCrash) -> str:
"""Get the task_id from the SubmissionEntry or TracedCrash."""
if isinstance(e, TracedCrash):
return e.crash.target.task_id # type: ignore[no-any-return]
elif isinstance(e, SubmissionEntry):
if isinstance(e, SubmissionEntry):
return e.crashes[0].crash.crash.target.task_id # type: ignore[no-any-return]
else:
raise ValueError(f"Unknown submission entry type: {type(e)}")
raise ValueError(f"Unknown submission entry type: {type(e)}")
def log_entry(
@@ -192,6 +192,7 @@ def _get_eligible_povs_for_submission(e: SubmissionEntry) -> list[CrashWithId]:
Returns:
List of CrashWithId objects that are eligible for submission
"""
return [
crash
@@ -219,8 +220,7 @@ def _find_matching_build_output(patch: SubmissionEntryPatch, build_output: Build
class CompetitionAPI:
"""
Simplified interface for the competition API.
"""Simplified interface for the competition API.
Handles submission formatting, error handling, and response parsing for:
- Vulnerability proofs (POVs)
@@ -231,12 +231,12 @@ class CompetitionAPI:
"""
def __init__(self, api_client: ApiClient, task_registry: TaskRegistry) -> None:
"""
Initialize with an API client.
"""Initialize with an API client.
Args:
api_client: Client for making HTTP requests
task_registry: Task registry for getting task metadata
"""
self.api_client = api_client
self.task_registry = task_registry
@@ -249,9 +249,8 @@ class CompetitionAPI:
"""
return dict(self.task_registry.get(task_id).metadata)
def submit_pov(self, crash: TracedCrash) -> Tuple[str | None, SubmissionResult]:
"""
Submit a vulnerability (POV) to the competition API.
def submit_pov(self, crash: TracedCrash) -> tuple[str | None, SubmissionResult]:
"""Submit a vulnerability (POV) to the competition API.
Reads crash input, encodes as base64, and submits with required metadata.
@@ -262,6 +261,7 @@ class CompetitionAPI:
Tuple[str | None, SubmissionResult]:
- POV ID if successful, None otherwise
- Submission status
"""
try:
# Read crash input file contents and encode as base64
@@ -304,7 +304,7 @@ class CompetitionAPI:
SubmissionResult.PASSED,
]:
logger.error(
f"[{crash.crash.target.task_id}] POV submission rejected (status: {response.status}) for harness: {crash.crash.harness_name}"
f"[{crash.crash.target.task_id}] POV submission rejected (status: {response.status}) for harness: {crash.crash.harness_name}",
)
span.set_status(Status(StatusCode.ERROR))
return None, mapped_status
@@ -316,8 +316,7 @@ class CompetitionAPI:
return None, SubmissionResult.ERRORED
def get_pov_status(self, task_id: str, pov_id: str) -> SubmissionResult:
"""
Get vulnerability submission status.
"""Get vulnerability submission status.
Args:
task_id: Task ID associated with the vulnerability
@@ -325,6 +324,7 @@ class CompetitionAPI:
Returns:
SubmissionResult: Current status (ACCEPTED, PASSED, FAILED, ERRORED, DEADLINE_EXCEEDED)
"""
assert task_id
assert pov_id
@@ -332,9 +332,8 @@ class CompetitionAPI:
response = PovApi(api_client=self.api_client).v1_task_task_id_pov_pov_id_get(task_id=task_id, pov_id=pov_id)
return _map_submission_status_to_result(response.status)
def submit_patch(self, task_id: str, patch: str) -> Tuple[str | None, SubmissionResult]:
"""
Submit a patch to the competition API.
def submit_patch(self, task_id: str, patch: str) -> tuple[str | None, SubmissionResult]:
"""Submit a patch to the competition API.
Args:
task_id: Task ID of the vulnerability being patched
@@ -344,6 +343,7 @@ class CompetitionAPI:
Tuple[str | None, SubmissionResult]:
- Patch ID if accepted, None otherwise
- Submission status
"""
assert task_id
assert patch
@@ -366,7 +366,8 @@ class CompetitionAPI:
logger.debug(f"[{task_id}] Submitting patch")
response = PatchApi(api_client=self.api_client).v1_task_task_id_patch_post(
task_id=task_id, payload=submission
task_id=task_id,
payload=submission,
)
logger.debug(f"[{task_id}] Patch submission response: {response}")
mapped_status = _map_submission_status_to_result(response.status)
@@ -382,8 +383,7 @@ class CompetitionAPI:
return (response.patch_id, mapped_status)
def get_patch_status(self, task_id: str, patch_id: str) -> SubmissionResult:
"""
Get patch submission status.
"""Get patch submission status.
Args:
task_id: Task ID associated with the patch
@@ -391,24 +391,29 @@ class CompetitionAPI:
Returns:
SubmissionResult: Current status (ACCEPTED, PASSED, FAILED, ERRORED, DEADLINE_EXCEEDED)
"""
assert task_id
assert patch_id
response = PatchApi(api_client=self.api_client).v1_task_task_id_patch_patch_id_get(
task_id=task_id, patch_id=patch_id
task_id=task_id,
patch_id=patch_id,
)
if response.functionality_tests_passing is not None:
logger.info(
f"[{task_id}] Patch {patch_id} functionality tests passing: {response.functionality_tests_passing}"
f"[{task_id}] Patch {patch_id} functionality tests passing: {response.functionality_tests_passing}",
)
return _map_submission_status_to_result(response.status)
def submit_bundle(
self, task_id: str, pov_id: str, patch_id: str, sarif_id: str
) -> Tuple[str | None, SubmissionResult]:
"""
Submit a bundle (vulnerability + patch).
self,
task_id: str,
pov_id: str,
patch_id: str,
sarif_id: str,
) -> tuple[str | None, SubmissionResult]:
"""Submit a bundle (vulnerability + patch).
Args:
task_id: Task ID for the submission
@@ -419,6 +424,7 @@ class CompetitionAPI:
Tuple[str | None, SubmissionResult]:
- Bundle ID if successful, None otherwise
- Submission status
"""
assert task_id
assert pov_id
@@ -444,7 +450,8 @@ class CompetitionAPI:
logger.debug(f"[{task_id}] Submitting bundle for harness: {pov_id} {patch_id} {sarif_id}")
response = BundleApi(api_client=self.api_client).v1_task_task_id_bundle_post(
task_id=task_id, payload=submission
task_id=task_id,
payload=submission,
)
logger.debug(f"[{task_id}] Bundle submission response: {response}")
mapped_status = _map_submission_status_to_result(response.status)
@@ -460,10 +467,14 @@ class CompetitionAPI:
return (response.bundle_id, mapped_status)
def patch_bundle(
self, task_id: str, bundle_id: str, pov_id: str, patch_id: str, sarif_id: str
) -> Tuple[bool, SubmissionResult]:
"""
Submit a bundle patch with SARIF association.
self,
task_id: str,
bundle_id: str,
pov_id: str,
patch_id: str,
sarif_id: str,
) -> tuple[bool, SubmissionResult]:
"""Submit a bundle patch with SARIF association.
Args:
task_id: Task ID for the submission
@@ -476,6 +487,7 @@ class CompetitionAPI:
Tuple[bool, SubmissionResult]:
- True if successful, False otherwise
- Submission status
"""
assert task_id
assert bundle_id
@@ -498,7 +510,9 @@ class CompetitionAPI:
task_metadata=self._get_task_metadata(task_id),
)
response = BundleApi(api_client=self.api_client).v1_task_task_id_bundle_bundle_id_patch(
task_id=task_id, bundle_id=bundle_id, payload=submission
task_id=task_id,
bundle_id=bundle_id,
payload=submission,
)
logger.debug(f"[{task_id}] Bundle patch submission response: {response}")
mapped_status = _map_submission_status_to_result(response.status)
@@ -507,7 +521,7 @@ class CompetitionAPI:
SubmissionResult.PASSED,
]:
logger.error(
f"[{task_id}] Bundle patch submission rejected (status: {response.status}) for harness: {pov_id} {patch_id} {sarif_id}"
f"[{task_id}] Bundle patch submission rejected (status: {response.status}) for harness: {pov_id} {patch_id} {sarif_id}",
)
span.set_status(Status(StatusCode.ERROR))
return (False, mapped_status)
@@ -516,8 +530,7 @@ class CompetitionAPI:
return (True, mapped_status)
def delete_bundle(self, task_id: str, bundle_id: str) -> bool:
"""
Delete a bundle by bundle_id.
"""Delete a bundle by bundle_id.
Args:
task_id: Task ID for the submission
@@ -525,6 +538,7 @@ class CompetitionAPI:
Returns:
bool: True if successful, False otherwise
"""
assert task_id
assert bundle_id
@@ -541,7 +555,8 @@ class CompetitionAPI:
try:
logger.debug(f"[{task_id}] Deleting bundle: {bundle_id}")
response = BundleApi(api_client=self.api_client).v1_task_task_id_bundle_bundle_id_delete(
task_id=task_id, bundle_id=bundle_id
task_id=task_id,
bundle_id=bundle_id,
)
logger.debug(f"[{task_id}] Bundle deletion response: {response}")
span.set_status(Status(StatusCode.OK))
@@ -552,9 +567,8 @@ class CompetitionAPI:
span.set_status(Status(StatusCode.ERROR))
return False
def submit_matching_sarif(self, task_id: str, sarif_id: str) -> Tuple[bool, SubmissionResult]:
"""
Submit a matching assessment for a SARIF report.
def submit_matching_sarif(self, task_id: str, sarif_id: str) -> tuple[bool, SubmissionResult]:
"""Submit a matching assessment for a SARIF report.
Used to claim overlap between our vulnerability findings and a broadcast SARIF.
@@ -566,6 +580,7 @@ class CompetitionAPI:
Tuple[bool, SubmissionResult]:
- True if successful, False otherwise
- Submission status
"""
assert task_id
assert sarif_id
@@ -587,9 +602,11 @@ class CompetitionAPI:
)
response = BroadcastSarifAssessmentApi(
api_client=self.api_client
api_client=self.api_client,
).v1_task_task_id_broadcast_sarif_assessment_broadcast_sarif_id_post(
task_id=task_id, broadcast_sarif_id=sarif_id, payload=submission
task_id=task_id,
broadcast_sarif_id=sarif_id,
payload=submission,
)
logger.debug(f"[{task_id}] Matching SARIF submission response: {response}")
mapped_status = _map_submission_status_to_result(response.status)
@@ -598,7 +615,7 @@ class CompetitionAPI:
SubmissionResult.PASSED,
]:
logger.error(
f"[{task_id}] Matching SARIF submission rejected (status: {response.status}) for sarif_id: {sarif_id}"
f"[{task_id}] Matching SARIF submission rejected (status: {response.status}) for sarif_id: {sarif_id}",
)
span.set_status(Status(StatusCode.ERROR))
return (False, mapped_status)
@@ -609,8 +626,7 @@ class CompetitionAPI:
@dataclass
class Submissions:
"""
Manages the complete lifecycle of vulnerability submissions to the competition API.
"""Manages the complete lifecycle of vulnerability submissions to the competition API.
This class implements a state machine that coordinates the submission of vulnerabilities (POVs),
patches, and bundles to a competition API. It handles deduplication, async patch generation,
@@ -683,15 +699,15 @@ class Submissions:
patch_submission_retry_limit: int = 60
patch_requests_per_vulnerability: int = 1
concurrent_patch_requests_per_task: int = 12
entries: List[SubmissionEntry] = field(init=False)
entries: list[SubmissionEntry] = field(init=False)
sarif_store: SARIFStore = field(init=False)
matched_sarifs: Set[str] = field(default_factory=set)
matched_sarifs: set[str] = field(default_factory=set)
build_requests_queue: ReliableQueue[BuildRequest] = field(init=False)
pov_reproduce_status: PoVReproduceStatus = field(init=False)
def __post_init__(self) -> None:
logger.info(
f"Initializing Submissions, patch_submission_retry_limit={self.patch_submission_retry_limit}, patch_requests_per_vulnerability={self.patch_requests_per_vulnerability}, concurrent_patch_requests_per_task={self.concurrent_patch_requests_per_task}"
f"Initializing Submissions, patch_submission_retry_limit={self.patch_submission_retry_limit}, patch_requests_per_vulnerability={self.patch_requests_per_vulnerability}, concurrent_patch_requests_per_task={self.concurrent_patch_requests_per_task}",
)
self.entries = self._get_stored_submissions()
self.sarif_store = SARIFStore(self.redis)
@@ -705,12 +721,12 @@ class Submissions:
self.matched_sarifs.add(sarif_id)
redis.sadd(self.MATCHED_SARIFS, sarif_id)
def _get_matched_sarifs(self, redis: Redis) -> Set[str]:
def _get_matched_sarifs(self, redis: Redis) -> set[str]:
"""Get all matched SARIF IDs from Redis."""
matched_sarifs = redis.smembers(self.MATCHED_SARIFS)
return set(matched_sarifs)
def _get_stored_submissions(self) -> List[SubmissionEntry]:
def _get_stored_submissions(self) -> list[SubmissionEntry]:
"""Get all stored submissions from Redis."""
submissions_data = self.redis.lrange(self.SUBMISSIONS, 0, -1)
return [SubmissionEntry.FromString(raw) for raw in submissions_data]
@@ -740,7 +756,7 @@ class Submissions:
continue
yield i, e
def _find_patch(self, internal_patch_id: str) -> Tuple[int, SubmissionEntry, SubmissionEntryPatch] | None:
def _find_patch(self, internal_patch_id: str) -> tuple[int, SubmissionEntry, SubmissionEntryPatch] | None:
"""Find a patch by its internal patch id."""
for i, e in self._enumerate_submissions():
for patch in e.patches:
@@ -749,14 +765,14 @@ class Submissions:
return None
def find_similar_entries(self, crash: TracedCrash) -> list[tuple[int, SubmissionEntry]]:
"""
Find existing submissions that have crashes similar to the given crash.
"""Find existing submissions that have crashes similar to the given crash.
Args:
crash: The traced crash to compare against existing submissions
Returns:
List of (index, SubmissionEntry) tuples for similar submissions
"""
crash_data = get_crash_data(crash.crash.stacktrace)
inst_key = get_inst_key(crash.crash.stacktrace)
@@ -786,8 +802,7 @@ class Submissions:
return similar_entries
def _add_to_similar_submission(self, crash: TracedCrash) -> bool:
"""
Check if the crash is similar to an existing submission and add it if so.
"""Check if the crash is similar to an existing submission and add it if so.
This method performs deduplication by comparing the crash data and instruction keys
of the new crash against existing submissions for the same task. If the crash is
@@ -800,21 +815,22 @@ class Submissions:
Returns:
True if the crash was added to an existing submission (duplicate found),
False if no similar submission was found
"""
similar_entries = self.find_similar_entries(crash)
if len(similar_entries) == 0:
# Unique PoV
return False
else:
# Similar submissions found - consolidate them (handles both single and multiple cases)
return self._consolidate_similar_submissions(crash, similar_entries)
# Similar submissions found - consolidate them (handles both single and multiple cases)
return self._consolidate_similar_submissions(crash, similar_entries)
def _consolidate_similar_submissions(
self, crash: TracedCrash | None, similar_entries: list[tuple[int, SubmissionEntry]]
self,
crash: TracedCrash | None,
similar_entries: list[tuple[int, SubmissionEntry]],
) -> bool:
"""
Consolidate multiple similar submissions into the first one.
"""Consolidate multiple similar submissions into the first one.
Args:
crash: The new crash that triggered the consolidation (if any)
@@ -822,6 +838,7 @@ class Submissions:
Returns:
True indicating the crash was handled (consolidated)
"""
# Use the first similar submission as the target
target_index, target_entry = similar_entries[0]
@@ -885,8 +902,7 @@ class Submissions:
return True
def submit_vulnerability(self, crash: TracedCrash) -> bool:
"""
Entry point for new vulnerability discoveries from fuzzers.
"""Entry point for new vulnerability discoveries from fuzzers.
Performs deduplication against existing submissions and either:
- Consolidates with similar existing submissions, or
@@ -900,6 +916,7 @@ class Submissions:
Returns:
True if the crash was handled (new submission or consolidated), False otherwise
"""
if self.task_registry.should_stop_processing(_task_id(crash)):
logger.info("Task is cancelled or expired, will not submit vulnerability.")
@@ -928,8 +945,7 @@ class Submissions:
return True
def _submit_pov_if_needed(self, i: int, e: SubmissionEntry, _redis: Redis) -> bool:
"""
Submit first eligible POV to competition API if none are pending/successful.
"""Submit first eligible POV to competition API if none are pending/successful.
Returns True if entry needs persistence, False otherwise.
"""
if _get_first_successful_pov(e):
@@ -949,18 +965,16 @@ class Submissions:
pov.result = status
log_entry(e, i=i, msg="Submitted POV")
return True
else:
logger.error(
f"[{_task_id(pov.crash)}] Failed to submit vulnerability. Competition API returned {status}. Will attempt the next PoV."
)
logger.debug(f"CrashInfo: {pov.crash}")
log_entry(e, i=i, msg="Failed to submit POV")
logger.error(
f"[{_task_id(pov.crash)}] Failed to submit vulnerability. Competition API returned {status}. Will attempt the next PoV.",
)
logger.debug(f"CrashInfo: {pov.crash}")
log_entry(e, i=i, msg="Failed to submit POV")
return False
def _update_pov_status(self, i: int, e: SubmissionEntry, _redis: Redis) -> bool:
"""
Poll competition API for status updates on pending POVs.
"""Poll competition API for status updates on pending POVs.
Returns True if any status changed and entry needs persistence.
"""
updated = False
@@ -973,9 +987,7 @@ class Submissions:
return updated
def _task_outstanding_patch_requests(self, task_id: str) -> int:
"""
Check the number of patch requests that have not been completed for the given task.
"""
"""Check the number of patch requests that have not been completed for the given task."""
n = 0
for _, e in self._enumerate_task_submissions(task_id):
maybe_patch = _current_patch(e)
@@ -984,12 +996,10 @@ class Submissions:
return n
def _request_patch_if_needed(self, i: int, e: SubmissionEntry, redis: Redis) -> bool:
"""
Request patch generation via queue if no current patch and conditions are met.
"""Request patch generation via queue if no current patch and conditions are met.
Respects concurrency limits and waits for potential cross-submission merging.
Returns True if patch request was made and entry needs persistence.
"""
# If we already have the "current patch" no need to request more.
if _current_patch(e):
return False
@@ -1026,10 +1036,7 @@ class Submissions:
return True
def _request_patched_builds_if_needed(self, i: int, e: SubmissionEntry, redis: Redis) -> bool:
"""
Make sure that builds are available for the current patch, if any.
"""
"""Make sure that builds are available for the current patch, if any."""
patch = _current_patch(e)
if not patch:
# No current patch, nothing to do.
@@ -1077,13 +1084,12 @@ class Submissions:
# Add the build output placeholder to the list of patch builds
patch.build_outputs.append(build_output)
logger.info(
f"[{task_id}] Pushed build request {BuildType.Name(build_req.build_type)} | {build_req.sanitizer} | {build_req.engine} | {build_req.apply_diff} | {build_req.internal_patch_id}"
f"[{task_id}] Pushed build request {BuildType.Name(build_req.build_type)} | {build_req.sanitizer} | {build_req.engine} | {build_req.apply_diff} | {build_req.internal_patch_id}",
)
return True
def record_patched_build(self, build_output: BuildOutput) -> bool:
"""
Entry point for completed patched builds from build system.
"""Entry point for completed patched builds from build system.
Updates the build output placeholder in the patch entry with the actual
build directory path. This enables POV reproduction testing to validate
@@ -1094,15 +1100,15 @@ class Submissions:
Returns:
True if build was recorded successfully
"""
"""
key = build_output.internal_patch_id
maybe_patch = self._find_patch(key)
if not maybe_patch:
# If we get here, the build output is not associated with any patch,
# it is possible that the task was cancelled or expired.
logger.error(
f"Build output {build_output.internal_patch_id} not found in any patch (task expired/cancelled?). Will discard."
f"Build output {build_output.internal_patch_id} not found in any patch (task expired/cancelled?). Will discard.",
)
return True
@@ -1112,7 +1118,7 @@ class Submissions:
if not bo:
# This should never happen, but just in case.
logger.error(
f"Build output {build_output.internal_patch_id} not found in patch {patch.internal_patch_id}. Will discard."
f"Build output {build_output.internal_patch_id} not found in patch {patch.internal_patch_id}. Will discard.",
)
return True
@@ -1120,14 +1126,14 @@ class Submissions:
if bo.task_dir:
# This could happen if the build takes longer than the build request timeout.
logger.warning(
f"Build output {build_output.internal_patch_id} already recorded for patch {patch.internal_patch_id}. Will discard."
f"Build output {build_output.internal_patch_id} already recorded for patch {patch.internal_patch_id}. Will discard.",
)
return True
if bo.task_id != build_output.task_id:
# This should never happen, but just in case.
logger.error(
f"Build output {build_output.internal_patch_id} has a different task id than the patch. Will discard."
f"Build output {build_output.internal_patch_id} has a different task id than the patch. Will discard.",
)
return True
@@ -1138,12 +1144,10 @@ class Submissions:
return True
def _submit_patch_if_good(self, i: int, e: SubmissionEntry, redis: Redis) -> bool:
"""
Test current patch effectiveness and submit to competition API if it mitigates all POVs.
"""Test current patch effectiveness and submit to competition API if it mitigates all POVs.
Advances to next patch if current one fails testing or submission retry limit exceeded.
Returns True if entry needs persistence.
"""
# Check that at least one POV has passed validation and has a competition_pov_id
if not _get_first_successful_pov_id(e):
return False
@@ -1200,7 +1204,9 @@ class Submissions:
# Bad patch (or undecided state), move on to next
_advance_patch_idx(e)
log_entry(
e, i=i, msg=f"Patch submission failed, advancing to next patch. Attempts={e.patch_submission_attempts}"
e,
i=i,
msg=f"Patch submission failed, advancing to next patch. Attempts={e.patch_submission_attempts}",
)
else:
# This is a status we won't do anything about, just record it. If it is a timeout, next iteration will probably
@@ -1210,9 +1216,7 @@ class Submissions:
return True
def _update_patch_status(self, i: int, e: SubmissionEntry, redis: Redis) -> bool:
"""
Update the status of any patch in the ACCEPTED state.
"""
"""Update the status of any patch in the ACCEPTED state."""
updated = False
for patch in _get_pending_patch_submissions(e):
status = self.competition_api.get_patch_status(_task_id(e), patch.competition_patch_id)
@@ -1242,12 +1246,10 @@ class Submissions:
return updated
def _ensure_single_bundle(self, i: int, e: SubmissionEntry, redis: Redis) -> bool:
"""
When SubmissionEntries are merged, we might end up having multiple bundles.
"""When SubmissionEntries are merged, we might end up having multiple bundles.
We want to avoid that so delete one bundle to get us closer to single bundle.
We only do one delete each iteration to prevent loosing too much data if something goes wrong.
"""
if len(e.bundles) <= 1:
# All good
return False
@@ -1278,7 +1280,6 @@ class Submissions:
when they are not set, for a new bundle we will set them to empty strings,
for an existing bundle we will not change them.
"""
nbundles = len(e.bundles)
if nbundles > 1:
# We only process when there is at most one bundle
@@ -1287,7 +1288,10 @@ class Submissions:
task_id = _task_id(e)
if nbundles == 0:
competition_bundle_id, status = self.competition_api.submit_bundle(
task_id, competition_pov_id, competition_patch_id or "", competition_sarif_id or ""
task_id,
competition_pov_id,
competition_patch_id or "",
competition_sarif_id or "",
)
if competition_bundle_id:
bundle = Bundle(
@@ -1305,57 +1309,54 @@ class Submissions:
msg=f"Submitted bundle {competition_bundle_id} for patch {competition_patch_id} and sarif {competition_sarif_id}",
)
return True
else:
log_entry(
e,
i=i,
msg=f"Failed to submit bundle, status: {status}. Will keep trying (not much else we can do).",
)
return False
else:
# We have a previous bundle, check if it is still valid
bundle = e.bundles[0]
bundle_needs_update = False
if competition_patch_id is not None:
# Want to make sure this is the value in the bundle, first check if it is already set
if bundle.competition_patch_id != competition_patch_id:
# If it is not set, set it
bundle.competition_patch_id = competition_patch_id
bundle_needs_update = True
if competition_sarif_id is not None:
# Want to make sure this is the value in the bundle, first check if it is already set
if bundle.competition_sarif_id != competition_sarif_id:
# If it is not set, set it
bundle.competition_sarif_id = competition_sarif_id
bundle_needs_update = True
if not bundle_needs_update:
# All good, no need to do anything
return False
log_entry(e, i=i, msg="Patching bundle")
# It bundles the wrong contents, patch the bundle using the correct contents.
success, status = self.competition_api.patch_bundle(
bundle.task_id,
bundle.bundle_id,
bundle.competition_pov_id,
bundle.competition_patch_id,
bundle.competition_sarif_id,
log_entry(
e,
i=i,
msg=f"Failed to submit bundle, status: {status}. Will keep trying (not much else we can do).",
)
if success:
log_entry(
e,
i=i,
msg=f"Patched bundle {bundle.bundle_id} with patch {competition_patch_id} and sarif {competition_sarif_id}",
)
return True
return False
# We have a previous bundle, check if it is still valid
bundle = e.bundles[0]
bundle_needs_update = False
if competition_patch_id is not None:
# Want to make sure this is the value in the bundle, first check if it is already set
if bundle.competition_patch_id != competition_patch_id:
# If it is not set, set it
bundle.competition_patch_id = competition_patch_id
bundle_needs_update = True
if competition_sarif_id is not None:
# Want to make sure this is the value in the bundle, first check if it is already set
if bundle.competition_sarif_id != competition_sarif_id:
# If it is not set, set it
bundle.competition_sarif_id = competition_sarif_id
bundle_needs_update = True
log_entry(e, i=i, msg=f"Failed to patch bundle {bundle.bundle_id}. Status: {status}. Will keep trying.")
if not bundle_needs_update:
# All good, no need to do anything
return False
log_entry(e, i=i, msg="Patching bundle")
# It bundles the wrong contents, patch the bundle using the correct contents.
success, status = self.competition_api.patch_bundle(
bundle.task_id,
bundle.bundle_id,
bundle.competition_pov_id,
bundle.competition_patch_id,
bundle.competition_sarif_id,
)
if success:
log_entry(
e,
i=i,
msg=f"Patched bundle {bundle.bundle_id} with patch {competition_patch_id} and sarif {competition_sarif_id}",
)
return True
log_entry(e, i=i, msg=f"Failed to patch bundle {bundle.bundle_id}. Status: {status}. Will keep trying.")
return False
def _ensure_patch_is_bundled(self, i: int, e: SubmissionEntry, redis: Redis) -> bool:
"""
Create or update bundle to include current passed patch with successful POV.
"""Create or update bundle to include current passed patch with successful POV.
Returns True if bundle was modified and entry needs persistence.
"""
current_patch = _current_patch(e)
@@ -1381,10 +1382,14 @@ class Submissions:
# Update the bundle with the current patch
return self._ensure_bundle_contents(
i, e, redis, competition_pov_id, competition_patch_id=current_patch.competition_patch_id
i,
e,
redis,
competition_pov_id,
competition_patch_id=current_patch.competition_patch_id,
)
def _get_available_sarifs_for_matching(self, task_id: str) -> List[SARIFBroadcastDetail]:
def _get_available_sarifs_for_matching(self, task_id: str) -> list[SARIFBroadcastDetail]:
"""Get SARIFs that are available for matching for the given task.
Returns SARIFs for the task that haven't been used in any existing bundles.
@@ -1394,6 +1399,7 @@ class Submissions:
Returns:
List of SARIF objects that are available for matching
"""
# Get SARIFs for the task, if none there is nothing to do.
sarifs = self.sarif_store.get_by_task_id(task_id)
@@ -1413,8 +1419,7 @@ class Submissions:
return [sarif for sarif in sarifs if sarif.sarif_id not in already_submitted_sarifs]
def _ensure_sarif_is_bundled(self, i: int, e: SubmissionEntry, redis: Redis) -> bool:
"""
Find external SARIF reports that match this entry's POVs and bundle them for additional scoring.
"""Find external SARIF reports that match this entry's POVs and bundle them for additional scoring.
Requires line-level matching for confidence. Returns True if bundle was created/updated.
"""
# If this already has a bundle with a SARIF no need to do anything
@@ -1447,13 +1452,16 @@ class Submissions:
fn=logging.info,
)
return self._ensure_bundle_contents(
i, e, redis, competition_pov_id, competition_sarif_id=sarif.sarif_id
i,
e,
redis,
competition_pov_id,
competition_sarif_id=sarif.sarif_id,
)
return False
def _confirm_matched_sarifs(self, i: int, e: SubmissionEntry, redis: Redis) -> bool:
"""Ensure the SARIF is submitted to the competition API"""
if len(e.bundles) != 1:
# Don't make any changes while we are fixing up bundles resulting from a merge
return False
@@ -1468,25 +1476,24 @@ class Submissions:
if success:
self._insert_matched_sarif(redis, bundle.competition_sarif_id)
return True
else:
log_entry(
e,
i=i,
msg=f"Failed to confirm SARIF {bundle.competition_sarif_id}. Status: {status}. Will keep trying.",
)
log_entry(
e,
i=i,
msg=f"Failed to confirm SARIF {bundle.competition_sarif_id}. Status: {status}. Will keep trying.",
)
# We have a bundle with a SARIF that has been confirmed, no need to do anything (or confirmation failed)
return True
def _reorder_patches_by_completion(self, e: SubmissionEntry) -> None:
"""
Reorder patches starting from patch_idx so that patches with content come before those without.
"""Reorder patches starting from patch_idx so that patches with content come before those without.
This ensures that completed patches are processed before outstanding patch requests,
regardless of the order in which they were received.
Args:
e: The submission entry to reorder patches for
"""
if not _current_patch(e):
return
@@ -1512,8 +1519,7 @@ class Submissions:
e.patches.append(patch)
def record_patch(self, patch: Patch) -> bool:
"""
Entry point for completed patches from patch generators.
"""Entry point for completed patches from patch generators.
Finds the submission entry associated with the patch's internal_patch_id and
records the patch content. Reorders patches to prioritize completed ones for
@@ -1527,6 +1533,7 @@ class Submissions:
Returns:
True if patch was recorded successfully
"""
key = patch.internal_patch_id
maybe_patch = self._find_patch(key)
@@ -1556,8 +1563,11 @@ class Submissions:
return True
def _pov_reproduce_patch_status(
self, patch: SubmissionEntryPatch, crashes: List[CrashWithId], task_id: str
) -> List[POVReproduceResponse | None]:
self,
patch: SubmissionEntryPatch,
crashes: list[CrashWithId],
task_id: str,
) -> list[POVReproduceResponse | None]:
result = []
for crash_with_id in crashes:
if crash_with_id.result in [
@@ -1580,19 +1590,19 @@ class Submissions:
return result
def _pov_reproduce_status_request(self, e: SubmissionEntry, patch_idx: int) -> List[POVReproduceResponse | None]:
def _pov_reproduce_status_request(self, e: SubmissionEntry, patch_idx: int) -> list[POVReproduceResponse | None]:
patch = e.patches[patch_idx]
task_id = _task_id(e)
return self._pov_reproduce_patch_status(patch, e.crashes, task_id)
def _check_all_povs_are_mitigated(self, i: int, e: SubmissionEntry, patch_idx: int) -> bool | None:
"""
Test if patch at patch_idx mitigates all POVs by running them against patched builds.
"""Test if patch at patch_idx mitigates all POVs by running them against patched builds.
Returns:
None: Some POV tests are still pending
True: All POVs are mitigated (patch is effective)
False: At least one POV still crashes (patch is ineffective)
"""
statuses = self._pov_reproduce_status_request(e, patch_idx)
n_pending = sum(1 for status in statuses if status is None)
@@ -1624,7 +1634,9 @@ class Submissions:
return SubmissionEntryPatch(internal_patch_id=str(uuid.uuid4()))
def _enqueue_patch_requests(
self, confirmed_vulnerability: ConfirmedVulnerability, q: ReliableQueue[ConfirmedVulnerability] | None
self,
confirmed_vulnerability: ConfirmedVulnerability,
q: ReliableQueue[ConfirmedVulnerability] | None,
) -> None:
"""Push N copies of vulnerability to queue for parallel patch generation."""
if q is None:
@@ -1634,8 +1646,7 @@ class Submissions:
q.push(confirmed_vulnerability)
def _should_wait_for_patch_mitigation_merge(self, i: int, e: SubmissionEntry) -> bool:
"""
Check if this submission's POVs are mitigated by patches from other submissions.
"""Check if this submission's POVs are mitigated by patches from other submissions.
If the POVs in this SubmissionEntry are mitigated by a patch in another SubmissionEntry,
we will merge the two entries later on. However, for now we need to check each of the
@@ -1648,6 +1659,7 @@ class Submissions:
Returns:
True if we should wait for patch mitigation evaluation or merge, False otherwise
"""
for j, e2 in self._enumerate_task_submissions(_task_id(e)):
if i == j:
@@ -1681,8 +1693,7 @@ class Submissions:
return False
def _merge_entries_by_patch_mitigation(self) -> None:
"""
Cross-submission optimization: merge entries when one submission's patch fixes another's POVs.
"""Cross-submission optimization: merge entries when one submission's patch fixes another's POVs.
This consolidates resources and avoids duplicate patch work by identifying when
a patch from submission A also mitigates POVs from submission B, then merging
@@ -1725,15 +1736,14 @@ class Submissions:
if len(to_merge) > 1:
merged_indices = [j for j, _ in to_merge[1:]] # Skip the first entry (i) since it's the target
logger.info(
f"[{i}:{_task_id(e)}] Merging {len(to_merge) - 1} similar submissions into this one. Merging indices: {', '.join(map(str, merged_indices))}"
f"[{i}:{_task_id(e)}] Merging {len(to_merge) - 1} similar submissions into this one. Merging indices: {', '.join(map(str, merged_indices))}",
)
self._consolidate_similar_submissions(crash=None, similar_entries=to_merge)
except Exception as err:
logger.error(f"[{i}:{_task_id(e)}] Error merging entries by patch mitigation: {err}")
def process_cycle(self) -> None:
"""
Main processing loop that advances all submission state machines.
"""Main processing loop that advances all submission state machines.
Called periodically by the scheduler to:
1. Submit POVs to competition API and poll for status updates
orchestrator/src/buttercup/orchestrator/scratch_cleaner/__cli__.py
+5 -3
View File
@@ -1,8 +1,10 @@
import logging
from redis import Redis
from buttercup.common.logger import setup_package_logger
from buttercup.orchestrator.scratch_cleaner.config import Settings
from buttercup.orchestrator.scratch_cleaner.scratch_cleaner import ScratchCleaner
from buttercup.common.logger import setup_package_logger
from redis import Redis
import logging
logger = logging.getLogger(__name__)
orchestrator/src/buttercup/orchestrator/scratch_cleaner/config.py
+6 -4
View File
@@ -1,7 +1,8 @@
from pydantic_settings import BaseSettings
from typing import Annotated
from pydantic import Field
from pathlib import Path
from typing import Annotated
from pydantic import Field
from pydantic_settings import BaseSettings
class Settings(BaseSettings):
@@ -10,7 +11,8 @@ class Settings(BaseSettings):
log_level: Annotated[str, Field(default="debug", description="Log level")]
sleep_time: Annotated[float, Field(default=60.0, description="Sleep time between checks in seconds")]
delete_old_tasks_scratch_delta_seconds: Annotated[
int, Field(default=1800, description="Time in seconds after which to delete old task directories")
int,
Field(default=1800, description="Time in seconds after which to delete old task directories"),
]
scratch_dir: Annotated[Path, Field(default=Path("/node-local/scratch"), description="Scratch directory")]
orchestrator/src/buttercup/orchestrator/scratch_cleaner/scratch_cleaner.py
+4 -2
View File
@@ -1,10 +1,12 @@
from __future__ import annotations
from dataclasses import dataclass, field
import logging
from pathlib import Path
import shutil
from dataclasses import dataclass, field
from pathlib import Path
from redis import Redis
from buttercup.common.task_registry import TaskRegistry
from buttercup.common.utils import serve_loop
orchestrator/src/buttercup/orchestrator/task_server/__cli__.py
+1
View File
@@ -1,4 +1,5 @@
import uvicorn
from buttercup.orchestrator.task_server.config import TaskServerSettings
orchestrator/src/buttercup/orchestrator/task_server/auth_tool.py
+14 -17
View File
@@ -1,6 +1,5 @@
#!/usr/bin/env python
"""
CLI tool for generating API keys and tokens for the CRS Task Server.
"""CLI tool for generating API keys and tokens for the CRS Task Server.
This tool generates UUID-based key IDs and secure tokens that can be used for authentication
with the CRS Task Server. The tokens are hashed using Argon2id for secure storage.
@@ -11,14 +10,12 @@ import secrets
import string
import sys
import uuid
from typing import Tuple
from argon2 import PasswordHasher, Type
from rich import print as rprint
from rich.console import Console
from rich.table import Table
# Create password hasher with Argon2id settings matching the server configuration
ph = PasswordHasher(
time_cost=3, # Number of iterations
@@ -35,35 +32,35 @@ TOKEN_LENGTH = 32
def generate_token(length: int = TOKEN_LENGTH) -> str:
"""
Generate a secure random token.
"""Generate a secure random token.
Args:
length: Length of the token to generate
Returns:
A secure random token string
"""
alphabet = string.ascii_letters + string.digits
return "".join(secrets.choice(alphabet) for _ in range(length))
def generate_key_id() -> str:
"""
Generate a UUID-based key ID.
"""Generate a UUID-based key ID.
Returns:
A UUID string to be used as key ID
"""
return str(uuid.uuid4())
def generate_api_key() -> Tuple[str, str, str]:
"""
Generate a complete API key (key ID and token) with hash.
def generate_api_key() -> tuple[str, str, str]:
"""Generate a complete API key (key ID and token) with hash.
Returns:
Tuple containing (key_id, token, token_hash)
"""
key_id = generate_key_id()
token = generate_token()
@@ -73,8 +70,7 @@ def generate_api_key() -> Tuple[str, str, str]:
def format_for_env(key_id: str, token_hash: str) -> str:
"""
Format a key ID and token hash for use in the CRS_API_TOKENS environment variable.
"""Format a key ID and token hash for use in the CRS_API_TOKENS environment variable.
Args:
key_id: The key ID
@@ -82,19 +78,20 @@ def format_for_env(key_id: str, token_hash: str) -> str:
Returns:
A formatted string for the environment variable
"""
return f"{key_id}:{token_hash}"
def print_api_key_info(key_id: str, token: str, token_hash: str, env_format: bool = False) -> None:
"""
Print API key information in a readable format.
"""Print API key information in a readable format.
Args:
key_id: The key ID
token: The plaintext token
token_hash: The hashed token
env_format: Whether to print in environment variable format
"""
console = Console()
@@ -120,11 +117,11 @@ def print_api_key_info(key_id: str, token: str, token_hash: str, env_format: boo
def main() -> int:
"""
Main entry point for the CLI tool.
"""Main entry point for the CLI tool.
Returns:
Exit code (0 for success, non-zero for error)
"""
parser = argparse.ArgumentParser(description="Generate API keys and tokens for the CRS Task Server")
parser.add_argument("--env", action="store_true", help="Format output as environment variables")
orchestrator/src/buttercup/orchestrator/task_server/backend.py
+23 -18
View File
@@ -1,24 +1,29 @@
import logging
import time
from uuid import UUID
from buttercup.orchestrator.task_server.models.types import (
Task,
TaskType,
SourceType,
StatusTasksState,
SARIFBroadcast,
from redis import Redis
from buttercup.common.datastructures.msg_pb2 import (
SourceDetail as SourceDetailProto,
)
from buttercup.common.datastructures.msg_pb2 import (
Task as TaskProto,
SourceDetail as SourceDetailProto,
)
from buttercup.common.datastructures.msg_pb2 import (
TaskDelete,
TaskDownload,
)
from buttercup.common.queues import ReliableQueue
from buttercup.common.sarif_store import SARIFStore
from buttercup.common.task_registry import TaskRegistry
from redis import Redis
import logging
from buttercup.orchestrator.task_server.models.types import (
SARIFBroadcast,
SourceType,
StatusTasksState,
Task,
TaskType,
)
logger = logging.getLogger(__name__)
@@ -77,8 +82,7 @@ def new_task(task: Task, tasks_queue: ReliableQueue) -> str:
def delete_task(task_id: UUID, delete_task_queue: ReliableQueue) -> str:
"""
Delete a task by pushing a delete request to the task deletion queue.
"""Delete a task by pushing a delete request to the task deletion queue.
Args:
task_id: The unique identifier of the task to delete
@@ -86,6 +90,7 @@ def delete_task(task_id: UUID, delete_task_queue: ReliableQueue) -> str:
Returns:
Empty string on successful deletion request
"""
task_delete = TaskDelete(task_id=str(task_id).lower(), received_at=time.time())
delete_task_queue.push(task_delete)
@@ -93,14 +98,14 @@ def delete_task(task_id: UUID, delete_task_queue: ReliableQueue) -> str:
def delete_all_tasks(delete_task_queue: ReliableQueue) -> str:
"""
Delete all tasks by pushing a delete request with the 'all' flag to the task deletion queue.
"""Delete all tasks by pushing a delete request with the 'all' flag to the task deletion queue.
Args:
delete_task_queue: Queue for processing task deletion requests
Returns:
Empty string on successful deletion request
"""
task_delete = TaskDelete(all=True, received_at=time.time())
delete_task_queue.push(task_delete)
@@ -108,8 +113,7 @@ def delete_all_tasks(delete_task_queue: ReliableQueue) -> str:
def store_sarif_broadcast(broadcast: SARIFBroadcast, sarif_store: SARIFStore) -> str:
"""
Store SARIF broadcast details in Redis.
"""Store SARIF broadcast details in Redis.
Args:
broadcast: The SARIFBroadcast containing a list of SARIFBroadcastDetail
@@ -117,6 +121,7 @@ def store_sarif_broadcast(broadcast: SARIFBroadcast, sarif_store: SARIFStore) ->
Returns:
Empty string on successful storage
"""
for sarif_detail in broadcast.broadcasts:
logger.info(f"Storing SARIF detail for task {sarif_detail.task_id}, SARIF ID: {sarif_detail.sarif_id}")
@@ -126,11 +131,11 @@ def store_sarif_broadcast(broadcast: SARIFBroadcast, sarif_store: SARIFStore) ->
def get_status_tasks_state(redis_url: str) -> StatusTasksState:
"""
Get the current state of tasks in the system.
"""Get the current state of tasks in the system.
Returns:
StatusTasksState: The current state of tasks in the system
"""
redis = Redis.from_url(redis_url)
registry = TaskRegistry(redis)
orchestrator/src/buttercup/orchestrator/task_server/config.py
+2 -2
View File
@@ -1,7 +1,7 @@
from pydantic_settings import BaseSettings
from typing import Annotated
from pydantic import Field
from pydantic_settings import CliImplicitFlag
from pydantic_settings import BaseSettings, CliImplicitFlag
class TaskServerSettings(BaseSettings):
orchestrator/src/buttercup/orchestrator/task_server/dependencies.py
+11 -9
View File
@@ -1,10 +1,12 @@
import logging
from functools import lru_cache
import redis
from redis import Redis
from functools import lru_cache
from buttercup.orchestrator.task_server.config import TaskServerSettings
from buttercup.common.queues import ReliableQueue, QueueNames, QueueFactory
from buttercup.common.queues import QueueFactory, QueueNames, ReliableQueue
from buttercup.common.sarif_store import SARIFStore
from buttercup.orchestrator.task_server.config import TaskServerSettings
logger = logging.getLogger(__name__)
@@ -22,11 +24,11 @@ def get_redis() -> Redis:
@lru_cache
def get_task_queue() -> ReliableQueue:
"""
Get a ReliableQueue instance for task messages.
"""Get a ReliableQueue instance for task messages.
Returns:
ReliableQueue: Queue for task messages
"""
logger.debug(f"Connecting to task queue at {QueueNames.DOWNLOAD_TASKS}")
return QueueFactory(get_redis()).create(QueueNames.DOWNLOAD_TASKS)
@@ -34,21 +36,21 @@ def get_task_queue() -> ReliableQueue:
@lru_cache
def get_delete_task_queue() -> ReliableQueue:
"""
Get a ReliableQueue instance for task deletion messages.
"""Get a ReliableQueue instance for task deletion messages.
Returns:
ReliableQueue: Queue for task deletion messages
"""
logger.debug(f"Connecting to delete task queue at {QueueNames.DELETE_TASK}")
return QueueFactory(get_redis()).create(QueueNames.DELETE_TASK)
def get_sarif_store() -> SARIFStore:
"""
Get a SARIFStore instance for SARIF storage and retrieval.
"""Get a SARIFStore instance for SARIF storage and retrieval.
Returns:
SARIFStore: Store for SARIF objects
"""
return SARIFStore(get_redis())
orchestrator/src/buttercup/orchestrator/task_server/server.py
+31 -39
View File
@@ -4,38 +4,38 @@
from __future__ import annotations
from typing import Annotated, Optional
from uuid import UUID
import importlib.metadata
from typing import Annotated
from uuid import UUID
from argon2 import PasswordHasher, Type
from argon2.exceptions import VerifyMismatchError
from fastapi import Depends, FastAPI, status, HTTPException
from fastapi import Depends, FastAPI, HTTPException, status
from fastapi.security import HTTPBasic, HTTPBasicCredentials
from pydantic import BaseModel, Field
from urllib3.exceptions import MaxRetryError, NewConnectionError
from buttercup.orchestrator.task_server.models.types import Status, Task, SARIFBroadcast, StatusState
from buttercup.orchestrator.task_server.backend import (
delete_task,
new_task,
delete_all_tasks,
store_sarif_broadcast,
get_status_tasks_state,
)
from buttercup.common.logger import setup_package_logger
from buttercup.orchestrator.task_server.dependencies import (
get_delete_task_queue,
get_task_queue,
get_settings,
get_sarif_store,
)
from buttercup.orchestrator.task_server.config import TaskServerSettings
from buttercup.common.queues import ReliableQueue
from buttercup.common.sarif_store import SARIFStore
from urllib3.exceptions import MaxRetryError, NewConnectionError
from buttercup.orchestrator.api_client_factory import create_api_client
from buttercup.orchestrator.competition_api_client.api.ping_api import PingApi
from buttercup.orchestrator.competition_api_client.models.types_ping_response import TypesPingResponse
from buttercup.orchestrator.task_server.backend import (
delete_all_tasks,
delete_task,
get_status_tasks_state,
new_task,
store_sarif_broadcast,
)
from buttercup.orchestrator.task_server.config import TaskServerSettings
from buttercup.orchestrator.task_server.dependencies import (
get_delete_task_queue,
get_sarif_store,
get_settings,
get_task_queue,
)
from buttercup.orchestrator.task_server.models.types import SARIFBroadcast, Status, StatusState, Task
settings = get_settings()
logger = setup_package_logger("task-server", __name__, settings.log_level, settings.log_max_line_length)
@@ -89,8 +89,7 @@ def check_auth(
credentials: Annotated[HTTPBasicCredentials, Depends(security)],
settings: Annotated[TaskServerSettings, Depends(get_settings)],
) -> str:
"""
Authenticate user with Argon2id-hashed token.
"""Authenticate user with Argon2id-hashed token.
Args:
credentials: HTTP Basic authentication credentials
@@ -101,6 +100,7 @@ def check_auth(
Raises:
HTTPException: If authentication fails
"""
# Get expected key ID and token hash from settings
expected_key_id = settings.api_key_id
@@ -110,7 +110,7 @@ def check_auth(
if not expected_key_id or not token_hash:
logger.error(
"Authentication configuration is missing. Configure BUTTERCUP_TASK_SERVER_API_KEY_ID and "
"BUTTERCUP_TASK_SERVER_API_TOKEN_HASH environment variables."
"BUTTERCUP_TASK_SERVER_API_TOKEN_HASH environment variables.",
)
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
@@ -148,9 +148,7 @@ def check_auth(
def get_status_(
credentials: Annotated[HTTPBasicCredentials, Depends(check_auth)],
) -> Status:
"""
CRS Status
"""
"""CRS Status"""
def is_competition_api_ready() -> bool:
is_ready = False
@@ -186,9 +184,7 @@ def get_status_(
def delete_status_(
credentials: Annotated[HTTPBasicCredentials, Depends(check_auth)],
) -> str:
"""
Reset status stats
"""
"""Reset status stats"""
logger.info("Resetting status")
return ""
@@ -199,9 +195,7 @@ def post_v1_sarif_(
body: SARIFBroadcast,
sarif_store: Annotated[SARIFStore, Depends(get_sarif_store)],
) -> str:
"""
Submit Sarif Broadcast
"""
"""Submit Sarif Broadcast"""
logger.info("Accepting Sarif Broadcast: %s", body)
return store_sarif_broadcast(body, sarif_store)
@@ -216,10 +210,8 @@ def post_v1_task_(
credentials: Annotated[HTTPBasicCredentials, Depends(check_auth)],
body: Task,
tasks_queue: Annotated[ReliableQueue, Depends(get_task_queue)],
) -> Optional[str]:
"""
Submit Task
"""
) -> str | None:
"""Submit Task"""
logger.debug("Accepting Task: %s", body)
return new_task(body, tasks_queue)
@@ -229,8 +221,7 @@ def delete_v1_task_(
credentials: Annotated[HTTPBasicCredentials, Depends(check_auth)],
delete_task_queue: Annotated[ReliableQueue, Depends(get_delete_task_queue)],
) -> str:
"""
Cancel All Tasks
"""Cancel All Tasks
This endpoint allows canceling all existing tasks in the system. All tasks will be marked for deletion.
@@ -243,6 +234,7 @@ def delete_v1_task_(
Raises:
HTTPException: If authentication fails
"""
logger.info("Deleting all tasks")
return delete_all_tasks(delete_task_queue)
@@ -254,8 +246,7 @@ def delete_v1_task_task_id_(
task_id: UUID,
delete_task_queue: Annotated[ReliableQueue, Depends(get_delete_task_queue)],
) -> str:
"""
Cancel a task by its ID.
"""Cancel a task by its ID.
This endpoint allows canceling an existing task by its unique identifier. The task will be marked for deletion
and removed from the system.
@@ -270,6 +261,7 @@ def delete_v1_task_task_id_(
Raises:
HTTPException: If authentication fails
"""
logger.info("Deleting task: %s", task_id)
return delete_task(task_id, delete_task_queue)
orchestrator/src/buttercup/orchestrator/ui/__cli__.py
+5 -3
View File
@@ -1,8 +1,10 @@
from buttercup.orchestrator.ui.config import Settings
import logging
import uvicorn
from buttercup.common.logger import setup_package_logger
from buttercup.common.telemetry import init_telemetry
import logging
import uvicorn
from buttercup.orchestrator.ui.config import Settings
# Import the generated FastAPI app
orchestrator/src/buttercup/orchestrator/ui/competition_api/main.py
+172 -196
View File
@@ -2,18 +2,22 @@
# filename: openapi.json
# timestamp: 2025-07-08T08:15:58+00:00
from __future__ import annotations
import uuid
import json
import base64
from datetime import datetime, timedelta
from typing import List, Dict, Any, Optional
from pathlib import Path
from fastapi import FastAPI, Depends, HTTPException
import base64
import json
import logging
import uuid
from datetime import datetime, timedelta
from functools import cache
from pathlib import Path
from typing import Any
from fastapi import Depends, FastAPI, HTTPException
from fastapi.responses import FileResponse, HTMLResponse, Response
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from buttercup.common.telemetry import crs_instance_id
from buttercup.orchestrator.ui.competition_api.models.types import (
BundleSubmission,
BundleSubmissionResponse,
@@ -37,10 +41,7 @@ from buttercup.orchestrator.ui.competition_api.models.types import (
)
from buttercup.orchestrator.ui.competition_api.services import ChallengeService, CRSClient
from buttercup.orchestrator.ui.config import Settings
from buttercup.orchestrator.ui.database import DatabaseManager, Task, Bundle, POV, Patch
from buttercup.common.telemetry import crs_instance_id
from functools import cache
import logging
from buttercup.orchestrator.ui.database import POV, Bundle, DatabaseManager, Patch, Task
logger = logging.getLogger(__name__)
@@ -62,26 +63,26 @@ dashboard_stats = {"activeTasks": 0, "totalPovs": 0, "totalPatches": 0, "totalBu
# Dashboard models
class TaskInfo(BaseModel):
task_id: str
name: Optional[str] = None
name: str | None = None
project_name: str
status: str # active, expired
duration: int
deadline: str
challenge_repo_url: Optional[str] = None
challenge_repo_head_ref: Optional[str] = None
challenge_repo_base_ref: Optional[str] = None
fuzz_tooling_url: Optional[str] = None
fuzz_tooling_ref: Optional[str] = None
povs: List[Dict[str, Any]] = []
patches: List[Dict[str, Any]] = []
bundles: List[Dict[str, Any]] = []
challenge_repo_url: str | None = None
challenge_repo_head_ref: str | None = None
challenge_repo_base_ref: str | None = None
fuzz_tooling_url: str | None = None
fuzz_tooling_ref: str | None = None
povs: list[dict[str, Any]] = []
patches: list[dict[str, Any]] = []
bundles: list[dict[str, Any]] = []
created_at: str
# CRS submission status and error information
crs_submission_status: Optional[str] = None # 'pending', 'success', 'failed'
crs_submission_error: Optional[str] = None
crs_error_details: Optional[str] = None # JSON string
crs_submission_timestamp: Optional[str] = None
crs_submission_status: str | None = None # 'pending', 'success', 'failed'
crs_submission_error: str | None = None
crs_error_details: str | None = None # JSON string
crs_submission_timestamp: str | None = None
class DashboardStats(BaseModel):
@@ -208,18 +209,17 @@ def get_artifact(task_id: str, artifact_type: str, artifact_id: str) -> Any:
if artifact_type == "bundles":
file_path = task_dir / f"{artifact_id}.json"
return json.load(file_path.open("r", encoding="utf-8"))
elif artifact_type == "patches":
if artifact_type == "patches":
file_path = task_dir / f"{artifact_id}.patch"
return file_path.read_text()
elif artifact_type == "povs":
if artifact_type == "povs":
file_path = task_dir / f"{artifact_id}.bin"
return file_path.read_bytes()
elif artifact_type == "sarifs":
if artifact_type == "sarifs":
file_path = task_dir / f"{artifact_id}.sarif"
return json.load(file_path.open("r", encoding="utf-8"))
else:
logger.error(f"Unknown artifact type: {artifact_type}")
return None
logger.error(f"Unknown artifact type: {artifact_type}")
return None
except Exception:
logger.exception(f"Failed to get {artifact_type} artifact {artifact_id} for task {task_id}")
return None
@@ -348,7 +348,7 @@ def update_dashboard_stats(database_manager: DatabaseManager) -> None:
"totalPatches": total_patches,
"totalBundles": total_bundles,
"failedTasks": failed_tasks_count,
}
},
)
@@ -448,7 +448,10 @@ def task_to_task_info(task: Task) -> TaskInfo:
def _create_task(
challenge: Challenge, challenge_service: ChallengeService, crs_client: CRSClient, database_manager: DatabaseManager
challenge: Challenge,
challenge_service: ChallengeService,
crs_client: CRSClient,
database_manager: DatabaseManager,
) -> Message | Error:
try:
# Create task for the challenge
@@ -495,42 +498,41 @@ def _create_task(
database_manager.update_task_crs_status(task_id=task_id, crs_submission_status="success")
logger.info(f"Task {task_id} CRS status updated to 'success'")
return Message(
message=f"Task {task_id} created and submitted successfully to CRS for challenge {challenge.name}"
message=f"Task {task_id} created and submitted successfully to CRS for challenge {challenge.name}",
)
else:
logger.error(f"Failed to submit task {task_id} to CRS")
logger.error(f"Failed to submit task {task_id} to CRS")
# Generate user-friendly error message
error_message = crs_response.get_user_friendly_error_message("Failed to submit task to CRS")
logger.error(f"Task {task_id} error message: {error_message}")
# Generate user-friendly error message
error_message = crs_response.get_user_friendly_error_message("Failed to submit task to CRS")
logger.error(f"Task {task_id} error message: {error_message}")
# Update the task with error information
logger.info(f"Updating task {task_id} CRS status to 'failed' with error details")
database_manager.update_task_crs_status(
task_id=task_id,
crs_submission_status="failed",
crs_submission_error=error_message,
crs_error_details=crs_response.error_details,
)
logger.info(f"Task {task_id} CRS status updated to 'failed' with error details")
# Update the task with error information
logger.info(f"Updating task {task_id} CRS status to 'failed' with error details")
database_manager.update_task_crs_status(
task_id=task_id,
crs_submission_status="failed",
crs_submission_error=error_message,
crs_error_details=crs_response.error_details,
)
logger.info(f"Task {task_id} CRS status updated to 'failed' with error details")
# Verify the task was updated correctly
with database_manager.get_task(task_id) as updated_task:
if updated_task:
logger.info(f"Verified task {task_id} in database:")
logger.info(f" - CRS status: {getattr(updated_task, 'crs_submission_status', 'N/A')}")
logger.info(f" - CRS error: {getattr(updated_task, 'crs_submission_error', 'N/A')}")
logger.info(f" - CRS error details: {getattr(updated_task, 'crs_error_details', 'N/A')}")
else:
logger.error(f"Task {task_id} not found in database after update!")
# Verify the task was updated correctly
with database_manager.get_task(task_id) as updated_task:
if updated_task:
logger.info(f"Verified task {task_id} in database:")
logger.info(f" - CRS status: {getattr(updated_task, 'crs_submission_status', 'N/A')}")
logger.info(f" - CRS error: {getattr(updated_task, 'crs_submission_error', 'N/A')}")
logger.info(f" - CRS error details: {getattr(updated_task, 'crs_error_details', 'N/A')}")
else:
logger.error(f"Task {task_id} not found in database after update!")
# Return success message since the task was created, but include the CRS error
return Message(
message=f"Task {task_id} created but failed to submit to CRS. "
f"Task is visible in the dashboard with error details. "
f"Error: {error_message}",
color="error",
)
# Return success message since the task was created, but include the CRS error
return Message(
message=f"Task {task_id} created but failed to submit to CRS. "
f"Task is visible in the dashboard with error details. "
f"Error: {error_message}",
color="error",
)
except Exception as e:
logger.error(f"Error creating task for challenge {challenge.name}: {e}")
@@ -559,7 +561,7 @@ def _create_task(
)
# Mark it as failed with the error details
error_message = f"Failed to create task: {str(e)}"
error_message = f"Failed to create task: {e!s}"
database_manager.update_task_crs_status(
task_id=failed_task_id,
crs_submission_status="failed",
@@ -580,15 +582,15 @@ def _create_task(
except Exception as db_error:
logger.error(f"Failed to create failed task in database: {db_error}")
# If we can't even create the failed task, return the original error
return Error(message=f"Failed to create task: {str(e)}")
return Error(message=f"Failed to create task: {e!s}")
def _create_sarif_broadcast(
body: dict[str, Any], challenge_service: ChallengeService, crs_client: CRSClient
body: dict[str, Any],
challenge_service: ChallengeService,
crs_client: CRSClient,
) -> Message | Error:
"""
Create a SARIF Broadcast
"""
"""Create a SARIF Broadcast"""
if "task_id" not in body:
return Error(message="Task ID is required")
task_id = body["task_id"]
@@ -603,32 +605,27 @@ def _create_sarif_broadcast(
if crs_response.success:
return Message(message=f"SARIF Broadcast for Task {task_id} created and submitted to CRS")
else:
# Use the helper method to generate user-friendly error message
error_message = crs_response.get_user_friendly_error_message(
f"Failed to submit SARIF Broadcast for Task {task_id} to CRS"
)
# Use the helper method to generate user-friendly error message
error_message = crs_response.get_user_friendly_error_message(
f"Failed to submit SARIF Broadcast for Task {task_id} to CRS",
)
# Return a message with the error details for better user experience
return Message(
message=f"SARIF Broadcast for Task {task_id} created but failed to submit to CRS. Error: {error_message}",
color="error",
)
# Return a message with the error details for better user experience
return Message(
message=f"SARIF Broadcast for Task {task_id} created but failed to submit to CRS. Error: {error_message}",
color="error",
)
@app.get("/v1/ping/", response_model=PingResponse, tags=["ping"])
def get_v1_ping_() -> PingResponse:
"""
Test authentication creds and network connectivity
"""
"""Test authentication creds and network connectivity"""
return PingResponse(status="pong")
@app.get("/v1/crs-ping/", tags=["ping"])
def get_v1_crs_ping_(crs_client: CRSClient = Depends(get_crs_client)) -> dict:
"""
Test connectivity to CRS
"""
"""Test connectivity to CRS"""
crs_ready = crs_client.ping()
return {
"crs_ready": crs_ready,
@@ -639,42 +636,33 @@ def get_v1_crs_ping_(crs_client: CRSClient = Depends(get_crs_client)) -> dict:
# Dashboard endpoints
@app.get("/", response_class=HTMLResponse, tags=["dashboard"])
def get_dashboard() -> HTMLResponse:
"""
Serve the main dashboard HTML page
"""
"""Serve the main dashboard HTML page"""
static_dir = Path(__file__).parent.parent / "static"
html_file = static_dir / "index.html"
if html_file.exists():
return HTMLResponse(content=html_file.read_text(), status_code=200)
else:
return HTMLResponse(content="<h1>Dashboard not found</h1>", status_code=404)
return HTMLResponse(content="<h1>Dashboard not found</h1>", status_code=404)
@app.get("/v1/dashboard/stats", response_model=DashboardStats, tags=["dashboard"])
def get_dashboard_stats(database_manager: DatabaseManager = Depends(get_database_manager)) -> DashboardStats:
"""
Get dashboard statistics
"""
"""Get dashboard statistics"""
update_dashboard_stats(database_manager)
return DashboardStats(**dashboard_stats)
@app.get("/v1/dashboard/config", tags=["dashboard"])
def get_dashboard_config() -> dict:
"""
Get dashboard configuration including instance ID
"""
"""Get dashboard configuration including instance ID"""
return {
"crs_instance_id": crs_instance_id,
}
@app.get("/v1/dashboard/tasks", response_model=List[TaskInfo], tags=["dashboard"])
def get_dashboard_tasks(database_manager: DatabaseManager = Depends(get_database_manager)) -> List[TaskInfo]:
"""
Get list of all tasks for dashboard
"""
@app.get("/v1/dashboard/tasks", response_model=list[TaskInfo], tags=["dashboard"])
def get_dashboard_tasks(database_manager: DatabaseManager = Depends(get_database_manager)) -> list[TaskInfo]:
"""Get list of all tasks for dashboard"""
update_dashboard_stats(database_manager)
with database_manager.get_all_tasks() as tasks:
tasks_list = []
@@ -686,11 +674,9 @@ def get_dashboard_tasks(database_manager: DatabaseManager = Depends(get_database
return tasks_list
@app.get("/v1/dashboard/tasks/failed", response_model=List[TaskInfo], tags=["dashboard"])
def get_failed_tasks(database_manager: DatabaseManager = Depends(get_database_manager)) -> List[TaskInfo]:
"""
Get list of failed tasks for debugging
"""
@app.get("/v1/dashboard/tasks/failed", response_model=list[TaskInfo], tags=["dashboard"])
def get_failed_tasks(database_manager: DatabaseManager = Depends(get_database_manager)) -> list[TaskInfo]:
"""Get list of failed tasks for debugging"""
logger.info("get_failed_tasks endpoint called")
try:
@@ -724,11 +710,9 @@ def get_failed_tasks(database_manager: DatabaseManager = Depends(get_database_ma
raise
@app.get("/v1/dashboard/tasks/active", response_model=List[TaskInfo], tags=["dashboard"])
def get_active_tasks(database_manager: DatabaseManager = Depends(get_database_manager)) -> List[TaskInfo]:
"""
Get list of active tasks
"""
@app.get("/v1/dashboard/tasks/active", response_model=list[TaskInfo], tags=["dashboard"])
def get_active_tasks(database_manager: DatabaseManager = Depends(get_database_manager)) -> list[TaskInfo]:
"""Get list of active tasks"""
active_tasks = database_manager.get_tasks_by_status("active")
tasks_list = []
for task in active_tasks:
@@ -739,11 +723,9 @@ def get_active_tasks(database_manager: DatabaseManager = Depends(get_database_ma
return tasks_list
@app.get("/v1/dashboard/tasks/expired", response_model=List[TaskInfo], tags=["dashboard"])
def get_expired_tasks(database_manager: DatabaseManager = Depends(get_database_manager)) -> List[TaskInfo]:
"""
Get list of expired tasks
"""
@app.get("/v1/dashboard/tasks/expired", response_model=list[TaskInfo], tags=["dashboard"])
def get_expired_tasks(database_manager: DatabaseManager = Depends(get_database_manager)) -> list[TaskInfo]:
"""Get list of expired tasks"""
expired_tasks = database_manager.get_tasks_by_status("expired")
tasks_list = []
for task in expired_tasks:
@@ -756,9 +738,7 @@ def get_expired_tasks(database_manager: DatabaseManager = Depends(get_database_m
@app.get("/v1/dashboard/tasks/{task_id}", response_model=TaskInfo, tags=["dashboard"])
def get_dashboard_task(task_id: str, database_manager: DatabaseManager = Depends(get_database_manager)) -> TaskInfo:
"""
Get detailed information about a specific task
"""
"""Get detailed information about a specific task"""
with database_manager.get_task(task_id) as task:
if not task:
raise HTTPException(status_code=404, detail="Task not found")
@@ -768,9 +748,7 @@ def get_dashboard_task(task_id: str, database_manager: DatabaseManager = Depends
@app.get("/v1/dashboard/tasks/{task_id}/crs-status", tags=["dashboard"])
def get_task_crs_status(task_id: str, database_manager: DatabaseManager = Depends(get_database_manager)) -> dict:
"""
Get detailed CRS submission status and error information for a specific task
"""
"""Get detailed CRS submission status and error information for a specific task"""
with database_manager.get_task(task_id) as task:
if not task:
raise HTTPException(status_code=404, detail="Task not found")
@@ -808,9 +786,7 @@ def get_task_crs_status(task_id: str, database_manager: DatabaseManager = Depend
tags=["request"],
)
def get_v1_request_list_() -> RequestListResponse | Error:
"""
Get a list of available challenges to task
"""
"""Get a list of available challenges to task"""
return RequestListResponse(challenges=[c.name for c in challenges if c.name is not None])
@@ -832,9 +808,7 @@ def post_v1_request_challenge_name(
crs_client: CRSClient = Depends(get_crs_client),
database_manager: DatabaseManager = Depends(get_database_manager),
) -> Message | Error:
"""
Send a task to the source of this request
"""
"""Send a task to the source of this request"""
if challenge_name not in [c.name for c in challenges]:
return Error(message=f"Challenge {challenge_name} not found")
@@ -861,13 +835,13 @@ def post_v1_request_challenge_name(
tags=["broadcast-sarif-assessment"],
)
def post_v1_task_task_id_broadcast_sarif_assessment_broadcast_sarif_id_(
task_id: str, broadcast_sarif_id: str, body: SarifAssessmentSubmission
task_id: str,
broadcast_sarif_id: str,
body: SarifAssessmentSubmission,
) -> SarifAssessmentResponse | Error:
"""
Submit a SARIF Assessment
"""
"""Submit a SARIF Assessment"""
logger.info(
f"SARIF Assessment submission - Task: {task_id}, Broadcast SARIF ID: {broadcast_sarif_id}, Assessment: {body.assessment}, Description: {body.description[:100]}..."
f"SARIF Assessment submission - Task: {task_id}, Broadcast SARIF ID: {broadcast_sarif_id}, Assessment: {body.assessment}, Description: {body.description[:100]}...",
)
return SarifAssessmentResponse(status=SubmissionStatus.SubmissionStatusAccepted)
@@ -884,11 +858,11 @@ def post_v1_task_task_id_broadcast_sarif_assessment_broadcast_sarif_id_(
tags=["bundle"],
)
def post_v1_task_task_id_bundle_(
task_id: str, body: BundleSubmission, database_manager: DatabaseManager = Depends(get_database_manager)
task_id: str,
body: BundleSubmission,
database_manager: DatabaseManager = Depends(get_database_manager),
) -> BundleSubmissionResponse | Error:
"""
Submit Bundle
"""
"""Submit Bundle"""
logger.info(f"Bundle submission - Task: {task_id}")
logger.debug(f"Bundle details: {json.dumps(body.model_dump(), indent=2)}")
@@ -919,11 +893,11 @@ def post_v1_task_task_id_bundle_(
tags=["bundle"],
)
def get_v1_task_task_id_bundle_bundle_id_(
task_id: str, bundle_id: str, database_manager: DatabaseManager = Depends(get_database_manager)
task_id: str,
bundle_id: str,
database_manager: DatabaseManager = Depends(get_database_manager),
) -> BundleSubmissionResponseVerbose | Error:
"""
Get Bundle
"""
"""Get Bundle"""
logger.info(f"Bundle retrieval - Task: {task_id}, Bundle ID: {bundle_id}")
with database_manager.get_bundle(bundle_id, task_id) as bundle:
if bundle is None:
@@ -953,11 +927,11 @@ def get_v1_task_task_id_bundle_bundle_id_(
tags=["bundle"],
)
def delete_v1_task_task_id_bundle_bundle_id_(
task_id: str, bundle_id: str, database_manager: DatabaseManager = Depends(get_database_manager)
task_id: str,
bundle_id: str,
database_manager: DatabaseManager = Depends(get_database_manager),
) -> str | Error | None:
"""
Delete Bundle
"""
"""Delete Bundle"""
logger.info(f"Bundle deletion - Task: {task_id}, Bundle ID: {bundle_id}")
with database_manager.get_bundle(bundle_id, task_id) as bundle:
if bundle is None:
@@ -986,9 +960,7 @@ def patch_v1_task_task_id_bundle_bundle_id_(
body: BundleSubmission,
database_manager: DatabaseManager = Depends(get_database_manager),
) -> BundleSubmissionResponseVerbose | Error:
"""
Update Bundle
"""
"""Update Bundle"""
logger.info(f"Bundle update - Task: {task_id}, Bundle ID: {bundle_id}")
logger.info(f"Updated bundle details: {json.dumps(body.model_dump(), indent=2)}")
@@ -1030,9 +1002,7 @@ def patch_v1_task_task_id_bundle_bundle_id_(
tags=["freeform"],
)
def post_v1_task_task_id_freeform_(task_id: str, body: FreeformSubmission) -> FreeformResponse | Error:
"""
Submit Freeform
"""
"""Submit Freeform"""
freeform_id = str(uuid.uuid4())
logger.info(f"Freeform submission - Task: {task_id}, Freeform ID: {freeform_id}")
logger.info(f"Freeform submission size: {len(body.submission)} bytes")
@@ -1051,11 +1021,11 @@ def post_v1_task_task_id_freeform_(task_id: str, body: FreeformSubmission) -> Fr
tags=["patch"],
)
def post_v1_task_task_id_patch_(
task_id: str, body: PatchSubmission, database_manager: DatabaseManager = Depends(get_database_manager)
task_id: str,
body: PatchSubmission,
database_manager: DatabaseManager = Depends(get_database_manager),
) -> PatchSubmissionResponse | Error:
"""
Submit Patch
"""
"""Submit Patch"""
logger.info(f"Patch submission - Task: {task_id}")
logger.debug(f"Patch details: {json.dumps(body.model_dump(), indent=2)}")
@@ -1066,7 +1036,9 @@ def post_v1_task_task_id_patch_(
save_patch(task_id, patch.patch_id, body.patch)
return PatchSubmissionResponse(
patch_id=patch.patch_id, status=SubmissionStatus.SubmissionStatusAccepted, functionality_tests_passing=None
patch_id=patch.patch_id,
status=SubmissionStatus.SubmissionStatusAccepted,
functionality_tests_passing=None,
)
@@ -1082,18 +1054,20 @@ def post_v1_task_task_id_patch_(
tags=["patch"],
)
def get_v1_task_task_id_patch_patch_id_(
task_id: str, patch_id: str, database_manager: DatabaseManager = Depends(get_database_manager)
task_id: str,
patch_id: str,
database_manager: DatabaseManager = Depends(get_database_manager),
) -> PatchSubmissionResponse | Error:
"""
Patch Status
"""
"""Patch Status"""
logger.info(f"Patch status check - Task: {task_id}, Patch ID: {patch_id}")
with database_manager.get_patch(patch_id, task_id) as patch:
if patch is None:
return Error(message=f"Patch {patch_id} not found")
return PatchSubmissionResponse(
patch_id=patch_id, status=SubmissionStatus.SubmissionStatusPassed, functionality_tests_passing=True
patch_id=patch_id,
status=SubmissionStatus.SubmissionStatusPassed,
functionality_tests_passing=True,
)
@@ -1109,14 +1083,14 @@ def get_v1_task_task_id_patch_patch_id_(
tags=["pov"],
)
def post_v1_task_task_id_pov_(
task_id: str, body: POVSubmission, database_manager: DatabaseManager = Depends(get_database_manager)
task_id: str,
body: POVSubmission,
database_manager: DatabaseManager = Depends(get_database_manager),
) -> POVSubmissionResponse | Error:
"""
Submit Vulnerability
"""
"""Submit Vulnerability"""
logger.info(f"POV submission - Task: {task_id}")
logger.debug(
f"POV details: architecture={body.architecture.value}, engine={body.engine.value}, fuzzer_name={body.fuzzer_name}, sanitizer={body.sanitizer}"
f"POV details: architecture={body.architecture.value}, engine={body.engine.value}, fuzzer_name={body.fuzzer_name}, sanitizer={body.sanitizer}",
)
with database_manager.get_task(task_id) as task:
@@ -1148,11 +1122,11 @@ def post_v1_task_task_id_pov_(
tags=["pov"],
)
def get_v1_task_task_id_pov_pov_id_(
task_id: str, pov_id: str, database_manager: DatabaseManager = Depends(get_database_manager)
task_id: str,
pov_id: str,
database_manager: DatabaseManager = Depends(get_database_manager),
) -> POVSubmissionResponse | Error:
"""
Vulnerability Status
"""
"""Vulnerability Status"""
logger.info(f"POV status check - Task: {task_id}, POV ID: {pov_id}")
with database_manager.get_pov(pov_id, task_id) as pov:
if pov is None:
@@ -1173,9 +1147,7 @@ def get_v1_task_task_id_pov_pov_id_(
tags=["submitted-sarif"],
)
def post_v1_task_task_id_submitted_sarif_(task_id: str, body: SARIFSubmission) -> SARIFSubmissionResponse | Error:
"""
Submit a CRS generated SARIF
"""
"""Submit a CRS generated SARIF"""
submitted_sarif_id = str(uuid.uuid4())
logger.info(f"SARIF submission - Task: {task_id}, Submitted SARIF ID: {submitted_sarif_id}")
logger.info(f"SARIF content: {json.dumps(body.sarif, indent=2)}")
@@ -1183,17 +1155,17 @@ def post_v1_task_task_id_submitted_sarif_(task_id: str, body: SARIFSubmission) -
save_sarif(task_id, submitted_sarif_id, body.sarif)
return SARIFSubmissionResponse(
submitted_sarif_id=submitted_sarif_id, status=SubmissionStatus.SubmissionStatusPassed
submitted_sarif_id=submitted_sarif_id,
status=SubmissionStatus.SubmissionStatusPassed,
)
@app.get("/files/{tarball_name}.tar.gz", tags=["files"])
def get_tarball(
tarball_name: str, challenge_service: ChallengeService = Depends(get_challenge_service)
tarball_name: str,
challenge_service: ChallengeService = Depends(get_challenge_service),
) -> FileResponse:
"""
Serve tarball files for CRS download
"""
"""Serve tarball files for CRS download"""
try:
return challenge_service.serve_tarball(tarball_name)
except Exception as e:
@@ -1208,9 +1180,7 @@ def trigger_task(
crs_client: CRSClient = Depends(get_crs_client),
database_manager: DatabaseManager = Depends(get_database_manager),
) -> Message | Error:
"""
Trigger a task
"""
"""Trigger a task"""
logger.info(f"Triggering task: {body.model_dump()}")
return _create_task(body, challenge_service, crs_client, database_manager)
@@ -1221,9 +1191,7 @@ def trigger_sarif(
challenge_service: ChallengeService = Depends(get_challenge_service),
crs_client: CRSClient = Depends(get_crs_client),
) -> Message | Error:
"""
Trigger a SARIF Broadcast
"""
"""Trigger a SARIF Broadcast"""
logger.info(f"Triggering SARIF Broadcast: {json.dumps(body, indent=2)}")
return _create_sarif_broadcast(body, challenge_service, crs_client)
@@ -1231,7 +1199,9 @@ def trigger_sarif(
# Download endpoints for PoVs, Patches, and Bundles
@app.get("/v1/dashboard/tasks/{task_id}/povs/{pov_id}/download", tags=["dashboard"])
def download_pov(
task_id: str, pov_id: str, database_manager: DatabaseManager = Depends(get_database_manager)
task_id: str,
pov_id: str,
database_manager: DatabaseManager = Depends(get_database_manager),
) -> Response:
"""Download a PoV testcase"""
with database_manager.get_pov(pov_id, task_id) as pov:
@@ -1247,7 +1217,9 @@ def download_pov(
@app.get("/v1/dashboard/tasks/{task_id}/patches/{patch_id}/download", tags=["dashboard"])
def download_patch(
task_id: str, patch_id: str, database_manager: DatabaseManager = Depends(get_database_manager)
task_id: str,
patch_id: str,
database_manager: DatabaseManager = Depends(get_database_manager),
) -> Response:
"""Download a patch"""
with database_manager.get_patch(patch_id, task_id) as patch:
@@ -1263,7 +1235,9 @@ def download_patch(
@app.get("/v1/dashboard/tasks/{task_id}/bundles/{bundle_id}/download", tags=["dashboard"])
def download_bundle(
task_id: str, bundle_id: str, database_manager: DatabaseManager = Depends(get_database_manager)
task_id: str,
bundle_id: str,
database_manager: DatabaseManager = Depends(get_database_manager),
) -> Response:
"""Download a bundle as JSON"""
with database_manager.get_bundle(bundle_id, task_id) as bundle:
@@ -1283,7 +1257,7 @@ def download_bundle(
# Detail view endpoints
@app.get("/v1/dashboard/povs/{pov_id}", tags=["dashboard"])
def get_pov_detail(pov_id: str, database_manager: DatabaseManager = Depends(get_database_manager)) -> Dict[str, Any]:
def get_pov_detail(pov_id: str, database_manager: DatabaseManager = Depends(get_database_manager)) -> dict[str, Any]:
"""Get detailed information about a specific PoV"""
with database_manager.get_pov(pov_id) as pov:
if pov is None:
@@ -1299,8 +1273,9 @@ def get_pov_detail(pov_id: str, database_manager: DatabaseManager = Depends(get_
@app.get("/v1/dashboard/patches/{patch_id}", tags=["dashboard"])
def get_patch_detail(
patch_id: str, database_manager: DatabaseManager = Depends(get_database_manager)
) -> Dict[str, Any]:
patch_id: str,
database_manager: DatabaseManager = Depends(get_database_manager),
) -> dict[str, Any]:
"""Get detailed information about a specific patch"""
with database_manager.get_patch(patch_id) as patch:
if patch is None:
@@ -1316,8 +1291,9 @@ def get_patch_detail(
@app.get("/v1/dashboard/bundles/{bundle_id}", tags=["dashboard"])
def get_bundle_detail(
bundle_id: str, database_manager: DatabaseManager = Depends(get_database_manager)
) -> Dict[str, Any]:
bundle_id: str,
database_manager: DatabaseManager = Depends(get_database_manager),
) -> dict[str, Any]:
"""Get detailed information about a specific bundle"""
with database_manager.get_bundle(bundle_id) as bundle:
if bundle is None:
@@ -1333,7 +1309,7 @@ def get_bundle_detail(
# List all PoVs and Patches across tasks
@app.get("/v1/dashboard/povs", tags=["dashboard"])
def get_all_povs(database_manager: DatabaseManager = Depends(get_database_manager)) -> List[Dict[str, Any]]:
def get_all_povs(database_manager: DatabaseManager = Depends(get_database_manager)) -> list[dict[str, Any]]:
"""Get all PoVs across all tasks"""
all_povs: list[dict[str, Any]] = []
@@ -1345,7 +1321,7 @@ def get_all_povs(database_manager: DatabaseManager = Depends(get_database_manage
"task_id": pov.task_id,
"task_name": pov.task.name or pov.task.project_name,
"pov": pov_dict,
}
},
)
all_povs.sort(key=lambda x: x["pov"].get("timestamp", ""), reverse=True)
@@ -1353,7 +1329,7 @@ def get_all_povs(database_manager: DatabaseManager = Depends(get_database_manage
@app.get("/v1/dashboard/patches", tags=["dashboard"])
def get_all_patches(database_manager: DatabaseManager = Depends(get_database_manager)) -> List[Dict[str, Any]]:
def get_all_patches(database_manager: DatabaseManager = Depends(get_database_manager)) -> list[dict[str, Any]]:
"""Get all patches across all tasks"""
all_patches: list[dict[str, Any]] = []
@@ -1365,7 +1341,7 @@ def get_all_patches(database_manager: DatabaseManager = Depends(get_database_man
"task_id": patch.task_id,
"task_name": patch.task.name or patch.task.project_name,
"patch": patch_dict,
}
},
)
# Sort by timestamp descending
orchestrator/src/buttercup/orchestrator/ui/competition_api/models/crs_types.py
+8 -7
View File
@@ -1,7 +1,8 @@
from __future__ import annotations
from enum import Enum
from typing import Any, Dict, List
from typing import Any
from pydantic import BaseModel
@@ -26,9 +27,9 @@ class TaskDetail(BaseModel):
deadline: int
focus: str
harnesses_included: bool
metadata: Dict[str, str]
metadata: dict[str, str]
project_name: str
source: List[SourceDetail]
source: list[SourceDetail]
task_id: str
type: TaskType
@@ -36,17 +37,17 @@ class TaskDetail(BaseModel):
class Task(BaseModel):
message_id: str
message_time: int
tasks: List[TaskDetail]
tasks: list[TaskDetail]
class SARIFBroadcastDetail(BaseModel):
metadata: Dict[str, Any]
sarif: Dict[str, Any]
metadata: dict[str, Any]
sarif: dict[str, Any]
sarif_id: str
task_id: str
class SARIFBroadcast(BaseModel):
broadcasts: List[SARIFBroadcastDetail]
broadcasts: list[SARIFBroadcastDetail]
message_id: str
message_time: int
orchestrator/src/buttercup/orchestrator/ui/competition_api/models/types.py
+21 -20
View File
@@ -5,7 +5,7 @@
from __future__ import annotations
from enum import Enum
from typing import Any, Dict, List, Optional
from typing import Any
from pydantic import BaseModel, Field, constr
@@ -21,19 +21,19 @@ class Assessment(Enum):
class BundleSubmission(BaseModel):
broadcast_sarif_id: Optional[str] = None
description: Optional[str] = Field(
broadcast_sarif_id: str | None = None
description: str | None = Field(
None,
description="optional plaintext description of the components of the bundle, such as would be found in a pull request description or bug report",
)
freeform_id: Optional[str] = None
patch_id: Optional[str] = None
pov_id: Optional[str] = None
submitted_sarif_id: Optional[str] = None
freeform_id: str | None = None
patch_id: str | None = None
pov_id: str | None = None
submitted_sarif_id: str | None = None
class Error(BaseModel):
fields: Optional[Dict[str, str]] = None
fields: dict[str, str] | None = None
message: str
@@ -50,14 +50,15 @@ class FuzzingEngine(Enum):
class Message(BaseModel):
message: str
color: Optional[str] = "default" # "default", "success", "warning", "error"
color: str | None = "default" # "default", "success", "warning", "error"
class POVSubmission(BaseModel):
architecture: Architecture
engine: FuzzingEngine
fuzzer_name: constr(max_length=4096) = Field(
..., description="Fuzz Tooling fuzzer that exercises this vuln\n\n4KiB max size"
...,
description="Fuzz Tooling fuzzer that exercises this vuln\n\n4KiB max size",
)
sanitizer: constr(max_length=4096) = Field(
...,
@@ -81,18 +82,18 @@ class PingResponse(BaseModel):
class RequestListResponse(BaseModel):
challenges: List[str] = Field(..., description="List of challenges that competitors may task themselves with")
challenges: list[str] = Field(..., description="List of challenges that competitors may task themselves with")
class RequestSubmission(BaseModel):
duration_secs: Optional[int] = Field(
duration_secs: int | None = Field(
None,
description="Time in seconds until a task should expire. If not provided, defaults to 3600.",
)
class SARIFSubmission(BaseModel):
sarif: Dict[str, Any] = Field(..., description="SARIF object compliant with the provided schema")
sarif: dict[str, Any] = Field(..., description="SARIF object compliant with the provided schema")
class SarifAssessmentSubmission(BaseModel):
@@ -121,17 +122,17 @@ class BundleSubmissionResponse(BaseModel):
class BundleSubmissionResponseVerbose(BaseModel):
broadcast_sarif_id: Optional[str] = None
broadcast_sarif_id: str | None = None
bundle_id: str
description: Optional[str] = None
freeform_id: Optional[str] = None
patch_id: Optional[str] = None
pov_id: Optional[str] = None
description: str | None = None
freeform_id: str | None = None
patch_id: str | None = None
pov_id: str | None = None
status: SubmissionStatus = Field(
...,
description="Schema-compliant submissions will only ever receive the statuses accepted or deadline_exceeded",
)
submitted_sarif_id: Optional[str] = None
submitted_sarif_id: str | None = None
class FreeformResponse(BaseModel):
@@ -148,7 +149,7 @@ class POVSubmissionResponse(BaseModel):
class PatchSubmissionResponse(BaseModel):
functionality_tests_passing: Optional[bool] = Field(None, description="null indicates the tests have not been run")
functionality_tests_passing: bool | None = Field(None, description="null indicates the tests have not been run")
patch_id: str
status: SubmissionStatus
orchestrator/src/buttercup/orchestrator/ui/competition_api/services/__init__.py
+1 -1
View File
@@ -1,4 +1,4 @@
from .challenge_service import ChallengeService
from .crs_client import CRSClient
__all__ = ["ChallengeService", "CRSClient"]
__all__ = ["CRSClient", "ChallengeService"]
orchestrator/src/buttercup/orchestrator/ui/competition_api/services/challenge_service.py
+16 -17
View File
@@ -8,21 +8,21 @@ import subprocess
import tarfile
import tempfile
import time
from pathlib import Path
from typing import Optional, Any
import uuid
from pathlib import Path
from typing import Any
from fastapi import HTTPException
from fastapi.responses import FileResponse
from buttercup.orchestrator.ui.competition_api.models.crs_types import (
SARIFBroadcast,
SARIFBroadcastDetail,
SourceDetail,
SourceType,
Task,
TaskDetail,
TaskType,
SARIFBroadcast,
SARIFBroadcastDetail,
)
logger = logging.getLogger(__name__)
@@ -39,11 +39,10 @@ class ChallengeService:
repo_url: str,
ref: str,
tarball_name: str,
exclude_dirs: Optional[list[str]] = None,
exclude_dirs: list[str] | None = None,
base_ref: str | None = None,
) -> tuple[str, str, str | None]:
"""
Clone a git repository, checkout the specified ref, and create a tarball.
"""Clone a git repository, checkout the specified ref, and create a tarball.
Args:
repo_url: Git repository URL
@@ -53,6 +52,7 @@ class ChallengeService:
Returns:
Tuple of (focus_dir, sha256_hash, diff_sha256_hash)
"""
if exclude_dirs is None:
exclude_dirs = [".git", ".aixcc"]
@@ -76,7 +76,8 @@ class ChallengeService:
# For GitHub repositories, use the PAT for authentication
# Convert https://github.com/owner/repo.git to https://username:pat@github.com/owner/repo.git
auth_url = repo_url.replace(
"https://github.com/", f"https://{github_username}:{github_pat}@github.com/"
"https://github.com/",
f"https://{github_username}:{github_pat}@github.com/",
)
logger.info("Using authenticated URL for private repository")
clone_url = auth_url
@@ -144,6 +145,7 @@ class ChallengeService:
# Create a git-diff file between the two directories (base_sub_path and ref_sub_path)
diff_result = subprocess.run(
["git", "diff", "--binary", "--no-index", base_sub_path.as_posix(), ref_sub_path.as_posix()],
check=False,
cwd=base_temp_path,
capture_output=True,
)
@@ -193,8 +195,7 @@ class ChallengeService:
def _extract_focus_dir(self, repo_url: str) -> str:
"""Extract the focus directory from the repository URL."""
focus_dir = repo_url.rstrip("/").split("/")[-1]
if focus_dir.endswith(".git"):
focus_dir = focus_dir[:-4] # Remove .git suffix
focus_dir = focus_dir.removesuffix(".git") # Remove .git suffix
return focus_dir
def create_task_for_challenge(
@@ -207,8 +208,7 @@ class ChallengeService:
fuzz_tooling_project_name: str,
duration_secs: int,
) -> Task:
"""
Create a task for a challenge by processing repositories and creating tarballs.
"""Create a task for a challenge by processing repositories and creating tarballs.
Args:
challenge_repo_url: URL of the challenge repository
@@ -221,6 +221,7 @@ class ChallengeService:
Returns:
Task object ready to be sent to CRS
"""
task_id = str(uuid.uuid4())
message_id = str(uuid.uuid4())
@@ -268,7 +269,7 @@ class ChallengeService:
sha256=diff_sha256,
type=SourceType.diff,
url=f"{self.base_url}/files/{diff_sha256}.tar.gz",
)
),
)
# Create task detail
@@ -300,9 +301,7 @@ class ChallengeService:
return task
def create_sarif_broadcast(self, task_id: str, sarif: dict[str, Any]) -> SARIFBroadcast:
"""
Create a SARIF Broadcast for a task
"""
"""Create a SARIF Broadcast for a task"""
sarif_id = str(uuid.uuid4())
message_id = str(uuid.uuid4())
message_time = int(time.time() * 1000)
@@ -313,7 +312,7 @@ class ChallengeService:
sarif=sarif,
sarif_id=sarif_id,
task_id=task_id,
)
),
],
message_id=message_id,
message_time=message_time,
orchestrator/src/buttercup/orchestrator/ui/competition_api/services/crs_client.py
+28 -26
View File
@@ -1,10 +1,11 @@
from __future__ import annotations
import logging
import requests
from typing import Optional, Dict, Any
from typing import Any
from buttercup.orchestrator.ui.competition_api.models.crs_types import Task, SARIFBroadcast
import requests
from buttercup.orchestrator.ui.competition_api.models.crs_types import SARIFBroadcast, Task
logger = logging.getLogger(__name__)
@@ -13,7 +14,11 @@ class CRSResponse:
"""Response from CRS API calls with detailed status and error information."""
def __init__(
self, success: bool, status_code: int, response_text: str = "", error_details: Optional[Dict[str, Any]] = None
self,
success: bool,
status_code: int,
response_text: str = "",
error_details: dict[str, Any] | None = None,
):
self.success = success
self.status_code = status_code
@@ -21,14 +26,14 @@ class CRSResponse:
self.error_details = error_details or {}
def get_user_friendly_error_message(self, base_message: str = "Operation failed") -> str:
"""
Generate a user-friendly error message from the response details.
"""Generate a user-friendly error message from the response details.
Args:
base_message: Base message to start with
Returns:
Formatted error message suitable for display to users
"""
if self.success:
return ""
@@ -65,12 +70,12 @@ class CRSResponse:
return error_message
def log_detailed_response(self, logger_instance: logging.Logger, operation: str = "CRS operation") -> None:
"""
Log detailed response information for debugging purposes.
"""Log detailed response information for debugging purposes.
Args:
logger_instance: Logger instance to use
operation: Description of the operation being logged
"""
if self.success:
logger_instance.info(f"{operation} completed successfully (HTTP {self.status_code})")
@@ -162,13 +167,12 @@ class CRSResponse:
"raw_response": response_data,
},
)
else:
# Non-JSON response, treat as success if status is 200/202
return cls(
success=response.status_code in (200, 202),
status_code=response.status_code,
response_text=response.text,
)
# Non-JSON response, treat as success if status is 200/202
return cls(
success=response.status_code in (200, 202),
status_code=response.status_code,
response_text=response.text,
)
except (ValueError, TypeError):
# Response is not valid JSON, treat as success if status is 200/202
@@ -186,14 +190,14 @@ class CRSClient:
self.password = password
def submit_task(self, task: Task) -> CRSResponse:
"""
Submit a task to the CRS via POST /v1/task endpoint
"""Submit a task to the CRS via POST /v1/task endpoint
Args:
task: Task object to submit
Returns:
CRSResponse object with detailed status and error information
"""
url = f"{self.crs_base_url}/v1/task/"
@@ -226,9 +230,7 @@ class CRSClient:
return CRSResponse(success=False, status_code=0, response_text=str(e), error_details={"exception": str(e)})
def submit_sarif_broadcast(self, broadcast: SARIFBroadcast) -> CRSResponse:
"""
Submit a SARIF Broadcast to the CRS via POST /v1/sarif/ endpoint
"""
"""Submit a SARIF Broadcast to the CRS via POST /v1/sarif/ endpoint"""
url = f"{self.crs_base_url}/v1/sarif/"
# Prepare authentication if provided
@@ -252,7 +254,8 @@ class CRSClient:
# Log detailed response information
crs_response.log_detailed_response(
logger, f"SARIF Broadcast submission for {len(broadcast.broadcasts)} tasks"
logger,
f"SARIF Broadcast submission for {len(broadcast.broadcasts)} tasks",
)
return crs_response
@@ -262,11 +265,11 @@ class CRSClient:
return CRSResponse(success=False, status_code=0, response_text=str(e), error_details={"exception": str(e)})
def ping(self) -> bool:
"""
Test connectivity to CRS via GET /status/ endpoint
"""Test connectivity to CRS via GET /status/ endpoint
Returns:
True if CRS is reachable and ready, False otherwise
"""
url = f"{self.crs_base_url}/status/"
@@ -289,9 +292,8 @@ class CRSClient:
ready = status_data.get("ready", False)
logger.info(f"CRS ping successful. Ready: {ready}")
return bool(ready)
else:
logger.error(f"CRS ping failed. Status: {response.status_code}, Response: {response.text}")
return False
logger.error(f"CRS ping failed. Status: {response.status_code}, Response: {response.text}")
return False
except Exception as e:
logger.error(f"Error pinging CRS: {e}")
orchestrator/src/buttercup/orchestrator/ui/config.py
+16 -10
View File
@@ -1,7 +1,8 @@
from pydantic_settings import BaseSettings
from typing import Annotated, Optional
from pydantic import Field
from pathlib import Path
from typing import Annotated
from pydantic import Field
from pydantic_settings import BaseSettings
class Settings(BaseSettings):
@@ -12,32 +13,37 @@ class Settings(BaseSettings):
# Competition API configuration
external_host: Annotated[
str, Field(default="localhost", description="External host, used to construct the files URLs")
str,
Field(default="localhost", description="External host, used to construct the files URLs"),
]
host: Annotated[str, Field(default="127.0.0.1", description="Host to bind the server to")]
port: Annotated[int, Field(default=1323, description="Port to bind the server to")]
# File server configuration
challenges_dir: Annotated[
Path, Field(default=Path("./challenges"), description="Directory containing challenge files")
Path,
Field(default=Path("./challenges"), description="Directory containing challenge files"),
]
# CRS configuration
crs_base_url: Annotated[str, Field(default="http://localhost:8000", description="CRS API base URL")]
crs_key_id: Annotated[Optional[str], Field(default=None, description="Key ID for CRS authentication")]
crs_key_token: Annotated[Optional[str], Field(default=None, description="Key token for CRS authentication")]
crs_key_id: Annotated[str | None, Field(default=None, description="Key ID for CRS authentication")]
crs_key_token: Annotated[str | None, Field(default=None, description="Key token for CRS authentication")]
# Storage configuration
storage_dir: Annotated[
Path, Field(default=Path("/tmp/buttercup-storage"), description="Directory for storing tarballs")
Path,
Field(default=Path("/tmp/buttercup-storage"), description="Directory for storing tarballs"),
]
run_data_dir: Annotated[
Path, Field(default=Path("/tmp/buttercup-run-data"), description="Directory for storing run data artifacts")
Path,
Field(default=Path("/tmp/buttercup-run-data"), description="Directory for storing run data artifacts"),
]
# Database configuration
database_url: Annotated[
str, Field(default="sqlite:///buttercup_ui.db", description="Database URL for storing submissions")
str,
Field(default="sqlite:///buttercup_ui.db", description="Database URL for storing submissions"),
]
class Config:
orchestrator/src/buttercup/orchestrator/ui/database.py
+12 -5
View File
@@ -1,22 +1,22 @@
"""Database models and operations for buttercup-ui."""
import logging
import uuid
from collections.abc import Iterator
from contextlib import contextmanager
from datetime import datetime
from typing import Iterator
import uuid
from sqlalchemy import (
BLOB,
DateTime,
ForeignKey,
Integer,
String,
Text,
create_engine,
func,
DateTime,
)
from sqlalchemy.orm import Session, relationship, sessionmaker, DeclarativeBase, Mapped, mapped_column
from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column, relationship, sessionmaker
logger = logging.getLogger(__name__)
@@ -313,7 +313,14 @@ class DatabaseManager:
# POV operations
def create_pov(
self, *, task_id: str, architecture: str, engine: str, fuzzer_name: str, sanitizer: str, testcase: bytes
self,
*,
task_id: str,
architecture: str,
engine: str,
fuzzer_name: str,
sanitizer: str,
testcase: bytes,
) -> POV:
"""Create a new POV."""
with self.get_session() as session:
orchestrator/src/buttercup/orchestrator/utils.py
+1
View File
@@ -1,5 +1,6 @@
import hashlib
from pathlib import Path
import requests
orchestrator/test/test_backend.py
+1
View File
@@ -1,5 +1,6 @@
import unittest
from unittest.mock import MagicMock, Mock, patch
from buttercup.orchestrator.task_server.backend import get_status_tasks_state
from buttercup.orchestrator.task_server.models.types import StatusTasksState
orchestrator/test/test_broadcast_sarif_assessment_api.py
+2 -6
View File
@@ -1,7 +1,4 @@
# coding: utf-8
"""
Example Competition API
"""Example Competition API
No description provided (generated by Openapi Generator https://github.com/openapitools/openapi-generator)
@@ -9,7 +6,7 @@ The version of the OpenAPI document: 1.4.0
Generated by OpenAPI Generator (https://openapi-generator.tech)
Do not edit the class manually.
""" # noqa: E501
"""
import unittest
@@ -30,7 +27,6 @@ class TestBroadcastSarifAssessmentApi(unittest.TestCase):
Submit a SARIF Assessment
"""
pass
if __name__ == "__main__":
orchestrator/test/test_bundle_api.py
+2 -9
View File
@@ -1,7 +1,4 @@
# coding: utf-8
"""
Example Competition API
"""Example Competition API
No description provided (generated by Openapi Generator https://github.com/openapitools/openapi-generator)
@@ -9,7 +6,7 @@ The version of the OpenAPI document: 1.4.0
Generated by OpenAPI Generator (https://openapi-generator.tech)
Do not edit the class manually.
""" # noqa: E501
"""
import unittest
@@ -30,28 +27,24 @@ class TestBundleApi(unittest.TestCase):
Delete Bundle
"""
pass
def test_v1_task_task_id_bundle_bundle_id_get(self) -> None:
"""Test case for v1_task_task_id_bundle_bundle_id_get
Get Bundle
"""
pass
def test_v1_task_task_id_bundle_bundle_id_patch(self) -> None:
"""Test case for v1_task_task_id_bundle_bundle_id_patch
Update Bundle
"""
pass
def test_v1_task_task_id_bundle_post(self) -> None:
"""Test case for v1_task_task_id_bundle_post
Submit Bundle
"""
pass
if __name__ == "__main__":
orchestrator/test/test_cancellation.py
+3 -3
View File
@@ -1,13 +1,13 @@
import time
from unittest.mock import Mock, patch
import pytest
from redis import Redis
from buttercup.common.datastructures.msg_pb2 import TaskDelete
from buttercup.common.datastructures.msg_pb2 import Task, TaskDelete
from buttercup.common.queues import ReliableQueue, RQItem
from buttercup.orchestrator.scheduler.cancellation import Cancellation
from buttercup.common.task_registry import TaskRegistry
from buttercup.common.datastructures.msg_pb2 import Task
from buttercup.orchestrator.scheduler.cancellation import Cancellation
@pytest.fixture
orchestrator/test/test_challenge_service.py
+24 -20
View File
@@ -1,20 +1,21 @@
import pytest
import subprocess
import tempfile
import tarfile
import tempfile
import time
from pathlib import Path
from unittest.mock import patch, MagicMock
from unittest.mock import MagicMock, patch
import pytest
from buttercup.orchestrator.ui.competition_api.services import ChallengeService
from buttercup.orchestrator.ui.competition_api.models.crs_types import (
SARIFBroadcast,
SARIFBroadcastDetail,
SourceType,
Task,
TaskDetail,
TaskType,
SARIFBroadcast,
SARIFBroadcastDetail,
)
import time
from buttercup.orchestrator.ui.competition_api.services import ChallengeService
class TestChallengeService:
@@ -65,7 +66,9 @@ class TestChallengeService:
mock_tarfile.return_value.__enter__.return_value = mock_tar
_, sha256_hash, _ = challenge_service.create_challenge_tarball(
repo_url="https://github.com/octocat/Hello-World", ref="main", tarball_name="test-repo"
repo_url="https://github.com/octocat/Hello-World",
ref="main",
tarball_name="test-repo",
)
# Verify git clone was called
@@ -102,7 +105,9 @@ class TestChallengeService:
with pytest.raises(subprocess.CalledProcessError):
challenge_service.create_challenge_tarball(
repo_url="https://github.com/invalid/repo", ref="main", tarball_name="test-repo"
repo_url="https://github.com/invalid/repo",
ref="main",
tarball_name="test-repo",
)
def test_serve_tarball_success(self, challenge_service):
@@ -261,8 +266,7 @@ class TestChallengeService:
(project_path / ".git" / "config").write_text("git config")
(project_path / "README.md").write_text("# Test Repository")
return MagicMock(returncode=0, stdout="", stderr="")
else:
return MagicMock(returncode=0, stdout="", stderr="")
return MagicMock(returncode=0, stdout="", stderr="")
mock_run.side_effect = git_clone_mock
@@ -330,12 +334,12 @@ class TestChallengeService:
"physicalLocation": {
"artifactLocation": {"uri": "src/main.c"},
"region": {"startLine": 10, "startColumn": 5},
}
}
},
},
],
}
},
],
}
},
],
}
@@ -394,7 +398,7 @@ class TestChallengeService:
"name": "complex-tool",
"version": "2.0.0",
"informationUri": "https://example.com/tool",
}
},
},
"results": [
{
@@ -409,14 +413,14 @@ class TestChallengeService:
"physicalLocation": {
"artifactLocation": {"uri": "src/vulnerable.c", "uriBaseId": "SRCROOT"},
"region": {"startLine": 25, "startColumn": 10, "endLine": 25, "endColumn": 15},
}
}
},
},
],
"properties": {"security-severity": "HIGH", "tags": ["buffer-overflow", "memory-safety"]},
}
},
],
"invocations": [{"executionSuccessful": True, "commandLine": "fuzzer --target vulnerable.c"}],
}
},
],
}
orchestrator/test/test_crs_client.py
+6 -5
View File
@@ -1,15 +1,16 @@
from unittest.mock import MagicMock, patch
import pytest
from unittest.mock import patch, MagicMock
import requests
from buttercup.orchestrator.ui.competition_api.services import CRSClient
from buttercup.orchestrator.ui.competition_api.models.crs_types import (
Task,
TaskDetail,
SourceDetail,
SourceType,
Task,
TaskDetail,
TaskType,
)
from buttercup.orchestrator.ui.competition_api.services import CRSClient
class TestCRSClient:
@@ -41,7 +42,7 @@ class TestCRSClient:
project_name="test-project",
source=[SourceDetail(sha256="a" * 64, type=SourceType.repo, url="/files/test-repo.tar.gz")],
type=TaskType.full,
)
),
],
)
orchestrator/test/test_downloader.py
+10 -9
View File
@@ -1,13 +1,15 @@
import pytest
import hashlib
import responses
import tarfile
import io
import tarfile
from pathlib import Path
from unittest.mock import patch
from buttercup.orchestrator.downloader.downloader import Downloader
from buttercup.common.datastructures.msg_pb2 import Task, SourceDetail
import pytest
import responses
from buttercup.common.datastructures.msg_pb2 import SourceDetail, Task
from buttercup.common.node_local import TmpDir
from buttercup.orchestrator.downloader.downloader import Downloader
@pytest.fixture
@@ -126,10 +128,9 @@ def mock_node_local(temp_download_dir):
# Return the destination path to simulate success
return dst
# Use os.rename for files
else:
os.rename(src, dst)
# Return the destination path to simulate success
return dst
os.rename(src, dst)
# Return the destination path to simulate success
return dst
except Exception as e:
print(f"Error in rename_atomically mock: {e}")
return None
orchestrator/test/test_feeform_api.py
+2 -6
View File
@@ -1,7 +1,4 @@
# coding: utf-8
"""
Example Competition API
"""Example Competition API
No description provided (generated by Openapi Generator https://github.com/openapitools/openapi-generator)
@@ -9,7 +6,7 @@ The version of the OpenAPI document: 1.1
Generated by OpenAPI Generator (https://openapi-generator.tech)
Do not edit the class manually.
""" # noqa: E501
"""
import unittest
@@ -30,7 +27,6 @@ class TestFeeformApi(unittest.TestCase):
Submit Freeform
"""
pass
if __name__ == "__main__":
orchestrator/test/test_freeform_api.py
+2 -6
View File
@@ -1,7 +1,4 @@
# coding: utf-8
"""
Example Competition API
"""Example Competition API
No description provided (generated by Openapi Generator https://github.com/openapitools/openapi-generator)
@@ -9,7 +6,7 @@ The version of the OpenAPI document: 1.4.0
Generated by OpenAPI Generator (https://openapi-generator.tech)
Do not edit the class manually.
""" # noqa: E501
"""
import unittest
@@ -30,7 +27,6 @@ class TestFreeformApi(unittest.TestCase):
Submit Freeform
"""
pass
if __name__ == "__main__":
orchestrator/test/test_ping_api.py
+2 -6
View File
@@ -1,7 +1,4 @@
# coding: utf-8
"""
Example Competition API
"""Example Competition API
No description provided (generated by Openapi Generator https://github.com/openapitools/openapi-generator)
@@ -9,7 +6,7 @@ The version of the OpenAPI document: 1.4.0
Generated by OpenAPI Generator (https://openapi-generator.tech)
Do not edit the class manually.
""" # noqa: E501
"""
import unittest
@@ -30,7 +27,6 @@ class TestPingApi(unittest.TestCase):
Test authentication creds and network connectivity
"""
pass
if __name__ == "__main__":
orchestrator/test/test_pov_api.py
+2 -7
View File
@@ -1,7 +1,4 @@
# coding: utf-8
"""
Example Competition API
"""Example Competition API
No description provided (generated by Openapi Generator https://github.com/openapitools/openapi-generator)
@@ -9,7 +6,7 @@ The version of the OpenAPI document: 1.4.0
Generated by OpenAPI Generator (https://openapi-generator.tech)
Do not edit the class manually.
""" # noqa: E501
"""
import unittest
@@ -30,14 +27,12 @@ class TestPovApi(unittest.TestCase):
Submit Vulnerability
"""
pass
def test_v1_task_task_id_pov_pov_id_get(self) -> None:
"""Test case for v1_task_task_id_pov_pov_id_get
Vulnerability Status
"""
pass
if __name__ == "__main__":
orchestrator/test/test_pov_reproducer.py
+7 -6
View File
@@ -1,13 +1,14 @@
import pytest
from unittest.mock import Mock, patch, MagicMock
from redis import Redis
from pathlib import Path
from unittest.mock import MagicMock, Mock, patch
import pytest
from redis import Redis
from buttercup.common.datastructures.msg_pb2 import BuildOutput, BuildType, POVReproduceRequest
from buttercup.common.task_registry import TaskRegistry
from buttercup.common.maps import BuildMap
from buttercup.common.challenge_task import ChallengeTask, ReproduceResult
from buttercup.common.datastructures.msg_pb2 import BuildOutput, BuildType, POVReproduceRequest
from buttercup.common.maps import BuildMap
from buttercup.common.sets import PoVReproduceStatus
from buttercup.common.task_registry import TaskRegistry
from buttercup.orchestrator.pov_reproducer.pov_reproducer import POVReproducer
orchestrator/test/test_request_api.py
+2 -7
View File
@@ -1,7 +1,4 @@
# coding: utf-8
"""
Example Competition API
"""Example Competition API
No description provided (generated by Openapi Generator https://github.com/openapitools/openapi-generator)
@@ -9,7 +6,7 @@ The version of the OpenAPI document: 1.4.0
Generated by OpenAPI Generator (https://openapi-generator.tech)
Do not edit the class manually.
""" # noqa: E501
"""
import unittest
@@ -30,14 +27,12 @@ class TestRequestApi(unittest.TestCase):
Send a task to the source of this request
"""
pass
def test_v1_request_list_get(self) -> None:
"""Test case for v1_request_list_get
Get a list of available challenges to task
"""
pass
if __name__ == "__main__":
orchestrator/test/test_sarif_matcher.py
+4 -4
View File
@@ -1,12 +1,12 @@
#!/usr/bin/env python3
import copy
import os
import unittest
import copy
from pathlib import Path
from buttercup.common.sarif_store import SARIFBroadcastDetail
from buttercup.common.datastructures.msg_pb2 import TracedCrash
from buttercup.common.sarif_store import SARIFBroadcastDetail
from buttercup.orchestrator.scheduler.sarif_matcher import match
@@ -22,7 +22,7 @@ class TestSarifMatcher(unittest.TestCase):
def load_sarif_broadcast(self) -> SARIFBroadcastDetail:
"""Test that we can load a SARIF broadcast from the file."""
with open(self.sarif_broadcast_path, "r") as f:
with open(self.sarif_broadcast_path) as f:
sarif_json = f.read()
sarif_broadcast = SARIFBroadcastDetail.model_validate_json(sarif_json)
@@ -131,7 +131,7 @@ class TestSarifMatcher(unittest.TestCase):
# Add a logical location with the exact function name
modified_sarif["runs"][0]["results"][0]["locations"][0]["logicalLocations"] = [
{"name": "OSS_FUZZ_png_handle_iCCP"}
{"name": "OSS_FUZZ_png_handle_iCCP"},
]
sarif_broadcast.sarif = modified_sarif

Some files were not shown because too many files have changed in this diff Show More