xref: /linux/scripts/make_fit.py (revision 3bf2237f59af842e546fa4b2c146198db95d1519) 
1#!/usr/bin/env python3
2# SPDX-License-Identifier: GPL-2.0+
3#
4# Copyright 2024 Google LLC
5# Written by Simon Glass <sjg@chromium.org>
6#
7
8"""Build a FIT containing a lot of devicetree files
9
10Usage:
11    make_fit.py -A arm64 -n 'Linux-6.6' -O linux
12        -o arch/arm64/boot/image.fit -k /tmp/kern/arch/arm64/boot/image.itk
13        -r /boot/initrd.img-6.14.0-27-generic @arch/arm64/boot/dts/dtbs-list
14        -E -c gzip
15
16Creates a FIT containing the supplied kernel, an optional ramdisk, and a set of
17devicetree files, either specified individually or listed in a file (with an
18'@' prefix).
19
20Use -r to specify an existing ramdisk/initrd file.
21
22Use -E to generate an external FIT (where the data is placed after the
23FIT data structure). This allows parsing of the data without loading
24the entire FIT.
25
26Use -c to compress the data, using bzip2, gzip, lz4, lzma, lzo and
27zstd algorithms.
28
29Use -D to decompose "composite" DTBs into their base components and
30deduplicate the resulting base DTBs and DTB overlays. This requires the
31DTBs to be sourced from the kernel build directory, as the implementation
32looks at the .cmd files produced by the kernel build.
33
34The resulting FIT can be booted by bootloaders which support FIT, such
35as U-Boot, Linuxboot, Tianocore, etc.
36"""
37
38import argparse
39import collections
40import multiprocessing
41import os
42import subprocess
43import sys
44import tempfile
45import time
46
47import libfdt
48
49
50# Tool extension and the name of the command-line tools
51CompTool = collections.namedtuple('CompTool', 'ext,tools')
52
53COMP_TOOLS = {
54    'bzip2': CompTool('.bz2', 'pbzip2,bzip2'),
55    'gzip': CompTool('.gz', 'pigz,gzip'),
56    'lz4': CompTool('.lz4', 'lz4'),
57    'lzma': CompTool('.lzma', 'lzma'),
58    'lzo': CompTool('.lzo', 'lzop'),
59    'xz': CompTool('.xz', 'xz'),
60    'zstd': CompTool('.zstd', 'zstd'),
61}
62
63
64def parse_args():
65    """Parse the program ArgumentParser
66
67    Returns:
68        Namespace object containing the arguments
69    """
70    epilog = 'Build a FIT from a directory tree containing .dtb files'
71    parser = argparse.ArgumentParser(epilog=epilog, fromfile_prefix_chars='@')
72    parser.add_argument('-A', '--arch', type=str, required=True,
73          help='Specifies the architecture')
74    parser.add_argument('-c', '--compress', type=str, default='none',
75          help='Specifies the compression')
76    parser.add_argument('-D', '--decompose-dtbs', action='store_true',
77          help='Decompose composite DTBs into base DTB and overlays')
78    parser.add_argument('-E', '--external', action='store_true',
79          help='Convert the FIT to use external data')
80    parser.add_argument('-n', '--name', type=str, required=True,
81          help='Specifies the name')
82    parser.add_argument('-o', '--output', type=str, required=True,
83          help='Specifies the output file (.fit)')
84    parser.add_argument('-O', '--os', type=str, required=True,
85          help='Specifies the operating system')
86    parser.add_argument('-k', '--kernel', type=str, required=True,
87          help='Specifies the (uncompressed) kernel input file (.itk)')
88    parser.add_argument('-r', '--ramdisk', type=str,
89          help='Specifies the ramdisk/initrd input file')
90    parser.add_argument('-v', '--verbose', action='store_true',
91                        help='Enable verbose output')
92    parser.add_argument('dtbs', type=str, nargs='*',
93          help='Specifies the devicetree files to process')
94
95    return parser.parse_args()
96
97
98def setup_fit(fsw, name):
99    """Make a start on writing the FIT
100
101    Outputs the root properties and the 'images' node
102
103    Args:
104        fsw (libfdt.FdtSw): Object to use for writing
105        name (str): Name of kernel image
106    """
107    fsw.INC_SIZE = 16 << 20
108    fsw.finish_reservemap()
109    fsw.begin_node('')
110    fsw.property_string('description', f'{name} with devicetree set')
111    fsw.property_u32('#address-cells', 1)
112
113    fsw.property_u32('timestamp', int(time.time()))
114    fsw.begin_node('images')
115
116
117def write_kernel(fsw, data, args):
118    """Write out the kernel image
119
120    Writes a kernel node along with the required properties
121
122    Args:
123        fsw (libfdt.FdtSw): Object to use for writing
124        data (bytes): Data to write (possibly compressed)
125        args (Namespace): Contains necessary strings:
126            arch: FIT architecture, e.g. 'arm64'
127            fit_os: Operating Systems, e.g. 'linux'
128            name: Name of OS, e.g. 'Linux-6.6.0-rc7'
129            compress: Compression algorithm to use, e.g. 'gzip'
130    """
131    with fsw.add_node('kernel'):
132        fsw.property_string('description', args.name)
133        fsw.property_string('type', 'kernel_noload')
134        fsw.property_string('arch', args.arch)
135        fsw.property_string('os', args.os)
136        fsw.property_string('compression', args.compress)
137        fsw.property('data', data)
138        fsw.property_u32('load', 0)
139        fsw.property_u32('entry', 0)
140
141
142def write_ramdisk(fsw, data, args):
143    """Write out the ramdisk image
144
145    Writes a ramdisk node along with the required properties
146
147    Args:
148        fsw (libfdt.FdtSw): Object to use for writing
149        data (bytes): Data to write (possibly compressed)
150        args (Namespace): Contains necessary strings:
151            arch: FIT architecture, e.g. 'arm64'
152            fit_os: Operating Systems, e.g. 'linux'
153    """
154    with fsw.add_node('ramdisk'):
155        fsw.property_string('description', 'Ramdisk')
156        fsw.property_string('type', 'ramdisk')
157        fsw.property_string('arch', args.arch)
158        fsw.property_string('compression', 'none')
159        fsw.property_string('os', args.os)
160        fsw.property('data', data)
161
162
163def finish_fit(fsw, entries, has_ramdisk=False):
164    """Finish the FIT ready for use
165
166    Writes the /configurations node and subnodes
167
168    Args:
169        fsw (libfdt.FdtSw): Object to use for writing
170        entries (list of tuple): List of configurations:
171            str: Description of model
172            str: Compatible stringlist
173        has_ramdisk (bool): True if a ramdisk is included in the FIT
174    """
175    fsw.end_node()
176    seq = 0
177    with fsw.add_node('configurations'):
178        for model, compat, files in entries:
179            seq += 1
180            with fsw.add_node(f'conf-{seq}'):
181                fsw.property('compatible', bytes(compat))
182                fsw.property_string('description', model)
183                fsw.property('fdt', bytes(''.join(f'fdt-{x}\x00' for x in files), "ascii"))
184                fsw.property_string('kernel', 'kernel')
185                if has_ramdisk:
186                    fsw.property_string('ramdisk', 'ramdisk')
187    fsw.end_node()
188
189
190def compress_data(inf, compress):
191    """Compress data using a selected algorithm
192
193    Args:
194        inf (IOBase): Filename containing the data to compress
195        compress (str): Compression algorithm, e.g. 'gzip'
196
197    Return:
198        bytes: Compressed data
199    """
200    if compress == 'none':
201        return inf.read()
202
203    comp = COMP_TOOLS.get(compress)
204    if not comp:
205        raise ValueError(f"Unknown compression algorithm '{compress}'")
206
207    with tempfile.NamedTemporaryFile() as comp_fname:
208        with open(comp_fname.name, 'wb') as outf:
209            done = False
210            for tool in comp.tools.split(','):
211                try:
212                    # Add parallel flags for tools that support them
213                    cmd = [tool]
214                    if tool in ('zstd', 'xz'):
215                        cmd.extend(['-T0'])  # Use all available cores
216                    cmd.append('-c')
217                    subprocess.call(cmd, stdin=inf, stdout=outf)
218                    done = True
219                    break
220                except FileNotFoundError:
221                    pass
222            if not done:
223                raise ValueError(f'Missing tool(s): {comp.tools}\n')
224            with open(comp_fname.name, 'rb') as compf:
225                comp_data = compf.read()
226    return comp_data
227
228
229def compress_dtb(fname, compress):
230    """Compress a single DTB file
231
232    Args:
233        fname (str): Filename containing the DTB
234        compress (str): Compression algorithm, e.g. 'gzip'
235
236    Returns:
237        tuple: (str: fname, bytes: compressed_data)
238    """
239    with open(fname, 'rb') as inf:
240        compressed = compress_data(inf, compress)
241    return fname, compressed
242
243
244def output_dtb(fsw, seq, fname, arch, compress, data=None):
245    """Write out a single devicetree to the FIT
246
247    Args:
248        fsw (libfdt.FdtSw): Object to use for writing
249        seq (int): Sequence number (1 for first)
250        fname (str): Filename containing the DTB
251        arch (str): FIT architecture, e.g. 'arm64'
252        compress (str): Compressed algorithm, e.g. 'gzip'
253        data (bytes): Pre-compressed data (optional)
254    """
255    with fsw.add_node(f'fdt-{seq}'):
256        fsw.property_string('description', os.path.basename(fname))
257        fsw.property_string('type', 'flat_dt')
258        fsw.property_string('arch', arch)
259        fsw.property_string('compression', compress)
260
261        if data is None:
262            with open(fname, 'rb') as inf:
263                data = compress_data(inf, compress)
264        fsw.property('data', data)
265
266
267def process_dtb(fname, args):
268    """Process an input DTB, decomposing it if requested and is possible
269
270    Args:
271        fname (str): Filename containing the DTB
272        args (Namespace): Program arguments
273    Returns:
274        tuple:
275            str: Model name string
276            str: Root compatible string
277            files: list of filenames corresponding to the DTB
278    """
279    # Get the compatible / model information
280    with open(fname, 'rb') as inf:
281        data = inf.read()
282    fdt = libfdt.FdtRo(data)
283    model = fdt.getprop(0, 'model').as_str()
284    compat = fdt.getprop(0, 'compatible')
285
286    if args.decompose_dtbs:
287        # Check if the DTB needs to be decomposed
288        path, basename = os.path.split(fname)
289        cmd_fname = os.path.join(path, f'.{basename}.cmd')
290        with open(cmd_fname, 'r', encoding='ascii') as inf:
291            cmd = inf.read()
292
293        if 'scripts/dtc/fdtoverlay' in cmd:
294            # This depends on the structure of the composite DTB command
295            # Ignore any subsequent commands, such as schema validation
296            files = cmd.split(';', 1)[0].split()
297            files = files[files.index('-i') + 1:]
298        else:
299            files = [fname]
300    else:
301        files = [fname]
302
303    return (model, compat, files)
304
305
306def _process_dtbs(args, fsw, entries, fdts):
307    """Process all DTB files and add them to the FIT
308
309    Args:
310        args: Program arguments
311        fsw: FIT writer object
312        entries: List to append entries to
313        fdts: Dictionary of processed DTBs
314
315    Returns:
316        tuple:
317            Number of files processed
318            Total size of files processed
319    """
320    seq = 0
321    size = 0
322
323    # First figure out the unique DTB files that need compression
324    todo = []
325    file_info = []  # List of (fname, model, compat, files) tuples
326
327    for fname in args.dtbs:
328        # Ignore non-DTB (*.dtb) files
329        if os.path.splitext(fname)[1] != '.dtb':
330            continue
331
332        try:
333            (model, compat, files) = process_dtb(fname, args)
334        except Exception as e:
335            sys.stderr.write(f'Error processing {fname}:\n')
336            raise e
337
338        file_info.append((fname, model, compat, files))
339        for fn in files:
340            if fn not in fdts and fn not in todo:
341                todo.append(fn)
342
343    # Compress all DTBs in parallel
344    cache = {}
345    if todo and args.compress != 'none':
346        if args.verbose:
347            print(f'Compressing {len(todo)} DTBs...')
348
349        with multiprocessing.Pool() as pool:
350            compress_args = [(fn, args.compress) for fn in todo]
351            # unpacks each tuple, calls compress_dtb(fn, compress) in parallel
352            results = pool.starmap(compress_dtb, compress_args)
353
354        cache = dict(results)
355
356    # Now write all DTBs to the FIT using pre-compressed data
357    for fname, model, compat, files in file_info:
358        for fn in files:
359            if fn not in fdts:
360                seq += 1
361                size += os.path.getsize(fn)
362                output_dtb(fsw, seq, fn, args.arch, args.compress,
363                           cache.get(fn))
364                fdts[fn] = seq
365
366        files_seq = [fdts[fn] for fn in files]
367        entries.append([model, compat, files_seq])
368
369    return seq, size
370
371
372def build_fit(args):
373    """Build the FIT from the provided files and arguments
374
375    Args:
376        args (Namespace): Program arguments
377
378    Returns:
379        tuple:
380            bytes: FIT data
381            int: Number of configurations generated
382            size: Total uncompressed size of data
383    """
384    size = 0
385    fsw = libfdt.FdtSw()
386    setup_fit(fsw, args.name)
387    entries = []
388    fdts = {}
389
390    # Handle the kernel
391    with open(args.kernel, 'rb') as inf:
392        comp_data = compress_data(inf, args.compress)
393    size += os.path.getsize(args.kernel)
394    write_kernel(fsw, comp_data, args)
395
396    # Handle the ramdisk if provided. Compression is not supported as it is
397    # already compressed.
398    if args.ramdisk:
399        with open(args.ramdisk, 'rb') as inf:
400            data = inf.read()
401        size += len(data)
402        write_ramdisk(fsw, data, args)
403
404    count, fdt_size = _process_dtbs(args, fsw, entries, fdts)
405    size += fdt_size
406
407    finish_fit(fsw, entries, bool(args.ramdisk))
408
409    # Include the kernel itself in the returned file count
410    fdt = fsw.as_fdt()
411    fdt.pack()
412    return fdt.as_bytearray(), count + 1 + bool(args.ramdisk), size
413
414
415def run_make_fit():
416    """Run the tool's main logic"""
417    args = parse_args()
418
419    out_data, count, size = build_fit(args)
420    with open(args.output, 'wb') as outf:
421        outf.write(out_data)
422
423    ext_fit_size = None
424    if args.external:
425        mkimage = os.environ.get('MKIMAGE', 'mkimage')
426        subprocess.check_call([mkimage, '-E', '-F', args.output],
427                              stdout=subprocess.DEVNULL)
428
429        with open(args.output, 'rb') as inf:
430            data = inf.read()
431        ext_fit = libfdt.FdtRo(data)
432        ext_fit_size = ext_fit.totalsize()
433
434    if args.verbose:
435        comp_size = len(out_data)
436        print(f'FIT size {comp_size:#x}/{comp_size / 1024 / 1024:.1f} MB',
437              end='')
438        if ext_fit_size:
439            print(f', header {ext_fit_size:#x}/{ext_fit_size / 1024:.1f} KB',
440                  end='')
441        print(f', {count} files, uncompressed {size / 1024 / 1024:.1f} MB')
442
443
444if __name__ == "__main__":
445    sys.exit(run_make_fit())
446