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