1#!/usr/bin/env python3
2#
3# Copyright (c) 2024 STMicroelectronics
4# SPDX-License-Identifier: Apache-2.0
5
6"""
7Script to prepare the LLEXT exports table of a Zephyr ELF
8
9This script performs compile-time processing of the LLEXT exports
10table for usage at runtime by the LLEXT subsystem code. The table
11is a special section filled with 'llext_const_symbol' structures
12generated by the EXPORT_SYMBOL macro.
13
14Currently, the preparatory work consists mostly of sorting the
15exports table to allow usage of binary search algorithms at runtime.
16If CONFIG_LLEXT_EXPORT_BUILTINS_BY_SLID option is enabled, SLIDs
17of all exported functions are also injected in the export table by
18this script. (In this case, the preparation process is destructive)
19"""
20
21import llext_slidlib
22
23from elftools.elf.elffile import ELFFile
24from elftools.elf.sections import Section
25
26import argparse
27import logging
28import pathlib
29import struct
30import sys
31
32#!!!!! WARNING !!!!!
33#
34#These constants MUST be kept in sync with the linker scripts
35#and the EXPORT_SYMBOL macro located in 'subsys/llext/llext.h'.
36#Otherwise, the LLEXT subsystem will be broken!
37#
38#!!!!! WARNING !!!!!
39
40LLEXT_EXPORT_TABLE_SECTION_NAME = "llext_const_symbol_area"
41LLEXT_EXPORT_NAMES_SECTION_NAME = "llext_exports_strtab"
42
43def _llext_const_symbol_struct(ptr_size: int, endianness: str):
44    """
45    ptr_size -- Platform pointer size in bytes
46    endianness -- Platform endianness ('little'/'big')
47    """
48    endspec = "<" if endianness == 'little' else ">"
49    if ptr_size == 4:
50        ptrspec = "I"
51    elif ptr_size == 8:
52        ptrspec = "Q"
53
54    # struct llext_const_symbol
55    # contains just two pointers.
56    lcs_spec = endspec + 2 * ptrspec
57    return struct.Struct(lcs_spec)
58
59#ELF Shdr flag applied to the export table section, to indicate
60#the section has already been prepared by this script. This is
61#mostly a security measure to prevent the script from running
62#twice on the same ELF file, which can result in catastrophic
63#failures if SLID-based linking is enabled (in this case, the
64#preparation process is destructive).
65#
66#This flag is part of the SHF_MASKOS mask, of which all bits
67#are "reserved for operating system-specific semantics".
68#See: https://refspecs.linuxbase.org/elf/gabi4+/ch4.sheader.html
69SHF_LLEXT_PREPARATION_DONE = 0x08000000
70
71class SectionDescriptor():
72    """ELF Section descriptor
73
74    This is a wrapper class around pyelftools' "Section" object.
75    """
76    def __init__(self, elffile, section_name):
77        self.name = section_name
78        self.section = elffile.get_section_by_name(section_name)
79        if not isinstance(self.section, Section):
80            raise KeyError(f"section {section_name} not found")
81
82        self.shdr_index = elffile.get_section_index(section_name)
83        self.shdr_offset = elffile['e_shoff'] + \
84            self.shdr_index * elffile['e_shentsize']
85        self.size = self.section['sh_size']
86        self.flags = self.section['sh_flags']
87        self.offset = self.section['sh_offset']
88
89class LLEXTExptabManipulator():
90    """Class used to wrap the LLEXT export table manipulation."""
91    def __init__(self, elf_fd, exptab_file_offset, lcs_struct, exports_count):
92        self.fd = elf_fd
93        self.exports_count = exports_count
94        self.base_offset = exptab_file_offset
95        self.lcs_struct = lcs_struct
96
97    def _seek_to_sym(self, index):
98        self.fd.seek(self.base_offset + index * self.lcs_struct.size)
99
100    def __getitem__(self, index):
101        if not isinstance(index, int):
102            raise TypeError(f"invalid type {type(index)} for index")
103
104        if index >= self.exports_count:
105            raise IndexError(f"index {index} is out of bounds (max {self.exports_count})")
106
107        self._seek_to_sym(index)
108        return self.lcs_struct.unpack(self.fd.read(self.lcs_struct.size))
109
110    def __setitem__(self, index, item):
111        if not isinstance(index, int):
112            raise TypeError(f"invalid type {type(index)} for index")
113
114        if index >= self.exports_count:
115            raise IndexError(f"index {index} is out of bounds (max {self.exports_count})")
116
117        (addr_or_slid, sym_addr) = item
118
119        self._seek_to_sym(index)
120        self.fd.write(self.lcs_struct.pack(addr_or_slid, sym_addr))
121
122class ZephyrElfExptabPreparator():
123    """Prepares the LLEXT export table of a Zephyr ELF.
124
125    Attributes:
126        elf_path: path to the Zephyr ELF to prepare
127        log: a logging.Logger object
128        slid_listing_path: path to the file where SLID listing should be saved
129    """
130    def __init__(self, elf_path: str, log: logging.Logger, slid_listing_path: str | None):
131        self.elf_path = elf_path
132        self.elf_fd = open(self.elf_path, 'rb+')
133        self.elf = ELFFile(self.elf_fd)
134        self.log = log
135
136        # Lazy-open the SLID listing file to ensure it is only created when necessary
137        self.slid_listing_path = slid_listing_path
138        self.slid_listing_fd = None
139
140    def _prepare_exptab_for_slid_linking(self):
141        """
142        IMPLEMENTATION NOTES:
143          In the linker script, we declare the export names table
144          as starting at address 0. Thanks to this, all "pointers"
145          to that section are equal to the offset inside the section.
146          Also note that symbol names are always NUL-terminated.
147
148          The export table is sorted by SLID in ASCENDING order.
149        """
150        def read_symbol_name(name_ptr):
151            raw_name = b''
152            self.elf_fd.seek(self.expstrtab_section.offset + name_ptr)
153
154            c = self.elf_fd.read(1)
155            while c != b'\0':
156                raw_name += c
157                c = self.elf_fd.read(1)
158
159            return raw_name.decode("utf-8")
160
161        #1) Load the export table
162        exports_list = []
163        for (name_ptr, export_address) in self.exptab_manipulator:
164            export_name = read_symbol_name(name_ptr)
165            exports_list.append((export_name, export_address))
166
167        #2) Generate the SLID for all exports
168        collided = False
169        sorted_exptab = dict()
170        for export_name, export_addr in exports_list:
171            slid = llext_slidlib.generate_slid(export_name, self.ptrsize)
172
173            collision = sorted_exptab.get(slid)
174            if collision:
175                #Don't abort immediately on collision: if there are others, we want to log them all.
176                self.log.error(f"SLID collision: {export_name} and {collision[0]} have the same SLID 0x{slid:X}")
177                collided = True
178            else:
179                sorted_exptab[slid] = (export_name, export_addr)
180
181        if collided:
182            return 1
183
184        #3) Sort the export table (order specified above)
185        sorted_exptab = dict(sorted(sorted_exptab.items()))
186
187        #4) Write the updated export table to ELF, and dump
188        #to SLID listing if requested by caller
189        if self.slid_listing_path:
190            self.slid_listing_fd = open(self.slid_listing_path, "w")
191
192        def slidlist_write(msg):
193            if self.slid_listing_fd:
194                self.slid_listing_fd.write(msg + "\n")
195
196        slidlist_write(f"/* SLID listing generated by {__file__} */")
197        slidlist_write("//")
198        slidlist_write("// This file contains the 'SLID -> name' mapping for all")
199        slidlist_write("// symbols exported to LLEXT by this Zephyr executable.")
200        slidlist_write("")
201
202        self.log.info("SLID -> export name mapping:")
203
204        i = 0
205        for (slid, name_and_symaddr) in sorted_exptab.items():
206            slid_as_str = llext_slidlib.format_slid(slid, self.ptrsize)
207            msg = f"{slid_as_str} -> {name_and_symaddr[0]}"
208            self.log.info(msg)
209            slidlist_write(msg)
210
211            self.exptab_manipulator[i] = (slid, name_and_symaddr[1])
212            i += 1
213
214        if self.slid_listing_fd:
215            self.slid_listing_fd.close()
216
217        return 0
218
219    def _prepare_exptab_for_str_linking(self):
220        #TODO: sort the export table by symbol
221        #      name to allow binary search too
222        #
223        # Plan of action:
224        #   1) Locate in which section the names are located
225        #   2) Load the export table and resolve names
226        #   3) Sort the exports by name
227        #       WARN: THIS MUST USE THE SAME SORTING RULES
228        #       AS LLEXT CODE OR DICHOTOMIC SEARCH WILL BREAK
229        #       Using a custom sorting function might be required.
230        #   4) Write back the updated export table
231        #
232        # N.B.: reusing part of the code in _prepare_elf_for_slid_linking
233        # might be possible and desireable.
234        #
235        # As of writing, this function will never be called as this script
236        # is only called if CONFIG_LLEXT_EXPORT_BUILTINS_BY_SLID is enabled,
237        # which makes _prepare_exptab_for_slid_linking be called instead.
238        #
239        self.log.warn(f"_prepare_exptab_for_str_linking: do nothing")
240        return 0
241
242    def _set_prep_done_shdr_flag(self):
243        #Offset and size of the 'sh_flags' member of
244        #the Elf_Shdr structure. The offset does not
245        #change between ELF32 and ELF64. Size in both
246        #is equal to pointer size (4 bytes for ELF32,
247        #8 bytes for ELF64).
248        SHF_OFFSET = 8
249        SHF_SIZE = self.ptrsize
250
251        off = self.exptab_section.shdr_offset + SHF_OFFSET
252
253        #Read existing sh_flags, set the PREPARATION_DONE flag
254        #and write back the new value.
255        self.elf_fd.seek(off)
256        sh_flags = int.from_bytes(self.elf_fd.read(SHF_SIZE), self.endianness)
257
258        sh_flags |= SHF_LLEXT_PREPARATION_DONE
259
260        self.elf_fd.seek(off)
261        self.elf_fd.write(int.to_bytes(sh_flags, self.ptrsize, self.endianness))
262
263    def _prepare_inner(self):
264        # Locate the export table section
265        try:
266            self.exptab_section = SectionDescriptor(
267                self.elf, LLEXT_EXPORT_TABLE_SECTION_NAME)
268        except KeyError as e:
269            self.log.error(e.args[0])
270            return 1
271
272        # Abort if the ELF has already been processed
273        if (self.exptab_section.flags & SHF_LLEXT_PREPARATION_DONE) != 0:
274            self.log.warning("exptab section flagged with LLEXT_PREPARATION_DONE "
275                "- not preparing again")
276            return 0
277
278        # Get the struct.Struct for export table entry
279        self.ptrsize = self.elf.elfclass // 8
280        self.endianness = 'little' if self.elf.little_endian else 'big'
281        self.lcs_struct = _llext_const_symbol_struct(self.ptrsize, self.endianness)
282
283        # Verify that the export table size is coherent
284        if (self.exptab_section.size % self.lcs_struct.size) != 0:
285            self.log.error(f"export table size (0x{self.exptab_section.size:X}) "
286                f"not aligned to 'llext_const_symbol' size (0x{self.lcs_struct.size:X})")
287            return 1
288
289        # Create the export table manipulator
290        num_exports = self.exptab_section.size // self.lcs_struct.size
291        self.exptab_manipulator = LLEXTExptabManipulator(
292            self.elf_fd, self.exptab_section.offset, self.lcs_struct, num_exports)
293
294        # Attempt to locate the export names section
295        try:
296            self.expstrtab_section = SectionDescriptor(
297                self.elf, LLEXT_EXPORT_NAMES_SECTION_NAME)
298        except KeyError:
299            self.expstrtab_section = None
300
301        self.log.debug(f"exports table section at file offset 0x{self.exptab_section.offset:X}")
302        if self.expstrtab_section:
303            self.log.debug(f"exports strtab section at file offset 0x{self.expstrtab_section.offset:X}")
304        else:
305            self.log.debug("no exports strtab section in ELF")
306        self.log.info(f"{num_exports} symbols are exported to LLEXTs by this ELF")
307
308        # Perform the export table preparation
309        if self.expstrtab_section:
310            res = self._prepare_exptab_for_slid_linking()
311        else:
312            res = self._prepare_exptab_for_str_linking()
313
314        if res == 0: # Add the "prepared" flag to export table section
315            self._set_prep_done_shdr_flag()
316
317    def prepare_elf(self):
318        res = self._prepare_inner()
319        self.elf_fd.close()
320        return res
321
322# pylint: disable=duplicate-code
323def _parse_args(argv):
324    """Parse the command line arguments."""
325    parser = argparse.ArgumentParser(
326        description=__doc__,
327        formatter_class=argparse.RawDescriptionHelpFormatter,
328        allow_abbrev=False)
329
330    parser.add_argument("-f", "--elf-file", default=pathlib.Path("build", "zephyr", "zephyr.elf"),
331                        help="ELF file to process")
332    parser.add_argument("-sl", "--slid-listing",
333                        help=("write the SLID listing to a file (only useful"
334                              "when CONFIG_LLEXT_EXPORT_BUILTINS_BY_SLID is enabled) "))
335    parser.add_argument("-v", "--verbose", action="count",
336                        help=("enable verbose output, can be used multiple times "
337                              "to increase verbosity level"))
338    parser.add_argument("--always-succeed", action="store_true",
339                        help="always exit with a return code of 0, used for testing")
340
341    return parser.parse_args(argv)
342
343def _init_log(verbose):
344    """Initialize a logger object."""
345    log = logging.getLogger(__file__)
346
347    console = logging.StreamHandler()
348    console.setFormatter(logging.Formatter("%(levelname)s: %(message)s"))
349    log.addHandler(console)
350
351    if verbose and verbose > 1:
352        log.setLevel(logging.DEBUG)
353    elif verbose and verbose > 0:
354        log.setLevel(logging.INFO)
355    else:
356        log.setLevel(logging.WARNING)
357
358    return log
359
360def main(argv=None):
361    args = _parse_args(argv)
362
363    log = _init_log(args.verbose)
364
365    log.info(f"prepare_llext_exptab: {args.elf_file}")
366
367    preparator = ZephyrElfExptabPreparator(args.elf_file, log, args.slid_listing)
368
369    res = preparator.prepare_elf()
370
371    if args.always_succeed:
372        return 0
373
374    return res
375
376if __name__ == "__main__":
377    sys.exit(main(sys.argv[1:]))
378