Swarming: Calculates CF zip url in preprocess to be used in main (#5235)

The setups scripts rely on a set of values that form the `GS storage
url` to download the latest CF .zip, these values either come from the
env vars, or the metadata server.

Since in swarming we don't have metadata server, we are calculating
these values beforehand and sending them as env vars in the swarming
request.

## Test Proofs
These changes have been in `dev` [since last
week](a609252).
We have successfully use them to fix this issue:
```
Downloading ClusterFuzz source code.
++ '[' '' = 1 ']'
++ gsutil cp gs:///linux-3.zip clusterfuzz_package.zip
```
Allowing the code to tgarget the correct URL:
```
Downloading ClusterFuzz source code.
++ '[' '' = 1 ']'
++ gsutil cp gs://deployment.clusterfuzz-development.appspot.com/linux-3.zip clusterfuzz_package.zip
``` 

Note: The failing e2e check is known to be failing right now so we can
ignore that

---------

Co-authored-by: Diego Jardon <37823380+jardondiego@users.noreply.github.com>

Author: IvanBM18

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)