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