Revert "Revert "[CIFuzz] Affected fuzzers (#3450)" (#3488)" (#3496)

This reverts commit e58ee49e2029a2ddaf3a801b1332b1677e6854f8.

2 files changed, 131 insertions, 11 deletions

diff --git a/infra/cifuzz/cifuzz.py b/infra/cifuzz/cifuzz.py
index d0919291..abf72237 100644
--- a/infra/cifuzz/cifuzz.py
+++ b/infra/cifuzz/cifuzz.py
@@ -163,6 +163,8 @@ def build_fuzzers(project_name,
   if helper.docker_run(command):
     logging.error('Building fuzzers failed.')
     return False
+  remove_unaffected_fuzzers(project_name, out_dir,
+                            build_repo_manager.get_git_diff(), src_in_docker)
   return True
 
 
@@ -293,23 +295,23 @@ def get_target_coverage_report(latest_cov_info, target_name):
   return get_json_from_url(target_url)
 
 
-def get_files_covered_by_target(latest_cov_info, target_name,
-                                oss_fuzz_project_base):
+def get_files_covered_by_target(latest_cov_info, target_name, src_in_docker):
   """Gets a list of files covered by the specific fuzz target.
 
   Args:
     latest_cov_info: A dict containing a project's latest cov report info.
     target_name: The name of the fuzz target whose coverage is requested.
-    oss_fuzz_project_base: The location where OSS-Fuzz project is cloned to for
-      the projects build.
+    src_in_docker: The location of the source dir in the docker image.
 
   Returns:
     A list of files that the fuzzer covers or None.
 
   Raises:
-    ValueError: When the oss_fuzz_project_base is not defined.
+    ValueError: When the src_in_docker is not defined.
   """
-  if not oss_fuzz_project_base:
+  if not src_in_docker:
+    logging.error('Project souce location in docker is not found.'
+                  'Can\'t get coverage information from OSS-Fuzz.')
     return None
   target_cov = get_target_coverage_report(latest_cov_info, target_name)
   if not target_cov:
@@ -319,25 +321,80 @@ def get_files_covered_by_target(latest_cov_info, target_name,
     logging.info('No files found in coverage report.')
     return None
 
-  # Cases like curl there is /src/curl and /src/curl_fuzzers/ are handled.
-  if not oss_fuzz_project_base.endswith('/'):
-    oss_fuzz_project_base += '/'
+  # Make sure cases like /src/curl and /src/curl/ are both handled.
+  norm_src_in_docker = os.path.normpath(src_in_docker)
+  if not norm_src_in_docker.endswith('/'):
+    norm_src_in_docker += '/'
 
   affected_file_list = []
   for file in coverage_per_file:
-    if not file['filename'].startswith(oss_fuzz_project_base):
+    norm_file_path = os.path.normpath(file['filename'])
+    if not norm_file_path.startswith(norm_src_in_docker):
       continue
     if not file['summary']['regions']['count']:
       # Don't consider a file affected if code in it is never executed.
       continue
 
-    relative_path = file['filename'].replace(oss_fuzz_project_base, '')
+    relative_path = file['filename'].replace(norm_src_in_docker, '')
     affected_file_list.append(relative_path)
   if not affected_file_list:
     return None
   return affected_file_list
 
 
+def remove_unaffected_fuzzers(project_name, out_dir, files_changed,
+                              src_in_docker):
+  """Removes all non affected fuzzers in the out directory.
+
+  Args:
+    project_name: The name of the relevant OSS-Fuzz project.
+    out_dir: The location of the fuzzer binaries.
+    files_changed: A list of files changed compared to HEAD.
+    src_in_docker: The location of the source dir in the docker image.
+  """
+  if not files_changed:
+    logging.info('No files changed compared to HEAD.')
+    return
+  fuzzer_paths = utils.get_fuzz_targets(out_dir)
+  if not fuzzer_paths:
+    logging.error('No fuzzers found in out dir.')
+    return
+
+  latest_cov_report_info = get_latest_cov_report_info(project_name)
+  if not latest_cov_report_info:
+    logging.error('Could not download latest coverage report.')
+    return
+  affected_fuzzers = []
+  for fuzzer in fuzzer_paths:
+    covered_files = get_files_covered_by_target(latest_cov_report_info,
+                                                os.path.basename(fuzzer),
+                                                src_in_docker)
+    if not covered_files:
+      # Assume a fuzzer is affected if we can't get its coverage from OSS-Fuzz.
+      affected_fuzzers.append(os.path.basename(fuzzer))
+      continue
+
+    for file in files_changed:
+      if file in covered_files:
+        affected_fuzzers.append(os.path.basename(fuzzer))
+
+  if not affected_fuzzers:
+    logging.info('No affected fuzzers detected, keeping all as fallback.')
+    return
+  logging.info('Using affected fuzzers.\n %s fuzzers affected by pull request',
+               ' '.join(affected_fuzzers))
+
+  all_fuzzer_names = map(os.path.basename, fuzzer_paths)
+
+  # Remove all the fuzzers that are not affected.
+  for fuzzer in all_fuzzer_names:
+    if fuzzer not in affected_fuzzers:
+      try:
+        os.remove(os.path.join(out_dir, fuzzer))
+      except OSError as error:
+        logging.error('%s occured while removing file %s', error, fuzzer)
+
+
 def get_json_from_url(url):
   """Gets a json object from a specified http url.
 
diff --git a/infra/cifuzz/cifuzz_test.py b/infra/cifuzz/cifuzz_test.py
index a73b94eb..47c9a6e1 100644
--- a/infra/cifuzz/cifuzz_test.py
+++ b/infra/cifuzz/cifuzz_test.py
@@ -361,5 +361,68 @@ class GetLatestCoverageReportUnitTest(unittest.TestCase):
     self.assertIsNone(cifuzz.get_latest_cov_report_info(''))
 
 
+class KeepAffectedFuzzersUnitTest(unittest.TestCase):
+  """Test the keep_affected_fuzzer method in the CIFuzz module."""
+
+  test_fuzzer_1 = os.path.join(TEST_FILES_PATH, 'out', 'example_crash_fuzzer')
+  test_fuzzer_2 = os.path.join(TEST_FILES_PATH, 'out', 'example_nocrash_fuzzer')
+  example_file_changed = 'test.txt'
+
+  def test_keeping_fuzzer_w_no_coverage(self):
+    """Tests that a specific fuzzer is kept if it is deemed affected."""
+    with tempfile.TemporaryDirectory() as tmp_dir, unittest.mock.patch.object(
+        cifuzz, 'get_latest_cov_report_info', return_value=1):
+      shutil.copy(self.test_fuzzer_1, tmp_dir)
+      shutil.copy(self.test_fuzzer_2, tmp_dir)
+      with unittest.mock.patch.object(cifuzz,
+                                      'get_files_covered_by_target',
+                                      side_effect=[[self.example_file_changed],
+                                                   None]):
+        cifuzz.remove_unaffected_fuzzers(EXAMPLE_PROJECT, tmp_dir,
+                                         [self.example_file_changed], '')
+        self.assertEqual(2, len(os.listdir(tmp_dir)))
+
+  def test_keeping_specific_fuzzer(self):
+    """Tests that a specific fuzzer is kept if it is deemed affected."""
+    with tempfile.TemporaryDirectory() as tmp_dir, unittest.mock.patch.object(
+        cifuzz, 'get_latest_cov_report_info', return_value=1):
+      shutil.copy(self.test_fuzzer_1, tmp_dir)
+      shutil.copy(self.test_fuzzer_2, tmp_dir)
+      with unittest.mock.patch.object(cifuzz,
+                                      'get_files_covered_by_target',
+                                      side_effect=[[self.example_file_changed],
+                                                   ['not/a/real/file']]):
+        cifuzz.remove_unaffected_fuzzers(EXAMPLE_PROJECT, tmp_dir,
+                                         [self.example_file_changed], '')
+        self.assertEqual(1, len(os.listdir(tmp_dir)))
+
+  def test_no_fuzzers_kept_fuzzer(self):
+    """Tests that if there is no affected then all fuzzers are kept."""
+    with tempfile.TemporaryDirectory() as tmp_dir, unittest.mock.patch.object(
+        cifuzz, 'get_latest_cov_report_info', return_value=1):
+      shutil.copy(self.test_fuzzer_1, tmp_dir)
+      shutil.copy(self.test_fuzzer_2, tmp_dir)
+      with unittest.mock.patch.object(cifuzz,
+                                      'get_files_covered_by_target',
+                                      side_effect=[None, None]):
+        cifuzz.remove_unaffected_fuzzers(EXAMPLE_PROJECT, tmp_dir,
+                                         [self.example_file_changed], '')
+        self.assertEqual(2, len(os.listdir(tmp_dir)))
+
+  def test_both_fuzzers_kept_fuzzer(self):
+    """Tests that if both fuzzers are affected then both fuzzers are kept."""
+    with tempfile.TemporaryDirectory() as tmp_dir, unittest.mock.patch.object(
+        cifuzz, 'get_latest_cov_report_info', return_value=1):
+      shutil.copy(self.test_fuzzer_1, tmp_dir)
+      shutil.copy(self.test_fuzzer_2, tmp_dir)
+      with unittest.mock.patch.object(
+          cifuzz,
+          'get_files_covered_by_target',
+          side_effect=[self.example_file_changed, self.example_file_changed]):
+        cifuzz.remove_unaffected_fuzzers(EXAMPLE_PROJECT, tmp_dir,
+                                         [self.example_file_changed], '')
+        self.assertEqual(2, len(os.listdir(tmp_dir)))
+
+
 if __name__ == '__main__':
   unittest.main()