1#!/usr/bin/env python3
2#
3# Copyright (c) 2022, CSIRO
4#
5# SPDX-License-Identifier: Apache-2.0
6
7import struct
8import sys
9
10import elftools
11from elftools.elf.elffile import ELFFile
12from elftools.elf.sections import SymbolTableSection
13from packaging import version
14
15if version.parse(elftools.__version__) < version.parse('0.24'):
16    sys.exit("pyelftools is out of date, need version 0.24 or later")
17
18
19class _Symbol:
20    """
21    Parent class for objects derived from an elf symbol.
22    """
23
24    def __init__(self, elf, sym):
25        self.elf = elf
26        self.sym = sym
27        self.data = self.elf.symbol_data(sym)
28
29    def __lt__(self, other):
30        return self.sym.entry.st_value < other.sym.entry.st_value
31
32    def _data_native_read(self, offset):
33        (format, size) = self.elf.native_struct_format
34        return struct.unpack(format, self.data[offset : offset + size])[0]
35
36
37class DevicePM(_Symbol):
38    """
39    Represents information about device PM capabilities.
40    """
41
42    required_ld_consts = ["_PM_DEVICE_STRUCT_FLAGS_OFFSET", "_PM_DEVICE_FLAG_PD"]
43
44    def __init__(self, elf, sym):
45        super().__init__(elf, sym)
46        self.flags = self._data_native_read(self.elf.ld_consts['_PM_DEVICE_STRUCT_FLAGS_OFFSET'])
47
48    @property
49    def is_power_domain(self):
50        return self.flags & (1 << self.elf.ld_consts["_PM_DEVICE_FLAG_PD"])
51
52
53class DeviceOrdinals(_Symbol):
54    """
55    Represents information about device dependencies.
56    """
57
58    DEVICE_HANDLE_SEP = -32768
59    DEVICE_HANDLE_ENDS = 32767
60    DEVICE_HANDLE_NULL = 0
61
62    def __init__(self, elf, sym):
63        super().__init__(elf, sym)
64        format = "<" if self.elf.little_endian else ">"
65        format += f"{len(self.data) // 2:d}h"
66        self._ordinals = struct.unpack(format, self.data)
67        self._ordinals_split = []
68
69        # Split ordinals on DEVICE_HANDLE_SEP
70        prev = 1
71        for idx, val in enumerate(self._ordinals, 1):
72            if val == self.DEVICE_HANDLE_SEP:
73                self._ordinals_split.append(self._ordinals[prev : idx - 1])
74                prev = idx
75        self._ordinals_split.append(self._ordinals[prev:])
76
77    @property
78    def self_ordinal(self):
79        return self._ordinals[0]
80
81    @property
82    def ordinals(self):
83        return self._ordinals_split
84
85
86class Device(_Symbol):
87    """
88    Represents information about a device object and its references to other objects.
89    """
90
91    required_ld_consts = ["_DEVICE_STRUCT_HANDLES_OFFSET", "_DEVICE_STRUCT_PM_OFFSET"]
92
93    def __init__(self, elf, sym):
94        super().__init__(elf, sym)
95        self.edt_node = None
96        self.handle = None
97        self.ordinals = None
98        self.pm = None
99
100        # Devicetree dependencies, injected dependencies, supported devices
101        self.devs_depends_on = set()
102        self.devs_depends_on_injected = set()
103        self.devs_supports = set()
104
105        # Point to the handles instance associated with the device;
106        # assigned by correlating the device struct handles pointer
107        # value with the addr of a Handles instance.
108        self.obj_ordinals = None
109        if '_DEVICE_STRUCT_HANDLES_OFFSET' in self.elf.ld_consts:
110            ordinal_offset = self.elf.ld_consts['_DEVICE_STRUCT_HANDLES_OFFSET']
111            self.obj_ordinals = self._data_native_read(ordinal_offset)
112
113        self.obj_pm = None
114        if '_DEVICE_STRUCT_PM_OFFSET' in self.elf.ld_consts:
115            pm_offset = self.elf.ld_consts['_DEVICE_STRUCT_PM_OFFSET']
116            self.obj_pm = self._data_native_read(pm_offset)
117
118    @property
119    def ordinal(self):
120        return self.ordinals.self_ordinal
121
122
123class ZephyrElf:
124    """
125    Represents information about devices in an elf file.
126    """
127
128    def __init__(self, kernel, edt, device_start_symbol):
129        self.elf = ELFFile(open(kernel, "rb"))  # noqa : SIM115
130        self.relocatable = self.elf['e_type'] == 'ET_REL'
131        self.edt = edt
132        self.devices = []
133        self.ld_consts = self._symbols_find_value(
134            set([device_start_symbol, *Device.required_ld_consts, *DevicePM.required_ld_consts])
135        )
136        self._device_parse_and_link()
137
138    @property
139    def little_endian(self):
140        """
141        True if the elf file is for a little-endian architecture.
142        """
143        return self.elf.little_endian
144
145    @property
146    def native_struct_format(self):
147        """
148        Get the struct format specifier and byte size of the native machine type.
149        """
150        format = "<" if self.little_endian else ">"
151        if self.elf.elfclass == 32:
152            format += "I"
153            size = 4
154        else:
155            format += "Q"
156            size = 8
157        return (format, size)
158
159    def symbol_data(self, sym):
160        """
161        Retrieve the raw bytes associated with a symbol from the elf file.
162        """
163        # Symbol data parameters
164        addr = sym.entry.st_value
165        length = sym.entry.st_size
166        # Section associated with the symbol
167        section = self.elf.get_section(sym.entry['st_shndx'])
168        data = section.data()
169        # Relocatable data does not appear to be shifted
170        offset = addr - (0 if self.relocatable else section['sh_addr'])
171        # Validate data extraction
172        assert offset + length <= len(data)
173        # Extract symbol bytes from section
174        return bytes(data[offset : offset + length])
175
176    def _symbols_find_value(self, names):
177        symbols = {}
178        for section in self.elf.iter_sections():
179            if isinstance(section, SymbolTableSection):
180                for sym in section.iter_symbols():
181                    if sym.name in names:
182                        symbols[sym.name] = sym.entry.st_value
183        return symbols
184
185    def _object_find_named(self, prefix, cb):
186        for section in self.elf.iter_sections():
187            if isinstance(section, SymbolTableSection):
188                for sym in section.iter_symbols():
189                    if sym.entry.st_info.type != 'STT_OBJECT':
190                        continue
191                    if sym.name.startswith(prefix):
192                        cb(sym)
193
194    def _link_devices(self, devices):
195        # Compute the dependency graph induced from the full graph restricted to the
196        # the nodes that exist in the application.  Note that the edges in the
197        # induced graph correspond to paths in the full graph.
198        root = self.edt.dep_ord2node[0]
199
200        for ord, dev in devices.items():
201            n = self.edt.dep_ord2node[ord]
202
203            deps = set(n.depends_on)
204            while len(deps) > 0:
205                dn = deps.pop()
206                if dn.dep_ordinal in devices:
207                    # this is used
208                    dev.devs_depends_on.add(devices[dn.dep_ordinal])
209                elif dn != root:
210                    # forward the dependency up one level
211                    for ddn in dn.depends_on:
212                        deps.add(ddn)
213
214            sups = set(n.required_by)
215            while len(sups) > 0:
216                sn = sups.pop()
217                if sn.dep_ordinal in devices:
218                    dev.devs_supports.add(devices[sn.dep_ordinal])
219                else:
220                    # forward the support down one level
221                    for ssn in sn.required_by:
222                        sups.add(ssn)
223
224    def _link_injected(self, devices):
225        for dev in devices.values():
226            injected = dev.ordinals.ordinals[1]
227            for inj in injected:
228                if inj in devices:
229                    dev.devs_depends_on_injected.add(devices[inj])
230                    devices[inj].devs_supports.add(dev)
231
232    def _device_parse_and_link(self):
233        # Find all PM structs
234        pm_structs = {}
235
236        def _on_pm(sym):
237            pm_structs[sym.entry.st_value] = DevicePM(self, sym)
238
239        self._object_find_named('__pm_device_', _on_pm)
240
241        # Find all ordinal arrays
242        ordinal_arrays = {}
243
244        def _on_ordinal(sym):
245            ordinal_arrays[sym.entry.st_value] = DeviceOrdinals(self, sym)
246
247        self._object_find_named('__devicedeps_', _on_ordinal)
248
249        # Find all device structs
250        def _on_device(sym):
251            self.devices.append(Device(self, sym))
252
253        self._object_find_named('__device_', _on_device)
254
255        # Sort the device array by address (st_value) for handle calculation
256        self.devices = sorted(self.devices)
257
258        # Assign handles to the devices
259        for idx, dev in enumerate(self.devices):
260            dev.handle = 1 + idx
261
262        # Link devices structs with PM and ordinals
263        for dev in self.devices:
264            if dev.obj_pm in pm_structs:
265                dev.pm = pm_structs[dev.obj_pm]
266            if dev.obj_ordinals in ordinal_arrays:
267                dev.ordinals = ordinal_arrays[dev.obj_ordinals]
268                if dev.ordinal != DeviceOrdinals.DEVICE_HANDLE_NULL:
269                    dev.edt_node = self.edt.dep_ord2node[dev.ordinal]
270
271        # Create mapping of ordinals to devices
272        devices_by_ord = {d.ordinal: d for d in self.devices if d.edt_node}
273
274        # Link devices to each other based on the EDT tree
275        self._link_devices(devices_by_ord)
276
277        # Link injected devices to each other
278        self._link_injected(devices_by_ord)
279
280    def device_dependency_graph(self, title, comment):
281        """
282        Construct a graphviz Digraph of the relationships between devices.
283        """
284        import graphviz
285
286        dot = graphviz.Digraph(title, comment=comment)
287        # Split iteration so nodes and edges are grouped in source
288        for dev in self.devices:
289            if dev.ordinal == DeviceOrdinals.DEVICE_HANDLE_NULL:
290                text = f'{dev.sym.name:s}\\nHandle: {dev.handle:d}'
291            else:
292                n = self.edt.dep_ord2node[dev.ordinal]
293                text = (
294                    f'{n.name:s}\\nOrdinal: {dev.ordinal:d} | Handle: {dev.handle:d}\\n{n.path:s}'
295                )
296            dot.node(str(dev.ordinal), text)
297        for dev in self.devices:
298            for sup in sorted(dev.devs_supports):
299                dot.edge(str(dev.ordinal), str(sup.ordinal))
300        return dot
301