1 : // Copyright 2013 Google Inc. All Rights Reserved.
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 : #include "syzygy/pe/transforms/coff_add_imports_transform.h"
16 :
17 : #include <windows.h>
18 : #include <map>
19 :
20 : #include "base/strings/string_piece.h"
21 : #include "base/strings/string_util.h"
22 : #include "base/strings/stringprintf.h"
23 : #include "syzygy/common/align.h"
24 :
25 : namespace pe {
26 : namespace transforms {
27 : namespace {
28 :
29 : using block_graph::BlockGraph;
30 : using block_graph::ConstTypedBlock;
31 : using block_graph::TypedBlock;
32 :
33 : static const BlockGraph::Offset kInvalidCoffSymbol = -1;
34 :
35 : } // namespace
36 :
37 : const char CoffAddImportsTransform::kTransformName[] =
38 : "CoffAddImportsTransform";
39 :
40 : bool CoffAddImportsTransform::TransformBlockGraph(
41 : const TransformPolicyInterface* policy,
42 : BlockGraph* block_graph,
43 E : BlockGraph::Block* headers_block) {
44 E : DCHECK_NE(reinterpret_cast<TransformPolicyInterface*>(NULL), policy);
45 E : DCHECK_NE(reinterpret_cast<BlockGraph*>(NULL), block_graph);
46 E : DCHECK_NE(reinterpret_cast<BlockGraph::Block*>(NULL), headers_block);
47 E : DCHECK_EQ(BlockGraph::COFF_IMAGE, block_graph->image_format());
48 :
49 : // Get file header.
50 E : TypedBlock<IMAGE_FILE_HEADER> file_header;
51 E : if (!file_header.Init(0, headers_block)) {
52 i : LOG(ERROR) << "Unable to cast COFF file header.";
53 i : return false;
54 : }
55 :
56 : // Get symbol and string tables.
57 E : BlockGraph::Block* symbols_block = NULL;
58 E : BlockGraph::Block* strings_block = NULL;
59 : if (!FindCoffSpecialBlocks(block_graph,
60 E : NULL, &symbols_block, &strings_block)) {
61 i : LOG(ERROR) << "Block graph is missing some COFF special blocks. "
62 : << "Not a COFF block graph?";
63 i : return false;
64 : }
65 E : DCHECK(symbols_block != NULL);
66 E : DCHECK(strings_block != NULL);
67 :
68 E : TypedBlock<IMAGE_SYMBOL> symbols;
69 E : if (!symbols.Init(0, symbols_block)) {
70 i : LOG(ERROR) << "Unable to cast symbol table.";
71 i : return false;
72 : }
73 E : DCHECK_EQ(file_header->NumberOfSymbols, symbols.ElementCount());
74 :
75 : // Read existing symbols.
76 E : CoffSymbolNameOffsetMap known_names;
77 : if (!BuildCoffSymbolNameOffsetMap(symbols_block, strings_block,
78 E : &known_names)) {
79 i : return false;
80 : }
81 :
82 : // Handle symbols from each library.
83 E : CoffSymbolNameOffsetMap names_to_add;
84 E : size_t string_len_to_add = 0;
85 E : for (size_t i = 0; i < imported_modules_.size(); ++i) {
86 : if (!FindAndCollectSymbolsFromModule(file_header, known_names,
87 : imported_modules_[i],
88 E : &names_to_add, &string_len_to_add))
89 i : return false;
90 E : }
91 :
92 : // Add symbols if necessary.
93 E : if (names_to_add.size() > 0) {
94 E : size_t old_symbols_block_size = symbols_block->size();
95 :
96 : // Update symbol and string blocks.
97 : symbols_block->InsertData(old_symbols_block_size,
98 : names_to_add.size() * sizeof(IMAGE_SYMBOL),
99 E : true);
100 E : symbols_block->ResizeData(symbols_block->size());
101 E : if (!symbols.Init(0, symbols_block)) {
102 i : LOG(ERROR) << "Unable to cast symbol table.";
103 i : return false;
104 : }
105 :
106 E : size_t string_cursor = strings_block->size();
107 E : strings_block->InsertData(string_cursor, string_len_to_add, true);
108 E : strings_block->ResizeData(strings_block->size());
109 :
110 E : CoffSymbolNameOffsetMap::iterator to_add_it = names_to_add.begin();
111 E : for (; to_add_it != names_to_add.end(); ++to_add_it) {
112 E : DCHECK_GT(strings_block->size(), string_cursor);
113 E : DCHECK_EQ(1u, to_add_it->second.size());
114 :
115 E : BlockGraph::Offset offset = *to_add_it->second.begin();
116 E : DCHECK_LE(old_symbols_block_size, static_cast<size_t>(offset));
117 E : size_t index = offset / sizeof(IMAGE_SYMBOL);
118 :
119 : std::memcpy(strings_block->GetMutableData() + string_cursor,
120 E : to_add_it->first.c_str(), to_add_it->first.size() + 1);
121 E : IMAGE_SYMBOL* symbol = &symbols[index];
122 E : symbol->N.Name.Short = 0;
123 E : symbol->N.Name.Long = string_cursor;
124 E : symbol->Type = IMAGE_SYM_DTYPE_FUNCTION << 4;
125 E : symbol->StorageClass = IMAGE_SYM_CLASS_EXTERNAL;
126 E : string_cursor += to_add_it->first.size() + 1;
127 E : }
128 E : DCHECK_EQ(strings_block->size(), string_cursor);
129 :
130 : // Update the file header.
131 E : file_header->NumberOfSymbols = symbols.ElementCount();
132 :
133 : // Update string table size.
134 E : TypedBlock<uint32> strings_size;
135 E : if (!strings_size.Init(0, strings_block)) {
136 i : LOG(ERROR) << "Unable to cast string table size prefix.";
137 i : return false;
138 : }
139 E : *strings_size = string_cursor;
140 : }
141 :
142 : // Update import module symbols.
143 E : for (size_t i = 0; i < imported_modules_.size(); ++i)
144 E : UpdateModuleReferences(symbols_block, imported_modules_[i]);
145 :
146 E : return true;
147 E : }
148 :
149 : bool CoffAddImportsTransform::FindAndCollectSymbolsFromModule(
150 : const TypedBlock<IMAGE_FILE_HEADER>& file_header,
151 : const CoffSymbolNameOffsetMap& known_names,
152 : ImportedModule* module,
153 : CoffSymbolNameOffsetMap* names_to_add,
154 E : size_t* string_len_to_add) {
155 E : DCHECK(module != NULL);
156 E : DCHECK(names_to_add != NULL);
157 E : DCHECK(string_len_to_add != NULL);
158 :
159 E : for (size_t i = 0; i < module->size(); ++i) {
160 E : BlockGraph::Offset symbol_import_offset = kInvalidCoffSymbol;
161 E : bool symbol_imported = false;
162 E : bool symbol_added = false;
163 :
164 E : std::string name(module->GetSymbolName(i));
165 E : CoffSymbolNameOffsetMap::const_iterator it = known_names.find(name);
166 E : if (it != known_names.end()) {
167 : // This symbol is already defined. Simply grab its offset which we can
168 : // use to draw a reference to it later. We grab the offset of its first
169 : // definition if it is multiply defined.
170 E : symbol_import_offset = *it->second.begin();
171 E : symbol_imported = true;
172 E : } else if (module->GetSymbolMode(i) == ImportedModule::kAlwaysImport) {
173 : // The symbol is not defined, but requested to be. Create it.
174 E : size_t new_index = file_header->NumberOfSymbols + names_to_add->size();
175 E : symbol_import_offset = new_index * sizeof(IMAGE_SYMBOL);
176 : CoffSymbolNameOffsetMap::iterator add_it =
177 E : names_to_add->insert(std::make_pair(name, CoffSymbolOffsets())).first;
178 :
179 E : if (add_it->second.empty()) {
180 : // When adding the symbol for the first time ensure to reserve room in
181 : // the string table.
182 E : add_it->second.insert(symbol_import_offset);
183 E : *string_len_to_add += name.size() + 1;
184 E : } else {
185 : // If this symbols is already to be added then reuse its offset.
186 i : symbol_import_offset = *add_it->second.begin();
187 : }
188 :
189 E : symbol_imported = true;
190 E : symbol_added = true;
191 : }
192 :
193 E : UpdateModuleSymbolInfo(i, symbol_imported, symbol_added, module);
194 E : symbols_added_ += symbol_added;
195 E : if (symbol_imported) {
196 : module_symbol_offset_map_.insert(
197 E : std::make_pair(std::make_pair(module, i), symbol_import_offset));
198 : }
199 E : }
200 :
201 : // All modules are considered imported in a COFF file, and none is ever
202 : // added by the transform.
203 E : UpdateModule(true, false, module);
204 :
205 E : return true;
206 E : }
207 :
208 : void CoffAddImportsTransform::UpdateModuleReferences(
209 : BlockGraph::Block* symbols_block,
210 E : ImportedModule* module) {
211 E : for (size_t i = 0; i < module->size(); ++i) {
212 : ModuleSymbolOffsetMap::const_iterator offset_it =
213 E : module_symbol_offset_map_.find(std::make_pair(module, i));
214 E : if (offset_it == module_symbol_offset_map_.end())
215 E : continue;
216 E : BlockGraph::Offset import_offset = offset_it->second;
217 E : DCHECK_NE(kInvalidCoffSymbol, import_offset);
218 : BlockGraph::Reference ref(BlockGraph::RELOC_ABSOLUTE_REF, sizeof(uint32),
219 E : symbols_block, import_offset, import_offset);
220 E : UpdateModuleSymbolReference(i, ref, false, module);
221 E : }
222 E : }
223 :
224 : } // namespace transforms
225 : } // namespace pe
|