Bug 1889223 - Submit all files to notarization before polling

This is the first behavior that consumes all files at once, so it's a bit odd how it's being handled in script.py

signingscript/src/signingscript/script.py

@@ -9,7 +9,7 @@
 import scriptworker.client
 
 from signingscript.exceptions import SigningScriptError
-from signingscript.task import build_filelist_dict, sign, task_cert_type, task_signing_formats
+from signingscript.task import apple_notarize_stacked, build_filelist_dict, sign, task_cert_type, task_signing_formats
 from signingscript.utils import copy_to_dir, load_apple_notarization_configs, load_autograph_configs
 
 log = logging.getLogger(__name__)
@@ -23,6 +23,7 @@ async def async_main(context):
         context (Context): the signing context.
 
     """
+    work_dir = context.config["work_dir"]
     async with aiohttp.ClientSession() as session:
         all_signing_formats = task_signing_formats(context)
         if "gpg" in all_signing_formats or "autograph_gpg" in all_signing_formats:
@@ -40,10 +41,23 @@ async def async_main(context):
                 raise Exception("Apple notarization is enabled but apple_notarization_configs is not defined")
             setup_apple_notarization_credentials(context)
 
+        # Only handle notarization in parallel if that's the only format requested
+        # TODO: This is the first behavior that consumes all files at once, so it's
+        #       here for now. Maybe this isn't the best place for this, maybe it
+        #       should be in task.py since most of the format handling is there
+        if len(all_signing_formats) == 1 and "apple_notarization" in all_signing_formats:
+            # Notarization handles ALL files at once - so we don't need to iterate over the filelist_dict
+            filelist_dict = build_filelist_dict(context)
+            output_files = await apple_notarize_stacked(context, filelist_dict)
+            for source in output_files:
+                source = os.path.relpath(source, work_dir)
+                copy_to_dir(os.path.join(work_dir, source), context.config["artifact_dir"], target=source)
+            log.info("Done!")
+            return
+
         context.session = session
         context.autograph_configs = load_autograph_configs(context.config["autograph_configs"])
 
-        work_dir = context.config["work_dir"]
         filelist_dict = build_filelist_dict(context)
         for path, path_dict in filelist_dict.items():
             copy_to_dir(path_dict["full_path"], context.config["work_dir"], target=path)

signingscript/src/signingscript/sign.py

@@ -1633,8 +1633,7 @@ async def apple_notarize(context, path, *args, **kwargs):
     """
     # Setup workdir
     notarization_workdir = os.path.join(context.config["work_dir"], "apple_notarize")
-    shutil.rmtree(notarization_workdir, ignore_errors=True)
-    utils.mkdir(notarization_workdir)
+    utils.mkdir(notarization_workdir, delete_before_create=True)
 
     _, extension = os.path.splitext(path)
     if extension == ".pkg":
@@ -1650,7 +1649,85 @@ async def apple_notarize_geckodriver(context, path, *args, **kwargs):
     """
     # Setup workdir
     notarization_workdir = os.path.join(context.config["work_dir"], "apple_notarize")
-    shutil.rmtree(notarization_workdir, ignore_errors=True)
-    utils.mkdir(notarization_workdir)
+    utils.mkdir(notarization_workdir, delete_before_create=True)
 
     return await _notarize_geckodriver(context, path, notarization_workdir)
+
+
+@time_async_function
+async def apple_notarize_stacked(context, filelist_dict):
+    """
+    Notarizes multiple packages using rcodesign.
+    Submits everything before polling for status.
+    """
+    ATTEMPTS = 5
+
+    # notarization submissions map (path -> submission_id)
+    submissions_map = {}
+    relpath_index_map = {}
+    task_index = 0
+    # Submit to notarization one by one
+    for relpath, path_dict in filelist_dict.items():
+        task_index += 1
+        relpath_index_map[relpath] = task_index
+        notarization_workdir = os.path.join(context.config["work_dir"], f"apple_notarize-{task_index}")
+        utils.mkdir(notarization_workdir, delete_before_create=True)
+        _, extension = os.path.splitext(relpath)
+        if extension == ".pkg":
+            path = os.path.join(notarization_workdir, relpath)
+            utils.copy_to_dir(path_dict["full_path"], notarization_workdir, target=relpath)
+            submissions_map[path] = await retry_async(
+                func=rcodesign_notarize,
+                args=(path, context.apple_credentials_path),
+                attempts=ATTEMPTS,
+                retry_exceptions=RCodesignError,
+            )
+        else:
+            await _extract_tarfile(context, path_dict["full_path"], extension, notarization_workdir)
+            workdir_files = os.listdir(notarization_workdir)
+            supported_files = [filename for filename in workdir_files if _can_notarize(filename, (".app", ".pkg"))]
+            if not supported_files:
+                raise SigningScriptError("No supported files found")
+            for file in supported_files:
+                path = os.path.join(notarization_workdir, file)
+                submissions_map[path] = await retry_async(
+                    func=rcodesign_notarize,
+                    args=(path, context.apple_credentials_path),
+                    attempts=ATTEMPTS,
+                    retry_exceptions=RCodesignError,
+                )
+
+    # Notary wait all files
+    for path, submission_id in submissions_map.items():
+        await retry_async(
+            func=rcodesign_notary_wait,
+            args=(submission_id, context.apple_credentials_path),
+            attempts=ATTEMPTS,
+            retry_exceptions=RCodesignError,
+        )
+
+    for path in submissions_map.keys():
+        await retry_async(
+            func=rcodesign_staple,
+            args=[path],
+            attempts=ATTEMPTS,
+            retry_exceptions=RCodesignError,
+        )
+
+    # Staple + create tarball where necessary
+    stapled_files = []
+    for relpath, path_dict in filelist_dict.items():
+        task_index = relpath_index_map[relpath]
+        notarization_workdir = os.path.join(context.config["work_dir"], f"apple_notarize-{task_index}")
+        target_path = os.path.join(context.config["work_dir"], relpath)
+        _, extension = os.path.splitext(relpath)
+        if extension == ".pkg":
+            utils.copy_to_dir(os.path.join(notarization_workdir, relpath), os.path.dirname(target_path))
+        else:
+            all_files = []
+            for root, _, files in os.walk(notarization_workdir):
+                for f in files:
+                    all_files.append(os.path.join(root, f))
+            await _create_tarfile(context, target_path, all_files, extension, notarization_workdir)
+        stapled_files.append(target_path)
+    return stapled_files

signingscript/src/signingscript/task.py

@@ -15,8 +15,8 @@
 from scriptworker.utils import get_single_item_from_sequence
 
 from signingscript.sign import (
-    apple_notarize,
     apple_notarize_geckodriver,
+    apple_notarize_stacked,
     sign_authenticode,
     sign_debian_pkg,
     sign_file,
@@ -53,7 +53,7 @@
         "privileged_webextension": sign_xpi,
         "system_addon": sign_xpi,
         "autograph_rsa": sign_file_detached,
-        "apple_notarization": apple_notarize,
+        "apple_notarization": apple_notarize_stacked,
         "apple_notarization_geckodriver": apple_notarize_geckodriver,
         "default": sign_file,
     }

signingscript/src/signingscript/utils.py

@@ -6,6 +6,7 @@
 import json
 import logging
 import os
+import shutil
 from asyncio.subprocess import PIPE, STDOUT
 from dataclasses import dataclass
 from shutil import copyfile
@@ -35,13 +36,15 @@ class AppleNotarization:
     private_key: str
 
 
-def mkdir(path):
+def mkdir(path, delete_before_create=False):
     """Equivalent to `mkdir -p`.
 
     Args:
         path (str): the path to mkdir
-
+        delete_before_create (bool, optional): whether to delete the path before creating it. Defaults to False
     """
+    if delete_before_create:
+        shutil.rmtree(path, ignore_errors=True)
     try:
         os.makedirs(path)
         log.info("mkdir {}".format(path))

signingscript/tests/test_script.py

@@ -30,13 +30,17 @@ async def fake_sign(_, val, *args, authenticode_comment=None):
             assert authenticode_comment == "Some authenticode comment"
         return [val]
 
+    async def fake_notarize_stacked(_, filelist_dict, *args, **kwargs):
+        return filelist_dict.keys()
+
     mocker.patch.object(script, "load_autograph_configs", new=noop_sync)
     mocker.patch.object(script, "load_apple_notarization_configs", new=noop_sync)
     mocker.patch.object(script, "setup_apple_notarization_credentials", new=noop_sync)
     # mocker.patch.object(script, "task_cert_type", new=noop_sync)
     mocker.patch.object(script, "task_signing_formats", return_value=formats)
     mocker.patch.object(script, "build_filelist_dict", new=fake_filelist_dict)
     mocker.patch.object(script, "sign", new=fake_sign)
+    mocker.patch.object(script, "apple_notarize_stacked", new=fake_notarize_stacked)
     context = mock.MagicMock()
     context.config = {"work_dir": tmpdir, "artifact_dir": tmpdir, "autograph_configs": {}, "apple_notarization_configs": "fake"}
     context.config.update(extra_config)

signingscript/tests/test_task.py

@@ -151,6 +151,7 @@ def fake_log(context, new_files, *args):
         ("widevine", stask.sign_widevine),
         ("autograph_authenticode", stask.sign_authenticode),
         ("autograph_authenticode_stub", stask.sign_authenticode),
+        ("apple_notarization", stask.apple_notarize_stacked),
         ("default", stask.sign_file),
         # Key id cases
         ("autograph_hash_only_mar384:firefox_20190321_dev", stask.sign_mar384_with_autograph_hash),