From 4cd36579ed20d28af36728ee6a84b2a04b3d9608 Mon Sep 17 00:00:00 2001
From: David Korczynski <david@adalogics.com>
Date: Sun, 13 Oct 2024 22:43:12 +0100
Subject: [PATCH] infra: enable building projects using cached images

Signed-off-by: David Korczynski <david@adalogics.com>
---
 infra/build/functions/build_project.py     | 17 ++++++++++++++---
 infra/build/functions/target_experiment.py | 10 +++++++---
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/infra/build/functions/build_project.py b/infra/build/functions/build_project.py
index 7c4ff7f142f1..a78f0598c59f 100755
--- a/infra/build/functions/build_project.py
+++ b/infra/build/functions/build_project.py
@@ -174,6 +174,12 @@ def image(self):
     """Returns the docker image for the project."""
     return f'gcr.io/{build_lib.IMAGE_PROJECT}/{self.name}'
 
+  def cached_image(self, coverage):
+    sanitizer_mapping = {'address': 'asan', 'coverage': 'cov'}
+    san_lookup = sanitizer_mapping[sanitizer]
+    return ('us-central1-docker.pkg.dev/oss-fuzz/oss-fuzz-gen/'
+            f'{self.name}-ofg-cached-{san_lookup}')
+
 
 def get_last_step_id(steps):
   """Returns the id of the last step in |steps|."""
@@ -308,10 +314,15 @@ def get_build_steps(  # pylint: disable=too-many-locals, too-many-statements, to
     project_yaml,
     dockerfile,
     config,
-    additional_env=None):
+    additional_env=None,
+    use_caching=False):
   """Returns build steps for project."""
 
   project = Project(project_name, project_yaml, dockerfile)
+  if use_caching:
+    project_image = project.cached_image()
+  else:
+    project_image = project.image
 
   if project.disabled:
     logging.info('Project "%s" is disabled.', project.name)
@@ -320,7 +331,7 @@ def get_build_steps(  # pylint: disable=too-many-locals, too-many-statements, to
   timestamp = get_datetime_now().strftime('%Y%m%d%H%M')
   build_steps = build_lib.get_project_image_steps(
       project.name,
-      project.image,
+      project_image,
       project.fuzzing_language,
       config=config,
       architectures=project.architectures,
@@ -357,7 +368,7 @@ def get_build_steps(  # pylint: disable=too-many-locals, too-many-statements, to
           # Report the build failure if it happened.
           build_steps.append({
               'name':
-                  project.image,
+                  project_image,
               'args': [
                   'bash', '-c',
                   f'cat {LOCAL_BUILD_LOG_PATH} && test -f {BUILD_SUCCESS_MARKER}'
diff --git a/infra/build/functions/target_experiment.py b/infra/build/functions/target_experiment.py
index 84edd93e69f3..b26d6ca86cd4 100644
--- a/infra/build/functions/target_experiment.py
+++ b/infra/build/functions/target_experiment.py
@@ -30,7 +30,7 @@
 
 def run_experiment(project_name, target_name, args, output_path, errlog_path,
                    build_output_path, upload_corpus_path, upload_coverage_path,
-                   experiment_name, upload_reproducer_path, tags):
+                   experiment_name, upload_reproducer_path, tags, use_cached_image):
   config = build_project.Config(testing=True,
                                 test_image_suffix='',
                                 repo=build_project.DEFAULT_OSS_FUZZ_REPO,
@@ -68,7 +68,8 @@ def run_experiment(project_name, target_name, args, output_path, errlog_path,
                                         project_yaml,
                                         dockerfile_contents,
                                         config,
-                                        additional_env=jcc_env)
+                                        additional_env=jcc_env,
+                                        use_caching=use_cached_image)
 
   build = build_project.Build('libfuzzer', 'address', 'x86_64')
   local_output_path = '/workspace/output.log'
@@ -330,12 +331,15 @@ def main():
                       nargs='*',
                       help='Tags for cloud build.',
                       default=[])
+  parser.add_argument('--use_cached_image',
+                      action='store_true',
+                      help='Use cached images post build.')
   args = parser.parse_args()
 
   run_experiment(args.project, args.target, args.args, args.upload_output_log,
                  args.upload_err_log, args.upload_build_log, args.upload_corpus,
                  args.upload_coverage, args.experiment_name,
-                 args.upload_reproducer, args.tags)
+                 args.upload_reproducer, args.tags, args.use_cached_image)
 
 
 if __name__ == '__main__':