[rearch][fuzz_task] Get data bundle directory in safe way. (#4202)

Related #3008

Author: jonathanmetzman

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

@@ -467,9 +467,9 @@ def _prepare_update_data_bundle(fuzzer, data_bundle):
   return data_bundle_directory
 
 
-def update_data_bundle(
+def _update_data_bundle(
     fuzzer: data_types.Fuzzer,
-    data_bundle_corpus: uworker_msg_pb2.DataBundleCorpus) -> bool:  # pylint: disable=no-member
+    data_bundle_corpus: uworker_msg_pb2.DataBundleCorpus) -> Optional[str]:  # pylint: disable=no-member
   """Updates a data bundle to the latest version."""
   data_bundle = uworker_io.entity_from_protobuf(data_bundle_corpus.data_bundle,
                                                 data_types.DataBundle)
@@ -656,7 +656,7 @@ def _set_up_data_bundles(update_input: uworker_msg_pb2.SetupInput):  # pylint: d
   fuzzer = uworker_io.entity_from_protobuf(update_input.fuzzer,
                                            data_types.Fuzzer)
   for data_bundle_corpus in update_input.data_bundle_corpuses:
-    if not update_data_bundle(fuzzer, data_bundle_corpus):
+    if not _update_data_bundle(fuzzer, data_bundle_corpus):
       return False
 
   return True
@@ -744,16 +744,6 @@ def _is_data_bundle_up_to_date(data_bundle, data_bundle_directory):
   return False
 
 
-def trusted_get_data_bundle_directory(fuzzer):
-  """For fuzz_task which doesn't get data bundles in an untrusted manner."""
-  # TODO(metzman): Delete this when fuzz_task is migrated.
-  # Check if we have a fuzzer-specific data bundle. Use it to calculate the
-  # data directory we will fetch our testcases from.
-  data_bundle = data_types.DataBundle.query(
-      data_types.DataBundle.name == fuzzer.data_bundle_name).get()
-  return get_data_bundle_directory(fuzzer, data_bundle)
-
-
 def get_data_bundle_directory(fuzzer, data_bundle):
   """Return data bundle data directory."""
   # Store corpora for built-in fuzzers like libFuzzer in the same directory

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

@@ -1783,7 +1783,9 @@ def run(self):
 
     # Data bundle directories can also have testcases which are kept in-place
     # because of dependencies.
-    self.data_directory = setup.trusted_get_data_bundle_directory(self.fuzzer)
+    self.data_directory = setup.get_data_bundle_directory(
+        self.fuzzer, self.uworker_input.setup_input.data_bundle_corpuses)
+
     if not self.data_directory:
       logs.error(
           'Unable to setup data bundle %s.' % self.fuzzer.data_bundle_name)

src/clusterfuzz/_internal/tests/core/bot/tasks/setup_test.py

@@ -214,7 +214,7 @@ def setUp(self):
         'clusterfuzz._internal.bot.tasks.task_types.is_remote_utask',
         'clusterfuzz._internal.bot.tasks.setup._update_fuzzer',
         'clusterfuzz._internal.bot.tasks.setup._clear_old_data_bundles_if_needed',
-        'clusterfuzz._internal.bot.tasks.setup.update_data_bundle',
+        'clusterfuzz._internal.bot.tasks.setup._update_data_bundle',
     ])
     self.mock.get_signed_upload_url.return_value = 'https://fake/upload'
     self.mock.get_signed_download_url.return_value = 'https://fake/download'
@@ -234,4 +234,4 @@ def test_data_bundles(self):
     setup_input = setup.preprocess_update_fuzzer_and_data_bundles(
         self.fuzzer_name)
     setup.update_fuzzer_and_data_bundles(setup_input)
-    self.assertEqual(self.mock.update_data_bundle.call_count, 2)
+    self.assertEqual(self.mock._update_data_bundle.call_count, 2)

src/clusterfuzz/_internal/tests/core/bot/untrusted_runner/untrusted_runner_integration_test.py

@@ -41,6 +41,8 @@
 from clusterfuzz._internal.system import shell
 from clusterfuzz._internal.tests.test_libs import untrusted_runner_helpers
 
+# pylint: disable=protected-access
+
 TEST_FILE_CONTENTS = (b'A' * config.FILE_TRANSFER_CHUNK_SIZE +
                       b'B' * config.FILE_TRANSFER_CHUNK_SIZE +
                       b'C' * (config.FILE_TRANSFER_CHUNK_SIZE // 2))
@@ -513,7 +515,7 @@ def test_symbolize(self):
     self.assertEqual(expected_symbolized_stacktrace, symbolized_stacktrace)
 
   def test_update_data_bundle(self):
-    """Test update_data_bundle."""
+    """Test _update_data_bundle."""
     self.mock.get_data_bundle_bucket_name.return_value = TEST_BUNDLE_BUCKET
 
     # Get a blobstore key for the fuzzer.
@@ -536,15 +538,15 @@ def test_update_data_bundle(self):
     data_bundle_corpus.data_bundle.CopyFrom(
         uworker_io.entity_to_protobuf(bundle))
     self.assertTrue(
-        setup.update_data_bundle(returned_fuzzer, data_bundle_corpus))
+        setup._update_data_bundle(returned_fuzzer, data_bundle_corpus))
 
     data_bundle_directory = file_host.rebase_to_worker_root(
         setup.get_data_bundle_directory(returned_fuzzer, bundle))
     self.assertTrue(os.path.exists(os.path.join(data_bundle_directory, 'a')))
     self.assertTrue(os.path.exists(os.path.join(data_bundle_directory, 'b')))
 
     self.assertTrue(
-        setup.update_data_bundle(returned_fuzzer, data_bundle_corpus))
+        setup._update_data_bundle(returned_fuzzer, data_bundle_corpus))
 
   def test_get_fuzz_targets(self):
     """Test get_fuzz_targets."""