From f4c82ace7ec8908cc5d6129e1436dacdf3332a77 Mon Sep 17 00:00:00 2001
From: jonathanmetzman <31354670+jonathanmetzman@users.noreply.github.com>
Date: Mon, 4 Nov 2024 11:47:05 -0500
Subject: [PATCH] Prepare to make analyze task in oss-fuzz run on uworkers.
 (#4368)

I'm going to migrate tasks in OSS-Fuzz one by one.
---
 configs/test/project.yaml                       |  3 +++
 src/clusterfuzz/_internal/base/task_utils.py    | 17 ++++++++++++++---
 .../_internal/bot/tasks/task_types.py           | 11 ++++++-----
 .../tests/core/base/task_utils_test.py          |  9 +++++++++
 4 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/configs/test/project.yaml b/configs/test/project.yaml
index 67a29800158..118018f333b 100644
--- a/configs/test/project.yaml
+++ b/configs/test/project.yaml
@@ -50,6 +50,9 @@ firebase:
   - google.com
   # - github.com
 
+uworker_tasks:
+  - analyze
+
 
 stacktrace:
   # Stack frames to ignore when determining the crash signature.
diff --git a/src/clusterfuzz/_internal/base/task_utils.py b/src/clusterfuzz/_internal/base/task_utils.py
index 0d697e2debe..4a87e2993d2 100644
--- a/src/clusterfuzz/_internal/base/task_utils.py
+++ b/src/clusterfuzz/_internal/base/task_utils.py
@@ -15,6 +15,7 @@
 any other module in tasks to prevent circular imports and issues with
 appengine."""
 
+from clusterfuzz._internal.config import local_config
 from clusterfuzz._internal.system import environment
 
 
@@ -25,11 +26,21 @@ def get_command_from_module(full_module_name: str) -> str:
   return module_name[:-len('_task')]
 
 
-def is_remotely_executing_utasks() -> bool:
+def is_remotely_executing_utasks(task=None) -> bool:
   """Returns True if the utask_main portions of utasks are being remotely
   executed on Google cloud batch."""
-  return bool(environment.is_production() and
-              environment.get_value('REMOTE_UTASK_EXECUTION'))
+  if bool(environment.is_production() and
+          environment.get_value('REMOTE_UTASK_EXECUTION')):
+    return True
+  if task is None:
+    return False
+  return is_task_opted_into_uworker_execution(task)
+
+
+def is_task_opted_into_uworker_execution(task):
+  # TODO(metzman): Remove this after OSS-Fuzz and Chrome are at parity.
+  uworker_tasks = local_config.ProjectConfig().get('uworker_tasks', [])
+  return task in uworker_tasks
 
 
 class UworkerMsgParseError(RuntimeError):
diff --git a/src/clusterfuzz/_internal/bot/tasks/task_types.py b/src/clusterfuzz/_internal/bot/tasks/task_types.py
index cdaf1f51071..d21d316112a 100644
--- a/src/clusterfuzz/_internal/bot/tasks/task_types.py
+++ b/src/clusterfuzz/_internal/bot/tasks/task_types.py
@@ -27,7 +27,8 @@ class BaseTask:
   """Base module for tasks."""
 
   @staticmethod
-  def is_execution_remote():
+  def is_execution_remote(command=None):
+    del command
     return False
 
   def __init__(self, module):
@@ -73,13 +74,13 @@ def preprocess(self, task_argument, job_type, uworker_env):
 
 
 def is_no_privilege_workload(command, job):
-  if not COMMAND_TYPES[command].is_execution_remote():
+  if not COMMAND_TYPES[command].is_execution_remote(command):
     return False
   return batch.is_no_privilege_workload(command, job)
 
 
 def is_remote_utask(command, job):
-  if not COMMAND_TYPES[command].is_execution_remote():
+  if not COMMAND_TYPES[command].is_execution_remote(command):
     return False
 
   if environment.is_uworker():
@@ -117,8 +118,8 @@ class UTask(BaseUTask):
   opted-in. Otherwise executes locally."""
 
   @staticmethod
-  def is_execution_remote():
-    return task_utils.is_remotely_executing_utasks()
+  def is_execution_remote(command=None):
+    return task_utils.is_remotely_executing_utasks(command)
 
   def execute(self, task_argument, job_type, uworker_env):
     """Executes a utask."""
diff --git a/src/clusterfuzz/_internal/tests/core/base/task_utils_test.py b/src/clusterfuzz/_internal/tests/core/base/task_utils_test.py
index 7a103e71d64..072ca62f683 100644
--- a/src/clusterfuzz/_internal/tests/core/base/task_utils_test.py
+++ b/src/clusterfuzz/_internal/tests/core/base/task_utils_test.py
@@ -33,3 +33,12 @@ def test_get_command_from_module(self):
       task_utils.get_command_from_module('postprocess')
     with self.assertRaises(ValueError):
       task_utils.get_command_from_module('uworker_main')
+
+
+class IsTaskOptedIntoUworkerExecution(unittest.TestCase):
+
+  def test_opt_in(self):
+    self.assertTrue(task_utils.is_task_opted_into_uworker_execution('analyze'))
+
+  def test_no_opt_in(self):
+    self.assertFalse(task_utils.is_task_opted_into_uworker_execution('fuzz'))