1 : // Copyright 2012 Google Inc.
2 : //
3 : // Licensed under the Apache License, Version 2.0 (the "License");
4 : // you may not use this file except in compliance with the License.
5 : // You may obtain a copy of the License at
6 : //
7 : // http://www.apache.org/licenses/LICENSE-2.0
8 : //
9 : // Unless required by applicable law or agreed to in writing, software
10 : // distributed under the License is distributed on an "AS IS" BASIS,
11 : // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 : // See the License for the specific language governing permissions and
13 : // limitations under the License.
14 : //
15 : // Defines a PE-specific block-graph transform that adds imports to a given
16 : // module. Multiple libraries may be specified, and multiple functions per
17 : // library. Any imports that don't exist in the image will be added.
18 : //
19 : // Use is as follows:
20 : //
21 : // ImportedModule foo_dll("foo.dll");
22 : // size_t foo_foo_index = foo_dll.AddSymbol("foo");
23 : // size_t foo_bar_index = foo_dll.AddSymbol("bar");
24 : //
25 : // AddImportsTransform add_imports_transform;
26 : // add_imports_transform.AddModule(&foo_dll);
27 : // add_imports_transform.TransformBlockGraph(block_graph, dos_header_block);
28 : //
29 : // // Create a reference to function 'bar' in 'foo.dll'.
30 : // BlockGraph::Reference foo_bar_ref;
31 : // CHECK(foo_dll.GetSymbolReference(foo_bar_index, &foo_bar_ref));
32 : // some_block->SetReference(some_offset, foo_bar_ref);
33 : //
34 : // NOTE: The references provided by GetSymbolReference are only valid
35 : // immediately after they are constructed. If the the import directory
36 : // entries are changed between creating the reference and adding it to a
37 : // block, than it may have been invalidated.
38 :
39 : #ifndef SYZYGY_PE_TRANSFORMS_ADD_IMPORTS_TRANSFORM_H_
40 : #define SYZYGY_PE_TRANSFORMS_ADD_IMPORTS_TRANSFORM_H_
41 :
42 : #include <windows.h>
43 :
44 : #include "syzygy/block_graph/typed_block.h"
45 : #include "syzygy/block_graph/transforms/named_transform.h"
46 :
47 : namespace pe {
48 : namespace transforms {
49 :
50 : using block_graph::transforms::NamedBlockGraphTransformImpl;
51 :
52 : // A transform for adding imported modules/symbols to a given block-graph.
53 : class AddImportsTransform
54 : : public NamedBlockGraphTransformImpl<AddImportsTransform> {
55 : public:
56 : typedef block_graph::BlockGraph BlockGraph;
57 :
58 : // Some forward declares.
59 : struct ImportedModule;
60 :
61 : AddImportsTransform();
62 :
63 : // Adds the given module and symbols to the list of modules and symbols to
64 : // import.
65 E : void AddModule(ImportedModule* imported_module) {
66 E : DCHECK(imported_module != NULL);
67 E : imported_modules_.push_back(imported_module);
68 E : }
69 :
70 : // Performs the transform. Adds entries for any missing modules and symbols,
71 : // returning references to their entries via the ImportedModule structures.
72 : //
73 : // @param block_graph the BlockGraph to populate.
74 : // @param dos_header_block the block containing the module's DOS header.
75 : // @returns true on success, false otherwise.
76 : virtual bool TransformBlockGraph(
77 : BlockGraph* block_graph, BlockGraph::Block* dos_header_block) OVERRIDE;
78 :
79 : // @returns the number of imported modules that were added to the image.
80 E : size_t modules_added() const { return modules_added_; }
81 :
82 : // @returns the number of imported symbols that were added to the image.
83 E : size_t symbols_added() const { return symbols_added_; }
84 :
85 : // @returns a pointer to the Block containing the Image Import Descriptor.
86 E : BlockGraph::Block* image_import_descriptor_block() {
87 E : return image_import_descriptor_block_;
88 E : }
89 :
90 : // @returns a pointer to the Block containing the Import Address Table.
91 E : BlockGraph::Block* import_address_table_block() {
92 E : return import_address_table_block_;
93 E : }
94 :
95 : // The name of this transform.
96 : static const char kTransformName[];
97 :
98 : protected:
99 : // A collection of modules (and symbols from them) to be imported. This
100 : // must be populated prior to calling the transform.
101 : std::vector<ImportedModule*> imported_modules_;
102 :
103 : // Statistics regarding the completed transform.
104 : size_t modules_added_;
105 : size_t symbols_added_;
106 :
107 : // We cache the blocks containing the IDT and IAT.
108 : BlockGraph::Block* image_import_descriptor_block_;
109 : BlockGraph::Block* import_address_table_block_;
110 : };
111 :
112 : // Describes a list of symbols to be imported from a module.
113 : struct AddImportsTransform::ImportedModule {
114 : ImportedModule() { }
115 :
116 : // @param module_name the name of the module to import.
117 E : explicit ImportedModule(const base::StringPiece& module_name)
118 : : name_(module_name.begin(), module_name.end()) {
119 E : }
120 :
121 : // Accesses the name of the module.
122 : // @returns the name of the module to import.
123 E : const std::string& name() const { return name_; }
124 :
125 : // Adds a symbol to be imported, returning its index.
126 : size_t AddSymbol(const base::StringPiece& symbol_name);
127 :
128 : // Returns the number of symbols that are to be imported from this module.
129 E : size_t size() const { return symbols_.size(); }
130 :
131 : // Accesses the name of the index'th symbol.
132 : //
133 : // @param index the index of the symbol to fetch.
134 : // @returns the name of the index'th symbol.
135 E : const std::string& GetSymbolName(size_t index) const {
136 E : DCHECK_LT(index, symbols_.size());
137 E : return symbols_[index].name;
138 E : }
139 :
140 : // Gets an absolute reference to the IAT entry of the ith symbol. Returns
141 : // true on success, false if this was not possible. This will fail if the
142 : // AddImportsTransform has not successfully run on this ImportedModule
143 : // object.
144 : //
145 : // The returned reference is only valid while the import data directory is
146 : // not modified. Once added to a block, the imports may be further modified
147 : // and reference tracking will ensure things are kept up to date; until this
148 : // time @p abs_reference is left dangling.
149 : //
150 : // @param index the index of the symbol to fetch.
151 : // @param abs_reference the reference to populate.
152 : // @returns true on success, false otherwise.
153 : bool GetSymbolReference(size_t index,
154 : BlockGraph::Reference* abs_reference) const;
155 :
156 : private:
157 : // The AddImportsTransform is a friend so that it may directly set certain
158 : // output structures.
159 : friend AddImportsTransform;
160 :
161 : // Represents a symbol imported from this library. Currently this only
162 : // supports importing by name, but we could always extend this to handle
163 : // ordinals.
164 : struct Symbol {
165 : // The name of the symbol to import.
166 : std::string name;
167 : // The index of the imported symbol in the module's Import Name Table.
168 : size_t index;
169 : };
170 :
171 : // Used to indicate that an ImportedSymbol hasn't yet been looked up.
172 : static const size_t kInvalidIndex;
173 :
174 : // The name of the module to be imported.
175 : std::string name_;
176 :
177 : // The image import descriptor associated with this module. This will refer
178 : // to a block in the block-graph provided to the AddImportsTransform, assuming
179 : // successful completion.
180 : block_graph::TypedBlock<IMAGE_IMPORT_DESCRIPTOR> import_descriptor_;
181 :
182 : // The list of symbols to be imported from this module.
183 : std::vector<Symbol> symbols_;
184 : };
185 :
186 : } // namespace transforms
187 : } // namespace pe
188 :
189 : #endif // SYZYGY_PE_TRANSFORMS_ADD_IMPORTS_TRANSFORM_H_
|