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