Consolidate Untrusted Tasks to Preprocess Queue & Support Platform Filtering in ENABLE_FUZZ_FOR_BOTS

src/clusterfuzz/_internal/base/tasks/__init__.py

@@ -25,6 +25,7 @@
 
 from google.cloud import monitoring_v3
 
+from clusterfuzz._internal import swarming
 from clusterfuzz._internal.base import external_tasks
 from clusterfuzz._internal.base import feature_flags
 from clusterfuzz._internal.base import memoize
@@ -52,6 +53,15 @@
 JOBS_TASKQUEUE = JOBS_PREFIX
 HIGH_END_JOBS_TASKQUEUE = HIGH_END_JOBS_PREFIX
 
+POSTPROCESS_QUEUE = 'postprocess'
+UTASK_MAIN_QUEUE = 'utask_main'
+PREPROCESS_QUEUE = 'preprocess'
+
+SWARMING_QUEUES = {
+    PREPROCESS_QUEUE: 'preprocess-swarming',
+    UTASK_MAIN_QUEUE: 'utask_main-swarming',
+}
+
 # Limits on number of tasks leased at once and in total.
 MAX_LEASED_TASKS_LIMIT = 1000
 MAX_TASKS_LIMIT = 100000
@@ -99,15 +109,6 @@ def get_task_duration(command):
 TASK_PAYLOAD_KEY = 'task_payload'
 TASK_END_TIME_KEY = 'task_end_time'
 
-POSTPROCESS_QUEUE = 'postprocess'
-UTASK_MAIN_QUEUE = 'utask_main'
-PREPROCESS_QUEUE = 'preprocess'
-
-SWARMING_QUEUES = {
-    PREPROCESS_QUEUE: 'preprocess-swarming',
-    UTASK_MAIN_QUEUE: 'utask_main-swarming',
-}
-
 # See https://github.com/google/clusterfuzz/issues/3347 for usage
 SUBQUEUE_IDENTIFIER = ':'
 
@@ -440,7 +441,14 @@ def get_task():
 
   logs.info(f'Could not get task from {regular_queue()}. Fuzzing.')
 
-  if not feature_flags.FeatureFlags.ENABLE_FUZZ_FOR_BOTS.enabled:
+  enable_fuzz_flag = feature_flags.FeatureFlags.ENABLE_FUZZ_FOR_BOTS
+  allowed_platforms = [
+      p.strip().lower()
+      for p in enable_fuzz_flag.string_value.split(',')
+      if p.strip()
+  ]
+  if (not enable_fuzz_flag.enabled or
+      environment.platform().lower() not in allowed_platforms):
     logs.warning('Fuzzing is disabled for long-lived bots.')
     return None
 
@@ -932,6 +940,18 @@ def add_task(command,
       external_tasks.add_external_task(command, argument, job)
       return
 
+  if queue is None:
+    from clusterfuzz._internal.bot.tasks import task_types
+    if task_types.is_untrusted_task(command):
+      if job_type != 'none' and swarming.is_swarming_task(job_type):
+        queue = SWARMING_QUEUES[PREPROCESS_QUEUE]
+      else:
+        queue = PREPROCESS_QUEUE
+    elif job_type != 'none':
+      queue = queue_for_job(job_type)
+    else:
+      queue = default_queue()
+
   # Add the task.
   eta = utils.utcnow() + datetime.timedelta(seconds=wait_time)
   extra_info = extra_info or {}
@@ -981,6 +1001,11 @@ def queue_for_job(job_name, is_high_end=False):
   if not job:
     raise Error('Job {} not found.'.format(job_name))
 
+  if full_utask_task_model():
+    if swarming.is_swarming_task(job_name, job):
+      return SWARMING_QUEUES[PREPROCESS_QUEUE]
+    return PREPROCESS_QUEUE
+
   return queue_for_platform(job.platform, is_high_end)
 
 
@@ -1066,45 +1091,20 @@ def redo_testcase(testcase, tasks, user_email):
   # If we are re-doing minimization, other tasks will be done automatically
   # after minimization completes. So, don't add those tasks.
   if minimize:
-    add_task(
-        'minimize',
-        testcase_id,
-        testcase.job_type,
-        queue_for_testcase(testcase),
-        wait_time=wait_time)
+    add_task('minimize', testcase_id, testcase.job_type, wait_time=wait_time)
     return
 
   if regression:
-    add_task(
-        'regression',
-        testcase_id,
-        testcase.job_type,
-        queue_for_testcase(testcase),
-        wait_time=wait_time)
+    add_task('regression', testcase_id, testcase.job_type, wait_time=wait_time)
 
   if progression:
-    add_task(
-        'progression',
-        testcase_id,
-        testcase.job_type,
-        queue_for_testcase(testcase),
-        wait_time=wait_time)
+    add_task('progression', testcase_id, testcase.job_type, wait_time=wait_time)
 
   if impact:
-    add_task(
-        'impact',
-        testcase_id,
-        testcase.job_type,
-        queue_for_testcase(testcase),
-        wait_time=wait_time)
+    add_task('impact', testcase_id, testcase.job_type, wait_time=wait_time)
 
   if blame:
-    add_task(
-        'blame',
-        testcase_id,
-        testcase.job_type,
-        queue_for_testcase(testcase),
-        wait_time=wait_time)
+    add_task('blame', testcase_id, testcase.job_type, wait_time=wait_time)
 
 
 def get_task_payload():

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

@@ -18,7 +18,6 @@
 from clusterfuzz._internal import swarming
 from clusterfuzz._internal.base import errors
 from clusterfuzz._internal.base import tasks
-from clusterfuzz._internal.base.tasks import pub_sub_task_queue
 from clusterfuzz._internal.base.tasks import task_utils
 from clusterfuzz._internal.batch import service as batch_service
 from clusterfuzz._internal.bot.tasks import utasks
@@ -156,6 +155,7 @@ def execute(self, task_argument, job_type, uworker_env):
       self.execute_locally(task_argument, job_type, uworker_env)
       return
 
+    from clusterfuzz._internal.base.tasks import pub_sub_task_queue
     queue = pub_sub_task_queue.UTASK_MAIN_QUEUE
     if swarming.is_swarming_task(job_type):
       queue = pub_sub_task_queue.SWARMING_UTASK_MAIN_QUEUE
@@ -245,7 +245,7 @@ def execute(self, task_argument, job_type, uworker_env):
     'analyze': UTask,
     'blame': TrustedTask,
     'corpus_pruning': UTask,
-    'fuzz': UTaskLocalExecutor,
+    'fuzz': UTask,
     'impact': TrustedTask,
     'minimize': UTask,
     'progression': UTask,
@@ -256,3 +256,15 @@ def execute(self, task_argument, job_type, uworker_env):
     'uworker_main': UworkerMainTask,
     'variant': UTask,
 }
+
+
+def is_trusted_task(task_name: str) -> bool:
+  """Returns True if the task is a trusted task."""
+  task_type = COMMAND_TYPES.get(task_name)
+  return task_type is not None and issubclass(task_type, TrustedTask)
+
+
+def is_untrusted_task(task_name: str) -> bool:
+  """Returns True if the task is an untrusted utask."""
+  task_type = COMMAND_TYPES.get(task_name)
+  return task_type is not None and issubclass(task_type, BaseUTask)

src/clusterfuzz/_internal/tests/core/base/tasks/tasks_test.py

@@ -15,6 +15,7 @@
 import unittest
 from unittest import mock
 
+from clusterfuzz._internal.base import feature_flags
 from clusterfuzz._internal.base import tasks
 from clusterfuzz._internal.base.tasks import pub_sub_task_queue
 from clusterfuzz._internal.datastore import data_types
@@ -589,3 +590,142 @@ def test_fallback_to_random_choice(self):
     self.assertEqual(tasks.tworker_get_task(override_queue=''), 'task2')
     self.mock.get_postprocess_task.assert_called_once()
     self.mock.get_preprocess_task.assert_not_called()
+
+
+class GetTaskFuzzingFeatureFlagTest(unittest.TestCase):
+  """Tests for ENABLE_FUZZ_FOR_BOTS feature flag in get_task."""
+
+  def setUp(self):
+    helpers.patch(self, [
+        'clusterfuzz._internal.base.tasks.get_command_override',
+        'clusterfuzz._internal.base.tasks.get_postprocess_task',
+        'clusterfuzz._internal.base.tasks.get_high_end_task',
+        'clusterfuzz._internal.base.tasks.get_regular_task',
+        'clusterfuzz._internal.base.tasks.get_fuzz_task',
+        'clusterfuzz._internal.system.environment.is_android',
+        'clusterfuzz._internal.system.environment.platform',
+    ])
+    self.mock.get_command_override.return_value = None
+    self.mock.get_postprocess_task.return_value = None
+    self.mock.get_high_end_task.return_value = None
+    self.mock.get_regular_task.return_value = None
+    self.mock.is_android.return_value = False
+    self.mock.platform.return_value = 'LINUX'
+    self.mock_fuzz_task = mock.Mock()
+    self.mock.get_fuzz_task.return_value = self.mock_fuzz_task
+
+  def test_flag_none(self):
+    """Test that fuzz task is not returned when feature flag is not set."""
+    with mock.patch.object(
+        feature_flags.FeatureFlags, 'flag',
+        new_callable=mock.PropertyMock) as mock_flag:
+      mock_flag.return_value = None
+      self.assertIsNone(tasks.get_task())
+      self.mock.get_fuzz_task.assert_not_called()
+
+  def test_flag_disabled(self):
+    """Test that fuzz task is not returned when feature flag is disabled."""
+    mock_flag_obj = mock.MagicMock(enabled=False, string_value='linux')
+    with mock.patch.object(
+        feature_flags.FeatureFlags, 'flag',
+        new_callable=mock.PropertyMock) as mock_flag:
+      mock_flag.return_value = mock_flag_obj
+      self.assertIsNone(tasks.get_task())
+      self.mock.get_fuzz_task.assert_not_called()
+
+  def test_platform_in_string_value(self):
+    """Test that fuzz task is returned when platform matches string value."""
+    mock_flag_obj = mock.MagicMock(enabled=True, string_value='linux,windows')
+    with mock.patch.object(
+        feature_flags.FeatureFlags, 'flag',
+        new_callable=mock.PropertyMock) as mock_flag:
+      mock_flag.return_value = mock_flag_obj
+      self.assertEqual(tasks.get_task(), self.mock_fuzz_task)
+      self.mock.get_fuzz_task.assert_called_once()
+
+  def test_platform_not_in_string_value(self):
+    """Test fuzz task is not returned when platform not in string value."""
+    mock_flag_obj = mock.MagicMock(enabled=True, string_value='windows,mac')
+    with mock.patch.object(
+        feature_flags.FeatureFlags, 'flag',
+        new_callable=mock.PropertyMock) as mock_flag:
+      mock_flag.return_value = mock_flag_obj
+      self.assertIsNone(tasks.get_task())
+      self.mock.get_fuzz_task.assert_not_called()
+
+  def test_empty_string_value(self):
+    """Test that fuzz task is not returned when string value is empty."""
+    mock_flag_obj = mock.MagicMock(enabled=True, string_value='')
+    with mock.patch.object(
+        feature_flags.FeatureFlags, 'flag',
+        new_callable=mock.PropertyMock) as mock_flag:
+      mock_flag.return_value = mock_flag_obj
+      self.assertIsNone(tasks.get_task())
+      self.mock.get_fuzz_task.assert_not_called()
+
+
+@mock.patch(
+    'clusterfuzz._internal.swarming.is_swarming_task', return_value=False)
+class TaskQueueResolutionTest(unittest.TestCase):
+  """Tests for task-aware queue resolution in add_task."""
+
+  def test_task_types_classification(self, _):
+    """Test that task_types correctly identifies untrusted and trusted tasks."""
+    from clusterfuzz._internal.bot.tasks import task_types
+    self.assertTrue(task_types.is_untrusted_task('fuzz'))
+    self.assertTrue(task_types.is_untrusted_task('minimize'))
+    self.assertTrue(task_types.is_untrusted_task('variant'))
+    self.assertTrue(task_types.is_untrusted_task('analyze'))
+    self.assertTrue(task_types.is_untrusted_task('progression'))
+    self.assertTrue(task_types.is_untrusted_task('regression'))
+    self.assertTrue(task_types.is_untrusted_task('symbolize'))
+    self.assertTrue(task_types.is_untrusted_task('corpus_pruning'))
+
+    self.assertTrue(task_types.is_trusted_task('blame'))
+    self.assertTrue(task_types.is_trusted_task('impact'))
+    self.assertTrue(task_types.is_trusted_task('unpack'))
+    self.assertFalse(task_types.is_untrusted_task('blame'))
+    self.assertFalse(task_types.is_trusted_task('fuzz'))
+
+  @mock.patch('clusterfuzz._internal.base.tasks.bulk_add_tasks')
+  @mock.patch('clusterfuzz._internal.base.tasks.data_types.Job.query')
+  def test_add_task_untrusted(self, mock_job_query, mock_bulk_add, _):
+    """Test that add_task automatically routes untrusted tasks to preprocess."""
+    mock_job = mock.MagicMock(platform='LINUX', base_os_version=None)
+    mock_job.is_external.return_value = False
+    mock_job_query.return_value.get.return_value = mock_job
+
+    tasks.add_task('minimize', '123', 'linux_asan_d8_dbg')
+    mock_bulk_add.assert_called_once()
+    self.assertEqual(mock_bulk_add.call_args[1]['queue'],
+                     tasks.PREPROCESS_QUEUE)
+
+  @mock.patch('clusterfuzz._internal.base.tasks.bulk_add_tasks')
+  @mock.patch('clusterfuzz._internal.base.tasks.data_types.Job.query')
+  def test_add_task_trusted_without_queue(self, mock_job_query, mock_bulk_add,
+                                          _):
+    """Test that add_task automatically routes trusted tasks to platform queue."""
+    mock_job = mock.MagicMock(platform='LINUX', base_os_version=None)
+    mock_job.is_external.return_value = False
+    mock_job_query.return_value.get.return_value = mock_job
+
+    tasks.add_task('blame', '123', 'linux_asan_d8_dbg')
+    mock_bulk_add.assert_called_once()
+    self.assertEqual(mock_bulk_add.call_args[1]['queue'], 'jobs-linux')
+
+  @mock.patch('clusterfuzz._internal.base.tasks.bulk_add_tasks')
+  @mock.patch('clusterfuzz._internal.base.tasks.data_types.Job.query')
+  def test_add_task_untrusted_with_explicit_queue(self, mock_job_query,
+                                                  mock_bulk_add, _):
+    """Test that add_task respects explicit queue for untrusted tasks (e.g. add_utask_main)."""
+    mock_job = mock.MagicMock(platform='LINUX', base_os_version=None)
+    mock_job.is_external.return_value = False
+    mock_job_query.return_value.get.return_value = mock_job
+
+    tasks.add_task(
+        'fuzz',
+        'https://storage.googleapis.com/uworker-input...',
+        'linux_asan_d8_dbg',
+        queue='utask_main')
+    mock_bulk_add.assert_called_once()
+    self.assertEqual(mock_bulk_add.call_args[1]['queue'], 'utask_main')