1 /*
2  * Copyright (C) 2015 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #ifndef ART_LIBDEXFILE_DEX_TEST_DEX_FILE_BUILDER_H_
18 #define ART_LIBDEXFILE_DEX_TEST_DEX_FILE_BUILDER_H_
19 
20 #include <zlib.h>
21 
22 #include <cstring>
23 #include <map>
24 #include <set>
25 #include <vector>
26 
27 #include <android-base/logging.h>
28 
29 #include "base/array_ref.h"
30 #include "base/bit_utils.h"
31 #include "dex/dex_file_loader.h"
32 #include "dex/standard_dex_file.h"
33 
34 namespace art {
35 
36 class TestDexFileBuilder {
37  public:
TestDexFileBuilder()38   TestDexFileBuilder()
39       : strings_(), types_(), fields_(), protos_() {
40   }
41 
AddString(const std::string & str)42   void AddString(const std::string& str) {
43     CHECK_LT(str.length(), 128u);  // Don't allow multi-byte length in uleb128.
44     strings_.emplace(str, IdxAndDataOffset());
45   }
46 
AddType(const std::string & descriptor)47   void AddType(const std::string& descriptor) {
48     AddString(descriptor);
49     types_.emplace(descriptor, 0u);
50   }
51 
AddField(const std::string & class_descriptor,const std::string & type,const std::string & name)52   void AddField(const std::string& class_descriptor,
53                 const std::string& type,
54                 const std::string& name) {
55     AddType(class_descriptor);
56     AddType(type);
57     AddString(name);
58     FieldKey key = { class_descriptor, type, name };
59     fields_.emplace(key, 0u);
60   }
61 
AddMethod(const std::string & class_descriptor,const std::string & signature,const std::string & name)62   void AddMethod(const std::string& class_descriptor,
63                  const std::string& signature,
64                  const std::string& name) {
65     AddType(class_descriptor);
66     AddString(name);
67 
68     ProtoKey proto_key = CreateProtoKey(signature);
69     AddString(proto_key.shorty);
70     AddType(proto_key.return_type);
71     for (const auto& arg_type : proto_key.args) {
72       AddType(arg_type);
73     }
74     auto it = protos_.emplace(proto_key, IdxAndDataOffset()).first;
75     const ProtoKey* proto = &it->first;  // Valid as long as the element remains in protos_.
76 
77     MethodKey method_key = {
78         class_descriptor, name, proto
79     };
80     methods_.emplace(method_key, 0u);
81   }
82 
83   std::unique_ptr<const DexFile> Build(const std::string& dex_location,
84                                        uint32_t location_checksum = 0u) {
85     union {
86       uint8_t data[sizeof(DexFile::Header)];
87       uint64_t force_alignment;
88     } header_data;
89     std::memset(header_data.data, 0, sizeof(header_data.data));
90     DexFile::Header* header = reinterpret_cast<DexFile::Header*>(&header_data.data);
91     std::copy_n(StandardDexFile::kDexMagic, 4u, header->magic_);
92     std::copy_n(StandardDexFile::kDexMagicVersions[0], 4u, header->magic_ + 4u);
93     header->header_size_ = sizeof(DexFile::Header);
94     header->endian_tag_ = DexFile::kDexEndianConstant;
95     header->link_size_ = 0u;  // Unused.
96     header->link_off_ = 0u;  // Unused.
97     header->map_off_ = 0u;  // Unused. TODO: This is wrong. Dex files created by this builder
98                             //               cannot be verified. b/26808512
99 
100     uint32_t data_section_size = 0u;
101 
102     uint32_t string_ids_offset = sizeof(DexFile::Header);
103     uint32_t string_idx = 0u;
104     for (auto& entry : strings_) {
105       entry.second.idx = string_idx;
106       string_idx += 1u;
107       entry.second.data_offset = data_section_size;
108       data_section_size += entry.first.length() + 1u /* length */ + 1u /* null-terminator */;
109     }
110     header->string_ids_size_ = strings_.size();
111     header->string_ids_off_ = strings_.empty() ? 0u : string_ids_offset;
112 
113     uint32_t type_ids_offset = string_ids_offset + strings_.size() * sizeof(dex::StringId);
114     uint32_t type_idx = 0u;
115     for (auto& entry : types_) {
116       entry.second = type_idx;
117       type_idx += 1u;
118     }
119     header->type_ids_size_ = types_.size();
120     header->type_ids_off_ = types_.empty() ? 0u : type_ids_offset;
121 
122     uint32_t proto_ids_offset = type_ids_offset + types_.size() * sizeof(dex::TypeId);
123     uint32_t proto_idx = 0u;
124     for (auto& entry : protos_) {
125       entry.second.idx = proto_idx;
126       proto_idx += 1u;
127       size_t num_args = entry.first.args.size();
128       if (num_args != 0u) {
129         entry.second.data_offset = RoundUp(data_section_size, 4u);
130         data_section_size = entry.second.data_offset + 4u + num_args * sizeof(dex::TypeItem);
131       } else {
132         entry.second.data_offset = 0u;
133       }
134     }
135     header->proto_ids_size_ = protos_.size();
136     header->proto_ids_off_ = protos_.empty() ? 0u : proto_ids_offset;
137 
138     uint32_t field_ids_offset = proto_ids_offset + protos_.size() * sizeof(dex::ProtoId);
139     uint32_t field_idx = 0u;
140     for (auto& entry : fields_) {
141       entry.second = field_idx;
142       field_idx += 1u;
143     }
144     header->field_ids_size_ = fields_.size();
145     header->field_ids_off_ = fields_.empty() ? 0u : field_ids_offset;
146 
147     uint32_t method_ids_offset = field_ids_offset + fields_.size() * sizeof(dex::FieldId);
148     uint32_t method_idx = 0u;
149     for (auto& entry : methods_) {
150       entry.second = method_idx;
151       method_idx += 1u;
152     }
153     header->method_ids_size_ = methods_.size();
154     header->method_ids_off_ = methods_.empty() ? 0u : method_ids_offset;
155 
156     // No class defs.
157     header->class_defs_size_ = 0u;
158     header->class_defs_off_ = 0u;
159 
160     uint32_t data_section_offset = method_ids_offset + methods_.size() * sizeof(dex::MethodId);
161     header->data_size_ = data_section_size;
162     header->data_off_ = (data_section_size != 0u) ? data_section_offset : 0u;
163 
164     uint32_t total_size = data_section_offset + data_section_size;
165     std::vector<uint8_t> dex_file_data(total_size, 0u);
166 
167     for (const auto& entry : strings_) {
168       CHECK_LT(entry.first.size(), 128u);
169       uint32_t raw_offset = data_section_offset + entry.second.data_offset;
170       dex_file_data[raw_offset] = static_cast<uint8_t>(entry.first.size());
171       std::memcpy(&dex_file_data[raw_offset + 1], entry.first.c_str(), entry.first.size() + 1);
172       Write32(dex_file_data,
173               string_ids_offset + entry.second.idx * sizeof(dex::StringId),
174               raw_offset);
175     }
176 
177     for (const auto& entry : types_) {
178       Write32(dex_file_data,
179               type_ids_offset + entry.second * sizeof(dex::TypeId),
180               GetStringIdx(entry.first));
181       ++type_idx;
182     }
183 
184     for (const auto& entry : protos_) {
185       size_t num_args = entry.first.args.size();
186       uint32_t type_list_offset =
187           (num_args != 0u) ? data_section_offset + entry.second.data_offset : 0u;
188       uint32_t raw_offset = proto_ids_offset + entry.second.idx * sizeof(dex::ProtoId);
189       Write32(dex_file_data, raw_offset + 0u, GetStringIdx(entry.first.shorty));
190       Write16(dex_file_data, raw_offset + 4u, GetTypeIdx(entry.first.return_type));
191       Write32(dex_file_data, raw_offset + 8u, type_list_offset);
192       if (num_args != 0u) {
193         CHECK_NE(entry.second.data_offset, 0u);
194         Write32(dex_file_data, type_list_offset, num_args);
195         for (size_t i = 0; i != num_args; ++i) {
196           Write16(dex_file_data,
197                   type_list_offset + 4u + i * sizeof(dex::TypeItem),
198                   GetTypeIdx(entry.first.args[i]));
199         }
200       }
201     }
202 
203     for (const auto& entry : fields_) {
204       uint32_t raw_offset = field_ids_offset + entry.second * sizeof(dex::FieldId);
205       Write16(dex_file_data, raw_offset + 0u, GetTypeIdx(entry.first.class_descriptor));
206       Write16(dex_file_data, raw_offset + 2u, GetTypeIdx(entry.first.type));
207       Write32(dex_file_data, raw_offset + 4u, GetStringIdx(entry.first.name));
208     }
209 
210     for (const auto& entry : methods_) {
211       uint32_t raw_offset = method_ids_offset + entry.second * sizeof(dex::MethodId);
212       Write16(dex_file_data, raw_offset + 0u, GetTypeIdx(entry.first.class_descriptor));
213       auto it = protos_.find(*entry.first.proto);
214       CHECK(it != protos_.end());
215       Write16(dex_file_data, raw_offset + 2u, it->second.idx);
216       Write32(dex_file_data, raw_offset + 4u, GetStringIdx(entry.first.name));
217     }
218 
219     // Leave signature as zeros.
220 
221     header->file_size_ = dex_file_data.size();
222 
223     // Write the complete header early, as part of it needs to be checksummed.
224     std::memcpy(&dex_file_data[0], header_data.data, sizeof(DexFile::Header));
225 
226     // Checksum starts after the checksum field.
227     size_t skip = sizeof(header->magic_) + sizeof(header->checksum_);
228     header->checksum_ = adler32(adler32(0L, Z_NULL, 0),
229                                 dex_file_data.data() + skip,
230                                 dex_file_data.size() - skip);
231 
232     // Write the complete header again, just simpler that way.
233     std::memcpy(&dex_file_data[0], header_data.data, sizeof(DexFile::Header));
234 
235     // Do not protect the final data from writing. Some tests need to modify it.
236 
237     static constexpr bool kVerify = false;
238     static constexpr bool kVerifyChecksum = false;
239     std::string error_msg;
240     std::unique_ptr<const DexFile> dex_file(DexFileLoader::Open(
241         dex_location,
242         location_checksum,
243         std::move(dex_file_data),
244         /*oat_dex_file=*/ nullptr,
245         kVerify,
246         kVerifyChecksum,
247         &error_msg));
248     CHECK(dex_file != nullptr) << error_msg;
249     return dex_file;
250   }
251 
GetStringIdx(const std::string & type)252   uint32_t GetStringIdx(const std::string& type) {
253     auto it = strings_.find(type);
254     CHECK(it != strings_.end());
255     return it->second.idx;
256   }
257 
GetTypeIdx(const std::string & type)258   uint32_t GetTypeIdx(const std::string& type) {
259     auto it = types_.find(type);
260     CHECK(it != types_.end());
261     return it->second;
262   }
263 
GetFieldIdx(const std::string & class_descriptor,const std::string & type,const std::string & name)264   uint32_t GetFieldIdx(const std::string& class_descriptor, const std::string& type,
265                        const std::string& name) {
266     FieldKey key = { class_descriptor, type, name };
267     auto it = fields_.find(key);
268     CHECK(it != fields_.end());
269     return it->second;
270   }
271 
GetMethodIdx(const std::string & class_descriptor,const std::string & signature,const std::string & name)272   uint32_t GetMethodIdx(const std::string& class_descriptor, const std::string& signature,
273                         const std::string& name) {
274     ProtoKey proto_key = CreateProtoKey(signature);
275     MethodKey method_key = { class_descriptor, name, &proto_key };
276     auto it = methods_.find(method_key);
277     CHECK(it != methods_.end());
278     return it->second;
279   }
280 
281  private:
282   struct IdxAndDataOffset {
283     uint32_t idx;
284     uint32_t data_offset;
285   };
286 
287   struct FieldKey {
288     const std::string class_descriptor;
289     const std::string type;
290     const std::string name;
291   };
292   struct FieldKeyComparator {
operatorFieldKeyComparator293     bool operator()(const FieldKey& lhs, const FieldKey& rhs) const {
294       if (lhs.class_descriptor != rhs.class_descriptor) {
295         return lhs.class_descriptor < rhs.class_descriptor;
296       }
297       if (lhs.name != rhs.name) {
298         return lhs.name < rhs.name;
299       }
300       return lhs.type < rhs.type;
301     }
302   };
303 
304   struct ProtoKey {
305     std::string shorty;
306     std::string return_type;
307     std::vector<std::string> args;
308   };
309   struct ProtoKeyComparator {
operatorProtoKeyComparator310     bool operator()(const ProtoKey& lhs, const ProtoKey& rhs) const {
311       if (lhs.return_type != rhs.return_type) {
312         return lhs.return_type < rhs.return_type;
313       }
314       size_t min_args = std::min(lhs.args.size(), rhs.args.size());
315       for (size_t i = 0; i != min_args; ++i) {
316         if (lhs.args[i] != rhs.args[i]) {
317           return lhs.args[i] < rhs.args[i];
318         }
319       }
320       return lhs.args.size() < rhs.args.size();
321     }
322   };
323 
324   struct MethodKey {
325     std::string class_descriptor;
326     std::string name;
327     const ProtoKey* proto;
328   };
329   struct MethodKeyComparator {
operatorMethodKeyComparator330     bool operator()(const MethodKey& lhs, const MethodKey& rhs) const {
331       if (lhs.class_descriptor != rhs.class_descriptor) {
332         return lhs.class_descriptor < rhs.class_descriptor;
333       }
334       if (lhs.name != rhs.name) {
335         return lhs.name < rhs.name;
336       }
337       return ProtoKeyComparator()(*lhs.proto, *rhs.proto);
338     }
339   };
340 
CreateProtoKey(const std::string & signature)341   ProtoKey CreateProtoKey(const std::string& signature) {
342     CHECK_EQ(signature[0], '(');
343     const char* args = signature.c_str() + 1;
344     const char* args_end = std::strchr(args, ')');
345     CHECK(args_end != nullptr);
346     const char* return_type = args_end + 1;
347 
348     ProtoKey key = {
349         std::string() + ((*return_type == '[') ? 'L' : *return_type),
350         return_type,
351         std::vector<std::string>()
352     };
353     while (args != args_end) {
354       key.shorty += (*args == '[') ? 'L' : *args;
355       const char* arg_start = args;
356       while (*args == '[') {
357         ++args;
358       }
359       if (*args == 'L') {
360         do {
361           ++args;
362           CHECK_NE(args, args_end);
363         } while (*args != ';');
364       }
365       ++args;
366       key.args.emplace_back(arg_start, args);
367     }
368     return key;
369   }
370 
Write32(std::vector<uint8_t> & dex_file_data,size_t offset,uint32_t value)371   static void Write32(std::vector<uint8_t>& dex_file_data, size_t offset, uint32_t value) {
372     CHECK_LE(offset + 4u, dex_file_data.size());
373     CHECK_EQ(dex_file_data[offset + 0], 0u);
374     CHECK_EQ(dex_file_data[offset + 1], 0u);
375     CHECK_EQ(dex_file_data[offset + 2], 0u);
376     CHECK_EQ(dex_file_data[offset + 3], 0u);
377     dex_file_data[offset + 0] = static_cast<uint8_t>(value >> 0);
378     dex_file_data[offset + 1] = static_cast<uint8_t>(value >> 8);
379     dex_file_data[offset + 2] = static_cast<uint8_t>(value >> 16);
380     dex_file_data[offset + 3] = static_cast<uint8_t>(value >> 24);
381   }
382 
Write16(std::vector<uint8_t> & dex_file_data,size_t offset,uint32_t value)383   static void Write16(std::vector<uint8_t>& dex_file_data, size_t offset, uint32_t value) {
384     CHECK_LE(value, 0xffffu);
385     CHECK_LE(offset + 2u, dex_file_data.size());
386     CHECK_EQ(dex_file_data[offset + 0], 0u);
387     CHECK_EQ(dex_file_data[offset + 1], 0u);
388     dex_file_data[offset + 0] = static_cast<uint8_t>(value >> 0);
389     dex_file_data[offset + 1] = static_cast<uint8_t>(value >> 8);
390   }
391 
392   std::map<std::string, IdxAndDataOffset> strings_;
393   std::map<std::string, uint32_t> types_;
394   std::map<FieldKey, uint32_t, FieldKeyComparator> fields_;
395   std::map<ProtoKey, IdxAndDataOffset, ProtoKeyComparator> protos_;
396   std::map<MethodKey, uint32_t, MethodKeyComparator> methods_;
397 };
398 
399 }  // namespace art
400 
401 #endif  // ART_LIBDEXFILE_DEX_TEST_DEX_FILE_BUILDER_H_
402