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:
@@ -35,17 +36,19 @@ async def async_main(context):
             if not context.config.get("widevine_cert"):
                 raise Exception("Widevine format is enabled, but widevine_cert is not defined")
 
-        if "apple_notarization" in all_signing_formats or "apple_notarization_geckodriver" in all_signing_formats:
+        if {"apple_notarization", "apple_notarization_geckodriver", "apple_notarization_stacked"}.intersection(all_signing_formats):
             if not context.config.get("apple_notarization_configs", False):
                 raise Exception("Apple notarization is enabled but apple_notarization_configs is not defined")
             setup_apple_notarization_credentials(context)
 
         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():
+            if path_dict["formats"] == ["apple_notarization_stacked"]:
+                # Skip if only format is notarization_stacked - handled below
+                continue
             copy_to_dir(path_dict["full_path"], context.config["work_dir"], target=path)
             log.info("signing %s", path)
             output_files = await sign(context, os.path.join(work_dir, path), path_dict["formats"], authenticode_comment=path_dict.get("comment"))
@@ -54,6 +57,15 @@ async def async_main(context):
                 copy_to_dir(os.path.join(work_dir, source), context.config["artifact_dir"], target=source)
             if "gpg" in path_dict["formats"] or "autograph_gpg" in path_dict["formats"]:
                 copy_to_dir(context.config["gpg_pubkey"], context.config["artifact_dir"], target="public/build/KEY")
+
+        # notarization_stacked is a special format that takes in all files at once instead of sequentially like other formats
+        notarization_dict = {path: path_dict for path, path_dict in filelist_dict.items() if "apple_notarization_stacked" in path_dict["formats"]}
+        if notarization_dict:
+            output_files = await apple_notarize_stacked(context, notarization_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!")
 
 

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

@@ -17,6 +17,7 @@
 from signingscript.sign import (
     apple_notarize,
     apple_notarize_geckodriver,
+    apple_notarize_stacked,  # noqa: F401
     sign_authenticode,
     sign_debian_pkg,
     sign_file,
@@ -32,6 +33,10 @@
 
 log = logging.getLogger(__name__)
 
+
+async def noop_sign(*args, **kwargs):
+    return []
+
 FORMAT_TO_SIGNING_FUNCTION = immutabledict(
     {
         "autograph_hash_only_mar384": sign_mar384_with_autograph_hash,
@@ -55,6 +60,9 @@
         "autograph_rsa": sign_file_detached,
         "apple_notarization": apple_notarize,
         "apple_notarization_geckodriver": apple_notarize_geckodriver,
+        # This format is handled in script.py
+        # "apple_notarization_stacked": apple_notarize_stacked,
+        "apple_notarization_stacked": noop_sign,
         "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)
@@ -99,6 +103,19 @@ async def test_async_main_apple_notarization(tmpdir, mocker):
     await async_main_helper(tmpdir, mocker, formats)
 
 
+@pytest.mark.asyncio
+async def test_async_main_apple_notarization_stacked(tmpdir, mocker):
+    formats = ["apple_notarization_stacked"]
+    mocker.patch.object(script, "copy_to_dir", new=noop_sync)
+    await async_main_helper(tmpdir, mocker, formats)
+
+
+@pytest.mark.asyncio
+async def test_async_main_apple_notarization_stacked_multi(tmpdir, mocker):
+    formats = ["autograph_mar", "apple_notarization_stacked"]
+    mocker.patch.object(script, "copy_to_dir", new=noop_sync)
+    await async_main_helper(tmpdir, mocker, formats)
+
 @pytest.mark.asyncio
 async def test_async_main_apple_notarization_no_config(tmpdir, mocker):
     formats = ["apple_notarization"]

signingscript/tests/test_sign.py

@@ -23,6 +23,7 @@
 
 import signingscript.sign as sign
 import signingscript.utils as utils
+import signingscript.rcodesign as rcodesign
 from signingscript.exceptions import SigningScriptError
 from signingscript.utils import get_hash
 
@@ -1481,3 +1482,29 @@ async def test_apple_notarize_geckodriver(mocker, context):
 
     await sign.apple_notarize_geckodriver(context, "/foo/bar.pkg")
     notarize_geckodriver.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_apple_notarize_stacked(mocker, context):
+    notarize = mock.AsyncMock()
+    mocker.patch.object(sign, "rcodesign_notarize", notarize)
+    wait = mock.AsyncMock()
+    mocker.patch.object(sign, "rcodesign_notary_wait", wait)
+    staple = mock.AsyncMock()
+    mocker.patch.object(sign, "rcodesign_staple", staple)
+
+    mocker.patch.object(sign, "_extract_tarfile", noop_async)
+    mocker.patch.object(sign, "_create_tarfile", noop_async)
+    mocker.patch.object(sign.os, "listdir", lambda *_: ["/foo.pkg", "/baz.app", "/foobar"])
+    mocker.patch.object(sign.shutil, "rmtree", noop_sync)
+    mocker.patch.object(sign.utils, "mkdir", noop_sync)
+    mocker.patch.object(sign.utils, "copy_to_dir", noop_sync)
+
+    await sign.apple_notarize_stacked(
+        context,
+        {"/app.tar.gz": {"full_path": "/app.tar.gz", "formats": ["apple_notarize_stacked"]}}
+        )
+    # one for each file format
+    assert notarize.await_count == 2
+    assert wait.await_count == 2
+    assert staple.await_count == 2

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),
         ("default", stask.sign_file),
         # Key id cases
         ("autograph_hash_only_mar384:firefox_20190321_dev", stask.sign_mar384_with_autograph_hash),