diff options
Diffstat (limited to 'scripts/make_fit.py')
| -rwxr-xr-x | scripts/make_fit.py | 179 |
1 files changed, 144 insertions, 35 deletions
diff --git a/scripts/make_fit.py b/scripts/make_fit.py index 1683e5ec6e67..e923cc8b05b7 100755 --- a/scripts/make_fit.py +++ b/scripts/make_fit.py @@ -10,10 +10,14 @@ Usage: make_fit.py -A arm64 -n 'Linux-6.6' -O linux -o arch/arm64/boot/image.fit -k /tmp/kern/arch/arm64/boot/image.itk - @arch/arm64/boot/dts/dtbs-list -E -c gzip + -r /boot/initrd.img-6.14.0-27-generic @arch/arm64/boot/dts/dtbs-list + -E -c gzip -Creates a FIT containing the supplied kernel and a set of devicetree files, -either specified individually or listed in a file (with an '@' prefix). +Creates a FIT containing the supplied kernel, an optional ramdisk, and a set of +devicetree files, either specified individually or listed in a file (with an +'@' prefix). + +Use -r to specify an existing ramdisk/initrd file. Use -E to generate an external FIT (where the data is placed after the FIT data structure). This allows parsing of the data without loading @@ -29,12 +33,11 @@ looks at the .cmd files produced by the kernel build. The resulting FIT can be booted by bootloaders which support FIT, such as U-Boot, Linuxboot, Tianocore, etc. - -Note that this tool does not yet support adding a ramdisk / initrd. """ import argparse import collections +import multiprocessing import os import subprocess import sys @@ -48,11 +51,12 @@ import libfdt CompTool = collections.namedtuple('CompTool', 'ext,tools') COMP_TOOLS = { - 'bzip2': CompTool('.bz2', 'bzip2'), + 'bzip2': CompTool('.bz2', 'pbzip2,bzip2'), 'gzip': CompTool('.gz', 'pigz,gzip'), 'lz4': CompTool('.lz4', 'lz4'), - 'lzma': CompTool('.lzma', 'lzma'), + 'lzma': CompTool('.lzma', 'plzip,lzma'), 'lzo': CompTool('.lzo', 'lzop'), + 'xz': CompTool('.xz', 'xz'), 'zstd': CompTool('.zstd', 'zstd'), } @@ -81,6 +85,8 @@ def parse_args(): help='Specifies the operating system') parser.add_argument('-k', '--kernel', type=str, required=True, help='Specifies the (uncompressed) kernel input file (.itk)') + parser.add_argument('-r', '--ramdisk', type=str, + help='Specifies the ramdisk/initrd input file') parser.add_argument('-v', '--verbose', action='store_true', help='Enable verbose output') parser.add_argument('dtbs', type=str, nargs='*', @@ -98,7 +104,7 @@ def setup_fit(fsw, name): fsw (libfdt.FdtSw): Object to use for writing name (str): Name of kernel image """ - fsw.INC_SIZE = 65536 + fsw.INC_SIZE = 16 << 20 fsw.finish_reservemap() fsw.begin_node('') fsw.property_string('description', f'{name} with devicetree set') @@ -133,7 +139,28 @@ def write_kernel(fsw, data, args): fsw.property_u32('entry', 0) -def finish_fit(fsw, entries): +def write_ramdisk(fsw, data, args): + """Write out the ramdisk image + + Writes a ramdisk node along with the required properties + + Args: + fsw (libfdt.FdtSw): Object to use for writing + data (bytes): Data to write (possibly compressed) + args (Namespace): Contains necessary strings: + arch: FIT architecture, e.g. 'arm64' + fit_os: Operating Systems, e.g. 'linux' + """ + with fsw.add_node('ramdisk'): + fsw.property_string('description', 'Ramdisk') + fsw.property_string('type', 'ramdisk') + fsw.property_string('arch', args.arch) + fsw.property_string('compression', 'none') + fsw.property_string('os', args.os) + fsw.property('data', data) + + +def finish_fit(fsw, entries, has_ramdisk=False): """Finish the FIT ready for use Writes the /configurations node and subnodes @@ -143,6 +170,7 @@ def finish_fit(fsw, entries): entries (list of tuple): List of configurations: str: Description of model str: Compatible stringlist + has_ramdisk (bool): True if a ramdisk is included in the FIT """ fsw.end_node() seq = 0 @@ -154,6 +182,8 @@ def finish_fit(fsw, entries): fsw.property_string('description', model) fsw.property('fdt', bytes(''.join(f'fdt-{x}\x00' for x in files), "ascii")) fsw.property_string('kernel', 'kernel') + if has_ramdisk: + fsw.property_string('ramdisk', 'ramdisk') fsw.end_node() @@ -179,7 +209,12 @@ def compress_data(inf, compress): done = False for tool in comp.tools.split(','): try: - subprocess.call([tool, '-c'], stdin=inf, stdout=outf) + # Add parallel flags for tools that support them + cmd = [tool] + if tool in ('zstd', 'xz'): + cmd.extend(['-T0']) # Use all available cores + cmd.append('-c') + subprocess.call(cmd, stdin=inf, stdout=outf) done = True break except FileNotFoundError: @@ -191,15 +226,31 @@ def compress_data(inf, compress): return comp_data -def output_dtb(fsw, seq, fname, arch, compress): +def compress_dtb(fname, compress): + """Compress a single DTB file + + Args: + fname (str): Filename containing the DTB + compress (str): Compression algorithm, e.g. 'gzip' + + Returns: + tuple: (str: fname, bytes: compressed_data) + """ + with open(fname, 'rb') as inf: + compressed = compress_data(inf, compress) + return fname, compressed + + +def output_dtb(fsw, seq, fname, arch, compress, data=None): """Write out a single devicetree to the FIT Args: fsw (libfdt.FdtSw): Object to use for writing seq (int): Sequence number (1 for first) fname (str): Filename containing the DTB - arch: FIT architecture, e.g. 'arm64' + arch (str): FIT architecture, e.g. 'arm64' compress (str): Compressed algorithm, e.g. 'gzip' + data (bytes): Pre-compressed data (optional) """ with fsw.add_node(f'fdt-{seq}'): fsw.property_string('description', os.path.basename(fname)) @@ -207,9 +258,10 @@ def output_dtb(fsw, seq, fname, arch, compress): fsw.property_string('arch', arch) fsw.property_string('compression', compress) - with open(fname, 'rb') as inf: - compressed = compress_data(inf, compress) - fsw.property('data', compressed) + if data is None: + with open(fname, 'rb') as inf: + data = compress_data(inf, compress) + fsw.property('data', data) def process_dtb(fname, args): @@ -249,30 +301,27 @@ def process_dtb(fname, args): return (model, compat, files) -def build_fit(args): - """Build the FIT from the provided files and arguments + +def _process_dtbs(args, fsw, entries, fdts): + """Process all DTB files and add them to the FIT Args: - args (Namespace): Program arguments + args: Program arguments + fsw: FIT writer object + entries: List to append entries to + fdts: Dictionary of processed DTBs Returns: tuple: - bytes: FIT data - int: Number of configurations generated - size: Total uncompressed size of data + Number of files processed + Total size of files processed """ seq = 0 size = 0 - fsw = libfdt.FdtSw() - setup_fit(fsw, args.name) - entries = [] - fdts = {} - # Handle the kernel - with open(args.kernel, 'rb') as inf: - comp_data = compress_data(inf, args.compress) - size += os.path.getsize(args.kernel) - write_kernel(fsw, comp_data, args) + # First figure out the unique DTB files that need compression + todo = [] + file_info = [] # List of (fname, model, compat, files) tuples for fname in args.dtbs: # Ignore non-DTB (*.dtb) files @@ -282,24 +331,84 @@ def build_fit(args): try: (model, compat, files) = process_dtb(fname, args) except Exception as e: - sys.stderr.write(f"Error processing {fname}:\n") + sys.stderr.write(f'Error processing {fname}:\n') raise e + file_info.append((fname, model, compat, files)) + for fn in files: + if fn not in fdts and fn not in todo: + todo.append(fn) + + # Compress all DTBs in parallel + cache = {} + if todo and args.compress != 'none': + if args.verbose: + print(f'Compressing {len(todo)} DTBs...') + + with multiprocessing.Pool() as pool: + compress_args = [(fn, args.compress) for fn in todo] + # unpacks each tuple, calls compress_dtb(fn, compress) in parallel + results = pool.starmap(compress_dtb, compress_args) + + cache = dict(results) + + # Now write all DTBs to the FIT using pre-compressed data + for fname, model, compat, files in file_info: for fn in files: if fn not in fdts: seq += 1 size += os.path.getsize(fn) - output_dtb(fsw, seq, fn, args.arch, args.compress) + output_dtb(fsw, seq, fn, args.arch, args.compress, + cache.get(fn)) fdts[fn] = seq files_seq = [fdts[fn] for fn in files] - entries.append([model, compat, files_seq]) - finish_fit(fsw, entries) + return seq, size + + +def build_fit(args): + """Build the FIT from the provided files and arguments + + Args: + args (Namespace): Program arguments + + Returns: + tuple: + bytes: FIT data + int: Number of configurations generated + size: Total uncompressed size of data + """ + size = 0 + fsw = libfdt.FdtSw() + setup_fit(fsw, args.name) + entries = [] + fdts = {} + + # Handle the kernel + with open(args.kernel, 'rb') as inf: + comp_data = compress_data(inf, args.compress) + size += os.path.getsize(args.kernel) + write_kernel(fsw, comp_data, args) + + # Handle the ramdisk if provided. Compression is not supported as it is + # already compressed. + if args.ramdisk: + with open(args.ramdisk, 'rb') as inf: + data = inf.read() + size += len(data) + write_ramdisk(fsw, data, args) + + count, fdt_size = _process_dtbs(args, fsw, entries, fdts) + size += fdt_size + + finish_fit(fsw, entries, bool(args.ramdisk)) # Include the kernel itself in the returned file count - return fsw.as_fdt().as_bytearray(), seq + 1, size + fdt = fsw.as_fdt() + fdt.pack() + return fdt.as_bytearray(), count + 1 + bool(args.ramdisk), size def run_make_fit(): |