diff --git a/README.md b/README.md index e36daee..7b2f2dc 100644 --- a/README.md +++ b/README.md @@ -43,18 +43,21 @@ Optionally, instruction context messages may also be left in the original Protob To run the test suite, use the following command: ```sh -solana-test-suite run-tests --input-dir --solana-target --target [--target ...] --output-dir --num-processes --chunk-size [--randomize-output-buffer] +solana-test-suite run-tests --input-dir --solana-target --target [--target ...] --output-dir --num-processes --chunk-size [--randomize-output-buffer] ``` +You can provide both `InstrContext` and `InstrFixture` within `--input-dir` - parsing is taken care of depending on the file extension `.bin` for `InstrContext` and `.fix` for `InstrFixture`. + | Argument | Description | |-----------------|-----------------------------------------------------------------------------------------------------| -| `--input-dir` | Input directory containing instruction context messages | +| `--input-dir` | Input directory containing instruction context or fixture messages | | `--solana-target` | Path to Solana Agave shared object (.so) target file | | `--target` | Additional shared object (.so) target file paths | | `--output-dir` | Log output directory for test results | | `--num-processes` | Number of processes to use | | `--randomize-output-buffer`| Randomizes bytes in output buffer before shared library execution | | `--chunk-size` | Number of test results per log file | +| `--verbose` | Verbose output: log failed test cases | **Note:** Each `.so` target file name should be unique. @@ -111,16 +114,22 @@ solana-test-suite minimize-tests --input-dir --solana-target --solana-target --output-dir --num-processes [--readable] +solana-test-suite create-fixtures --input-dir --solana-target --target [--target ...] --output-dir --num-processes [--readable] [--keep-passing] [--group-by-program] ``` +You have an additional option to produce fixtures for only passing test cases (makes it easier to produce fixtures from larger batches of new-passing mismatches). + + | Argument | Description | |-----------------|-----------------------------------------------------------------------------------------------------| | `--input-dir` | Input directory containing instruction context messages | | `--solana-target` | Path to Solana Agave shared object (.so) target file | +| `--target` | Shared object (.so) target file paths (pairs with `--keep-passing`) | `--output-dir` | Instruction fixtures dumping directory | | `--num-processes` | Number of processes to use | | `--readable` | Output fixtures in human-readable format | +| `--keep-passing` | Only keep passing test cases | +| `--group-by-program` | Group fixture output by program type | ### Create Instruction Context from Fixtures @@ -137,23 +146,6 @@ solana-test-suite instr-from-fixtures --input-dir --solana-target --target [--target ...] --output-dir --num-iterations --num-processes [--randomize-output-buffer] -``` - -| Argument | Description | -|----------------------------|--------------------------------------------------------------------------------------------------------------------------| -| `--input-dir` | Input directory containing instruction context messages | -| `--target` | Additional shared object (.so) target file paths | -| `--output-dir` | Log output directory for test results | -| `--num-iterations` | Number of consistency iterations to run for each library | -| `--num-processes` | Number of processes to use | -| `--randomize-output-buffer`| Randomizes bytes in output buffer before shared library execution | - ### Uninstalling diff --git a/src/test_suite/constants.py b/src/test_suite/constants.py index cd1a0cc..39f42ef 100644 --- a/src/test_suite/constants.py +++ b/src/test_suite/constants.py @@ -3,3 +3,17 @@ # Output buffer size OUTPUT_BUFFER_SIZE = 100 * 1024 * 1024 + +# Native program mappings +NATIVE_PROGRAM_MAPPING = { + "11111111111111111111111111111111": "system", + "Config1111111111111111111111111111111111111": "config", + "ComputeBudget111111111111111111111111111111": "compute-budget", + "Stake11111111111111111111111111111111111111": "stake", + "Vote111111111111111111111111111111111111111": "vote", + "AddressLookupTab1e1111111111111111111111111": "address-lookup-table", + "BPFLoader1111111111111111111111111111111111": "bpf-loader-v1", + "BPFLoader2111111111111111111111111111111111": "bpf-loader-v2", + "BPFLoaderUpgradeab1e11111111111111111111111": "bpf-loader-upgradeable-v1", + "LoaderV411111111111111111111111111111111111": "bpf-loader-v4", +} diff --git a/src/test_suite/fixture_utils.py b/src/test_suite/fixture_utils.py index 49a4875..8b32221 100644 --- a/src/test_suite/fixture_utils.py +++ b/src/test_suite/fixture_utils.py @@ -1,57 +1,59 @@ +import fd58 from test_suite.codec_utils import encode_input, encode_output -from test_suite.multiprocessing_utils import prune_execution_result +from test_suite.constants import NATIVE_PROGRAM_MAPPING +from test_suite.multiprocessing_utils import ( + build_test_results, + read_instr, + process_single_test_case, + prune_execution_result, +) import test_suite.globals as globals import test_suite.invoke_pb2 as pb from google.protobuf import text_format from pathlib import Path -def create_fixture( - file_serialized_instruction_context: tuple[str, dict], - file_serialized_instruction_effects: tuple[str, dict[str, str | None]], -) -> tuple[str, str | None]: +def create_fixture(test_file: Path) -> int: """ Create instruction fixture for an instruction context and effects. Args: - - file_serialized_instruction_context (tuple[str, str]): Tuple of file stem and serialized instruction context. - - file_serialized_instruction_effects (tuple[str, dict[str, str | None]]): Tuple of file stem and dictionary of target library names and serialized instruction effects. + - test_file (Path): Path to the file containing serialized instruction contexts Returns: - - tuple[str, str | None]: Tuple of file stem and instruction fixture. + - int: 1 on success, 0 on failure """ + serialized_instr_context = read_instr(test_file) + results = process_single_test_case(serialized_instr_context) + pruned_results = prune_execution_result(serialized_instr_context, results) - file_stem, serialized_instruction_context = file_serialized_instruction_context - file_stem_2, serialized_instruction_effects = file_serialized_instruction_effects - - assert file_stem == file_stem_2, f"{file_stem} != {file_stem_2}" - - # Both instruction context and instruction effects should not be None - if serialized_instruction_context is None or serialized_instruction_effects is None: - return file_stem, None + # This is only relevant when you gather results for multiple targets + if globals.only_keep_passing: + status, _ = build_test_results(pruned_results) + if status != 1: + return 0 - _, targets_to_serialized_pruned_instruction_effects = prune_execution_result( - file_serialized_instruction_context, file_serialized_instruction_effects - ) + if pruned_results is None: + return 0 - pruned_instruction_effects = targets_to_serialized_pruned_instruction_effects[ - globals.solana_shared_library - ] + serialized_instr_effects = pruned_results[globals.solana_shared_library] - if pruned_instruction_effects is None: - return file_stem, None + if serialized_instr_context is None or serialized_instr_effects is None: + return 0 # Create instruction fixture instr_context = pb.InstrContext() - instr_context.ParseFromString(serialized_instruction_context) + instr_context.ParseFromString(serialized_instr_context) instr_effects = pb.InstrEffects() - instr_effects.ParseFromString(pruned_instruction_effects) + instr_effects.ParseFromString(serialized_instr_effects) fixture = pb.InstrFixture() fixture.input.MergeFrom(instr_context) fixture.output.MergeFrom(instr_effects) - return file_stem, fixture.SerializeToString(deterministic=True) + return write_fixture_to_disk( + test_file.stem, fixture.SerializeToString(deterministic=True) + ) def write_fixture_to_disk(file_stem: str, serialized_instruction_fixture: str) -> int: @@ -61,7 +63,6 @@ def write_fixture_to_disk(file_stem: str, serialized_instruction_fixture: str) - Args: - file_stem (str): File stem - - serialized_instruction_fixture (str): Serialized instruction fixture Returns: - int: 0 on failure, 1 on success @@ -69,6 +70,15 @@ def write_fixture_to_disk(file_stem: str, serialized_instruction_fixture: str) - if serialized_instruction_fixture is None: return 0 + output_dir = globals.output_dir + + if globals.organize_fixture_dir: + instr_fixture = pb.InstrFixture() + instr_fixture.ParseFromString(serialized_instruction_fixture) + program_type = get_program_type(instr_fixture) + output_dir = output_dir / program_type + output_dir.mkdir(parents=True, exist_ok=True) + if globals.readable: # Deserialize fixture instr_fixture = pb.InstrFixture() @@ -86,12 +96,12 @@ def write_fixture_to_disk(file_stem: str, serialized_instruction_fixture: str) - instr_fixture.input.CopyFrom(instr_context) instr_fixture.output.CopyFrom(instr_effects) - with open(globals.output_dir / (file_stem + ".fix.txt"), "w") as f: + with open(output_dir / (file_stem + ".fix.txt"), "w") as f: f.write( text_format.MessageToString(instr_fixture, print_unknown_fields=False) ) else: - with open(f"{globals.output_dir}/{file_stem}.fix", "wb") as f: + with open(output_dir / (file_stem + ".fix"), "wb") as f: f.write(serialized_instruction_fixture) return 1 @@ -118,3 +128,33 @@ def extract_instr_context_from_fixture(fixture_file: Path): return 0 return 1 + + +def get_program_type(instr_fixture: pb.InstrFixture) -> str: + """ + Get the program type based on the program / loader id. + + Args: + - fixture (pb.InstrFixture): Instruction fixture + + Returns: + - str | None: Program type (unknown if not found) + """ + # Check if the program type can be deduced from program_id + program_id = fd58.enc32(instr_fixture.input.program_id).decode() + + program_type = NATIVE_PROGRAM_MAPPING.get(program_id, None) + if program_type: + return program_type + + # Use the program_id owner instead (loader_id may not be reliable) + for account_state in instr_fixture.input.accounts: + if account_state.address == instr_fixture.input.program_id: + program_type = NATIVE_PROGRAM_MAPPING.get( + fd58.enc32(account_state.owner).decode(), "unknown" + ) + if program_type != "unknown": + program_type += "-programs" + return program_type + + return "unknown" diff --git a/src/test_suite/globals.py b/src/test_suite/globals.py index 2a43013..587483c 100644 --- a/src/test_suite/globals.py +++ b/src/test_suite/globals.py @@ -21,3 +21,9 @@ # (For fixtures) Whether to output in human-readable format readable = False + +# (For fixtures) Whether to organize fixtures by program type +organize_fixture_dir = False + +# (For fixtures) Whether to only keep passing tests +only_keep_passing = False diff --git a/src/test_suite/minimize_utils.py b/src/test_suite/minimize_utils.py index e463759..bcc7cc3 100644 --- a/src/test_suite/minimize_utils.py +++ b/src/test_suite/minimize_utils.py @@ -2,7 +2,7 @@ import test_suite.invoke_pb2 as pb import test_suite.globals as globals from test_suite.multiprocessing_utils import ( - generate_test_case, + read_instr, process_instruction, ) @@ -18,7 +18,7 @@ def minimize_single_test_case(test_file: Path) -> int: Returns: int: 0 on failure, 1 on success """ - _, serialized_instruction_context = generate_test_case(test_file) + _, serialized_instruction_context = read_instr(test_file) # Skip if input is invalid if serialized_instruction_context is None: diff --git a/src/test_suite/multiprocessing_utils.py b/src/test_suite/multiprocessing_utils.py index 96a47e5..4915966 100644 --- a/src/test_suite/multiprocessing_utils.py +++ b/src/test_suite/multiprocessing_utils.py @@ -2,7 +2,7 @@ from test_suite.constants import OUTPUT_BUFFER_SIZE import test_suite.invoke_pb2 as pb from test_suite.codec_utils import encode_input, encode_output, decode_input -from test_suite.validation_utils import check_account_unchanged, is_valid +from test_suite.validation_utils import check_account_unchanged import ctypes from ctypes import c_uint64, c_int, POINTER, Structure from pathlib import Path @@ -11,10 +11,6 @@ import os -def lazy_starmap(args, function): - return function(*args) - - def process_instruction( library: ctypes.CDLL, serialized_instruction_context: str ) -> pb.InstrEffects | None: @@ -61,15 +57,15 @@ def process_instruction( return output_object -def generate_test_case(test_file: Path) -> tuple[str, str | None]: +def read_instr(test_file: Path) -> str | None: """ Reads in test files and generates an InstrContext Protobuf object for a test case. Args: - - test_file (Path): Path to the file containing serialized instruction contexts. + - test_file (Path): Path to the instruction context message. Returns: - - tuple[str, str | None]: Tuple of file stem and serialized instruction context, if exists. + - str | None: Serialized instruction context, or None if reading failed. """ # Try to read in first as binary-encoded Protobuf messages try: @@ -91,13 +87,46 @@ def generate_test_case(test_file: Path) -> tuple[str, str | None]: if instruction_context is None: # Unreadable file, skip it - return test_file.stem, None + return None # Discard unknown fields instruction_context.DiscardUnknownFields() # Serialize instruction context to string (pickleable) - return test_file.stem, instruction_context.SerializeToString(deterministic=True) + return instruction_context.SerializeToString(deterministic=True) + + +def read_fixture(fixture_file: Path) -> str | None: + """ + Same as read_instr, but for InstrFixture protobuf messages. + + DOES NOT SUPPORT HUMAN READABLE MESSAGES!!! + + Args: + - fixture_file (Path): Path to the instruction fixture message. + + Returns: + - str | None: Serialized instruction fixture, or None if reading failed. + """ + # Try to read in first as binary-encoded Protobuf messages + try: + # Read in binary Protobuf messages + with open(fixture_file, "rb") as f: + instruction_fixture = pb.InstrFixture() + instruction_fixture.ParseFromString(f.read()) + except: + # Unable to read message, skip and continue + instruction_fixture = None + + if instruction_fixture is None: + # Unreadable file, skip it + return None + + # Discard unknown fields + instruction_fixture.DiscardUnknownFields() + + # Serialize instruction fixture to string (pickleable) + return instruction_fixture.SerializeToString(deterministic=True) def decode_single_test_case(test_file: Path) -> int: @@ -110,7 +139,7 @@ def decode_single_test_case(test_file: Path) -> int: Returns: - int: 1 if successfully decoded and written, 0 if skipped. """ - _, serialized_instruction_context = generate_test_case(test_file) + serialized_instruction_context = read_instr(test_file) # Skip if input is invalid if serialized_instruction_context is None: @@ -129,24 +158,21 @@ def decode_single_test_case(test_file: Path) -> int: def process_single_test_case( - file_stem: str, serialized_instruction_context: str | None -) -> tuple[str, dict[str, str | None] | None]: + serialized_instruction_context: str | None, +) -> dict[str, str | None] | None: """ Process a single execution context (file, serialized instruction context) through - all target libraries and returns serialized instruction effects. This - function is called by processes. + all target libraries and returns serialized instruction effects. Args: - - file_stem (str): Stem of file containing serialized instruction context. - serialized_instruction_context (str | None): Serialized instruction context. Returns: - - tuple[str, dict[str, str | None] | None]: Tuple of file stem and dictionary of target library names - and instruction effects. + - dict[str, str | None] | None: Dictionary of target library names and instruction effects. """ # Mark as skipped if instruction context doesn't exist if serialized_instruction_context is None: - return file_stem, None + return None # Execute test case on each target library results = {} @@ -161,7 +187,7 @@ def process_single_test_case( ) results[target] = result - return file_stem, results + return results def merge_results_over_iterations(results: tuple) -> tuple[str, dict]: @@ -193,31 +219,25 @@ def merge_results_over_iterations(results: tuple) -> tuple[str, dict]: def prune_execution_result( - file_serialized_instruction_context: tuple[str, str], - file_serialized_instruction_effects: tuple[str, dict[str, str | None]], -) -> tuple[str, dict]: + serialized_instruction_context: str, + targets_to_serialized_instruction_effects: dict[str, str | None], +) -> dict[str, str | None] | None: """ Prune execution result to only include actually modified accounts. Args: - - file_serialized_instruction_context (tuple[str, str]): Tuple of file stem and serialized instruction context. - - file_serialized_instruction_effects (tuple[str, dict[str, str | None]]): Tuple of file stem and dictionary of target library names and serialized instruction effects. + - serialized_instruction_context (str): Serialized instruction context. + - serialized_instruction_effects (dict[str, str | None]): Dictionary of target library names and serialized instruction effects. Returns: - - tuple[str, dict]: Tuple of file stem and serialized pruned instruction effects for each target. + - dict[str, str | None] | None: Serialized pruned instruction effects for each target. """ - file_stem, serialized_instruction_context = file_serialized_instruction_context if serialized_instruction_context is None: - return file_stem, None + return None instruction_context = pb.InstrContext() instruction_context.ParseFromString(serialized_instruction_context) - file_stem_2, targets_to_serialized_instruction_effects = ( - file_serialized_instruction_effects - ) - assert file_stem == file_stem_2, f"{file_stem}, {file_stem_2}" - targets_to_serialized_pruned_instruction_effects = {} for ( target, @@ -249,7 +269,7 @@ def prune_execution_result( instruction_effects.SerializeToString(deterministic=True) ) - return file_stem, targets_to_serialized_pruned_instruction_effects + return targets_to_serialized_pruned_instruction_effects def check_consistency_in_results(file_stem: str, results: dict) -> dict[str, bool]: @@ -302,26 +322,22 @@ def check_consistency_in_results(file_stem: str, results: dict) -> dict[str, boo return results_per_target -def build_test_results( - file_stem: str, results: dict[str, str | None] -) -> tuple[str, int, dict | None]: +def build_test_results(results: dict[str, str | None]) -> tuple[int, dict | None]: """ Build a single result of single test execution and returns whether the test passed or failed. Args: - - file_stem (str): File stem of the test case. - results (dict[str, str | None]): Dictionary of target library names and serialized instruction effects. Returns: - - tuple[str, int, dict | None]: Tuple of: - - File stem + - tuple[int, dict | None]: Tuple of: - 1 if passed, -1 if failed, 0 if skipped - Dictionary of target library names and file-dumpable serialized instruction effects """ # If no results or Agave rejects input, mark case as skipped if results is None: # Mark as skipped (0) - return file_stem, 0, None + return 0, None outputs = {target: "None\n" for target in results} @@ -345,7 +361,7 @@ def build_test_results( ) # 1 = passed, -1 = failed - return file_stem, 1 if test_case_passed else -1, outputs + return 1 if test_case_passed else -1, outputs def initialize_process_output_buffers(randomize_output_buffer=False): @@ -410,7 +426,14 @@ def run_test(test_file: Path) -> tuple[str, int, dict | None]: - 1 if passed, -1 if failed, 0 if skipped - Dictionary of target library names and file-dumpable serialized instruction effects """ - test_case = generate_test_case(test_file) - results = process_single_test_case(*test_case) - pruned_results = prune_execution_result(test_case, results) - return build_test_results(*pruned_results) + # Process fixtures through this entrypoint as well + if test_file.suffix == ".fix": + fixture = pb.InstrFixture() + serialized_fixture = read_fixture(test_file) + fixture.MergeFromString(serialized_fixture) + serialized_instr_context = fixture.input.SerializeToString(deterministic=True) + else: + serialized_instr_context = read_instr(test_file) + results = process_single_test_case(serialized_instr_context) + pruned_results = prune_execution_result(serialized_instr_context, results) + return test_file.stem, *build_test_results(pruned_results) diff --git a/src/test_suite/test_suite.py b/src/test_suite/test_suite.py index c0d4920..bbfd116 100644 --- a/src/test_suite/test_suite.py +++ b/src/test_suite/test_suite.py @@ -6,24 +6,19 @@ import ctypes from multiprocessing import Pool from pathlib import Path -from test_suite.constants import LOG_FILE_SEPARATOR_LENGTH +from test_suite.constants import LOG_FILE_SEPARATOR_LENGTH, NATIVE_PROGRAM_MAPPING from test_suite.fixture_utils import ( create_fixture, extract_instr_context_from_fixture, - write_fixture_to_disk, ) import test_suite.invoke_pb2 as pb from test_suite.codec_utils import encode_output from test_suite.minimize_utils import minimize_single_test_case from test_suite.multiprocessing_utils import ( - check_consistency_in_results, decode_single_test_case, - generate_test_case, + read_instr, initialize_process_output_buffers, - lazy_starmap, - merge_results_over_iterations, process_instruction, - process_single_test_case, prune_execution_result, get_feature_pool, run_test, @@ -55,7 +50,7 @@ def exec_instr( help="Randomizes bytes in output buffer before shared library execution", ), ): - _, instruction_context = generate_test_case(file) + instruction_context = read_instr(file) assert instruction_context is not None, f"Unable to read {file.name}" # Initialize output buffers and shared library @@ -73,9 +68,9 @@ def exec_instr( instruction_effects = instruction_effects.SerializeToString(deterministic=True) # Prune execution results - _, pruned_instruction_effects = prune_execution_result( - (file.stem, instruction_context), - (file.stem, {shared_library: instruction_effects}), + pruned_instruction_effects = prune_execution_result( + instruction_context, + {shared_library: instruction_effects}, ) parsed_instruction_effects = pb.InstrEffects() parsed_instruction_effects.ParseFromString( @@ -108,131 +103,11 @@ def debug_instr( print(f"Processing {file.name}...") # Decode the file and pass it into GDB - _, instruction_context = generate_test_case(file) + instruction_context = read_instr(file) assert instruction_context is not None, f"Unable to read {file.name}" debug_host(shared_library, instruction_context, gdb=debugger) -@app.command() -def check_consistency( - input_dir: Path = typer.Option( - Path("corpus8"), - "--input-dir", - "-i", - help="Input directory containing instruction context messages", - ), - shared_libraries: List[Path] = typer.Option( - [], "--target", "-t", help="Shared object (.so) target file paths" - ), - output_dir: Path = typer.Option( - Path("consistency_results"), - "--output-dir", - "-o", - help="Output directory for test results", - ), - num_iterations: int = typer.Option( - 2, - "--num-iterations", - "-n", - help="Number of consistency iterations to run for each library", - ), - num_processes: int = typer.Option( - 4, "--num-processes", "-p", help="Number of processes to use" - ), - randomize_output_buffer: bool = typer.Option( - False, - "--randomize-output-buffer", - "-r", - help="Randomizes bytes in output buffer before shared library execution", - ), -): - # Initialize globals - globals.output_dir = output_dir - globals.n_iterations = num_iterations - - # Create the output directory, if necessary - if globals.output_dir.exists(): - shutil.rmtree(globals.output_dir) - globals.output_dir.mkdir(parents=True, exist_ok=True) - - # Generate the test cases in parallel from files on disk - print("Reading test files...") - with Pool(processes=num_processes) as pool: - execution_contexts = pool.map(generate_test_case, input_dir.iterdir()) - - results_per_iteration = [] - for iteration in range(globals.n_iterations): - print(f"Starting iteration {iteration}...") - - # Use the target libraries global map to store shared libraries - for target in shared_libraries: - lib = ctypes.CDLL(target) - lib.sol_compat_init() - globals.target_libraries[target] = lib - - # Initialize the libraries for each iteration - for iteration in range(globals.n_iterations): - # Make output directory - (globals.output_dir / target.stem / str(iteration)).mkdir( - parents=True, exist_ok=True - ) - - # Process the test cases in parallel through shared libraries for n interations - print("Executing tests...") - with Pool( - processes=num_processes, - initializer=initialize_process_output_buffers, - initargs=(randomize_output_buffer,), - ) as pool: - execution_results = pool.starmap( - process_single_test_case, execution_contexts - ) - results_per_iteration.append(execution_results) - - print("Cleaning up...") - for target in shared_libraries: - globals.target_libraries[target].sol_compat_fini() - - # Build the results properly - with Pool(processes=num_processes) as pool: - execution_results = pool.map( - merge_results_over_iterations, zip(*results_per_iteration) - ) - - # Process the test results in parallel - print("Building test results...") - with Pool(processes=num_processes) as pool: - test_case_results = pool.starmap( - check_consistency_in_results, execution_results - ) - - # Compute per-library results - library_results = {} - for library in globals.target_libraries: - library_results[library] = {"passed": 0, "failed": 0, "skipped": 0} - - # Build the results - for result in test_case_results: - for library, outcome in result.items(): - library_results[library]["passed"] += outcome == 1 - library_results[library]["failed"] += outcome == -1 - library_results[library]["skipped"] += outcome == 0 - - peak_memory_usage_kb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss - print(f"Peak Memory Usage: {peak_memory_usage_kb / 1024} MB") - - print("-" * LOG_FILE_SEPARATOR_LENGTH) - - for library in globals.target_libraries: - results = library_results[library] - print(f"{library} results") - print(f"Total test cases: {sum(results.values())}") - print( - f"Passed: {results['passed']}, Failed: {results['failed']}, Skipped: {results['skipped']}" - ) - print("-" * LOG_FILE_SEPARATOR_LENGTH) - - @app.command() def minimize_tests( input_dir: Path = typer.Option( @@ -317,20 +192,21 @@ def instr_from_fixtures( shutil.rmtree(globals.output_dir) globals.output_dir.mkdir(parents=True, exist_ok=True) - num_test_cases = len(list(input_dir.iterdir())) + test_cases = list(input_dir.iterdir()) + num_test_cases = len(test_cases) print("Converting to InstrContext...") - execution_contexts = [] + results = [] with Pool(processes=num_processes) as pool: for result in tqdm.tqdm( - pool.imap(extract_instr_context_from_fixture, input_dir.iterdir()), + pool.imap(extract_instr_context_from_fixture, test_cases), total=num_test_cases, ): - execution_contexts.append(result) + results.append(result) print("-" * LOG_FILE_SEPARATOR_LENGTH) - print(f"{len(execution_contexts)} total files seen") - print(f"{sum(execution_contexts)} files successfully written") + print(f"{len(results)} total files seen") + print(f"{sum(results)} files successfully written") @app.command() @@ -347,6 +223,12 @@ def create_fixtures( "-s", help="Solana (or ground truth) shared object (.so) target file path", ), + shared_libraries: List[Path] = typer.Option( + [], + "--target", + "-t", + help="Shared object (.so) target file paths (pairs with --keep-passing)", + ), output_dir: Path = typer.Option( Path("test_fixtures"), "--output-dir", @@ -359,11 +241,22 @@ def create_fixtures( readable: bool = typer.Option( False, "--readable", "-r", help="Output fixtures in human-readable format" ), + only_keep_passing: bool = typer.Option( + False, "--keep-passing", "-k", help="Only keep passing test cases" + ), + organize_fixture_dir: bool = typer.Option( + False, "--group-by-program", "-g", help="Group fixture output by program type" + ), ): + # Add Solana library to shared libraries + shared_libraries = [solana_shared_library] + shared_libraries + # Specify globals globals.output_dir = output_dir globals.solana_shared_library = solana_shared_library globals.readable = readable + globals.only_keep_passing = only_keep_passing + globals.organize_fixture_dir = organize_fixture_dir # Create the output directory, if necessary if globals.output_dir.exists(): @@ -371,58 +264,25 @@ def create_fixtures( globals.output_dir.mkdir(parents=True, exist_ok=True) # Initialize shared library - globals.solana_shared_library = solana_shared_library - lib = ctypes.CDLL(solana_shared_library) - lib.sol_compat_init() - globals.target_libraries[solana_shared_library] = lib + for target in shared_libraries: + # Load in and initialize shared libraries + lib = ctypes.CDLL(target) + lib.sol_compat_init() + globals.target_libraries[target] = lib - num_test_cases = len(list(input_dir.iterdir())) + test_cases = list(input_dir.iterdir()) + num_test_cases = len(test_cases) # Generate the test cases in parallel from files on disk - print("Reading test files...") - execution_contexts = [] - with Pool(processes=num_processes) as pool: - for result in tqdm.tqdm( - pool.imap(generate_test_case, input_dir.iterdir()), total=num_test_cases - ): - execution_contexts.append(result) - - # Process the test cases in parallel through shared libraries - print("Executing tests...") - execution_results = [] + print("Creating fixtures...") + write_results = [] with Pool( processes=num_processes, initializer=initialize_process_output_buffers ) as pool: for result in tqdm.tqdm( pool.imap( - functools.partial(lazy_starmap, function=process_single_test_case), - execution_contexts, - ), - total=num_test_cases, - ): - execution_results.append(result) - - # Prune effects and create fixtures - print("Creating fixtures...") - execution_fixtures = [] - with Pool(processes=num_processes) as pool: - for result in tqdm.tqdm( - pool.imap( - functools.partial(lazy_starmap, function=create_fixture), - zip(execution_contexts, execution_results), - ), - total=num_test_cases, - ): - execution_fixtures.append(result) - - # Write fixtures to disk - print("Writing results to disk...") - write_results = [] - with Pool(processes=num_processes) as pool: - for result in tqdm.tqdm( - pool.imap( - functools.partial(lazy_starmap, function=write_fixture_to_disk), - execution_fixtures, + create_fixture, + test_cases, ), total=num_test_cases, ): @@ -430,7 +290,8 @@ def create_fixtures( # Clean up print("Cleaning up...") - lib.sol_compat_fini() + for target in shared_libraries: + globals.target_libraries[target].sol_compat_fini() print("-" * LOG_FILE_SEPARATOR_LENGTH) print(f"{len(write_results)} total files seen") @@ -443,7 +304,7 @@ def run_tests( Path("corpus8"), "--input-dir", "-i", - help="Input directory containing instruction context messages", + help="Input directory containing instruction context or fixture messages", ), solana_shared_library: Path = typer.Option( Path("impl/lib/libsolfuzz_agave_v2.0.so"), @@ -502,7 +363,8 @@ def run_tests( log_dir = globals.output_dir / target.stem log_dir.mkdir(parents=True, exist_ok=True) - num_test_cases = len(list(input_dir.iterdir())) + test_cases = list(input_dir.iterdir()) + num_test_cases = len(test_cases) # Process the test results in parallel print("Running tests...") @@ -513,7 +375,7 @@ def run_tests( initargs=(randomize_output_buffer,), ) as pool: for result in tqdm.tqdm( - pool.imap(run_test, input_dir.iterdir()), + pool.imap(run_test, test_cases), total=num_test_cases, ): test_case_results.append(result)