[Centipede] Implement corpus prune for centipede on corpus prune task (#4707)

It implements corpus pruning for centipede fuzzer.

---------

Signed-off-by: Javan Lacerda <javanlacerda@google.com>
Co-authored-by: Ali Hijazi <ahijazi@google.com>
Co-authored-by: Jonathan Metzman <metzman@chromium.org>

Author: 3 people

src/clusterfuzz/_internal/bot/fuzzers/centipede/engine.py

@@ -551,8 +551,8 @@ def minimize_corpus(self,
       # Only remove this directory if it was created in this method.
       shutil.rmtree(full_corpus_workdir)
 
-    return engine.ReproduceResult(result.command, result.return_code,
-                                  result.time_executed, result.output)
+    return engine.FuzzResult(result.output, result.command, [], None,
+                             result.time_executed, result.timed_out)
 
   def _get_smallest_crasher(self, workdir_path):
     """Returns the path to the smallest crash in Centipede's |workdir_path|."""

src/clusterfuzz/_internal/bot/tasks/utasks/corpus_pruning_task.py

@@ -19,7 +19,6 @@
 import os
 import random
 import shutil
-from typing import Dict
 from typing import List
 import zipfile
 
@@ -301,8 +300,8 @@ def _cross_pollinate_other_fuzzer_corpuses(self):
             'Failed to unpack corpus backup from url %s.' % corpus_backup_url)
 
 
-class Runner:
-  """Runner for libFuzzer."""
+class BaseRunner:
+  """Base Runner"""
 
   def __init__(self, build_directory, context):
     self.build_directory = build_directory
@@ -312,12 +311,44 @@ def __init__(self, build_directory, context):
         self.build_directory, self.context.fuzz_target.binary)
     if not self.target_path:
       raise CorpusPruningError(
-          'Failed to get fuzzer path for %s.' % self.context.fuzz_target.binary)
-
+          f'Failed to get fuzzer path for {self.context.fuzz_target.binary}')
     self.fuzzer_options = options.get_fuzz_target_options(self.target_path)
 
-  def get_libfuzzer_flags(self):
-    """Get default libFuzzer options."""
+  def get_fuzzer_flags(self):
+    return []
+
+  def process_sanitizer_options(self):
+    """Process sanitizer options overrides."""
+    if not self.fuzzer_options:
+      return
+
+    # Only need to look as ASan, as that's what we prune with.
+    overrides = self.fuzzer_options.get_asan_options()
+    if not overrides:
+      return
+
+    asan_options = environment.get_memory_tool_options('ASAN_OPTIONS')
+    if not asan_options:
+      return
+    asan_options.update(overrides)
+    environment.set_memory_tool_options('ASAN_OPTIONS', asan_options)
+
+  def reproduce(self, input_path, arguments, max_time):
+    return self.context.engine.reproduce(self.target_path, input_path,
+                                         arguments, max_time)
+
+  def minimize_corpus(self, arguments, input_dirs, output_dir, reproducers_dir,
+                      max_time):
+    return self.context.engine.minimize_corpus(self.target_path, arguments,
+                                               input_dirs, output_dir,
+                                               reproducers_dir, max_time)
+
+
+class LibFuzzerRunner(BaseRunner):
+  """Runner for libFuzzer."""
+
+  def get_fuzzer_flags(self):
+    """Get default libFuzzer options for pruning."""
     rss_limit = RSS_LIMIT
     max_len = engine_common.CORPUS_INPUT_SIZE_LIMIT
     detect_leaks = 1
@@ -352,22 +383,6 @@ def get_libfuzzer_flags(self):
 
     return arguments.list()
 
-  def process_sanitizer_options(self):
-    """Process sanitizer options overrides."""
-    if not self.fuzzer_options:
-      return
-
-    # Only need to look at ASan, as that's what we prune with.
-    overrides = self.fuzzer_options.get_asan_options()
-    if not overrides:
-      return
-
-    asan_options = environment.get_memory_tool_options('ASAN_OPTIONS')
-    if not asan_options:
-      return
-    asan_options.update(overrides)
-    environment.set_memory_tool_options('ASAN_OPTIONS', asan_options)
-
   def reproduce(self, input_path, arguments, max_time):
     return self.context.engine.reproduce(self.target_path, input_path,
                                          arguments, max_time)
@@ -379,33 +394,81 @@ def minimize_corpus(self, arguments, input_dirs, output_dir, reproducers_dir,
                                                reproducers_dir, max_time)
 
 
-class CorpusPruner:
-  """Class that handles corpus pruning."""
+class CentipedeRunner(BaseRunner):
+  """Runner implementation for Centipede fuzzing engine."""
+
+
+class CorpusPrunerBase:
+  """Base class for corpus pruning that is engine‐agnostic."""
 
   def __init__(self, runner):
     self.runner = runner
-    self.context = self.runner.context
+    self.context = runner.context
+
+  def run(self, initial_corpus_path, minimized_corpus_path, bad_units_path):
+    """Running generic corpus prunning"""
+    if not shell.get_directory_file_count(initial_corpus_path):
+      # Empty corpus, nothing to do.
+      return None
+
+    # Unpack seed corpus if needed.
+    engine_common.unpack_seed_corpus_if_needed(
+        self.runner.target_path, initial_corpus_path, force_unpack=True)
+
+    environment.reset_current_memory_tool_options(
+        redzone_size=MIN_REDZONE, leaks=True)
+    self.runner.process_sanitizer_options()
+
+    additional_args = self.runner.get_fuzzer_flags()
+    logs.info('Running merge...')
+    try:
+      result = self.runner.minimize_corpus(
+          additional_args, [initial_corpus_path], minimized_corpus_path,
+          bad_units_path, CORPUS_PRUNING_TIMEOUT)
+    except TimeoutError as e:
+      raise CorpusPruningError(
+          'Corpus pruning timed out while minimizing corpus\n' + repr(e))
+    except engine.Error as e:
+      raise CorpusPruningError('Corpus pruning failed to minimize corpus\n' +
+                               repr(e))
+
+    symbolized_output = stack_symbolizer.symbolize_stacktrace(result.logs)
+
+    if not shell.get_directory_file_count(minimized_corpus_path):
+      raise CorpusPruningError('Corpus pruning failed to minimize corpus\n' +
+                               symbolized_output)
+
+    logs.info('Corpus merge finished successfully.', output=symbolized_output)
+    return result.stats
+
+  def process_bad_units(self, bad_units_path, quarantine_corpus_path):
+    del bad_units_path
+    del quarantine_corpus_path
+    return {}
+
+
+class LibFuzzerPruner(CorpusPrunerBase):
+  """
+  LibFuzzerPruner is a specialized pruner for libFuzzer that handles
+  quarantining of problematic units and related special cases.
+  """
 
   def _run_single_unit(self, unit_path):
-    """Run a single unit, and return the result."""
-    arguments = self.runner.get_libfuzzer_flags()
+    arguments = self.runner.get_fuzzer_flags()  # Expect libFuzzer flags.
     return self.runner.reproduce(unit_path, arguments, SINGLE_UNIT_TIMEOUT)
 
   def _quarantine_unit(self, unit_path, quarantine_corpus_path):
-    """Moves the given unit to the quarantine, and returns the path to the unit
-    in the quarantine."""
     quarantined_unit_path = os.path.join(quarantine_corpus_path,
                                          os.path.basename(unit_path))
     shutil.move(unit_path, quarantined_unit_path)
-
     return quarantined_unit_path
 
-  def process_bad_units(self, bad_units_path, quarantine_corpus_path
-                       ) -> Dict[str, uworker_msg_pb2.CrashInfo]:  # pylint: disable=no-member
-    """Process bad units found during merge."""
-    # TODO(ochang): A lot of this function is similar to parts of fuzz_task.
-    # Ideally fuzz_task can be refactored in a way that lets us share the common
-    # code.
+  def process_bad_units(self, bad_units_path, quarantine_corpus_path):
+    """
+    Process bad units by running each test case individually,
+    quarantining those that timeout, OOM, or crash due to memory sanitizer
+    errors.
+    """
     crashes = {}
 
     environment.reset_current_memory_tool_options(redzone_size=DEFAULT_REDZONE)
@@ -415,41 +478,36 @@ def process_bad_units(self, bad_units_path, quarantine_corpus_path
     corpus_file_paths = _get_corpus_file_paths(bad_units_path)
     num_bad_units = 0
 
-    # Run each corpus item individually.
     for i, unit_path in enumerate(corpus_file_paths, 1):
       if i % 100 == 0:
         logs.info('Up to %d' % i)
 
       unit_name = os.path.basename(unit_path)
       if unit_name.startswith('timeout-') or unit_name.startswith('oom-'):
-        # Don't waste time re-running timeout or oom testcases.
+        # Immediately quarantine timeouts/oom testcases.
         self._quarantine_unit(unit_path, quarantine_corpus_path)
         num_bad_units += 1
         continue
 
       try:
         result = self._run_single_unit(unit_path)
       except TimeoutError:
-        # Slow unit. Quarantine it.
         self._quarantine_unit(unit_path, quarantine_corpus_path)
         num_bad_units += 1
         continue
 
       if not crash_analyzer.is_memory_tool_crash(result.output):
-        # Didn't crash.
         continue
 
-      # Get memory tool crash information.
       state = stack_analyzer.get_crash_data(result.output, symbolize_flag=True)
 
-      # Crashed or caused a leak. Quarantine it.
+      # Quarantine the crashing unit.
       unit_path = self._quarantine_unit(unit_path, quarantine_corpus_path)
       num_bad_units += 1
 
       if crash_analyzer.ignore_stacktrace(state.crash_stacktrace):
         continue
 
-      # Local de-duplication.
       if state.crash_state not in crashes:
         security_flag = crash_analyzer.is_security_issue(
             state.crash_stacktrace, state.crash_type, state.crash_address)
@@ -460,49 +518,13 @@ def process_bad_units(self, bad_units_path, quarantine_corpus_path
             crash_stacktrace=state.crash_stacktrace,
             unit_path=unit_path,
             security_flag=security_flag)
-
     logs.info(
         f'Found {num_bad_units} bad units, {len(crashes)} unique crashes.')
     return crashes
 
-  def run(self, initial_corpus_path, minimized_corpus_path, bad_units_path):
-    """Run corpus pruning. Output result to directory."""
-    if not shell.get_directory_file_count(initial_corpus_path):
-      logs.info('Empty corpus, nothing to do.')
-      return None
 
-    # Set memory tool options and fuzzer arguments.
-    engine_common.unpack_seed_corpus_if_needed(
-        self.runner.target_path, initial_corpus_path, force_unpack=True)
-
-    environment.reset_current_memory_tool_options(
-        redzone_size=MIN_REDZONE, leaks=True)
-    self.runner.process_sanitizer_options()
-    additional_args = self.runner.get_libfuzzer_flags()
-
-    # Execute fuzzer with arguments for corpus pruning.
-    logs.info('Running merge...')
-    try:
-      result = self.runner.minimize_corpus(
-          additional_args, [initial_corpus_path], minimized_corpus_path,
-          bad_units_path, CORPUS_PRUNING_TIMEOUT)
-    except TimeoutError as e:
-      raise CorpusPruningError(
-          'Corpus pruning timed out while minimizing corpus\n' + repr(e))
-    except engine.Error as e:
-      raise CorpusPruningError('Corpus pruning failed to minimize corpus\n' +
-                               repr(e))
-
-    symbolized_output = stack_symbolizer.symbolize_stacktrace(result.logs)
-
-    # Sanity check that there are files in minimized corpus after merging.
-    if not shell.get_directory_file_count(minimized_corpus_path):
-      raise CorpusPruningError('Corpus pruning failed to minimize corpus\n' +
-                               symbolized_output)
-
-    logs.info('Corpus merge finished successfully.', output=symbolized_output)
-
-    return result.stats
+class CentipedePruner(CorpusPrunerBase):
+  """Centipede pruner."""
 
 
 class CrossPollinator:
@@ -597,6 +619,22 @@ def _record_cross_pollination_stats(output):
   client.insert([big_query.Insert(row=bigquery_row, insert_id=None)])
 
 
+def _get_pruner_and_runner(context):
+  """Get pruner and runner object acording with the FuzzTarget into the context
+  """
+  build_directory = environment.get_value('BUILD_DIR')
+  if context.fuzz_target.engine.lower() == 'libfuzzer':
+    runner = LibFuzzerRunner(build_directory, context)
+    pruner = LibFuzzerPruner(runner)
+  elif context.fuzz_target.engine.lower() == 'centipede':
+    runner = CentipedeRunner(build_directory, context)
+    pruner = CentipedePruner(runner)
+  else:
+    raise CorpusPruningError(
+        'Corpus pruner task does not support the given engine.')
+  return pruner, runner
+
+
 def do_corpus_pruning(uworker_input, context, revision) -> CorpusPruningResult:
   """Run corpus pruning."""
   # Set |FUZZ_TARGET| environment variable to help with unarchiving only fuzz
@@ -611,10 +649,8 @@ def do_corpus_pruning(uworker_input, context, revision) -> CorpusPruningResult:
       revision=revision, fuzz_target=context.fuzz_target.binary):
     raise CorpusPruningError('Failed to setup build.')
 
-  build_directory = environment.get_value('BUILD_DIR')
   start_time = datetime.datetime.utcnow()
-  runner = Runner(build_directory, context)
-  pruner = CorpusPruner(runner)
+  pruner, runner = _get_pruner_and_runner(context)
   fuzzer_binary_name = os.path.basename(runner.target_path)
 
   logs.info('Getting the initial corpus to process from GCS.')

src/clusterfuzz/_internal/tests/core/bot/fuzzers/centipede/centipede_engine_test.py

@@ -383,7 +383,6 @@ def test_crash_timeout(self):
         'slo',
         timeout_per_input=_TIMEOUT_PER_INPUT_TEST)
 
-  @unittest.skip('This test is failing, blocking deploy.')
   def test_minimize_corpus(self):
     """Tests minimizing a corpus."""
     unminimized_corpus = setup_testcase('unmin_corpus', self.test_paths)
@@ -402,7 +401,9 @@ def test_minimize_corpus(self):
     self.assertEqual(len(os.listdir(crash_corpus)), 1)
     crasher = os.path.join(crash_corpus, os.listdir(crash_corpus)[0])
     with open(crasher) as crasher_file:
-      self.assertEqual(crasher_file.read(), '?f???u???z?')
+      crash_content = crasher_file.read()
+      crasher_file.close()
+    self.assertEqual(crash_content, '?f???u???z?')
 
   def test_minimize_testcase(self):
     """Tests minimizing a testcase."""