1 /*
2  * Copyright (C) 2017 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 #pragma once
18 
19 #include "common.h"
20 #include "buffer.h"
21 #include "arrayview.h"
22 #include "dex_format.h"
23 #include "dex_ir.h"
24 
25 #include <map>
26 #include <memory>
27 #include <vector>
28 
29 namespace dex {
30 
31 // Specialized buffer for creating a .dex image section
32 // (tracking the section offset, section type, ...)
33 class Section : public slicer::Buffer {
34  public:
Section(dex::u2 mapEntryType)35   explicit Section(dex::u2 mapEntryType) : map_entry_type_(mapEntryType) {}
36   ~Section() = default;
37 
38   Section(const Section&) = delete;
39   Section& operator=(const Section&) = delete;
40 
SetOffset(dex::u4 offset)41   void SetOffset(dex::u4 offset) {
42     SLICER_CHECK(offset > 0 && offset % 4 == 0);
43     offset_ = offset;
44   }
45 
SectionOffset()46   dex::u4 SectionOffset() const {
47     SLICER_CHECK(offset_ > 0 && offset_ % 4 == 0);
48     return ItemsCount() > 0 ? offset_ : 0;
49   }
50 
AbsoluteOffset(dex::u4 itemOffset)51   dex::u4 AbsoluteOffset(dex::u4 itemOffset) const {
52     SLICER_CHECK(offset_ > 0);
53     SLICER_CHECK(itemOffset < size());
54     return offset_ + itemOffset;
55   }
56 
57   // TODO: return absolute offsets?
58   dex::u4 AddItem(dex::u4 alignment = 1) {
59     ++count_;
60     Align(alignment);
61     return size();
62   }
63 
ItemsCount()64   dex::u4 ItemsCount() const { return count_; }
65 
MapEntryType()66   dex::u2 MapEntryType() const { return map_entry_type_; }
67 
68  private:
69   dex::u4 offset_ = 0;
70   dex::u4 count_ = 0;
71   const dex::u2 map_entry_type_;
72 };
73 
74 // A specialized container for an .dex index section
75 // (strings, types, fields, methods, ...)
76 template <class T>
77 class Index {
78  public:
Index(dex::u2 mapEntryType)79   explicit Index(dex::u2 mapEntryType) : map_entry_type_(mapEntryType) {}
80   ~Index() = default;
81 
82   Index(const Index&) = delete;
83   Index& operator=(const Index&) = delete;
84 
Init(dex::u4 offset,dex::u4 count)85   dex::u4 Init(dex::u4 offset, dex::u4 count) {
86     values_.reset(new T[count]);
87     offset_ = offset;
88     count_ = count;
89     return size();
90   }
91 
Free()92   void Free() {
93     values_.reset();
94     offset_ = 0;
95     count_ = 0;
96   }
97 
SectionOffset()98   dex::u4 SectionOffset() const {
99     SLICER_CHECK(offset_ > 0 && offset_ % 4 == 0);
100     return ItemsCount() > 0 ? offset_ : 0;
101   }
102 
begin()103   T* begin() { return values_.get(); }
end()104   T* end() { return begin() + count_; }
105 
empty()106   bool empty() const { return count_ == 0; }
107 
ItemsCount()108   dex::u4 ItemsCount() const { return count_; }
data()109   const T* data() const { return values_.get(); }
size()110   dex::u4 size() const { return count_ * sizeof(T); }
111 
112   T& operator[](int i) {
113     SLICER_CHECK(i >= 0 && i < count_);
114     return values_[i];
115   }
116 
MapEntryType()117   dex::u2 MapEntryType() const { return map_entry_type_; }
118 
119  private:
120   dex::u4 offset_ = 0;
121   dex::u4 count_ = 0;
122   std::unique_ptr<T[]> values_;
123   const dex::u2 map_entry_type_;
124 };
125 
126 // Creates an in-memory .dex image from a .dex IR
127 class Writer {
128   // The container for the individual sections in a .dex image
129   // (factored out from Writer for a more granular lifetime control)
130   struct DexImage {
DexImageDexImage131     DexImage()
132         : string_ids(dex::kStringIdItem),
133           type_ids(dex::kTypeIdItem),
134           proto_ids(dex::kProtoIdItem),
135           field_ids(dex::kFieldIdItem),
136           method_ids(dex::kMethodIdItem),
137           class_defs(dex::kClassDefItem),
138           string_data(dex::kStringDataItem),
139           type_lists(dex::kTypeList),
140           debug_info(dex::kDebugInfoItem),
141           encoded_arrays(dex::kEncodedArrayItem),
142           code(dex::kCodeItem),
143           class_data(dex::kClassDataItem),
144           ann_directories(dex::kAnnotationsDirectoryItem),
145           ann_set_ref_lists(dex::kAnnotationSetRefList),
146           ann_sets(dex::kAnnotationSetItem),
147           ann_items(dex::kAnnotationItem),
148           map_list(dex::kMapList) {}
149 
150     Index<dex::StringId> string_ids;
151     Index<dex::TypeId> type_ids;
152     Index<dex::ProtoId> proto_ids;
153     Index<dex::FieldId> field_ids;
154     Index<dex::MethodId> method_ids;
155     Index<dex::ClassDef> class_defs;
156 
157     Section string_data;
158     Section type_lists;
159     Section debug_info;
160     Section encoded_arrays;
161     Section code;
162     Section class_data;
163     Section ann_directories;
164     Section ann_set_ref_lists;
165     Section ann_sets;
166     Section ann_items;
167     Section map_list;
168   };
169 
170  public:
171   // interface for allocating the final in-memory image
172   struct Allocator {
173     virtual void* Allocate(size_t size) = 0;
174     virtual void Free(void* ptr) = 0;
175     virtual ~Allocator() = default;
176   };
177 
178  public:
Writer(std::shared_ptr<ir::DexFile> dex_ir)179   explicit Writer(std::shared_ptr<ir::DexFile> dex_ir) : dex_ir_(dex_ir) {}
180   ~Writer() = default;
181 
182   Writer(const Writer&) = delete;
183   Writer& operator=(const Writer&) = delete;
184 
185   // .dex image creation
186   dex::u1* CreateImage(Allocator* allocator, size_t* new_image_size);
187 
188  private:
189   // helpers for creating various .dex sections
190   dex::u4 CreateStringDataSection(dex::u4 section_offset);
191   dex::u4 CreateMapSection(dex::u4 section_offset);
192   dex::u4 CreateAnnItemSection(dex::u4 section_offset);
193   dex::u4 CreateAnnSetsSection(dex::u4 section_offset);
194   dex::u4 CreateAnnSetRefListsSection(dex::u4 section_offset);
195   dex::u4 CreateTypeListsSection(dex::u4 section_offset);
196   dex::u4 CreateCodeItemSection(dex::u4 section_offset);
197   dex::u4 CreateDebugInfoSection(dex::u4 section_offset);
198   dex::u4 CreateClassDataSection(dex::u4 section_offset);
199   dex::u4 CreateAnnDirectoriesSection(dex::u4 section_offset);
200   dex::u4 CreateEncodedArrayItemSection(dex::u4 section_offset);
201 
202   // back-fill the indexes
203   void FillTypes();
204   void FillProtos();
205   void FillFields();
206   void FillMethods();
207   void FillClassDefs();
208 
209   // helpers for writing .dex structures
210   dex::u4 WriteTypeList(const std::vector<ir::Type*>& types);
211   dex::u4 WriteAnnotationItem(const ir::Annotation* ir_annotation);
212   dex::u4 WriteAnnotationSet(const ir::AnnotationSet* ir_annotation_set);
213   dex::u4 WriteAnnotationSetRefList(const ir::AnnotationSetRefList* ir_annotation_set_ref_list);
214   dex::u4 WriteClassAnnotations(const ir::Class* ir_class);
215   dex::u4 WriteDebugInfo(const ir::DebugInfo* ir_debug_info);
216   dex::u4 WriteCode(const ir::Code* ir_code);
217   dex::u4 WriteClassData(const ir::Class* ir_class);
218   dex::u4 WriteClassStaticValues(const ir::Class* ir_class);
219 
220   // Map indexes from the original .dex to the
221   // corresponding index in the new image
222   dex::u4 MapStringIndex(dex::u4 index) const;
223   dex::u4 MapTypeIndex(dex::u4 index) const;
224   dex::u4 MapFieldIndex(dex::u4 index) const;
225   dex::u4 MapMethodIndex(dex::u4 index) const;
226 
227   // writing parts of a class definition
228   void WriteInstructions(slicer::ArrayView<const dex::u2> instructions);
229   void WriteTryBlocks(const ir::Code* ir_code);
230   void WriteEncodedField(const ir::EncodedField* irEncodedField, dex::u4* base_index);
231   void WriteEncodedMethod(const ir::EncodedMethod* irEncodedMethod, dex::u4* base_index);
232 
233   dex::u4 FilePointer(const ir::Node* ir_node) const;
234 
235  private:
236   std::shared_ptr<ir::DexFile> dex_ir_;
237   std::unique_ptr<DexImage> dex_;
238 
239   // CONSIDER: we can have multiple maps per IR node type
240   //  (that's what the reader does)
241   std::map<const ir::Node*, dex::u4> node_offset_;
242 };
243 
244 }  // namespace dex
245