Bug 1889223 - Submit all files to notarization before polling

signingscript/src/signingscript/script.py

@@ -10,7 +10,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__)
@@ -24,6 +24,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:
@@ -36,17 +37,24 @@ 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"]
+        # TODO: Make task.sign take in the whole filelist_dict and return a dict of output files.
+        #       That would likely mean changing all behaviors to accept and deal with multiple files at once.
+
         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
+            if "apple_notarization_stacked" in path_dict["formats"]:
+                raise SigningScriptError("apple_notarization_stacked cannot be mixed with other signing types")
             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"))
@@ -55,6 +63,16 @@ 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
+        # Should be fixed in https://github.com/mozilla-releng/scriptworker-scripts/issues/980
+        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

@@ -1604,7 +1604,11 @@ async def _notarize_geckodriver(context, path, workdir):
 
 
 async def _notarize_all(context, path, workdir):
-    """Notarizes all files in a tarball"""
+    """
+    Notarizes all files in a tarball
+
+    @Deprecated: This function is deprecated and will be removed in the future. Use apple_notarize_stacked instead.
+    """
     _, extension = os.path.splitext(path)
     # Attempt extracting
     await _extract_tarfile(context, path, extension, tmp_dir=workdir)
@@ -1634,10 +1638,11 @@ async def _notarize_all(context, path, workdir):
 async def apple_notarize(context, path, *args, **kwargs):
     """
     Notarizes given package(s) using rcodesign.
+
+    @Deprecated: This function is deprecated and will be removed in the future. Use apple_notarize_stacked instead.
     """
     # Setup workdir
     notarization_workdir = os.path.join(context.config["work_dir"], "apple_notarize")
-    shutil.rmtree(notarization_workdir, ignore_errors=True)
     utils.mkdir(notarization_workdir)
 
     _, extension = os.path.splitext(path)
@@ -1658,3 +1663,88 @@ async def apple_notarize_geckodriver(context, path, *args, **kwargs):
     utils.mkdir(notarization_workdir)
 
     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
+
+    relpath_index_map = {}
+    paths_to_notarize = []
+    task_index = 0
+
+    # Create list of files to be notarized + check for potential problems
+    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}")
+        shutil.rmtree(notarization_workdir, ignore_errors=True)
+        utils.mkdir(notarization_workdir)
+        _, 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)
+            paths_to_notarize.append(path)
+        elif extension == ".gz":
+            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(f"No supported files found for file {relpath}")
+            for file in supported_files:
+                path = os.path.join(notarization_workdir, file)
+                paths_to_notarize.append(path)
+        else:
+            raise SigningScriptError(f"Unsupported file extension: {extension} for file {relpath}")
+
+    # notarization submissions map (path -> submission_id)
+    submissions_map = {}
+    # Submit to notarization one by one
+    for path in paths_to_notarize:
+        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,
+        )
+
+    # Staple files
+    for path in submissions_map.keys():
+        await retry_async(
+            func=rcodesign_staple,
+            args=[path],
+            attempts=ATTEMPTS,
+            retry_exceptions=RCodesignError,
+        )
+
+    # Wrap up
+    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)
+        # Pkgs don't need to be tarred
+        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

@@ -18,6 +18,7 @@
 from signingscript.sign import (
     apple_notarize,
     apple_notarize_geckodriver,
+    apple_notarize_stacked,  # noqa: F401
     sign_authenticode,
     sign_debian_pkg,
     sign_file,
@@ -33,6 +34,7 @@
 
 log = logging.getLogger(__name__)
 
+
 FORMAT_TO_SIGNING_FUNCTION = immutabledict(
     {
         "autograph_hash_only_mar384": sign_mar384_with_autograph_hash,
@@ -58,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
+        # Should be refactored in https://github.com/mozilla-releng/scriptworker-scripts/issues/980
+        # "apple_notarization_stacked": apple_notarize_stacked,
         "default": sign_file,
     }
 )

signingscript/src/signingscript/utils.py

@@ -40,7 +40,6 @@ def mkdir(path):
 
     Args:
         path (str): the path to mkdir
-
     """
     try:
         os.makedirs(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,21 @@ 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_mixed_fail(tmpdir, mocker):
+    formats = ["autograph_mar", "apple_notarization_stacked"]
+    mocker.patch.object(script, "copy_to_dir", new=noop_sync)
+    with pytest.raises(SigningScriptError):
+        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

@@ -1481,3 +1481,64 @@ 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.os, "walk", lambda *_: [("/", None, ["foo.pkg", "baz.app"])])
+    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"]},
+            "/app2.pkg": {"full_path": "/app2.pkg", "formats": ["apple_notarize_stacked"]},
+        },
+    )
+    # one for each file format
+    assert notarize.await_count == 3
+    assert wait.await_count == 3
+    assert staple.await_count == 3
+
+
+@pytest.mark.asyncio
+async def test_apple_notarize_stacked_unsupported(mocker, context):
+    """Test unsupported file extensions"""
+
+    mocker.patch.object(sign, "_extract_tarfile", noop_async)
+    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)
+
+    # Returns unsupported file formats
+    mocker.patch.object(sign.os, "listdir", lambda *_: ["/foo.aaa", "/baz.bbb", "/foobar"])
+
+    with pytest.raises(SigningScriptError):
+        # Main file is supported, contents uses the above os.listdir
+        await sign.apple_notarize_stacked(
+            context,
+            {
+                "/app.tar.gz": {"full_path": "/app.tar.gz", "formats": ["apple_notarize_stacked"]},
+            },
+        )
+
+    with pytest.raises(SigningScriptError):
+        # Main file extension is unsupported
+        await sign.apple_notarize_stacked(
+            context,
+            {
+                "/app.bbb": {"full_path": "/app.bbb", "formats": ["apple_notarize_stacked"]},
+            },
+        )

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),