ukify: allow multiple initrds

If given, multiple initrds are concatenated into a temporary file which then
becomes the .initrd section.

It is also possible to give no initrd. After all, some machines boot without an
initrd, and it should be possible to use the stub without requiring an initrd.
(The stub might not like this, but this is something to fix there.)

2 files changed, 38 insertions, 11 deletions

diff --git a/src/ukify/test/test_ukify.py b/src/ukify/test/test_ukify.py
index 48ffc6d495..34701402e5 100755
--- a/src/ukify/test/test_ukify.py
+++ b/src/ukify/test/test_ukify.py
@@ -49,13 +49,13 @@ def test_round_up():
 def test_parse_args_minimal():
     opts = ukify.parse_args('arg1 arg2'.split())
     assert opts.linux == pathlib.Path('arg1')
-    assert opts.initrd == pathlib.Path('arg2')
+    assert opts.initrd == [pathlib.Path('arg2')]
     assert opts.os_release in (pathlib.Path('/etc/os-release'),
                                pathlib.Path('/usr/lib/os-release'))
 
 def test_parse_args_many():
     opts = ukify.parse_args(
-        ['/ARG1', '///ARG2',
+        ['/ARG1', '///ARG2', '/ARG3 WITH SPACE',
          '--cmdline=a b c',
          '--os-release=K1=V1\nK2=V2',
          '--devicetree=DDDDTTTT',
@@ -77,7 +77,7 @@ def test_parse_args_many():
          '--no-measure',
          ])
     assert opts.linux == pathlib.Path('/ARG1')
-    assert opts.initrd == pathlib.Path('/ARG2')
+    assert opts.initrd == [pathlib.Path('/ARG2'), pathlib.Path('/ARG3 WITH SPACE')]
     assert opts.os_release == 'K1=V1\nK2=V2'
     assert opts.devicetree == pathlib.Path('DDDDTTTT')
     assert opts.splash == pathlib.Path('splash')
@@ -103,7 +103,7 @@ def test_parse_sections():
          ])
 
     assert opts.linux == pathlib.Path('/ARG1')
-    assert opts.initrd == pathlib.Path('/ARG2')
+    assert opts.initrd == [pathlib.Path('/ARG2')]
     assert len(opts.sections) == 2
 
     assert opts.sections[0].name == 'test'
@@ -334,9 +334,13 @@ def test_pcr_signing2(kernel_initrd, tmpdir):
     pub2 = unbase64(ourdir / 'example.tpm2-pcr-public2.pem.base64')
     priv2 = unbase64(ourdir / 'example.tpm2-pcr-private2.pem.base64')
 
+    # simulate a microcode file
+    with open(f'{tmpdir}/microcode', 'wb') as microcode:
+        microcode.write(b'1234567890')
+
     output = f'{tmpdir}/signed.efi'
     opts = ukify.parse_args([
-        *kernel_initrd,
+        kernel_initrd[0], microcode.name, kernel_initrd[1],
         f'--output={output}',
         '--uname=1.2.3',
         '--cmdline=ARG1 ARG2 ARG3',
@@ -367,7 +371,7 @@ def test_pcr_signing2(kernel_initrd, tmpdir):
     subprocess.check_call([
         'objcopy',
         *(f'--dump-section=.{n}={tmpdir}/out.{n}' for n in (
-            'pcrpkey', 'pcrsig', 'osrel', 'uname', 'cmdline')),
+            'pcrpkey', 'pcrsig', 'osrel', 'uname', 'cmdline', 'initrd')),
         output,
         tmpdir / 'dummy',
     ],
@@ -377,6 +381,8 @@ def test_pcr_signing2(kernel_initrd, tmpdir):
     assert open(tmpdir / 'out.osrel').read() == 'ID=foobar\n'
     assert open(tmpdir / 'out.uname').read() == '1.2.3'
     assert open(tmpdir / 'out.cmdline').read() == 'ARG1 ARG2 ARG3'
+    assert open(tmpdir / 'out.initrd', 'rb').read(10) == b'1234567890'
+
     sig = open(tmpdir / 'out.pcrsig').read()
     sig = json.loads(sig)
     assert list(sig.keys()) == ['sha1']
diff --git a/src/ukify/ukify.py b/src/ukify/ukify.py
index 83423fc720..e9e5d13d13 100755
--- a/src/ukify/ukify.py
+++ b/src/ukify/ukify.py
@@ -206,8 +206,9 @@ class Section:
 
     @classmethod
     def create(cls, name, contents, flags=None, measure=False):
-        if isinstance(contents, str):
-            tmp = tempfile.NamedTemporaryFile(mode='wt', prefix=f'tmp{name}')
+        if isinstance(contents, str | bytes):
+            mode = 'wt' if isinstance(contents, str) else 'wb'
+            tmp = tempfile.NamedTemporaryFile(mode=mode, prefix=f'tmp{name}')
             tmp.write(contents)
             tmp.flush()
             contents = pathlib.Path(tmp.name)
@@ -404,6 +405,24 @@ def call_systemd_measure(uki, linux, opts):
         uki.add_section(Section.create('.pcrsig', combined))
 
 
+def join_initrds(initrds):
+    match initrds:
+        case []:
+            return None
+        case [initrd]:
+            return initrd
+        case multiple:
+            seq = []
+            for file in multiple:
+                initrd = file.read_bytes()
+                padding = b'\0' * round_up(len(initrd), 4)  # pad to 32 bit alignment
+                seq += [initrd, padding]
+
+            return b''.join(seq)
+
+    assert False
+
+
 def make_uki(opts):
     # kernel payload signing
 
@@ -455,6 +474,7 @@ def make_uki(opts):
         opts.uname = Uname.scrape(opts.linux, opts=opts)
 
     uki = UKI(opts.stub)
+    initrd = join_initrds(opts.initrd)
 
     # TODO: derive public key from from opts.pcr_private_keys?
     pcrpkey = opts.pcrpkey
@@ -469,7 +489,7 @@ def make_uki(opts):
         ('.dtb',     opts.devicetree, True ),
         ('.splash',  opts.splash,     True ),
         ('.pcrpkey', pcrpkey,         True ),
-        ('.initrd',  opts.initrd,     True ),
+        ('.initrd',  initrd,          True ),
         ('.uname',   opts.uname,      False),
 
         # linux shall be last to leave breathing room for decompression.
@@ -541,7 +561,7 @@ def parse_args(args=None):
         description='Build and sign Unified Kernel Images',
         allow_abbrev=False,
         usage='''\
-usage: ukify [options…] linux initrd
+usage: ukify [options…] linux initrd…
        ukify -h | --help
 ''')
 
@@ -553,7 +573,8 @@ usage: ukify [options…] linux initrd
                    help='vmlinuz file [.linux section]')
     p.add_argument('initrd',
                    type=pathlib.Path,
-                   help='initrd file [.initrd section]')
+                   nargs='*',
+                   help='initrd files [.initrd section]')
 
     p.add_argument('--cmdline',
                    metavar='TEXT|@PATH',