From 88dba120a8b331b224c01345399530cab8117ffc Mon Sep 17 00:00:00 2001
From: jonathanmetzman <31354670+jonathanmetzman@users.noreply.github.com>
Date: Wed, 4 Aug 2021 14:19:22 -0700
Subject: [PATCH] [CIFuzz][NFC] Add function for setting env vars in docker
 command. (#6162)

---
 infra/cifuzz/docker.py                   | 19 +++++++++++++++----
 infra/cifuzz/fuzz_target.py              | 12 ++++++------
 infra/cifuzz/generate_coverage_report.py | 16 ++++++++++------
 3 files changed, 31 insertions(+), 16 deletions(-)

diff --git a/infra/cifuzz/docker.py b/infra/cifuzz/docker.py
index 6397fa874..15858b458 100644
--- a/infra/cifuzz/docker.py
+++ b/infra/cifuzz/docker.py
@@ -43,6 +43,15 @@ _DEFAULT_DOCKER_RUN_COMMAND = [
 ]
 
 
+def get_docker_env_vars(env_mapping):
+  """Returns a list of docker arguments that sets each key in |env_mapping| as
+  an env var and the value of that key in |env_mapping| as the value."""
+  env_var_args = []
+  for env_var, env_var_val in env_mapping.items():
+    env_var_args.extend(['-e', f'{env_var}={env_var_val}'])
+  return env_var_args
+
+
 def get_project_image_name(project):
   """Returns the name of the project builder image for |project_name|."""
   # TODO(ochang): We may need unique names to support parallel fuzzing.
@@ -65,10 +74,12 @@ def get_base_docker_run_args(workspace,
   """Returns arguments that should be passed to every invocation of 'docker
   run'."""
   docker_args = _DEFAULT_DOCKER_RUN_ARGS.copy()
-  docker_args += [
-      '-e', f'SANITIZER={sanitizer}', '-e', f'FUZZING_LANGUAGE={language}',
-      '-e', 'OUT=' + workspace.out
-  ]
+  env_mapping = {
+      'SANITIZER': sanitizer,
+      'FUZZING_LANGUAGE': language,
+      'OUT': workspace.out
+  }
+  docker_args += get_docker_env_vars(env_mapping)
   docker_container = utils.get_container_name()
   logging.info('Docker container: %s.', docker_container)
   if docker_container:
diff --git a/infra/cifuzz/fuzz_target.py b/infra/cifuzz/fuzz_target.py
index 19dec4022..e972b55b5 100644
--- a/infra/cifuzz/fuzz_target.py
+++ b/infra/cifuzz/fuzz_target.py
@@ -99,12 +99,12 @@ class FuzzTarget:  # pylint: disable=too-many-instance-attributes
     # If corpus can be downloaded use it for fuzzing.
     self.latest_corpus_path = self.clusterfuzz_deployment.download_corpus(
         self.target_name)
-    command += ['-e', 'CORPUS_DIR=' + self.latest_corpus_path]
+    command += docker.get_docker_env_vars({
+        'CORPUS_DIR': self.latest_corpus_path,
+        'RUN_FUZZER_MODE': 'interactive'
+    })
 
-    command += [
-        '-e', 'RUN_FUZZER_MODE=interactive', docker.BASE_RUNNER_TAG, 'bash',
-        '-c'
-    ]
+    command += [docker.BASE_RUNNER_TAG, 'bash', '-c']
 
     options = LIBFUZZER_OPTIONS.copy() + [
         f'-max_total_time={self.duration}',
@@ -195,7 +195,7 @@ class FuzzTarget:  # pylint: disable=too-many-instance-attributes
     command, container = docker.get_base_docker_run_command(
         self.workspace, self.config.sanitizer, self.config.language)
     if container:
-      command += ['-e', f'TESTCASE={testcase}']
+      command += docker.get_docker_env_vars({'TESTCASE': testcase})
     else:
       command += ['-v', f'{testcase}:/testcase']
 
diff --git a/infra/cifuzz/generate_coverage_report.py b/infra/cifuzz/generate_coverage_report.py
index fa4fc8339..1d392c15a 100644
--- a/infra/cifuzz/generate_coverage_report.py
+++ b/infra/cifuzz/generate_coverage_report.py
@@ -22,12 +22,16 @@ def run_coverage_command(workspace, config):
   """Runs the coverage command in base-runner to generate a coverage report."""
   docker_args, _ = docker.get_base_docker_run_args(workspace, config.sanitizer,
                                                    config.language)
-  docker_args += [
-      '-e', 'COVERAGE_EXTRA_ARGS=', '-e', 'HTTP_PORT=', '-e',
-      f'CORPUS_DIR={workspace.corpora}', '-e',
-      f'COVERAGE_OUTPUT_DIR={workspace.coverage_report}', '-t',
-      docker.BASE_RUNNER_TAG, 'coverage'
-  ]
+  env_mapping = {
+      'COVERAGE_EXTRA_ARGS': '',
+      'HTTP_PORT': '',
+      'CORPUS_DIR': workspace.corpora,
+      'COVERAGE_OUTPUT_DIR': workspace.coverage_report
+  }
+  docker_args += docker.get_docker_env_vars(env_mapping)
+
+  docker_args += ['-t', docker.BASE_RUNNER_TAG, 'coverage']
+
   return helper.docker_run(docker_args)