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/pe_coff_add_imports_transform.h"
  16    :  
  17    :  namespace pe {
  18    :  namespace transforms {
  19    :  
  20    :  using block_graph::BlockGraph;
  21    :  
  22    :  size_t ImportedModule::AddSymbol(const base::StringPiece& symbol_name,
  23  E :                                   TransformMode mode) {
  24  E :    Symbol symbol = { symbol_name.as_string(),
  25  E :                      symbols_by_index_.size(),
  26  E :                      kInvalidImportIndex,
  27  E :                      mode };
  28  E :    std::pair<SymbolSet::iterator, bool> result = symbols_by_name_.insert(symbol);
  29    :  
  30    :    // We can safely cast away constness because we are not changing the key
  31    :    // portion of the Symbol (it's name). This is a bit of a hack that allows us
  32    :    // to use a std::set rather than a std::map and some hoop jumping.
  33  E :    Symbol* inserted_symbol = const_cast<Symbol*>(&(*result.first));
  34    :  
  35  E :    if (!result.second) {
  36    :      // Upgrade the mode to always-import if the symbol was previously inserted
  37    :      // as find-only.
  38    :      if (mode == ImportedModule::kAlwaysImport &&
  39  E :          inserted_symbol->mode == ImportedModule::kFindOnly) {
  40  E :        inserted_symbol->mode = ImportedModule::kAlwaysImport;
  41    :      }
  42  E :    } else {
  43    :      // This symbol was newly added. Insert it into the reverse lookup array.
  44  E :      DCHECK_EQ(symbol.symbol_index, inserted_symbol->symbol_index);
  45  E :      symbols_by_index_.push_back(inserted_symbol);
  46    :    }
  47    :  
  48    :    // Keep track of whether all symbols in this module are kFindOnly; if at
  49    :    // least one is not, the whole module is considered kAlwaysImport.
  50  E :    if (mode != ImportedModule::kFindOnly)
  51  E :      mode_ = ImportedModule::kAlwaysImport;
  52    :  
  53    :    // Return the index of the symbol.
  54  E :    return inserted_symbol->symbol_index;
  55  E :  }
  56    :  
  57    :  bool ImportedModule::GetSymbolReference(size_t index,
  58    :                                          BlockGraph::Reference* ref,
  59  E :                                          bool* is_ptr) const {
  60  E :    DCHECK_GT(symbols_by_index_.size(), index);
  61  E :    DCHECK(ref != NULL);
  62  E :    DCHECK(is_ptr != NULL);
  63    :  
  64  E :    Symbol* symbol = symbols_by_index_[index];
  65  E :    if (symbol->import_index == kInvalidImportIndex)
  66  i :      return false;
  67    :  
  68  E :    *ref = symbol->ref;
  69  E :    *is_ptr = symbol->is_ptr;
  70  E :    return true;
  71  E :  }
  72    :  
  73    :  void PECoffAddImportsTransform::UpdateModule(bool imported,
  74    :                                               bool added,
  75  E :                                               ImportedModule* imported_module) {
  76  E :    DCHECK(imported_module != NULL);
  77  E :    imported_module->imported_ = imported;
  78  E :    imported_module->added_ = added;
  79  E :  }
  80    :  
  81    :  void PECoffAddImportsTransform::UpdateModuleSymbolIndex(
  82    :      size_t index,
  83    :      size_t import_index,
  84    :      bool added,
  85  E :      ImportedModule* imported_module) {
  86  E :    DCHECK(imported_module != NULL);
  87  E :    DCHECK_GT(imported_module->symbols_by_index_.size(), index);
  88  E :    ImportedModule::Symbol* symbol = imported_module->symbols_by_index_[index];
  89  E :    symbol->import_index = import_index;
  90  E :    symbol->added = added;
  91  E :  }
  92    :  
  93    :  void PECoffAddImportsTransform::UpdateModuleSymbolReference(
  94    :      size_t index,
  95    :      BlockGraph::Reference ref,
  96    :      bool is_ptr,
  97  E :      ImportedModule* imported_module) {
  98  E :    DCHECK(imported_module != NULL);
  99  E :    DCHECK_GT(imported_module->symbols_by_index_.size(), index);
 100  E :    ImportedModule::Symbol* symbol = imported_module->symbols_by_index_[index];
 101  E :    symbol->ref = ref;
 102  E :    symbol->is_ptr = is_ptr;
 103  E :  }
 104    :  
 105    :  }  // namespace transforms
 106    :  }  // namespace pe