1 //===- InstrumentationMap.cpp - XRay Instrumentation Map ------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // Implementation of the InstrumentationMap type for XRay sleds.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "llvm/XRay/InstrumentationMap.h"
14 #include "llvm/ADT/DenseMap.h"
15 #include "llvm/ADT/STLExtras.h"
16 #include "llvm/ADT/StringRef.h"
17 #include "llvm/ADT/Twine.h"
18 #include "llvm/Object/Binary.h"
19 #include "llvm/Object/ELFObjectFile.h"
20 #include "llvm/Object/ObjectFile.h"
21 #include "llvm/Object/RelocationResolver.h"
22 #include "llvm/Support/DataExtractor.h"
23 #include "llvm/Support/Error.h"
24 #include "llvm/Support/FileSystem.h"
25 #include "llvm/Support/YAMLTraits.h"
26 #include "llvm/TargetParser/Triple.h"
27 #include <algorithm>
28 #include <cstddef>
29 #include <cstdint>
30 #include <system_error>
31 #include <vector>
32 
33 using namespace llvm;
34 using namespace xray;
35 
36 std::optional<int32_t> InstrumentationMap::getFunctionId(uint64_t Addr) const {
37   auto I = FunctionIds.find(Addr);
38   if (I != FunctionIds.end())
39     return I->second;
40   return std::nullopt;
41 }
42 
43 std::optional<uint64_t>
44 InstrumentationMap::getFunctionAddr(int32_t FuncId) const {
45   auto I = FunctionAddresses.find(FuncId);
46   if (I != FunctionAddresses.end())
47     return I->second;
48   return std::nullopt;
49 }
50 
51 using RelocMap = DenseMap<uint64_t, uint64_t>;
52 
53 static Error
54 loadObj(StringRef Filename, object::OwningBinary<object::ObjectFile> &ObjFile,
55         InstrumentationMap::SledContainer &Sleds,
56         InstrumentationMap::FunctionAddressMap &FunctionAddresses,
57         InstrumentationMap::FunctionAddressReverseMap &FunctionIds) {
58   InstrumentationMap Map;
59 
60   // Find the section named "xray_instr_map".
61   if ((!ObjFile.getBinary()->isELF() && !ObjFile.getBinary()->isMachO()) ||
62       !(ObjFile.getBinary()->getArch() == Triple::x86_64 ||
63         ObjFile.getBinary()->getArch() == Triple::loongarch64 ||
64         ObjFile.getBinary()->getArch() == Triple::ppc64le ||
65         ObjFile.getBinary()->getArch() == Triple::arm ||
66         ObjFile.getBinary()->getArch() == Triple::aarch64))
67     return make_error<StringError>(
68         "File format not supported (only does ELF and Mach-O little endian "
69         "64-bit).",
70         std::make_error_code(std::errc::not_supported));
71 
72   StringRef Contents = "";
73   const auto &Sections = ObjFile.getBinary()->sections();
74   uint64_t Address = 0;
75   auto I = llvm::find_if(Sections, [&](object::SectionRef Section) {
76     Expected<StringRef> NameOrErr = Section.getName();
77     if (NameOrErr) {
78       Address = Section.getAddress();
79       return *NameOrErr == "xray_instr_map";
80     }
81     consumeError(NameOrErr.takeError());
82     return false;
83   });
84 
85   if (I == Sections.end())
86     return make_error<StringError>(
87         "Failed to find XRay instrumentation map.",
88         std::make_error_code(std::errc::executable_format_error));
89 
90   if (Error E = I->getContents().moveInto(Contents))
91     return E;
92 
93   RelocMap Relocs;
94   if (ObjFile.getBinary()->isELF()) {
95     uint32_t RelativeRelocation = [](object::ObjectFile *ObjFile) {
96       if (const auto *ELFObj = dyn_cast<object::ELF32LEObjectFile>(ObjFile))
97         return ELFObj->getELFFile().getRelativeRelocationType();
98       else if (const auto *ELFObj =
99                    dyn_cast<object::ELF32BEObjectFile>(ObjFile))
100         return ELFObj->getELFFile().getRelativeRelocationType();
101       else if (const auto *ELFObj =
102                    dyn_cast<object::ELF64LEObjectFile>(ObjFile))
103         return ELFObj->getELFFile().getRelativeRelocationType();
104       else if (const auto *ELFObj =
105                    dyn_cast<object::ELF64BEObjectFile>(ObjFile))
106         return ELFObj->getELFFile().getRelativeRelocationType();
107       else
108         return static_cast<uint32_t>(0);
109     }(ObjFile.getBinary());
110 
111     object::SupportsRelocation Supports;
112     object::RelocationResolver Resolver;
113     std::tie(Supports, Resolver) =
114         object::getRelocationResolver(*ObjFile.getBinary());
115 
116     for (const object::SectionRef &Section : Sections) {
117       for (const object::RelocationRef &Reloc : Section.relocations()) {
118         if (ObjFile.getBinary()->getArch() == Triple::arm) {
119           if (Supports && Supports(Reloc.getType())) {
120             Expected<uint64_t> ValueOrErr = Reloc.getSymbol()->getValue();
121             if (!ValueOrErr)
122               return ValueOrErr.takeError();
123             Relocs.insert(
124                 {Reloc.getOffset(),
125                  object::resolveRelocation(Resolver, Reloc, *ValueOrErr, 0)});
126           }
127         } else if (Supports && Supports(Reloc.getType())) {
128           auto AddendOrErr = object::ELFRelocationRef(Reloc).getAddend();
129           auto A = AddendOrErr ? *AddendOrErr : 0;
130           Expected<uint64_t> ValueOrErr = Reloc.getSymbol()->getValue();
131           if (!ValueOrErr)
132             // TODO: Test this error.
133             return ValueOrErr.takeError();
134           Relocs.insert(
135               {Reloc.getOffset(),
136                object::resolveRelocation(Resolver, Reloc, *ValueOrErr, A)});
137         } else if (Reloc.getType() == RelativeRelocation) {
138           if (auto AddendOrErr = object::ELFRelocationRef(Reloc).getAddend())
139             Relocs.insert({Reloc.getOffset(), *AddendOrErr});
140         }
141       }
142     }
143   }
144 
145   // Copy the instrumentation map data into the Sleds data structure.
146   auto C = Contents.bytes_begin();
147   bool Is32Bit = ObjFile.getBinary()->makeTriple().isArch32Bit();
148   size_t ELFSledEntrySize = Is32Bit ? 16 : 32;
149 
150   if ((C - Contents.bytes_end()) % ELFSledEntrySize != 0)
151     return make_error<StringError>(
152         Twine("Instrumentation map entries not evenly divisible by size of "
153               "an XRay sled entry."),
154         std::make_error_code(std::errc::executable_format_error));
155 
156   auto RelocateOrElse = [&](uint64_t Offset, uint64_t Address) {
157     if (!Address) {
158       uint64_t A = I->getAddress() + C - Contents.bytes_begin() + Offset;
159       RelocMap::const_iterator R = Relocs.find(A);
160       if (R != Relocs.end())
161         return R->second;
162     }
163     return Address;
164   };
165 
166   const int WordSize = Is32Bit ? 4 : 8;
167   int32_t FuncId = 1;
168   uint64_t CurFn = 0;
169   for (; C != Contents.bytes_end(); C += ELFSledEntrySize) {
170     DataExtractor Extractor(
171         StringRef(reinterpret_cast<const char *>(C), ELFSledEntrySize), true,
172         8);
173     Sleds.push_back({});
174     auto &Entry = Sleds.back();
175     uint64_t OffsetPtr = 0;
176     uint64_t AddrOff = OffsetPtr;
177     if (Is32Bit)
178       Entry.Address = RelocateOrElse(AddrOff, Extractor.getU32(&OffsetPtr));
179     else
180       Entry.Address = RelocateOrElse(AddrOff, Extractor.getU64(&OffsetPtr));
181     uint64_t FuncOff = OffsetPtr;
182     if (Is32Bit)
183       Entry.Function = RelocateOrElse(FuncOff, Extractor.getU32(&OffsetPtr));
184     else
185       Entry.Function = RelocateOrElse(FuncOff, Extractor.getU64(&OffsetPtr));
186     auto Kind = Extractor.getU8(&OffsetPtr);
187     static constexpr SledEntry::FunctionKinds Kinds[] = {
188         SledEntry::FunctionKinds::ENTRY, SledEntry::FunctionKinds::EXIT,
189         SledEntry::FunctionKinds::TAIL,
190         SledEntry::FunctionKinds::LOG_ARGS_ENTER,
191         SledEntry::FunctionKinds::CUSTOM_EVENT};
192     if (Kind >= std::size(Kinds))
193       return errorCodeToError(
194           std::make_error_code(std::errc::executable_format_error));
195     Entry.Kind = Kinds[Kind];
196     Entry.AlwaysInstrument = Extractor.getU8(&OffsetPtr) != 0;
197     Entry.Version = Extractor.getU8(&OffsetPtr);
198     if (Entry.Version >= 2) {
199       Entry.Address += C - Contents.bytes_begin() + Address;
200       Entry.Function += C - Contents.bytes_begin() + WordSize + Address;
201     }
202 
203     // We do replicate the function id generation scheme implemented in the
204     // XRay runtime.
205     // FIXME: Figure out how to keep this consistent with the XRay runtime.
206     if (CurFn == 0) {
207       CurFn = Entry.Function;
208       FunctionAddresses[FuncId] = Entry.Function;
209       FunctionIds[Entry.Function] = FuncId;
210     }
211     if (Entry.Function != CurFn) {
212       ++FuncId;
213       CurFn = Entry.Function;
214       FunctionAddresses[FuncId] = Entry.Function;
215       FunctionIds[Entry.Function] = FuncId;
216     }
217   }
218   return Error::success();
219 }
220 
221 static Error
222 loadYAML(sys::fs::file_t Fd, size_t FileSize, StringRef Filename,
223          InstrumentationMap::SledContainer &Sleds,
224          InstrumentationMap::FunctionAddressMap &FunctionAddresses,
225          InstrumentationMap::FunctionAddressReverseMap &FunctionIds) {
226   std::error_code EC;
227   sys::fs::mapped_file_region MappedFile(
228       Fd, sys::fs::mapped_file_region::mapmode::readonly, FileSize, 0, EC);
229   sys::fs::closeFile(Fd);
230   if (EC)
231     return make_error<StringError>(
232         Twine("Failed memory-mapping file '") + Filename + "'.", EC);
233 
234   std::vector<YAMLXRaySledEntry> YAMLSleds;
235   yaml::Input In(StringRef(MappedFile.data(), MappedFile.size()));
236   In >> YAMLSleds;
237   if (In.error())
238     return make_error<StringError>(
239         Twine("Failed loading YAML document from '") + Filename + "'.",
240         In.error());
241 
242   Sleds.reserve(YAMLSleds.size());
243   for (const auto &Y : YAMLSleds) {
244     FunctionAddresses[Y.FuncId] = Y.Function;
245     FunctionIds[Y.Function] = Y.FuncId;
246     Sleds.push_back(SledEntry{Y.Address, Y.Function, Y.Kind, Y.AlwaysInstrument,
247                               Y.Version});
248   }
249   return Error::success();
250 }
251 
252 // FIXME: Create error types that encapsulate a bit more information than what
253 // StringError instances contain.
254 Expected<InstrumentationMap>
255 llvm::xray::loadInstrumentationMap(StringRef Filename) {
256   // At this point we assume the file is an object file -- and if that doesn't
257   // work, we treat it as YAML.
258   // FIXME: Extend to support non-ELF and non-x86_64 binaries.
259 
260   InstrumentationMap Map;
261   auto ObjectFileOrError = object::ObjectFile::createObjectFile(Filename);
262   if (!ObjectFileOrError) {
263     auto E = ObjectFileOrError.takeError();
264     // We try to load it as YAML if the ELF load didn't work.
265     Expected<sys::fs::file_t> FdOrErr =
266         sys::fs::openNativeFileForRead(Filename);
267     if (!FdOrErr) {
268       // Report the ELF load error if YAML failed.
269       consumeError(FdOrErr.takeError());
270       return std::move(E);
271     }
272 
273     uint64_t FileSize;
274     if (sys::fs::file_size(Filename, FileSize))
275       return std::move(E);
276 
277     // If the file is empty, we return the original error.
278     if (FileSize == 0)
279       return std::move(E);
280 
281     // From this point on the errors will be only for the YAML parts, so we
282     // consume the errors at this point.
283     consumeError(std::move(E));
284     if (auto E = loadYAML(*FdOrErr, FileSize, Filename, Map.Sleds,
285                           Map.FunctionAddresses, Map.FunctionIds))
286       return std::move(E);
287   } else if (auto E = loadObj(Filename, *ObjectFileOrError, Map.Sleds,
288                               Map.FunctionAddresses, Map.FunctionIds)) {
289     return std::move(E);
290   }
291   return Map;
292 }
293