Swarming: Calculates CF zip url in preprocess to be used in main

src/clusterfuzz/_internal/swarming/__init__.py

@@ -25,6 +25,7 @@
 from clusterfuzz._internal.base.feature_flags import FeatureFlags
 from clusterfuzz._internal.config import local_config
 from clusterfuzz._internal.datastore import data_types
+from clusterfuzz._internal.google_cloud_utils import compute_metadata
 from clusterfuzz._internal.google_cloud_utils import credentials
 from clusterfuzz._internal.metrics import logs
 from clusterfuzz._internal.protos import swarming_pb2
@@ -112,6 +113,68 @@ def _get_task_dimensions(job: data_types.Job, platform_specific_dimensions: list
   return task_dimensions
 
 
+def _append_metadata_env_var(
+    env_vars: list[swarming_pb2.StringPair],  # pylint: disable=no-member
+    env_var_name: str,
+    metadata_path: str) -> None:
+  """Attempts to get a variable from the environment or metadata and appends it.
+
+  Args:
+    env_vars: The list of string pairs to append the retrieved value to.
+    env_var_name: The name of the environment variable to check and set.
+    metadata_path: The path in the compute metadata to check if not set in env.
+  """
+  value = environment.get_value(env_var_name)
+  if not value:
+    try:
+      value = compute_metadata.get(metadata_path)
+    except Exception:
+      pass
+
+  if value:
+    env_vars.append(
+        swarming_pb2.StringPair(  # pylint: disable=no-member
+            key=env_var_name, value=str(value)))
+  else:
+    logs.warning(f'{env_var_name} is not set or cannot be fetched.')
+
+
+def _get_env_vars(logs_project_id: str,
+                  instance_spec: dict) -> list[swarming_pb2.StringPair]:  # pylint: disable=no-member
+  """Retrieve required environment variables from metadata and config."""
+  default_task_environment = [
+      swarming_pb2.StringPair(key='UWORKER', value='True'),  # pylint: disable=no-member
+      swarming_pb2.StringPair(key='SWARMING_BOT', value='True'),  # pylint: disable=no-member
+      swarming_pb2.StringPair(key='LOG_TO_GCP', value='True'),  # pylint: disable=no-member
+      swarming_pb2.StringPair(key='IS_K8S_ENV', value='True'),  # pylint: disable=no-member
+      swarming_pb2.StringPair(  # pylint: disable=no-member
+          key='LOGGING_CLOUD_PROJECT_ID',
+          value=logs_project_id or ''),
+  ]
+
+  _append_metadata_env_var(default_task_environment, 'DEPLOYMENT_BUCKET',
+                           'project/attributes/deployment-bucket')
+  _append_metadata_env_var(default_task_environment, 'HOST_JOB_SELECTION',
+                           'instance/attributes/host-job-selection')
+  _append_metadata_env_var(default_task_environment, 'DEPLOYMENT_ZIP',
+                           'project/attributes/deployment-zip')
+
+  env_vars = []
+  env_vars.append(
+      swarming_pb2.StringPair(  # pylint: disable=no-member
+          key='DOCKER_IMAGE',
+          value=instance_spec.get('docker_image', '')))
+
+  platform_specific_env = instance_spec.get('env', [])
+  for var in platform_specific_env:
+    env_vars.append(swarming_pb2.StringPair(key=var['key'], value=var['value']))  # pylint: disable=no-member
+
+  env_vars.append(_env_vars_to_json(default_task_environment))
+  env_vars.extend(default_task_environment)
+
+  return env_vars
+
+
 def _env_vars_to_json(
     env_vars: list[swarming_pb2.StringPair]) -> swarming_pb2.StringPair:  # pylint: disable=no-member
   """
@@ -166,27 +229,7 @@ def create_new_task_request(command: str, job_name: str, download_url: str
   # env_prefixes allows the modification of existing environment variables by
   # adding the values as prefixes to the env variable.
   env_prefixes = instance_spec.get('env_prefixes', {})
-  default_task_environment = [
-      swarming_pb2.StringPair(key='UWORKER', value='True'),  # pylint: disable=no-member
-      swarming_pb2.StringPair(key='SWARMING_BOT', value='True'),  # pylint: disable=no-member
-      swarming_pb2.StringPair(key='LOG_TO_GCP', value='True'),  # pylint: disable=no-member
-      swarming_pb2.StringPair(key='IS_K8S_ENV', value='True'),  # pylint: disable=no-member
-      swarming_pb2.StringPair(  # pylint: disable=no-member
-          key='LOGGING_CLOUD_PROJECT_ID',
-          value=logs_project_id),
-  ]
-
-  platform_specific_env = instance_spec.get('env', [])
-  swarming_bot_environment = []
-  swarming_bot_environment.append(
-      swarming_pb2.StringPair(  # pylint: disable=no-member
-          key='DOCKER_IMAGE',
-          value=instance_spec.get('docker_image', '')))
-  for var in platform_specific_env:
-    swarming_bot_environment.append(
-        swarming_pb2.StringPair(key=var['key'], value=var['value']))  # pylint: disable=no-member
-  swarming_bot_environment.append(_env_vars_to_json(default_task_environment))
-  swarming_bot_environment.extend(default_task_environment)
+  swarming_bot_environment = _get_env_vars(logs_project_id, instance_spec)
   dimensions = instance_spec.get('dimensions', [])
   cas_input_root = instance_spec.get('cas_input_root', {})
 

src/clusterfuzz/_internal/tests/core/swarming/service_test.py

@@ -31,9 +31,11 @@ def setUp(self):
         'clusterfuzz._internal.swarming.create_new_task_request',
         'clusterfuzz._internal.base.tasks.task_utils.get_command_from_module',
         'clusterfuzz._internal.metrics.logs.error',
+        'clusterfuzz._internal.google_cloud_utils.compute_metadata.get',
     ])
     self.service = service.SwarmingService()
     self.mock.create_new_task_request.return_value = 'fake_request'
+    self.mock.get.return_value = None
 
   def test_create_utask_main_job_success(self):
     """Test creating a single task successfully."""

src/clusterfuzz/_internal/tests/core/swarming/swarming_config_error_test.py

@@ -29,9 +29,11 @@ class SwarmingConfigErrorTest(unittest.TestCase):
   def setUp(self):
     helpers.patch(self, [
         'clusterfuzz._internal.swarming.FeatureFlags',
+        'clusterfuzz._internal.google_cloud_utils.compute_metadata.get',
     ])
     helpers.patch_environ(self)
     self.mock.FeatureFlags.SWARMING_REMOTE_EXECUTION.enabled = True
+    self.mock.get.return_value = None
 
   def test_is_swarming_task_bad_config(self):
     """Tests that is_swarming_task returns False when there's a BadConfigError."""

src/clusterfuzz/_internal/tests/core/swarming/swarming_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Swarming tests."""
 import base64
+import os
 import unittest
 from unittest import mock
 
@@ -38,11 +39,17 @@ def setUp(self):
         'clusterfuzz._internal.google_cloud_utils.credentials.get_scoped_service_account_credentials',
         'google.auth.transport.requests.Request',
         'clusterfuzz._internal.swarming.FeatureFlags',
+        'clusterfuzz._internal.google_cloud_utils.compute_metadata.get',
     ])
     helpers.patch_environ(self)
     self.mock._get_task_name.return_value = 'task_name'  # pylint: disable=protected-access
     self.mock.FeatureFlags.SWARMING_REMOTE_EXECUTION.enabled = True
+    self.mock.get.return_value = None
     self.maxDiff = None
+    os.environ.pop('DEPLOYMENT_ZIP', None)
+    os.environ.pop('DEPLOYMENT_BUCKET', None)
+    os.environ.pop('PROJECT_NAME', None)
+    os.environ.pop('HOST_JOB_SELECTION', None)
 
   def test_get_spec_from_config_with_docker_image(self):
     """Tests that create_new_task_request works as expected."""
@@ -430,3 +437,37 @@ def test_get_task_dimensions_job_precedence(self):
         swarming_pb2.StringPair(key='key2', value='value2'),
     ]
     self.assertCountEqual(dimensions, expected_dimensions)
+
+  def test_get_env_vars_with_metadata_server(self):
+    """Tests that _get_env_vars uses values from the metadata server when available."""
+
+    def metadata_get(path):
+      if path == 'project/attributes/deployment-bucket':
+        return 'test-bucket-from-metadata'
+      return None
+
+    self.mock.get.side_effect = metadata_get
+    instance_spec = {
+        "docker_image": "gcr.io/clusterfuzz-images/base:a2f4dd6-202202070654"
+    }
+    env = swarming._get_env_vars('project_id', instance_spec)  # pylint: disable=protected-access
+
+    expected_env = [
+        swarming_pb2.StringPair(
+            key='DOCKER_IMAGE',
+            value='gcr.io/clusterfuzz-images/base:a2f4dd6-202202070654'),
+        swarming_pb2.StringPair(
+            key='DOCKER_ENV_VARS',
+            value=
+            '{"UWORKER": "True", "SWARMING_BOT": "True", "LOG_TO_GCP": "True", "IS_K8S_ENV": "True", "LOGGING_CLOUD_PROJECT_ID": "project_id", "DEPLOYMENT_BUCKET": "test-bucket-from-metadata"}'
+        ),
+        swarming_pb2.StringPair(key='UWORKER', value='True'),
+        swarming_pb2.StringPair(key='SWARMING_BOT', value='True'),
+        swarming_pb2.StringPair(key='LOG_TO_GCP', value='True'),
+        swarming_pb2.StringPair(key='IS_K8S_ENV', value='True'),
+        swarming_pb2.StringPair(
+            key='LOGGING_CLOUD_PROJECT_ID', value='project_id'),
+        swarming_pb2.StringPair(
+            key='DEPLOYMENT_BUCKET', value='test-bucket-from-metadata'),
+    ]
+    self.assertEqual(env, expected_env)