Skip to content

Commit 6121550

Browse files
authored
Refactor reproduce (#4988)
This PR is a follow-up to #4976, addressing the feedback and discussions from the original review. The primary focus is on refactoring and code quality improvements. *(Note: This PR involves a significant amount of code being moved to more appropriate modules.)* --- Quick summary: - Moved local setup functions into the `setup.py` module - Updated imports to be module-level, adhering to style guidelines. - Added a `--local-logging` flag to `butler` to enable convenient local log output during development. - Reduced the reliance on mocks in tests - Implemented logic to correctly handle `minimized_keys` during setup. - Improved error handling to be more Pythonic. - Removed unnecessary type hints and enhanced overall test reliability.
1 parent 06491e2 commit 6121550

File tree

6 files changed

+449
-405
lines changed

6 files changed

+449
-405
lines changed

butler.py

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -97,11 +97,13 @@ def _setup_args_for_remote(parser):
9797

9898
subparsers.add_parser('reboot', help='Reboot with `sudo reboot`.')
9999

100+
100101
def _add_integration_tests_subparsers(toplevel_subparsers):
101102
"""Adds a parser for the `integration_tests` command."""
102103
toplevel_subparsers.add_parser(
103104
'integration_tests', help='Run end-to-end integration tests.')
104-
105+
106+
105107
def _add_weights_fuzzer_subparser(weights_subparsers):
106108
"""Adds a parser for the `weights fuzzer` command."""
107109
parser = weights_subparsers.add_parser(
@@ -236,19 +238,29 @@ def _add_weights_subparser(toplevel_subparsers):
236238
_add_weights_batches_subparser(subparsers)
237239
_add_weights_target_subparser(subparsers)
238240

241+
239242
def _add_reproduce_subparser(toplevel_subparsers):
240243
"""Adds a parser for the `reproduce` command."""
241244
parser = toplevel_subparsers.add_parser(
242245
'reproduce', help='Reproduce a testcase locally.')
243246
parser.add_argument(
244247
'-c', '--config-dir', required=True, help='Path to application config.')
245248
parser.add_argument(
246-
'-t', '--testcase-id', required=True, help='The testcase ID to reproduce.')
249+
'-t',
250+
'--testcase-id',
251+
required=True,
252+
help='The testcase ID to reproduce.')
253+
247254

248255
def main():
249256
"""Parse the command-line args and invoke the right command."""
250257
parser = _ArgumentParser(
251258
description='Butler is here to help you with command-line tasks.')
259+
parser.add_argument(
260+
'--local-logging',
261+
action='store_true',
262+
default=False,
263+
help='Force logs to be local-only.')
252264
subparsers = parser.add_subparsers(dest='command')
253265

254266
subparsers.add_parser(
@@ -424,12 +436,12 @@ def main():
424436
parser.print_help()
425437
return 0
426438

427-
_setup()
439+
_setup(args)
428440
command = importlib.import_module(f'local.butler.{args.command}')
429441
return command.execute(args)
430442

431443

432-
def _setup():
444+
def _setup(args):
433445
"""Set up configs and import paths."""
434446
os.environ['ROOT_DIR'] = os.path.abspath('.')
435447
os.environ['PYTHONIOENCODING'] = 'UTF-8'
@@ -438,6 +450,10 @@ def _setup():
438450
from clusterfuzz._internal.base import modules
439451
modules.fix_module_search_paths()
440452

453+
if args.local_logging:
454+
from clusterfuzz._internal.system import environment
455+
environment.set_local_log_only()
456+
441457

442458
if __name__ == '__main__':
443459
sys.exit(main())

src/clusterfuzz/_internal/bot/tasks/setup.py

Lines changed: 120 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@
4949
_DATA_BUNDLE_SYNC_INTERVAL_IN_SECONDS = 6 * 60 * 60
5050
_SYNC_FILENAME = '.sync'
5151
_TESTCASE_ARCHIVE_EXTENSION = '.zip'
52+
_EXECUTABLE_PERMISSIONS = 0o750
5253

5354

5455
def _set_timeout_value_from_user_upload(testcase_id, uworker_env):
@@ -352,10 +353,12 @@ def _get_testcase_key_and_archive_status(testcase):
352353

353354

354355
def _is_testcase_minimized(testcase):
356+
"""Returns True if the testcase is minimized."""
355357
return testcase.minimized_keys and testcase.minimized_keys != 'NA'
356358

357359

358360
def download_testcase(testcase_download_url, dst):
361+
"""Downloads a testcase from a signed url"""
359362
return storage.download_signed_url_to_file(testcase_download_url, dst)
360363

361364

@@ -651,7 +654,7 @@ def _update_fuzzer(
651654
return False
652655

653656
# Make fuzzer executable.
654-
os.chmod(fuzzer_path, 0o750)
657+
os.chmod(fuzzer_path, _EXECUTABLE_PERMISSIONS)
655658

656659
# Cleanup unneeded archive.
657660
shell.remove_file(archive_path)
@@ -890,3 +893,119 @@ def archive_testcase_and_dependencies_in_gcs(resource_list, testcase_path: str,
890893
shell.remove_file(zip_path)
891894

892895
return archived, absolute_filename, zip_filename
896+
897+
898+
def setup_local_testcase(testcase: data_types.Testcase) -> str | None:
899+
"""Sets up the testcase file locally.
900+
901+
Args:
902+
testcase: The Testcase object.
903+
904+
Returns:
905+
- str | None : The local file path to the testcase, or None on failure.
906+
"""
907+
try:
908+
shell.clear_testcase_directories()
909+
except Exception as e:
910+
logs.error(f'Error clearing testcase directories: {e}')
911+
return None
912+
913+
try:
914+
_, testcase_file_path = _get_testcase_file_and_path(testcase)
915+
if _is_testcase_minimized(testcase):
916+
blob_key = testcase.minimized_keys
917+
else:
918+
blob_key = testcase.fuzzed_keys
919+
if not blobs.read_blob_to_disk(blob_key, testcase_file_path):
920+
logs.error(f'Failed to download testcase from blobstore: {blob_key}')
921+
# Returning None for path when download fails
922+
return None
923+
prepare_environment_for_testcase(testcase)
924+
except Exception as e:
925+
logs.error(f'Error setting up testcase locally: {e}')
926+
return None
927+
928+
return testcase_file_path
929+
930+
931+
def setup_local_fuzzer(fuzzer_name: str) -> bool:
932+
"""Sets up the fuzzer binaries and environment.
933+
934+
Args:
935+
fuzzer_name: The name of the fuzzer to set up.
936+
937+
Returns:
938+
True if setup was successful, False otherwise.
939+
"""
940+
fuzzer = data_types.Fuzzer.query(data_types.Fuzzer.name == fuzzer_name).get()
941+
if not fuzzer:
942+
logs.error(f'Fuzzer {fuzzer_name} not found.')
943+
return False
944+
945+
if fuzzer.untrusted_content:
946+
logs.warning('You are about to run an untrusted fuzzer locally. '
947+
'This can be dangerous.')
948+
949+
environment.set_value('UNTRUSTED_CONTENT', fuzzer.untrusted_content)
950+
951+
if fuzzer.data_bundle_name:
952+
logs.info(
953+
f'Fuzzer {fuzzer_name} uses data bundle: {fuzzer.data_bundle_name}')
954+
955+
if fuzzer.launcher_script:
956+
logs.error('Fuzzers with launch script not supported yet.')
957+
return False
958+
959+
if fuzzer.builtin:
960+
logs.info(f'Fuzzer {fuzzer_name} is builtin, no setup required.')
961+
return True
962+
963+
fuzzer_directory: str = get_fuzzer_directory(fuzzer.name)
964+
965+
try:
966+
if not shell.remove_directory(fuzzer_directory, recreate=True):
967+
logs.error(f'Failed to clear fuzzer directory: {fuzzer_directory}')
968+
return False
969+
except Exception as e:
970+
logs.error(f'Error clearing fuzzer directory {fuzzer_directory}: {e}')
971+
return False
972+
973+
archive_path = os.path.join(fuzzer_directory, fuzzer.filename)
974+
if not blobs.read_blob_to_disk(fuzzer.blobstore_key, archive_path):
975+
logs.error('Failed to download fuzzer archive from blobstore: '
976+
f'{fuzzer.blobstore_key}')
977+
return False
978+
979+
try:
980+
with archive.open(archive_path) as reader:
981+
reader.extract_all(fuzzer_directory)
982+
except archive.ArchiveError as e:
983+
logs.error(f'Failed to unpack fuzzer archive {fuzzer.filename}: {e}')
984+
return False
985+
except Exception as e:
986+
logs.error(
987+
f'Unexpected error unpacking fuzzer archive {fuzzer.filename}: {e}')
988+
return False
989+
finally:
990+
if os.path.exists(archive_path):
991+
try:
992+
shell.remove_file(archive_path)
993+
except Exception as e:
994+
logs.warning(
995+
f'Failed to remove temporary archive file {archive_path}: {e}')
996+
997+
fuzzer_path = os.path.join(fuzzer_directory, fuzzer.executable_path)
998+
if not os.path.exists(fuzzer_path):
999+
logs.error(
1000+
f'Fuzzer executable {fuzzer.executable_path} not found in archive. '
1001+
'Check fuzzer configuration.')
1002+
return False
1003+
1004+
try:
1005+
os.chmod(fuzzer_path, _EXECUTABLE_PERMISSIONS)
1006+
except OSError as e:
1007+
logs.error(
1008+
f'Failed to set permissions on fuzzer executable {fuzzer_path}: {e}')
1009+
return False
1010+
1011+
return True

src/clusterfuzz/_internal/system/environment.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1001,6 +1001,16 @@ def set_bot_environment():
10011001
return True
10021002

10031003

1004+
def set_local_log_only():
1005+
"""Force logs to be local-only."""
1006+
1007+
# We set this to an empty string because currently the logs
1008+
# module does not correctly evaluate env vars
1009+
# (i.e., 'false' is evaluated to true).
1010+
set_value('LOG_TO_GCP', '')
1011+
set_value('LOG_TO_CONSOLE', 'True')
1012+
1013+
10041014
def set_tsan_max_history_size():
10051015
"""Sets maximum history size for TSAN tool."""
10061016
tsan_options = get_value('TSAN_OPTIONS')

0 commit comments

Comments
 (0)