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