projects/linux·blob uaccess-cleanups · refs

444 lines · 14.1 KiB

1#!/usr/bin/env python3
2# SPDX-License-Identifier: GPL-2.0+
3#
4# Copyright 2024 Google LLC
5# Written by Simon Glass <[email protected]>
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 files = cmd.split()
296 files = files[files.index('-i') + 1:]
297 else:
298 files = [fname]
299 else:
300 files = [fname]
301
302 return (model, compat, files)
303
304
305def _process_dtbs(args, fsw, entries, fdts):
306 """Process all DTB files and add them to the FIT
307
308 Args:
309 args: Program arguments
310 fsw: FIT writer object
311 entries: List to append entries to
312 fdts: Dictionary of processed DTBs
313
314 Returns:
315 tuple:
316 Number of files processed
317 Total size of files processed
318 """
319 seq = 0
320 size = 0
321
322 # First figure out the unique DTB files that need compression
323 todo = []
324 file_info = [] # List of (fname, model, compat, files) tuples
325
326 for fname in args.dtbs:
327 # Ignore non-DTB (*.dtb) files
328 if os.path.splitext(fname)[1] != '.dtb':
329 continue
330
331 try:
332 (model, compat, files) = process_dtb(fname, args)
333 except Exception as e:
334 sys.stderr.write(f'Error processing {fname}:\n')
335 raise e
336
337 file_info.append((fname, model, compat, files))
338 for fn in files:
339 if fn not in fdts and fn not in todo:
340 todo.append(fn)
341
342 # Compress all DTBs in parallel
343 cache = {}
344 if todo and args.compress != 'none':
345 if args.verbose:
346 print(f'Compressing {len(todo)} DTBs...')
347
348 with multiprocessing.Pool() as pool:
349 compress_args = [(fn, args.compress) for fn in todo]
350 # unpacks each tuple, calls compress_dtb(fn, compress) in parallel
351 results = pool.starmap(compress_dtb, compress_args)
352
353 cache = dict(results)
354
355 # Now write all DTBs to the FIT using pre-compressed data
356 for fname, model, compat, files in file_info:
357 for fn in files:
358 if fn not in fdts:
359 seq += 1
360 size += os.path.getsize(fn)
361 output_dtb(fsw, seq, fn, args.arch, args.compress,
362 cache.get(fn))
363 fdts[fn] = seq
364
365 files_seq = [fdts[fn] for fn in files]
366 entries.append([model, compat, files_seq])
367
368 return seq, size
369
370
371def build_fit(args):
372 """Build the FIT from the provided files and arguments
373
374 Args:
375 args (Namespace): Program arguments
376
377 Returns:
378 tuple:
379 bytes: FIT data
380 int: Number of configurations generated
381 size: Total uncompressed size of data
382 """
383 size = 0
384 fsw = libfdt.FdtSw()
385 setup_fit(fsw, args.name)
386 entries = []
387 fdts = {}
388
389 # Handle the kernel
390 with open(args.kernel, 'rb') as inf:
391 comp_data = compress_data(inf, args.compress)
392 size += os.path.getsize(args.kernel)
393 write_kernel(fsw, comp_data, args)
394
395 # Handle the ramdisk if provided. Compression is not supported as it is
396 # already compressed.
397 if args.ramdisk:
398 with open(args.ramdisk, 'rb') as inf:
399 data = inf.read()
400 size += len(data)
401 write_ramdisk(fsw, data, args)
402
403 count, fdt_size = _process_dtbs(args, fsw, entries, fdts)
404 size += fdt_size
405
406 finish_fit(fsw, entries, bool(args.ramdisk))
407
408 # Include the kernel itself in the returned file count
409 fdt = fsw.as_fdt()
410 fdt.pack()
411 return fdt.as_bytearray(), count + 1 + bool(args.ramdisk), size
412
413
414def run_make_fit():
415 """Run the tool's main logic"""
416 args = parse_args()
417
418 out_data, count, size = build_fit(args)
419 with open(args.output, 'wb') as outf:
420 outf.write(out_data)
421
422 ext_fit_size = None
423 if args.external:
424 mkimage = os.environ.get('MKIMAGE', 'mkimage')
425 subprocess.check_call([mkimage, '-E', '-F', args.output],
426 stdout=subprocess.DEVNULL)
427
428 with open(args.output, 'rb') as inf:
429 data = inf.read()
430 ext_fit = libfdt.FdtRo(data)
431 ext_fit_size = ext_fit.totalsize()
432
433 if args.verbose:
434 comp_size = len(out_data)
435 print(f'FIT size {comp_size:#x}/{comp_size / 1024 / 1024:.1f} MB',
436 end='')
437 if ext_fit_size:
438 print(f', header {ext_fit_size:#x}/{ext_fit_size / 1024:.1f} KB',
439 end='')
440 print(f', {count} files, uncompressed {size / 1024 / 1024:.1f} MB')
441
442
443if __name__ == "__main__":
444 sys.exit(run_make_fit())