[cifuzz] Add pruning task (#6188)

Fixes: #6064

1 files changed, 35 insertions, 8 deletions

diff --git a/infra/cifuzz/run_fuzzers.py b/infra/cifuzz/run_fuzzers.py
index 93a30a37..79cf432c 100644
--- a/infra/cifuzz/run_fuzzers.py
+++ b/infra/cifuzz/run_fuzzers.py
@@ -166,6 +166,26 @@ class BaseFuzzTargetRunner:
     return bug_found
 
 
+class PruneTargetRunner(BaseFuzzTargetRunner):
+  """Runner that prunes corpora."""
+
+  @property
+  def quit_on_bug_found(self):
+    return False
+
+  def run_fuzz_target(self, fuzz_target_obj):
+    """Prunes with |fuzz_target_obj| and returns the result."""
+    result = fuzz_target_obj.prune()
+    logging.debug('Corpus path contents: %s.', os.listdir(result.corpus_path))
+    self.clusterfuzz_deployment.upload_corpus(fuzz_target_obj.target_name,
+                                              result.corpus_path)
+    return result
+
+  def cleanup_after_fuzz_target_run(self, fuzz_target_obj):  # pylint: disable=no-self-use
+    """Cleans up after pruning with |fuzz_target_obj|."""
+    fuzz_target_obj.free_disk_if_needed()
+
+
 class CoverageTargetRunner(BaseFuzzTargetRunner):
   """Runner that runs the 'coverage' command."""
 
@@ -224,8 +244,9 @@ class BatchFuzzTargetRunner(BaseFuzzTargetRunner):
   def run_fuzz_target(self, fuzz_target_obj):
     """Fuzzes with |fuzz_target_obj| and returns the result."""
     result = fuzz_target_obj.fuzz()
-    logging.debug('corpus_path: %s', os.listdir(result.corpus_path))
-    self.clusterfuzz_deployment.upload_corpus(fuzz_target_obj.target_name)
+    logging.debug('Corpus path contents: %s.', os.listdir(result.corpus_path))
+    self.clusterfuzz_deployment.upload_corpus(fuzz_target_obj.target_name,
+                                              result.corpus_path)
     return result
 
   def cleanup_after_fuzz_target_run(self, fuzz_target_obj):
@@ -255,15 +276,21 @@ class BatchFuzzTargetRunner(BaseFuzzTargetRunner):
     return result
 
 
+_RUN_FUZZERS_MODE_RUNNER_MAPPING = {
+    'batch': BatchFuzzTargetRunner,
+    'coverage': CoverageTargetRunner,
+    'prune': PruneTargetRunner,
+    'ci': CiFuzzTargetRunner,
+}
+
+
 def get_fuzz_target_runner(config):
   """Returns a fuzz target runner object based on the run_fuzzers_mode of
   |config|."""
-  logging.info('RUN_FUZZERS_MODE is: %s', config.run_fuzzers_mode)
-  if config.run_fuzzers_mode == 'batch':
-    return BatchFuzzTargetRunner(config)
-  if config.run_fuzzers_mode == 'coverage':
-    return CoverageTargetRunner(config)
-  return CiFuzzTargetRunner(config)
+  runner = _RUN_FUZZERS_MODE_RUNNER_MAPPING[config.run_fuzzers_mode](config)
+  logging.info('RUN_FUZZERS_MODE is: %s. Runner: %s.', config.run_fuzzers_mode,
+               runner)
+  return runner
 
 
 def run_fuzzers(config):  # pylint: disable=too-many-locals