mirror of
https://github.com/trailofbits/buttercup
synced 2026-06-21 14:11:39 +00:00
Merge SubmissionEntries based on similarity and on patch mitigation (#899)
* Use an internal_patch_id instead of indices This is the first step in being able to merge sets of PoVs and test patches against all PoVs within a task. * Discard redundant builds * Initial PoV-merging strategy Still not optimal in terms of SARIF-matching/bundling etc * Appears to be working version of merging including bundle and sarif handling * Make tests pass * Update integration test steps * Fixes and cleanup from review * Removed additional request for patched builds * Refactored some loops into find-style functions to simplify * Inline small function used once * Refactors for increased robustness and readability including additional testing * SARIF matching - additional tests and refactor * Add enumerate_task_submissions * Refactor and simplify tests using a builder Cleanup unused code * Fix read_submissions to use CrashWithId * Improvements based on review * Cache final states of PoV reproduce (#909) As these never change we can limit the load on redis by caching the results. * Merge SubmissionEntries based on patches (#910) * Cache final states of PoV reproduce As these never change we can limit the load on redis by caching the results. * Merge SubmissionEntries based on patches If a PoV in another entry is mitigated by the current entry's patch, merge the entries as athey should be considered the same ChallengeVulnerability. * Add positional argument * Hold of submitting a patch while evaluating Check each already submitted patch before submitting a new one for the same task. If any of the already submitted patches mitigates any PoV in the current SubmissionEntry - do not submit this. It will be merged later on. * Additional logging, truncate ids * Only request patch if no submitted patch mitigates Before we request a new patch, we check each of the already submitted patches to see if any of them already mitigates the PoVs in the current SubmissionEntry. If they do, this will be merged at a later stage. * PR feedback
This commit is contained in:
Henrik Brodin
@@ -148,13 +148,13 @@ jobs:
|
||||
- name: Wait for vuln to be found
|
||||
timeout-minutes: ${{ fromJSON(env.VULN_TIMEOUT) }}
|
||||
run: |
|
||||
while ! kubectl logs -n $BUTTERCUP_NAMESPACE -l app=scheduler --tail=-1 | grep "] SUBMIT_PATCH_REQUEST pov_id=" ; do
|
||||
while ! kubectl logs -n $BUTTERCUP_NAMESPACE -l app=scheduler --tail=-1 | grep "] POV submission response: pov_id=" ; do
|
||||
sleep 60
|
||||
done
|
||||
- name: Wait for vuln to enter the passed state in competition api
|
||||
timeout-minutes: ${{ fromJSON(env.VULN_TIMEOUT) }}
|
||||
run: |
|
||||
while ! kubectl logs -n $BUTTERCUP_NAMESPACE -l app=scheduler --tail=-1 | grep "WAIT_POV_PASS -> SUBMIT_PATCH" ; do
|
||||
while ! kubectl logs -n $BUTTERCUP_NAMESPACE -l app=scheduler --tail=-1 | grep "Updated POV status. New status PASSED" ; do
|
||||
sleep 60
|
||||
done
|
||||
|
||||
@@ -165,30 +165,30 @@ jobs:
|
||||
sleep 60
|
||||
done
|
||||
|
||||
- name: Wait for patch to be submitted
|
||||
- name: Wait for patch to be recorded
|
||||
timeout-minutes: ${{ fromJSON(env.PATCH_TIMEOUT) }}
|
||||
run: |
|
||||
while ! kubectl logs -n $BUTTERCUP_NAMESPACE -l app=scheduler --tail=-1 | grep "SUBMIT_PATCH -> WAIT_PATCH_PASS"; do
|
||||
while ! kubectl logs -n $BUTTERCUP_NAMESPACE -l app=scheduler --tail=-1 | grep "Appending patch for task"; do
|
||||
sleep 60
|
||||
done
|
||||
- name: Wait for patch to enter the passed state in competition api
|
||||
timeout-minutes: ${{ fromJSON(env.PATCH_TIMEOUT) }}
|
||||
run: |
|
||||
while ! kubectl logs -n $BUTTERCUP_NAMESPACE -l app=scheduler --tail=-1 | grep "WAIT_PATCH_PASS -> SUBMIT_BUNDLE"; do
|
||||
while ! kubectl logs -n $BUTTERCUP_NAMESPACE -l app=scheduler --tail=-1 | grep "Patch passed"; do
|
||||
sleep 60
|
||||
done
|
||||
|
||||
- name: Wait for bundle to be submitted
|
||||
timeout-minutes: ${{ fromJSON(env.BUNDLE_TIMEOUT) }}
|
||||
run: |
|
||||
while ! kubectl logs -n $BUTTERCUP_NAMESPACE -l app=scheduler --tail=-1 | grep "SUBMIT_BUNDLE -> SUBMIT_MATCHING_SARIF"; do
|
||||
while ! kubectl logs -n $BUTTERCUP_NAMESPACE -l app=scheduler --tail=-1 | grep "Bundle submission response: bundle_id="; do
|
||||
sleep 60
|
||||
done
|
||||
|
||||
- name: Send a SARIF broadcast
|
||||
timeout-minutes: ${{ fromJSON(env.SARIF_TIMEOUT) }}
|
||||
run: |
|
||||
TASK_ID=$(kubectl logs -n $BUTTERCUP_NAMESPACE -l app=scheduler --tail=-1 | grep "SUBMIT_BUNDLE -> SUBMIT_MATCHING_SARIF" | grep -o "\[[^]]*\]" | grep -o "[^[]*$" | cut -d: -f3 | tr -d ']')
|
||||
TASK_ID=$(kubectl logs -n $BUTTERCUP_NAMESPACE -l app=scheduler --tail=-1 | grep "Submitting bundle for harness" | grep -o "\[[^]]*\]" | grep -o "[^[]*$" | cut -d: -f3 | tr -d ']')
|
||||
echo "Task ID: $TASK_ID"
|
||||
while true; do
|
||||
# Kill any existing port-forward processes
|
||||
@@ -209,16 +209,16 @@ jobs:
|
||||
# Failed - wait a bit before retrying
|
||||
sleep 10
|
||||
done
|
||||
- name: Wait for SARIF to be submitted as correct
|
||||
timeout-minutes: ${{ fromJSON(env.SARIF_TIMEOUT) }}
|
||||
run: |
|
||||
while ! kubectl logs -n $BUTTERCUP_NAMESPACE -l app=scheduler --tail=-1 | grep "SUBMIT_MATCHING_SARIF -> SUBMIT_BUNDLE_PATCH"; do
|
||||
sleep 60
|
||||
done
|
||||
- name: Wait for Bundle to be patched to include the SARIF
|
||||
timeout-minutes: ${{ fromJSON(env.SARIF_TIMEOUT) }}
|
||||
run: |
|
||||
while ! kubectl logs -n $BUTTERCUP_NAMESPACE -l app=scheduler --tail=-1 | grep "SUBMIT_BUNDLE_PATCH -> STOP"; do
|
||||
while ! kubectl logs -n $BUTTERCUP_NAMESPACE -l app=scheduler --tail=-1 | grep "Bundle patch submission response: broadcast_sarif_id="; do
|
||||
sleep 60
|
||||
done
|
||||
- name: Wait for SARIF to be submitted as correct
|
||||
timeout-minutes: ${{ fromJSON(env.SARIF_TIMEOUT) }}
|
||||
run: |
|
||||
while ! kubectl logs -n $BUTTERCUP_NAMESPACE -l app=scheduler --tail=-1 | grep "Matching SARIF submission response"; do
|
||||
sleep 60
|
||||
done
|
||||
|
||||
|
||||
+34
-21
@@ -137,40 +137,53 @@ message FunctionCoverage {
|
||||
int32 covered_lines = 4;
|
||||
}
|
||||
|
||||
enum SubmissionResult {
|
||||
NONE = 0; // Default/unset state - no submission result yet
|
||||
ACCEPTED = 1;
|
||||
PASSED = 2;
|
||||
FAILED = 3;
|
||||
DEADLINE_EXCEEDED = 4;
|
||||
ERRORED = 5;
|
||||
INCONCLUSIVE = 6;
|
||||
}
|
||||
|
||||
message SubmissionEntryPatch {
|
||||
string patch = 1;
|
||||
string internal_patch_id = 2;
|
||||
string competition_patch_id = 3;
|
||||
repeated BuildOutput build_outputs = 4;
|
||||
optional SubmissionResult result = 5;
|
||||
}
|
||||
|
||||
message SubmissionEntry {
|
||||
enum SubmissionState {
|
||||
SUBMIT_PATCH_REQUEST = 0;
|
||||
WAIT_POV_PASS = 1;
|
||||
SUBMIT_PATCH = 2;
|
||||
WAIT_PATCH_PASS = 3;
|
||||
SUBMIT_BUNDLE = 4;
|
||||
SUBMIT_MATCHING_SARIF = 5;
|
||||
SUBMIT_BUNDLE_PATCH = 6;
|
||||
STOP = 7;
|
||||
}
|
||||
|
||||
SubmissionState state = 1;
|
||||
|
||||
repeated TracedCrash crashes = 2;
|
||||
|
||||
string pov_id = 3;
|
||||
string competition_patch_id = 4;
|
||||
message Bundle {
|
||||
string task_id = 1;
|
||||
string competition_pov_id = 2;
|
||||
string competition_patch_id = 3;
|
||||
string competition_sarif_id = 4;
|
||||
string bundle_id = 5;
|
||||
string sarif_id = 6;
|
||||
}
|
||||
|
||||
message CrashWithId {
|
||||
TracedCrash crash = 1;
|
||||
string competition_pov_id = 2;
|
||||
optional SubmissionResult result = 3;
|
||||
}
|
||||
|
||||
// This is our model of a Challenge Vulnerability
|
||||
message SubmissionEntry {
|
||||
bool stop = 1;
|
||||
|
||||
repeated CrashWithId crashes = 2;
|
||||
|
||||
repeated Bundle bundles = 3;
|
||||
|
||||
// Content of patches targeting this vulnerability
|
||||
repeated SubmissionEntryPatch patches = 7;
|
||||
repeated SubmissionEntryPatch patches = 4;
|
||||
|
||||
// Current patch being considered
|
||||
int32 patch_idx = 8;
|
||||
int32 patch_submission_attempt = 9;
|
||||
int32 patch_idx = 5;
|
||||
int32 patch_submission_attempts = 6;
|
||||
}
|
||||
|
||||
message POVReproduceRequest {
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -4,9 +4,16 @@ from google.protobuf import descriptor as _descriptor
|
||||
from google.protobuf import message as _message
|
||||
from typing import ClassVar as _ClassVar, Iterable as _Iterable, Mapping as _Mapping, Optional as _Optional, Union as _Union
|
||||
|
||||
ACCEPTED: SubmissionResult
|
||||
COVERAGE: BuildType
|
||||
DEADLINE_EXCEEDED: SubmissionResult
|
||||
DESCRIPTOR: _descriptor.FileDescriptor
|
||||
ERRORED: SubmissionResult
|
||||
FAILED: SubmissionResult
|
||||
FUZZER: BuildType
|
||||
INCONCLUSIVE: SubmissionResult
|
||||
NONE: SubmissionResult
|
||||
PASSED: SubmissionResult
|
||||
PATCH: BuildType
|
||||
TRACER_NO_DIFF: BuildType
|
||||
|
||||
@@ -48,6 +55,20 @@ class BuildRequest(_message.Message):
|
||||
task_id: str
|
||||
def __init__(self, engine: _Optional[str] = ..., sanitizer: _Optional[str] = ..., task_dir: _Optional[str] = ..., task_id: _Optional[str] = ..., build_type: _Optional[_Union[BuildType, str]] = ..., apply_diff: bool = ..., patch: _Optional[str] = ..., internal_patch_id: _Optional[str] = ...) -> None: ...
|
||||
|
||||
class Bundle(_message.Message):
|
||||
__slots__ = ["bundle_id", "competition_patch_id", "competition_pov_id", "competition_sarif_id", "task_id"]
|
||||
BUNDLE_ID_FIELD_NUMBER: _ClassVar[int]
|
||||
COMPETITION_PATCH_ID_FIELD_NUMBER: _ClassVar[int]
|
||||
COMPETITION_POV_ID_FIELD_NUMBER: _ClassVar[int]
|
||||
COMPETITION_SARIF_ID_FIELD_NUMBER: _ClassVar[int]
|
||||
TASK_ID_FIELD_NUMBER: _ClassVar[int]
|
||||
bundle_id: str
|
||||
competition_patch_id: str
|
||||
competition_pov_id: str
|
||||
competition_sarif_id: str
|
||||
task_id: str
|
||||
def __init__(self, task_id: _Optional[str] = ..., competition_pov_id: _Optional[str] = ..., competition_patch_id: _Optional[str] = ..., competition_sarif_id: _Optional[str] = ..., bundle_id: _Optional[str] = ...) -> None: ...
|
||||
|
||||
class ConfirmedVulnerability(_message.Message):
|
||||
__slots__ = ["crashes", "internal_patch_id"]
|
||||
CRASHES_FIELD_NUMBER: _ClassVar[int]
|
||||
@@ -70,6 +91,16 @@ class Crash(_message.Message):
|
||||
target: BuildOutput
|
||||
def __init__(self, target: _Optional[_Union[BuildOutput, _Mapping]] = ..., harness_name: _Optional[str] = ..., crash_input_path: _Optional[str] = ..., stacktrace: _Optional[str] = ..., crash_token: _Optional[str] = ...) -> None: ...
|
||||
|
||||
class CrashWithId(_message.Message):
|
||||
__slots__ = ["competition_pov_id", "crash", "result"]
|
||||
COMPETITION_POV_ID_FIELD_NUMBER: _ClassVar[int]
|
||||
CRASH_FIELD_NUMBER: _ClassVar[int]
|
||||
RESULT_FIELD_NUMBER: _ClassVar[int]
|
||||
competition_pov_id: str
|
||||
crash: TracedCrash
|
||||
result: SubmissionResult
|
||||
def __init__(self, crash: _Optional[_Union[TracedCrash, _Mapping]] = ..., competition_pov_id: _Optional[str] = ..., result: _Optional[_Union[SubmissionResult, str]] = ...) -> None: ...
|
||||
|
||||
class FunctionCoverage(_message.Message):
|
||||
__slots__ = ["covered_lines", "function_name", "function_paths", "total_lines"]
|
||||
COVERED_LINES_FIELD_NUMBER: _ClassVar[int]
|
||||
@@ -158,48 +189,34 @@ class SourceDetail(_message.Message):
|
||||
def __init__(self, sha256: _Optional[str] = ..., source_type: _Optional[_Union[SourceDetail.SourceType, str]] = ..., url: _Optional[str] = ...) -> None: ...
|
||||
|
||||
class SubmissionEntry(_message.Message):
|
||||
__slots__ = ["bundle_id", "competition_patch_id", "crashes", "patch_idx", "patch_submission_attempt", "patches", "pov_id", "sarif_id", "state"]
|
||||
class SubmissionState(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
|
||||
__slots__ = []
|
||||
BUNDLE_ID_FIELD_NUMBER: _ClassVar[int]
|
||||
COMPETITION_PATCH_ID_FIELD_NUMBER: _ClassVar[int]
|
||||
__slots__ = ["bundles", "crashes", "patch_idx", "patch_submission_attempts", "patches", "stop"]
|
||||
BUNDLES_FIELD_NUMBER: _ClassVar[int]
|
||||
CRASHES_FIELD_NUMBER: _ClassVar[int]
|
||||
PATCHES_FIELD_NUMBER: _ClassVar[int]
|
||||
PATCH_IDX_FIELD_NUMBER: _ClassVar[int]
|
||||
PATCH_SUBMISSION_ATTEMPT_FIELD_NUMBER: _ClassVar[int]
|
||||
POV_ID_FIELD_NUMBER: _ClassVar[int]
|
||||
SARIF_ID_FIELD_NUMBER: _ClassVar[int]
|
||||
STATE_FIELD_NUMBER: _ClassVar[int]
|
||||
STOP: SubmissionEntry.SubmissionState
|
||||
SUBMIT_BUNDLE: SubmissionEntry.SubmissionState
|
||||
SUBMIT_BUNDLE_PATCH: SubmissionEntry.SubmissionState
|
||||
SUBMIT_MATCHING_SARIF: SubmissionEntry.SubmissionState
|
||||
SUBMIT_PATCH: SubmissionEntry.SubmissionState
|
||||
SUBMIT_PATCH_REQUEST: SubmissionEntry.SubmissionState
|
||||
WAIT_PATCH_PASS: SubmissionEntry.SubmissionState
|
||||
WAIT_POV_PASS: SubmissionEntry.SubmissionState
|
||||
bundle_id: str
|
||||
competition_patch_id: str
|
||||
crashes: _containers.RepeatedCompositeFieldContainer[TracedCrash]
|
||||
PATCH_SUBMISSION_ATTEMPTS_FIELD_NUMBER: _ClassVar[int]
|
||||
STOP_FIELD_NUMBER: _ClassVar[int]
|
||||
bundles: _containers.RepeatedCompositeFieldContainer[Bundle]
|
||||
crashes: _containers.RepeatedCompositeFieldContainer[CrashWithId]
|
||||
patch_idx: int
|
||||
patch_submission_attempt: int
|
||||
patch_submission_attempts: int
|
||||
patches: _containers.RepeatedCompositeFieldContainer[SubmissionEntryPatch]
|
||||
pov_id: str
|
||||
sarif_id: str
|
||||
state: SubmissionEntry.SubmissionState
|
||||
def __init__(self, state: _Optional[_Union[SubmissionEntry.SubmissionState, str]] = ..., crashes: _Optional[_Iterable[_Union[TracedCrash, _Mapping]]] = ..., pov_id: _Optional[str] = ..., competition_patch_id: _Optional[str] = ..., bundle_id: _Optional[str] = ..., sarif_id: _Optional[str] = ..., patches: _Optional[_Iterable[_Union[SubmissionEntryPatch, _Mapping]]] = ..., patch_idx: _Optional[int] = ..., patch_submission_attempt: _Optional[int] = ...) -> None: ...
|
||||
stop: bool
|
||||
def __init__(self, stop: bool = ..., crashes: _Optional[_Iterable[_Union[CrashWithId, _Mapping]]] = ..., bundles: _Optional[_Iterable[_Union[Bundle, _Mapping]]] = ..., patches: _Optional[_Iterable[_Union[SubmissionEntryPatch, _Mapping]]] = ..., patch_idx: _Optional[int] = ..., patch_submission_attempts: _Optional[int] = ...) -> None: ...
|
||||
|
||||
class SubmissionEntryPatch(_message.Message):
|
||||
__slots__ = ["build_outputs", "competition_patch_id", "internal_patch_id", "patch"]
|
||||
__slots__ = ["build_outputs", "competition_patch_id", "internal_patch_id", "patch", "result"]
|
||||
BUILD_OUTPUTS_FIELD_NUMBER: _ClassVar[int]
|
||||
COMPETITION_PATCH_ID_FIELD_NUMBER: _ClassVar[int]
|
||||
INTERNAL_PATCH_ID_FIELD_NUMBER: _ClassVar[int]
|
||||
PATCH_FIELD_NUMBER: _ClassVar[int]
|
||||
RESULT_FIELD_NUMBER: _ClassVar[int]
|
||||
build_outputs: _containers.RepeatedCompositeFieldContainer[BuildOutput]
|
||||
competition_patch_id: str
|
||||
internal_patch_id: str
|
||||
patch: str
|
||||
def __init__(self, patch: _Optional[str] = ..., internal_patch_id: _Optional[str] = ..., competition_patch_id: _Optional[str] = ..., build_outputs: _Optional[_Iterable[_Union[BuildOutput, _Mapping]]] = ...) -> None: ...
|
||||
result: SubmissionResult
|
||||
def __init__(self, patch: _Optional[str] = ..., internal_patch_id: _Optional[str] = ..., competition_patch_id: _Optional[str] = ..., build_outputs: _Optional[_Iterable[_Union[BuildOutput, _Mapping]]] = ..., result: _Optional[_Union[SubmissionResult, str]] = ...) -> None: ...
|
||||
|
||||
class Task(_message.Message):
|
||||
__slots__ = ["cancelled", "deadline", "focus", "message_id", "message_time", "metadata", "project_name", "sources", "task_id", "task_type"]
|
||||
@@ -280,3 +297,6 @@ class WeightedHarness(_message.Message):
|
||||
|
||||
class BuildType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
|
||||
__slots__ = []
|
||||
|
||||
class SubmissionResult(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
|
||||
__slots__ = []
|
||||
|
||||
@@ -5,6 +5,7 @@ from bson.json_util import dumps, CANONICAL_JSON_OPTIONS
|
||||
from contextlib import contextmanager
|
||||
import random
|
||||
import json
|
||||
from functools import lru_cache
|
||||
|
||||
# Import POVReproduceRequest for the refactored PoVReproduceStatus
|
||||
from buttercup.common.datastructures.msg_pb2 import POVReproduceRequest, POVReproduceResponse
|
||||
@@ -105,6 +106,28 @@ class PoVReproduceStatus:
|
||||
def __init__(self, redis: Redis):
|
||||
self.redis = redis
|
||||
|
||||
@lru_cache(maxsize=1000)
|
||||
def _did_crash(self, key: str) -> bool | None:
|
||||
"""Check if POV crashed (is in final states) and return crash status.
|
||||
|
||||
Args:
|
||||
key: The serialized key for the POV reproduction request
|
||||
|
||||
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)
|
||||
pipeline.sismember(POV_REPRODUCE_NON_MITIGATED_SET_NAME, key)
|
||||
result = pipeline.execute()
|
||||
|
||||
if result[0]:
|
||||
return False # Mitigated - didn't crash
|
||||
elif result[1]:
|
||||
return True # Non-mitigated - did crash
|
||||
else:
|
||||
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."""
|
||||
return dumps(
|
||||
@@ -122,11 +145,19 @@ class PoVReproduceStatus:
|
||||
None if pending, POVReproduceResponse if completed
|
||||
"""
|
||||
key = self._make_key(request)
|
||||
|
||||
# First check cache for final states only
|
||||
did_crash = self._did_crash(key)
|
||||
if did_crash is not None:
|
||||
return POVReproduceResponse(request=request, did_crash=did_crash)
|
||||
|
||||
# If not in final states, do the regular logic including pending check
|
||||
pipeline = self.redis.pipeline()
|
||||
pipeline.sismember(POV_REPRODUCE_PENDING_SET_NAME, key)
|
||||
pipeline.sismember(POV_REPRODUCE_MITIGATED_SET_NAME, key)
|
||||
pipeline.sismember(POV_REPRODUCE_NON_MITIGATED_SET_NAME, key)
|
||||
result = pipeline.execute()
|
||||
|
||||
if result[0]:
|
||||
return None # Pending
|
||||
elif result[1]:
|
||||
|
||||
@@ -29,7 +29,8 @@ def truncate_stacktraces(submission: SubmissionEntry, max_length: int = 80) -> S
|
||||
text_format.Parse(submission_text, truncated_submission)
|
||||
|
||||
# Now truncate the stacktraces and crash token
|
||||
for crash in truncated_submission.crashes:
|
||||
for crash_with_id in truncated_submission.crashes:
|
||||
crash = crash_with_id.crash
|
||||
if crash.crash.stacktrace and len(crash.crash.stacktrace) > max_length:
|
||||
crash.crash.stacktrace = crash.crash.stacktrace[:max_length] + "... (truncated)"
|
||||
|
||||
|
||||
@@ -2,6 +2,7 @@ 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
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -558,3 +559,339 @@ class TestPoVReproduceStatus:
|
||||
# Should have 2 pending items now (requests[4] and requests[5])
|
||||
pending1 = pov_status.get_one_pending()
|
||||
assert pending1.task_id in [requests[4].task_id, requests[5].task_id]
|
||||
|
||||
|
||||
class TestPoVReproduceStatusCaching:
|
||||
"""Test suite for PoVReproduceStatus caching functionality."""
|
||||
|
||||
def _create_mock_redis(self):
|
||||
"""Helper to create properly mocked Redis client with pipeline support."""
|
||||
mock_redis = MagicMock()
|
||||
mock_pipeline = MagicMock()
|
||||
mock_redis.pipeline.return_value = mock_pipeline
|
||||
mock_pipeline.__enter__.return_value = mock_pipeline
|
||||
mock_pipeline.__exit__.return_value = None
|
||||
return mock_redis, mock_pipeline
|
||||
|
||||
def test_cache_hit_for_mitigated_status(self):
|
||||
"""Test that repeated requests for mitigated status hit cache, not Redis."""
|
||||
# Create mock Redis client
|
||||
mock_redis, mock_pipeline = self._create_mock_redis()
|
||||
|
||||
# Mock the final states check to return mitigated
|
||||
mock_pipeline.execute.return_value = [True, False] # mitigated=True, non_mitigated=False
|
||||
|
||||
pov_status = PoVReproduceStatus(mock_redis)
|
||||
|
||||
# Create sample request
|
||||
request = POVReproduceRequest()
|
||||
request.task_id = "test-task"
|
||||
request.internal_patch_id = "0"
|
||||
request.pov_path = "/test/path"
|
||||
request.sanitizer = "asan"
|
||||
request.harness_name = "test_harness"
|
||||
|
||||
# First call should hit Redis
|
||||
result1 = pov_status.request_status(request)
|
||||
assert isinstance(result1, POVReproduceResponse)
|
||||
assert result1.did_crash is False
|
||||
|
||||
# Second call should hit cache, not Redis
|
||||
result2 = pov_status.request_status(request)
|
||||
assert isinstance(result2, POVReproduceResponse)
|
||||
assert result2.did_crash is False
|
||||
|
||||
# Third call should also hit cache
|
||||
result3 = pov_status.request_status(request)
|
||||
assert isinstance(result3, POVReproduceResponse)
|
||||
assert result3.did_crash is False
|
||||
|
||||
# Redis should only be called once (for the first request)
|
||||
assert mock_redis.pipeline.call_count == 1
|
||||
|
||||
def test_cache_hit_for_non_mitigated_status(self):
|
||||
"""Test that repeated requests for non-mitigated status hit cache, not Redis."""
|
||||
# Create mock Redis client
|
||||
mock_redis, mock_pipeline = self._create_mock_redis()
|
||||
|
||||
# Mock the final states check to return non-mitigated
|
||||
mock_pipeline.execute.return_value = [False, True] # mitigated=False, non_mitigated=True
|
||||
|
||||
pov_status = PoVReproduceStatus(mock_redis)
|
||||
|
||||
# Create sample request
|
||||
request = POVReproduceRequest()
|
||||
request.task_id = "test-task"
|
||||
request.internal_patch_id = "0"
|
||||
request.pov_path = "/test/path"
|
||||
request.sanitizer = "asan"
|
||||
request.harness_name = "test_harness"
|
||||
|
||||
# First call should hit Redis
|
||||
result1 = pov_status.request_status(request)
|
||||
assert isinstance(result1, POVReproduceResponse)
|
||||
assert result1.did_crash is True
|
||||
|
||||
# Second call should hit cache, not Redis
|
||||
result2 = pov_status.request_status(request)
|
||||
assert isinstance(result2, POVReproduceResponse)
|
||||
assert result2.did_crash is True
|
||||
|
||||
# Redis should only be called once (for the first request)
|
||||
assert mock_redis.pipeline.call_count == 1
|
||||
|
||||
def test_cache_miss_for_pending_status(self):
|
||||
"""Test that pending status doesn't get cached and always hits Redis."""
|
||||
# Create mock Redis client
|
||||
mock_redis, mock_pipeline = self._create_mock_redis()
|
||||
mock_redis.sadd.return_value = 1 # Mock successful add to pending set
|
||||
|
||||
# Mock responses:
|
||||
# First call: _did_crash returns [False, False], then full check returns [True, False, False]
|
||||
# Second call: _did_crash hits cache (no Redis call), then full check returns [True, False, False]
|
||||
mock_pipeline.execute.side_effect = [
|
||||
[False, False], # First request: _did_crash: not in final states
|
||||
[
|
||||
True,
|
||||
False,
|
||||
False,
|
||||
], # First request: request_status full check: pending=True, mitigated=False, non_mitigated=False
|
||||
[True, False, False], # Second request: request_status full check: still pending (cache hit for _did_crash)
|
||||
]
|
||||
|
||||
pov_status = PoVReproduceStatus(mock_redis)
|
||||
|
||||
# Create sample request
|
||||
request = POVReproduceRequest()
|
||||
request.task_id = "test-task"
|
||||
request.internal_patch_id = "0"
|
||||
request.pov_path = "/test/path"
|
||||
request.sanitizer = "asan"
|
||||
request.harness_name = "test_harness"
|
||||
|
||||
# First call should return None (pending)
|
||||
result1 = pov_status.request_status(request)
|
||||
assert result1 is None
|
||||
|
||||
# Second call should also return None (pending) and hit Redis again for full check
|
||||
# but _did_crash should hit cache
|
||||
result2 = pov_status.request_status(request)
|
||||
assert result2 is None
|
||||
|
||||
# Redis should be called 3 times total: 1 for first _did_crash + 2 for full checks
|
||||
# (second _did_crash hits cache)
|
||||
assert mock_redis.pipeline.call_count == 3
|
||||
|
||||
def test_cache_works_with_different_requests(self):
|
||||
"""Test that cache works correctly for different requests."""
|
||||
# Create mock Redis client
|
||||
mock_redis, mock_pipeline = self._create_mock_redis()
|
||||
|
||||
# Mock different responses for different requests
|
||||
mock_pipeline.execute.side_effect = [
|
||||
[True, False], # request1: mitigated
|
||||
[False, True], # request2: non-mitigated
|
||||
# Subsequent calls should hit cache
|
||||
]
|
||||
|
||||
pov_status = PoVReproduceStatus(mock_redis)
|
||||
|
||||
# Create two different requests
|
||||
request1 = POVReproduceRequest()
|
||||
request1.task_id = "test-task-1"
|
||||
request1.internal_patch_id = "0"
|
||||
request1.pov_path = "/test/path1"
|
||||
request1.sanitizer = "asan"
|
||||
request1.harness_name = "test_harness"
|
||||
|
||||
request2 = POVReproduceRequest()
|
||||
request2.task_id = "test-task-2"
|
||||
request2.internal_patch_id = "0"
|
||||
request2.pov_path = "/test/path2"
|
||||
request2.sanitizer = "msan"
|
||||
request2.harness_name = "test_harness"
|
||||
|
||||
# First calls should hit Redis
|
||||
result1 = pov_status.request_status(request1)
|
||||
assert isinstance(result1, POVReproduceResponse)
|
||||
assert result1.did_crash is False # mitigated
|
||||
|
||||
result2 = pov_status.request_status(request2)
|
||||
assert isinstance(result2, POVReproduceResponse)
|
||||
assert result2.did_crash is True # non-mitigated
|
||||
|
||||
# Repeat calls should hit cache
|
||||
result1_cached = pov_status.request_status(request1)
|
||||
assert isinstance(result1_cached, POVReproduceResponse)
|
||||
assert result1_cached.did_crash is False
|
||||
|
||||
result2_cached = pov_status.request_status(request2)
|
||||
assert isinstance(result2_cached, POVReproduceResponse)
|
||||
assert result2_cached.did_crash is True
|
||||
|
||||
# Redis should only be called twice (once for each unique request)
|
||||
assert mock_redis.pipeline.call_count == 2
|
||||
|
||||
def test_cache_respects_maxsize(self):
|
||||
"""Test that cache behavior is consistent (simplified test without maxsize override)."""
|
||||
# Create mock Redis client
|
||||
mock_redis, mock_pipeline = self._create_mock_redis()
|
||||
|
||||
# Mock all responses as mitigated
|
||||
mock_pipeline.execute.return_value = [True, False] # mitigated=True, non_mitigated=False
|
||||
|
||||
pov_status = PoVReproduceStatus(mock_redis)
|
||||
|
||||
# Create multiple different requests
|
||||
requests = []
|
||||
for i in range(3):
|
||||
request = POVReproduceRequest()
|
||||
request.task_id = f"test-task-{i}"
|
||||
request.internal_patch_id = "0"
|
||||
request.pov_path = f"/test/path{i}"
|
||||
request.sanitizer = "asan"
|
||||
request.harness_name = "test_harness"
|
||||
requests.append(request)
|
||||
|
||||
# First requests should hit Redis
|
||||
result1 = pov_status.request_status(requests[0])
|
||||
result2 = pov_status.request_status(requests[1])
|
||||
result3 = pov_status.request_status(requests[2])
|
||||
assert isinstance(result1, POVReproduceResponse)
|
||||
assert isinstance(result2, POVReproduceResponse)
|
||||
assert isinstance(result3, POVReproduceResponse)
|
||||
assert mock_redis.pipeline.call_count == 3
|
||||
|
||||
# Repeated requests should hit cache
|
||||
result1_cached = pov_status.request_status(requests[0])
|
||||
result2_cached = pov_status.request_status(requests[1])
|
||||
result3_cached = pov_status.request_status(requests[2])
|
||||
assert isinstance(result1_cached, POVReproduceResponse)
|
||||
assert isinstance(result2_cached, POVReproduceResponse)
|
||||
assert isinstance(result3_cached, POVReproduceResponse)
|
||||
# Should not have increased - all cache hits
|
||||
assert mock_redis.pipeline.call_count == 3
|
||||
|
||||
def test_separate_instances_have_separate_caches(self):
|
||||
"""Test that different PoVReproduceStatus instances have separate caches."""
|
||||
# Create two mock Redis clients
|
||||
mock_redis1 = MagicMock()
|
||||
mock_pipeline1 = MagicMock()
|
||||
mock_redis1.pipeline.return_value = mock_pipeline1
|
||||
mock_pipeline1.__enter__.return_value = mock_pipeline1
|
||||
mock_pipeline1.__exit__.return_value = None
|
||||
mock_pipeline1.execute.return_value = [True, False] # mitigated
|
||||
|
||||
mock_redis2 = MagicMock()
|
||||
mock_pipeline2 = MagicMock()
|
||||
mock_redis2.pipeline.return_value = mock_pipeline2
|
||||
mock_pipeline2.__enter__.return_value = mock_pipeline2
|
||||
mock_pipeline2.__exit__.return_value = None
|
||||
mock_pipeline2.execute.return_value = [False, True] # non-mitigated
|
||||
|
||||
# Create two instances
|
||||
pov_status1 = PoVReproduceStatus(mock_redis1)
|
||||
pov_status2 = PoVReproduceStatus(mock_redis2)
|
||||
|
||||
# Create same request for both
|
||||
request = POVReproduceRequest()
|
||||
request.task_id = "test-task"
|
||||
request.internal_patch_id = "0"
|
||||
request.pov_path = "/test/path"
|
||||
request.sanitizer = "asan"
|
||||
request.harness_name = "test_harness"
|
||||
|
||||
# First calls should hit respective Redis instances
|
||||
result1 = pov_status1.request_status(request)
|
||||
result2 = pov_status2.request_status(request)
|
||||
|
||||
assert isinstance(result1, POVReproduceResponse)
|
||||
assert result1.did_crash is False # mitigated
|
||||
assert isinstance(result2, POVReproduceResponse)
|
||||
assert result2.did_crash is True # non-mitigated
|
||||
|
||||
# Each instance should have hit its Redis once
|
||||
assert mock_redis1.pipeline.call_count == 1
|
||||
assert mock_redis2.pipeline.call_count == 1
|
||||
|
||||
# Repeat calls should hit respective caches
|
||||
result1_cached = pov_status1.request_status(request)
|
||||
result2_cached = pov_status2.request_status(request)
|
||||
|
||||
assert isinstance(result1_cached, POVReproduceResponse)
|
||||
assert result1_cached.did_crash is False
|
||||
assert isinstance(result2_cached, POVReproduceResponse)
|
||||
assert result2_cached.did_crash is True
|
||||
|
||||
# Redis call counts should not increase (cache hits)
|
||||
assert mock_redis1.pipeline.call_count == 1
|
||||
assert mock_redis2.pipeline.call_count == 1
|
||||
|
||||
def test_cache_integration_with_real_workflow(self, pov_status, sample_request):
|
||||
"""Test that caching works correctly in a real workflow with Redis."""
|
||||
# This test uses real Redis and validates the cache behavior
|
||||
|
||||
# Initial request should return None (pending)
|
||||
result1 = pov_status.request_status(sample_request)
|
||||
assert result1 is None
|
||||
|
||||
# Mark as mitigated
|
||||
mark_result = pov_status.mark_mitigated(sample_request)
|
||||
assert mark_result is True
|
||||
|
||||
# Now test caching: multiple requests should return same result
|
||||
# We can't easily test Redis call count with real Redis, but we can
|
||||
# verify that the behavior is consistent
|
||||
results = []
|
||||
for _ in range(5):
|
||||
result = pov_status.request_status(sample_request)
|
||||
results.append(result)
|
||||
|
||||
# All results should be identical POVReproduceResponse objects
|
||||
for result in results:
|
||||
assert isinstance(result, POVReproduceResponse)
|
||||
assert result.did_crash is False # mitigated
|
||||
assert result.request.task_id == sample_request.task_id
|
||||
assert result.request.internal_patch_id == sample_request.internal_patch_id
|
||||
assert result.request.pov_path == sample_request.pov_path
|
||||
assert result.request.sanitizer == sample_request.sanitizer
|
||||
assert result.request.harness_name == sample_request.harness_name
|
||||
|
||||
def test_cache_clears_across_different_final_states(self):
|
||||
"""Test that moving between final states works correctly with caching."""
|
||||
# Create mock Redis client
|
||||
mock_redis, mock_pipeline = self._create_mock_redis()
|
||||
|
||||
pov_status = PoVReproduceStatus(mock_redis)
|
||||
|
||||
# Create sample request
|
||||
request = POVReproduceRequest()
|
||||
request.task_id = "test-task"
|
||||
request.internal_patch_id = "0"
|
||||
request.pov_path = "/test/path"
|
||||
request.sanitizer = "asan"
|
||||
request.harness_name = "test_harness"
|
||||
|
||||
# First: mock as mitigated
|
||||
mock_pipeline.execute.return_value = [True, False]
|
||||
result1 = pov_status.request_status(request)
|
||||
assert isinstance(result1, POVReproduceResponse)
|
||||
assert result1.did_crash is False
|
||||
|
||||
# Second call should hit cache
|
||||
result2 = pov_status.request_status(request)
|
||||
assert isinstance(result2, POVReproduceResponse)
|
||||
assert result2.did_crash is False
|
||||
assert mock_redis.pipeline.call_count == 1 # Only first call hit Redis
|
||||
|
||||
# Now simulate the state changing in Redis (e.g., due to another process)
|
||||
# This would require cache invalidation in a real system, but our current
|
||||
# implementation doesn't handle this case - the cache will return stale data
|
||||
# This test documents the current behavior
|
||||
mock_pipeline.execute.return_value = [False, True] # Now non-mitigated
|
||||
|
||||
# This call will still return cached (stale) mitigated result
|
||||
result3 = pov_status.request_status(request)
|
||||
assert isinstance(result3, POVReproduceResponse)
|
||||
assert result3.did_crash is False # Still cached mitigated result
|
||||
assert mock_redis.pipeline.call_count == 1 # No new Redis calls
|
||||
|
||||
@@ -378,7 +378,7 @@ class Scheduler:
|
||||
vuln_item: RQItem[TracedCrash] | None = self.traced_vulnerabilities_queue.pop()
|
||||
if vuln_item is not None:
|
||||
crash: TracedCrash = vuln_item.deserialized
|
||||
logger.info(f"Submitting vulnerability for task {crash.crash.target.task_id}")
|
||||
logger.info(f"Recording vulnerability for task {crash.crash.target.task_id}")
|
||||
if self.submissions.submit_vulnerability(crash):
|
||||
self.traced_vulnerabilities_queue.ack_item(vuln_item.item_id)
|
||||
collected_item = True
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
+3781
-422
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user