Download databundles for non engine jobs (#5085)

This fixes b/446692241 by downloading databundles on blacbox fuzzer
runs.
It does so without hurting the performance of non-blackbox batch tasks.
Since downloading the default databundle takes ~5 minutes. It also
optimizes the unzip by about 50% by using the linux unzip binary instead
of Python's slow implementation.

---------

Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>
Co-authored-by: jonathanmetzman <[email protected]>
Co-authored-by: Jonathan Metzman <[email protected]>

Author: google-labs-jules[bot]

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

@@ -16,6 +16,8 @@
 import datetime
 import os
 import platform
+import shutil
+import subprocess
 import sys
 import time
 
@@ -288,14 +290,15 @@ def update_source_code():
             f'(release = {utils.get_clusterfuzz_release()}).')
 
 
-def update_tests_if_needed():
+def update_tests_if_needed(tests_url=None):
   """Updates layout tests every day."""
   data_directory = environment.get_value('FUZZ_DATA')
   error_occured = False
   expected_task_duration = 60 * 60  # 1 hour.
   retry_limit = environment.get_value('FAIL_RETRIES')
   temp_archive = os.path.join(data_directory, 'temp.zip')
-  tests_url = environment.get_value('WEB_TESTS_URL')
+  if not tests_url:
+    tests_url = environment.get_value('WEB_TESTS_URL')
 
   # Check if we have a valid tests url.
   if not tests_url:
@@ -323,9 +326,19 @@ def update_tests_if_needed():
   for _ in range(retry_limit):
     try:
       shell.remove_directory(data_directory, recreate=True)
-      storage.copy_file_from(tests_url, temp_archive)
-      with archive.open(temp_archive) as reader:
-        reader.extract_all(data_directory, trusted=True)
+      if tests_url.startswith('http'):
+        storage.download_signed_url_to_file(tests_url, temp_archive)
+      else:
+        storage.copy_file_from(tests_url, temp_archive)
+
+      if shutil.which('unzip'):
+        subprocess.run(
+            ['unzip', '-q', '-o', temp_archive, '-d', data_directory],
+            check=True,
+            capture_output=True)
+      else:
+        with archive.open(temp_archive) as reader:
+          reader.extract_all(data_directory, trusted=True)
       shell.remove_file(temp_archive)
       error_occured = False
       break

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

@@ -38,6 +38,7 @@
 from clusterfuzz._internal.bot.tasks import setup
 from clusterfuzz._internal.bot.tasks import task_creation
 from clusterfuzz._internal.bot.tasks import trials
+from clusterfuzz._internal.bot.tasks import update_task
 from clusterfuzz._internal.bot.tasks.utasks import fuzz_task_knobs
 from clusterfuzz._internal.bot.tasks.utasks import uworker_handle_errors
 from clusterfuzz._internal.bot.tasks.utasks import uworker_io
@@ -2097,6 +2098,10 @@ def _get_fuzz_target(uworker_input):
 
 def utask_main(uworker_input):
   """Runs the given fuzzer for one round."""
+  if not engine.get(uworker_input.fuzzer_name):
+    update_task.update_tests_if_needed(
+        uworker_input.fuzz_task_input.web_tests_url)
+
   # Sets fuzzing logs context before running the fuzzer.
   fuzz_target = _get_fuzz_target(uworker_input)
   with logs.fuzzer_log_context(uworker_input.fuzzer_name,
@@ -2191,6 +2196,11 @@ def _utask_preprocess(fuzzer_name, job_type, uworker_env):
             use_backup=True).serialize())
 
   fuzz_task_input.trials.extend(trials.preprocess_get_db_trials())
+  web_tests_url = environment.get_value('WEB_TESTS_URL')
+  if web_tests_url:
+    fuzz_task_input.web_tests_url = storage.get_signed_download_url(
+        web_tests_url)
+
   for _ in range(MAX_CRASHES_UPLOADED):
     url = fuzz_task_input.crash_upload_urls.add()
     url.key = blobs.generate_new_blob_name()

src/clusterfuzz/_internal/protos/uworker_msg.proto

@@ -53,6 +53,7 @@ message FuzzTaskInput {
   repeated string global_blacklisted_functions = 6;
   repeated BlobUploadUrl crash_upload_urls = 7;
   repeated google.protobuf.Any trials = 8;
+  optional string web_tests_url = 9;
 }
 
 message DataBundleCorpus {

src/clusterfuzz/_internal/protos/uworker_msg_pb2.py

@@ -251,6 +251,7 @@ class FuzzTaskInput(google.protobuf.message.Message):
     GLOBAL_BLACKLISTED_FUNCTIONS_FIELD_NUMBER: builtins.int
     CRASH_UPLOAD_URLS_FIELD_NUMBER: builtins.int
     TRIALS_FIELD_NUMBER: builtins.int
+    WEB_TESTS_URL_FIELD_NUMBER: builtins.int
     sample_testcase_upload_key: builtins.str
     sample_testcase_upload_url: builtins.str
     script_log_upload_url: builtins.str
@@ -265,6 +266,7 @@ class FuzzTaskInput(google.protobuf.message.Message):
     def crash_upload_urls(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___BlobUploadUrl]: ...
     @property
     def trials(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[google.protobuf.any_pb2.Any]: ...
+    web_tests_url: builtins.str
     def __init__(
         self,
         *,
@@ -276,9 +278,10 @@ class FuzzTaskInput(google.protobuf.message.Message):
         global_blacklisted_functions: collections.abc.Iterable[builtins.str] | None = ...,
         crash_upload_urls: collections.abc.Iterable[global___BlobUploadUrl] | None = ...,
         trials: collections.abc.Iterable[google.protobuf.any_pb2.Any] | None = ...,
+        web_tests_url: builtins.str | None = ...,
     ) -> None: ...
-    def HasField(self, field_name: typing_extensions.Literal["_corpus", b"_corpus", "_fuzz_target", b"_fuzz_target", "_sample_testcase_upload_key", b"_sample_testcase_upload_key", "_sample_testcase_upload_url", b"_sample_testcase_upload_url", "_script_log_upload_url", b"_script_log_upload_url", "corpus", b"corpus", "fuzz_target", b"fuzz_target", "sample_testcase_upload_key", b"sample_testcase_upload_key", "sample_testcase_upload_url", b"sample_testcase_upload_url", "script_log_upload_url", b"script_log_upload_url"]) -> builtins.bool: ...
-    def ClearField(self, field_name: typing_extensions.Literal["_corpus", b"_corpus", "_fuzz_target", b"_fuzz_target", "_sample_testcase_upload_key", b"_sample_testcase_upload_key", "_sample_testcase_upload_url", b"_sample_testcase_upload_url", "_script_log_upload_url", b"_script_log_upload_url", "corpus", b"corpus", "crash_upload_urls", b"crash_upload_urls", "fuzz_target", b"fuzz_target", "global_blacklisted_functions", b"global_blacklisted_functions", "sample_testcase_upload_key", b"sample_testcase_upload_key", "sample_testcase_upload_url", b"sample_testcase_upload_url", "script_log_upload_url", b"script_log_upload_url", "trials", b"trials"]) -> None: ...
+    def HasField(self, field_name: typing_extensions.Literal["_corpus", b"_corpus", "_fuzz_target", b"_fuzz_target", "_sample_testcase_upload_key", b"_sample_testcase_upload_key", "_sample_testcase_upload_url", b"_sample_testcase_upload_url", "_script_log_upload_url", b"_script_log_upload_url", "_web_tests_url", b"_web_tests_url", "corpus", b"corpus", "fuzz_target", b"fuzz_target", "sample_testcase_upload_key", b"sample_testcase_upload_key", "sample_testcase_upload_url", b"sample_testcase_upload_url", "script_log_upload_url", b"script_log_upload_url", "web_tests_url", b"web_tests_url"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing_extensions.Literal["_corpus", b"_corpus", "_fuzz_target", b"_fuzz_target", "_sample_testcase_upload_key", b"_sample_testcase_upload_key", "_sample_testcase_upload_url", b"_sample_testcase_upload_url", "_script_log_upload_url", b"_script_log_upload_url", "_web_tests_url", b"_web_tests_url", "corpus", b"corpus", "crash_upload_urls", b"crash_upload_urls", "fuzz_target", b"fuzz_target", "global_blacklisted_functions", b"global_blacklisted_functions", "sample_testcase_upload_key", b"sample_testcase_upload_key", "sample_testcase_upload_url", b"sample_testcase_upload_url", "script_log_upload_url", b"script_log_upload_url", "trials", b"trials", "web_tests_url", b"web_tests_url"]) -> None: ...
     @typing.overload
     def WhichOneof(self, oneof_group: typing_extensions.Literal["_corpus", b"_corpus"]) -> typing_extensions.Literal["corpus"] | None: ...
     @typing.overload
@@ -289,6 +292,8 @@ class FuzzTaskInput(google.protobuf.message.Message):
     def WhichOneof(self, oneof_group: typing_extensions.Literal["_sample_testcase_upload_url", b"_sample_testcase_upload_url"]) -> typing_extensions.Literal["sample_testcase_upload_url"] | None: ...
     @typing.overload
     def WhichOneof(self, oneof_group: typing_extensions.Literal["_script_log_upload_url", b"_script_log_upload_url"]) -> typing_extensions.Literal["script_log_upload_url"] | None: ...
+    @typing.overload
+    def WhichOneof(self, oneof_group: typing_extensions.Literal["_web_tests_url", b"_web_tests_url"]) -> typing_extensions.Literal["web_tests_url"] | None: ...
 
 global___FuzzTaskInput = FuzzTaskInput