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
30namespace art {
31
32class TestDexFileBuilder {
33 public:
34  TestDexFileBuilder()
35      : strings_(), types_(), fields_(), protos_(), dex_file_data_() {
36  }
37
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
44  void AddType(const std::string& descriptor) {
45    CHECK(dex_file_data_.empty());
46    AddString(descriptor);
47    types_.emplace(descriptor, 0u);
48  }
49
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
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.
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    std::string error_msg;
231    std::unique_ptr<const DexFile> dex_file(DexFile::Open(
232        &dex_file_data_[0], dex_file_data_.size(), dex_location, 0u, nullptr, false, &error_msg));
233    CHECK(dex_file != nullptr) << error_msg;
234    return dex_file;
235  }
236
237  uint32_t GetStringIdx(const std::string& type) {
238    auto it = strings_.find(type);
239    CHECK(it != strings_.end());
240    return it->second.idx;
241  }
242
243  uint32_t GetTypeIdx(const std::string& type) {
244    auto it = types_.find(type);
245    CHECK(it != types_.end());
246    return it->second;
247  }
248
249  uint32_t GetFieldIdx(const std::string& class_descriptor, const std::string& type,
250                       const std::string& name) {
251    FieldKey key = { class_descriptor, type, name };
252    auto it = fields_.find(key);
253    CHECK(it != fields_.end());
254    return it->second;
255  }
256
257  uint32_t GetMethodIdx(const std::string& class_descriptor, const std::string& signature,
258                        const std::string& name) {
259    ProtoKey proto_key = CreateProtoKey(signature);
260    MethodKey method_key = { class_descriptor, name, &proto_key };
261    auto it = methods_.find(method_key);
262    CHECK(it != methods_.end());
263    return it->second;
264  }
265
266 private:
267  struct IdxAndDataOffset {
268    uint32_t idx;
269    uint32_t data_offset;
270  };
271
272  struct FieldKey {
273    const std::string class_descriptor;
274    const std::string type;
275    const std::string name;
276  };
277  struct FieldKeyComparator {
278    bool operator()(const FieldKey& lhs, const FieldKey& rhs) const {
279      if (lhs.class_descriptor != rhs.class_descriptor) {
280        return lhs.class_descriptor < rhs.class_descriptor;
281      }
282      if (lhs.name != rhs.name) {
283        return lhs.name < rhs.name;
284      }
285      return lhs.type < rhs.type;
286    }
287  };
288
289  struct ProtoKey {
290    std::string shorty;
291    std::string return_type;
292    std::vector<std::string> args;
293  };
294  struct ProtoKeyComparator {
295    bool operator()(const ProtoKey& lhs, const ProtoKey& rhs) const {
296      if (lhs.return_type != rhs.return_type) {
297        return lhs.return_type < rhs.return_type;
298      }
299      size_t min_args = std::min(lhs.args.size(), rhs.args.size());
300      for (size_t i = 0; i != min_args; ++i) {
301        if (lhs.args[i] != rhs.args[i]) {
302          return lhs.args[i] < rhs.args[i];
303        }
304      }
305      return lhs.args.size() < rhs.args.size();
306    }
307  };
308
309  struct MethodKey {
310    std::string class_descriptor;
311    std::string name;
312    const ProtoKey* proto;
313  };
314  struct MethodKeyComparator {
315    bool operator()(const MethodKey& lhs, const MethodKey& rhs) const {
316      if (lhs.class_descriptor != rhs.class_descriptor) {
317        return lhs.class_descriptor < rhs.class_descriptor;
318      }
319      if (lhs.name != rhs.name) {
320        return lhs.name < rhs.name;
321      }
322      return ProtoKeyComparator()(*lhs.proto, *rhs.proto);
323    }
324  };
325
326  ProtoKey CreateProtoKey(const std::string& signature) {
327    CHECK_EQ(signature[0], '(');
328    const char* args = signature.c_str() + 1;
329    const char* args_end = std::strchr(args, ')');
330    CHECK(args_end != nullptr);
331    const char* return_type = args_end + 1;
332
333    ProtoKey key = {
334        std::string() + ((*return_type == '[') ? 'L' : *return_type),
335        return_type,
336        std::vector<std::string>()
337    };
338    while (args != args_end) {
339      key.shorty += (*args == '[') ? 'L' : *args;
340      const char* arg_start = args;
341      while (*args == '[') {
342        ++args;
343      }
344      if (*args == 'L') {
345        do {
346          ++args;
347          CHECK_NE(args, args_end);
348        } while (*args != ';');
349      }
350      ++args;
351      key.args.emplace_back(arg_start, args);
352    }
353    return key;
354  }
355
356  void Write32(size_t offset, uint32_t value) {
357    CHECK_LE(offset + 4u, dex_file_data_.size());
358    CHECK_EQ(dex_file_data_[offset + 0], 0u);
359    CHECK_EQ(dex_file_data_[offset + 1], 0u);
360    CHECK_EQ(dex_file_data_[offset + 2], 0u);
361    CHECK_EQ(dex_file_data_[offset + 3], 0u);
362    dex_file_data_[offset + 0] = static_cast<uint8_t>(value >> 0);
363    dex_file_data_[offset + 1] = static_cast<uint8_t>(value >> 8);
364    dex_file_data_[offset + 2] = static_cast<uint8_t>(value >> 16);
365    dex_file_data_[offset + 3] = static_cast<uint8_t>(value >> 24);
366  }
367
368  void Write16(size_t offset, uint32_t value) {
369    CHECK_LE(value, 0xffffu);
370    CHECK_LE(offset + 2u, dex_file_data_.size());
371    CHECK_EQ(dex_file_data_[offset + 0], 0u);
372    CHECK_EQ(dex_file_data_[offset + 1], 0u);
373    dex_file_data_[offset + 0] = static_cast<uint8_t>(value >> 0);
374    dex_file_data_[offset + 1] = static_cast<uint8_t>(value >> 8);
375  }
376
377  std::map<std::string, IdxAndDataOffset> strings_;
378  std::map<std::string, uint32_t> types_;
379  std::map<FieldKey, uint32_t, FieldKeyComparator> fields_;
380  std::map<ProtoKey, IdxAndDataOffset, ProtoKeyComparator> protos_;
381  std::map<MethodKey, uint32_t, MethodKeyComparator> methods_;
382
383  std::vector<uint8_t> dex_file_data_;
384};
385
386}  // namespace art
387
388#endif  // ART_COMPILER_UTILS_TEST_DEX_FILE_BUILDER_H_
389