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