1#
2# Copyright (c) 2023, Arm Limited. All rights reserved.
3#
4# SPDX-License-Identifier: BSD-3-Clause
5#
6
7import re
8from dataclasses import asdict, dataclass
9from typing import BinaryIO
10
11from elftools.elf.elffile import ELFFile
12
13
14@dataclass(frozen=True)
15class TfaMemObject:
16    name: str
17    start: int
18    end: int
19    size: int
20    children: list
21
22
23class TfaElfParser:
24    """A class representing an ELF file built for TF-A.
25
26    Provides a basic interface for reading the symbol table and other
27    attributes of an ELF file. The constructor accepts a file-like object with
28    the contents an ELF file.
29    """
30
31    def __init__(self, elf_file: BinaryIO):
32        self._segments = {}
33        self._memory_layout = {}
34
35        elf = ELFFile(elf_file)
36
37        self._symbols = {
38            sym.name: sym.entry["st_value"]
39            for sym in elf.get_section_by_name(".symtab").iter_symbols()
40        }
41
42        self.set_segment_section_map(elf.iter_segments(), elf.iter_sections())
43        self._memory_layout = self.get_memory_layout_from_symbols()
44        self._start = elf["e_entry"]
45        self._size, self._free = self._get_mem_usage()
46        self._end = self._start + self._size
47
48    @property
49    def symbols(self):
50        return self._symbols.items()
51
52    @staticmethod
53    def tfa_mem_obj_factory(elf_obj, name=None, children=None, segment=False):
54        """Converts a pyelfparser Segment or Section to a TfaMemObject."""
55        # Ensure each segment is provided a name since they aren't in the
56        # program header.
57        assert not (
58            segment and name is None
59        ), "Attempting to make segment without a name"
60
61        if children is None:
62            children = list()
63
64        # Segment and sections header keys have different prefixes.
65        vaddr = "p_vaddr" if segment else "sh_addr"
66        size = "p_memsz" if segment else "sh_size"
67
68        # TODO figure out how to handle free space for sections and segments
69        return TfaMemObject(
70            name if segment else elf_obj.name,
71            elf_obj[vaddr],
72            elf_obj[vaddr] + elf_obj[size],
73            elf_obj[size],
74            [] if not children else children,
75        )
76
77    def _get_mem_usage(self) -> (int, int):
78        """Get total size and free space for this component."""
79        size = free = 0
80
81        # Use information encoded in the segment header if we can't get a
82        # memory configuration.
83        if not self._memory_layout:
84            return sum(s.size for s in self._segments.values()), 0
85
86        for v in self._memory_layout.values():
87            size += v["length"]
88            free += v["start"] + v["length"] - v["end"]
89
90        return size, free
91
92    def set_segment_section_map(self, segments, sections):
93        """Set segment to section mappings."""
94        segments = list(
95            filter(lambda seg: seg["p_type"] == "PT_LOAD", segments)
96        )
97
98        for sec in sections:
99            for n, seg in enumerate(segments):
100                if seg.section_in_segment(sec):
101                    if n not in self._segments.keys():
102                        self._segments[n] = self.tfa_mem_obj_factory(
103                            seg, name=f"{n:#02}", segment=True
104                        )
105
106                    self._segments[n].children.append(
107                        self.tfa_mem_obj_factory(sec)
108                    )
109
110    def get_memory_layout_from_symbols(self, expr=None) -> dict:
111        """Retrieve information about the memory configuration from the symbol
112        table.
113        """
114        assert len(self._symbols), "Symbol table is empty!"
115
116        expr = r".*(.?R.M)_REGION.*(START|END|LENGTH)" if not expr else expr
117        region_symbols = filter(lambda s: re.match(expr, s), self._symbols)
118        memory_layout = {}
119
120        for symbol in region_symbols:
121            region, _, attr = tuple(symbol.lower().strip("__").split("_"))
122            if region not in memory_layout:
123                memory_layout[region] = {}
124
125            # Retrieve the value of the symbol using the symbol as the key.
126            memory_layout[region][attr] = self._symbols[symbol]
127
128        return memory_layout
129
130    def get_seg_map_as_dict(self):
131        """Get a dictionary of segments and their section mappings."""
132        return [asdict(v) for k, v in self._segments.items()]
133
134    def get_memory_layout(self):
135        """Get the total memory consumed by this module from the memory
136        configuration.
137            {"rom": {"start": 0x0, "end": 0xFF, "length": ... }
138        """
139        mem_dict = {}
140
141        for mem, attrs in self._memory_layout.items():
142            limit = attrs["start"] + attrs["length"]
143            mem_dict[mem] = {
144                "start": attrs["start"],
145                "limit": limit,
146                "size": attrs["end"] - attrs["start"],
147                "free": limit - attrs["end"],
148                "total": attrs["length"],
149            }
150        return mem_dict
151
152    def get_mod_mem_usage_dict(self):
153        """Get the total memory consumed by the module, this combines the
154        information in the memory configuration.
155        """
156        return {
157            "start": self._start,
158            "end": self._end,
159            "size": self._size,
160            "free": self._free,
161        }
162